import matplotlib.pyplot as plt
import numpy as np
import os
from scipy.stats import ranksums

# ============================================================
# GRAPH CONFIGURATION - THIN LINES & SMALL FONT
# ============================================================
plt.style.use('seaborn-v0_8-whitegrid')
fontsize = 8
linewidth = 0.8

# palette accessible daltonien
colors = ['#4477AA', '#CCBB44', '#AA3377', '#EE6677', '#228833', '#66CCEE']

params = {
    'font.size': fontsize,
    'axes.labelsize': fontsize + 1,
    'axes.titlesize': fontsize + 2,
    'axes.linewidth': linewidth,
    'lines.linewidth': linewidth,
    'lines.markersize': 2,
    'xtick.labelsize': fontsize,
    'ytick.labelsize': fontsize,
    'xtick.major.width': linewidth,
    'ytick.major.width': linewidth,
    'xtick.major.size': 3,
    'ytick.major.size': 3,
    'figure.dpi': 300,
    'svg.fonttype': 'path',
    'font.family': 'sans-serif',
    'mathtext.default': 'regular',
    'legend.fontsize': fontsize - 1,
    'legend.frameon': True,
    'legend.framealpha': 0.8
}
plt.rcParams.update(params)

# ============================================================
# SYSTEMS / PATHS
# ============================================================
systems = {
    'CNP007': ('CNP0075206.3/CNP007', colors[0]),
    'CNP018': ('CNP0188167.1/CNP018', colors[1]),
    'CNP019': ('CNP0196376.4/CNP019', colors[2])
}

# ============================================================
# FONCTION D'EXPORT MULTI-FORMAT
# ============================================================
def save_plot_multiformat(fig, base_name):
    """Sauvegarde le graphique en SVG et PNG"""
    # SVG (qualité vectorielle)
    svg_name = f"{base_name}.svg"
    fig.savefig(svg_name, format='svg', bbox_inches='tight', dpi=300)
    
    # PNG (pour publication/web)
    png_name = f"{base_name}.png"
    fig.savefig(png_name, format='png', bbox_inches='tight', dpi=300)
    
    return svg_name, png_name

# ============================================================
# FONCTIONS DE LECTURE .xvg AVEC DÉTECTION D'UNITÉS
# ============================================================
def detect_time_units(times):
    """
    Détecte si les temps sont en ps ou ns basé sur les valeurs
    Retourne: 'ps' ou 'ns'
    """
    if len(times) == 0:
        return 'ns'
    
    # Prendre le dernier point de temps pour la détection
    max_time = np.max(times)
    
    # Si le temps max a 6 chiffres ou plus -> probablement en ps
    if max_time >= 100000:  # 100 ns = 100000 ps
        return 'ps'
    # Si le temps max a 3-4 chiffres -> probablement en ns  
    elif max_time <= 10000:  # 10000 ns = 10 µs (peu probable)
        return 'ns'
    else:
        # Valeur intermédiaire, vérifier le pattern
        if max_time > 1000 and max_time < 100000:
            # Regarder plusieurs points pour décider
            sample_times = times[:min(100, len(times))]
            decimal_parts = sample_times - np.floor(sample_times)
            # Si beaucoup de décimales -> probablement ns
            if np.mean(decimal_parts > 0) > 0.5:
                return 'ns'
            else:
                return 'ps'
        return 'ns'

def read_xvg_file(filename):
    """Lit un .xvg et renvoie (time, value) avec conversion automatique ps->ns"""
    times = []
    vals = []
    if not os.path.exists(filename):
        return np.array([]), np.array([])
    
    with open(filename, 'r') as f:
        for line in f:
            if line.startswith('#') or line.startswith('@') or not line.strip():
                continue
            parts = line.strip().split()
            try:
                if len(parts) >= 2:
                    t = float(parts[0])
                    v = float(parts[1])
                    times.append(t)
                    vals.append(v)
                elif len(parts) == 1:
                    v = float(parts[0])
                    vals.append(v)
            except ValueError:
                continue
    
    if len(times) == 0 and len(vals) > 0:
        times = np.arange(len(vals))
    
    times = np.array(times)
    vals = np.array(vals)
    
    # Détection automatique des unités et conversion si nécessaire
    if len(times) > 0:
        time_units = detect_time_units(times)
        if time_units == 'ps':
            print(f"   ⚠️  Conversion ps→ns détectée pour {os.path.basename(filename)} "
                  f"(t_max={np.max(times):.1f} ps → {np.max(times)/1000:.1f} ns)")
            times = times / 1000.0  # Conversion ps -> ns
        else:
            print(f"   ✓ Unités ns confirmées pour {os.path.basename(filename)} "
                  f"(t_max={np.max(times):.1f} ns)")
    
    return times, vals

def read_pca_columns_only(filename):
    """
    Lit seulement les colonnes PC1 et PC2 (sans temps)
    Retourne: (pc1, pc2)
    """
    pc1 = []
    pc2 = []
    if not os.path.exists(filename):
        return np.array([]), np.array([])
    
    with open(filename, 'r') as f:
        for line in f:
            if line.startswith('#') or line.startswith('@') or not line.strip():
                continue
            parts = line.strip().split()
            try:
                if len(parts) >= 2:
                    pc1.append(float(parts[0]))
                    pc2.append(float(parts[1]))
            except ValueError:
                continue
    
    return np.array(pc1), np.array(pc2)

def read_eigenval_file(filename):
    """Lit un fichier eigenval.xvg et renvoie les valeurs propres"""
    eigenvals = []
    if not os.path.exists(filename):
        return np.array([])
    with open(filename, 'r') as f:
        for line in f:
            if line.startswith('#') or line.startswith('@') or not line.strip():
                continue
            parts = line.strip().split()
            try:
                if len(parts) >= 2:
                    eigenvals.append(float(parts[1]))
                elif len(parts) == 1:
                    eigenvals.append(float(parts[0]))
            except ValueError:
                continue
    return np.array(eigenvals)

# ============================================================
# ALIGNEMENT RMSD <-> ENERGY (interpolation)
# ============================================================
def align_rmsd_energy_data(time_rmsd, rmsd, time_energy, energy):
    """Interpole pour aligner RMSD et Energy sur une grille de temps commune"""
    if len(time_rmsd) == 0 or len(time_energy) == 0:
        return np.array([]), np.array([])
    
    if len(time_energy) >= len(time_rmsd):
        target_time = time_energy
        rmsd_aligned = np.interp(target_time, time_rmsd, rmsd)
        energy_aligned = energy
    else:
        target_time = time_rmsd
        rmsd_aligned = rmsd
        energy_aligned = np.interp(target_time, time_energy, energy)
    
    return rmsd_aligned, energy_aligned

# ============================================================
# PLOTTING HELPERS
# ============================================================
def style_axes(ax):
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.grid(True, alpha=0.3, linewidth=0.5)

# ============================================================
# FONCTIONS POUR TESTS STATISTIQUES
# ============================================================
def perform_wilcoxon_tests(data_dict):
    """
    Effectue des tests de Wilcoxon rank-sum entre toutes les paires de systèmes
    Retourne: dict avec p-values
    """
    systems_names = list(data_dict.keys())
    n_systems = len(systems_names)
    results = {}
    
    for i in range(n_systems):
        for j in range(i+1, n_systems):
            sys1 = systems_names[i]
            sys2 = systems_names[j]
            
            data1 = data_dict[sys1]
            data2 = data_dict[sys2]
            
            if len(data1) > 0 and len(data2) > 0:
                stat, p_value = ranksums(data1, data2)
                results[f"{sys1}_vs_{sys2}"] = p_value
            else:
                results[f"{sys1}_vs_{sys2}"] = np.nan
    
    return results

def add_statistical_annotation(ax, data_dict, y_positions, y_range):
    """
    Ajoute les annotations statistiques sur le plot
    """
    results = perform_wilcoxon_tests(data_dict)
    systems_names = list(data_dict.keys())
    
    # Déterminer la hauteur des lignes d'annotation
    y_range_total = y_range[1] - y_range[0]
    line_height = y_range_total * 0.05
    text_height = y_range_total * 0.08
    
    current_y = y_range[1] + line_height
    
    for i, (comparison, p_value) in enumerate(results.items()):
        if np.isnan(p_value):
            continue
            
        sys1, sys2 = comparison.split('_vs_')
        idx1 = systems_names.index(sys1)
        idx2 = systems_names.index(sys2)
        
        # Dessiner la ligne horizontale
        x1 = idx1 + 1
        x2 = idx2 + 1
        ax.plot([x1, x1, x2, x2], [current_y, current_y + line_height, current_y + line_height, current_y], 
                'k-', linewidth=0.8)
        
        # Ajouter la significativité
        if p_value < 0.001:
            sig_symbol = '***'
        elif p_value < 0.01:
            sig_symbol = '**'
        elif p_value < 0.05:
            sig_symbol = '*'
        else:
            sig_symbol = 'ns'
        
        # Afficher la p-value avec format scientifique si très petite
        if p_value < 0.001:
            p_text = f"p = {p_value:.2e}"
        else:
            p_text = f"p = {p_value:.3f}"
        
        ax.text((x1 + x2) / 2, current_y + text_height, 
                f"{sig_symbol}\n{p_text}", 
                ha='center', va='bottom', fontsize=fontsize-2, 
                linespacing=0.8)
        
        current_y += text_height * 2.5
    
    # Ajuster les limites pour accommoder les annotations
    ax.set_ylim([y_range[0], current_y + text_height])

# ============================================================
# ANALYSES SECONDAIRES
# ============================================================
def plot_rmsd_protein_boxplot():
    print("Generating Protein RMSD boxplot...")
    fig, ax = plt.subplots(figsize=(6,4))
    data_dict = {}
    labels = []
    
    for sys_name, (prefix, color) in systems.items():
        t, v = read_xvg_file(f"{prefix}_rmsd_prot.xvg")
        if len(v) > 0:
            data_dict[sys_name] = v
            labels.append(sys_name)
    
    if data_dict:
        data = [data_dict[label] for label in labels]
        bp = ax.boxplot(data, labels=labels, patch_artist=True)
        
        for i, box in enumerate(bp['boxes']):
            box.set_facecolor(colors[i])
            box.set_alpha(0.7)
        for median in bp['medians']:
            median.set_color('black')
            median.set_linewidth(1)
        
        # Ajouter les tests statistiques
        y_range = ax.get_ylim()
        add_statistical_annotation(ax, data_dict, range(1, len(labels)+1), y_range)
        
        # Afficher les résultats dans la console
        print("\n📊 Wilcoxon Rank-Sum Test Results - Protein RMSD:")
        results = perform_wilcoxon_tests(data_dict)
        for comparison, p_value in results.items():
            if not np.isnan(p_value):
                sig_status = "SIGNIFICANT" if p_value < 0.05 else "not significant"
                print(f"   {comparison}: p = {p_value:.4f} ({sig_status})")
    
    ax.set_ylabel(r'RMSD ($\AA$)')
    ax.set_title('Protein RMSD Distribution with Wilcoxon Tests')
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'RMSD_Protein_Boxplot')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_rmsd_ligand_boxplot():
    print("Generating Ligand RMSD boxplot...")
    fig, ax = plt.subplots(figsize=(6,4))
    data_dict = {}
    labels = []
    
    for sys_name, (prefix, color) in systems.items():
        t, v = read_xvg_file(f"{prefix}_rmsd_lig.xvg")
        if len(v) > 0:
            data_dict[sys_name] = v
            labels.append(sys_name)
    
    if data_dict:
        data = [data_dict[label] for label in labels]
        bp = ax.boxplot(data, labels=labels, patch_artist=True)
        
        for i, box in enumerate(bp['boxes']):
            box.set_facecolor(colors[i])
            box.set_alpha(0.7)
        for median in bp['medians']:
            median.set_color('black')
            median.set_linewidth(1)
        
        # Ajouter les tests statistiques
        y_range = ax.get_ylim()
        add_statistical_annotation(ax, data_dict, range(1, len(labels)+1), y_range)
        
        # Afficher les résultats dans la console
        print("\n📊 Wilcoxon Rank-Sum Test Results - Ligand RMSD:")
        results = perform_wilcoxon_tests(data_dict)
        for comparison, p_value in results.items():
            if not np.isnan(p_value):
                sig_status = "SIGNIFICANT" if p_value < 0.05 else "not significant"
                print(f"   {comparison}: p = {p_value:.4f} ({sig_status})")
    
    ax.set_ylabel(r'RMSD ($\AA$)')
    ax.set_title('Ligand RMSD Distribution with Wilcoxon Tests')
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'RMSD_Ligand_Boxplot')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_rmsd_vs_energy_scatter():
    print("Generating RMSD vs Energy scatter plot...")
    fig, ax = plt.subplots(figsize=(8,6))
    for sys_name, (prefix, color) in systems.items():
        t_r, rmsd = read_xvg_file(f"{prefix}_rmsd_prot.xvg")
        t_e, energy = read_xvg_file(f"{prefix}_energy.xvg")
        if len(rmsd)>0 and len(energy)>0:
            rmsd_a, energy_a = align_rmsd_energy_data(t_r, rmsd, t_e, energy)
            if len(rmsd_a)>0:
                ax.scatter(rmsd_a, energy_a, color=color, label=sys_name, 
                          alpha=0.6, s=20, edgecolors='black', linewidth=0.2)
    ax.set_xlabel(r'Protein RMSD ($\AA$)')
    ax.set_ylabel('Energy (kJ/mol)')
    ax.set_title('Protein RMSD vs Potential Energy')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'RMSD_vs_Energy_Scatter')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_pca_colored_by_rmsd_time():
    """
    PCA colorée par le temps EXTRACT du fichier RMSD
    Utilise le temps réel de la dynamique MD
    """
    print("Generating PCA colored by RMSD-extracted time...")
    
    for sys_name, (prefix, color) in systems.items():
        # 1. Lire RMSD pour obtenir le TEMPS RÉEL
        time_rmsd, rmsd_values = read_xvg_file(f"{prefix}_rmsd_prot.xvg")
        
        # 2. Lire PCA (sans temps)
        pc1, pc2 = read_pca_columns_only(f"{prefix}_pca.xvg")
        
        if len(pc1) == 0 or len(time_rmsd) == 0:
            print(f"  [WARN] {sys_name}: PCA or RMSD data missing -> skipping")
            continue
        
        # 3. VÉRIFICATION : même nombre de points ?
        if len(pc1) != len(time_rmsd):
            print(f"  ⚠️  {sys_name}: PCA frames ({len(pc1)}) ≠ RMSD frames ({len(time_rmsd)})")
            # Prendre le minimum des deux
            n_points = min(len(pc1), len(time_rmsd))
            time_rmsd = time_rmsd[:n_points]
            pc1 = pc1[:n_points]
            pc2 = pc2[:n_points]
            rmsd_values = rmsd_values[:n_points]
            print(f"  → Utilisation de {n_points} points communs")
        
        if len(time_rmsd) == 0:
            print(f"  [WARN] {sys_name}: No data available -> skipping")
            continue
        
        # 4. Créer le plot PCA coloré par temps uniquement
        fig, ax = plt.subplots(figsize=(8, 6))
        
        # PCA colorée par TEMPS
        sc = ax.scatter(pc1, pc2, c=time_rmsd, cmap='viridis', 
                       alpha=0.7, s=30, edgecolors='black', linewidth=0.3)
        cbar = plt.colorbar(sc, ax=ax)
        cbar.set_label('Time (ns)')
        ax.set_xlabel('PC1')
        ax.set_ylabel('PC2')
        ax.set_title(f'{sys_name} - PCA colored by Time')
        
        # Style
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.grid(True, alpha=0.3, linewidth=0.5)
        
        plt.tight_layout()
        
        # Sauvegarde en double format
        base_name = f"PCA_TimeColored_{sys_name}"
        svg_file, png_file = save_plot_multiformat(fig, base_name)
        plt.close()
        
        print(f"  ✓ {sys_name}: PCA time-colored saved (SVG: {svg_file}, PNG: {png_file})")

def plot_eigenvalues():
    """Plot des valeurs propres pour chaque système"""
    print("Generating PCA eigenvalues plot...")
    
    fig, ax = plt.subplots(figsize=(8,4))
    
    for sys_name, (prefix, color) in systems.items():
        eigenvals = read_eigenval_file(f"{prefix}_eigenval.xvg")
        if len(eigenvals) > 0:
            # Limiter à 20 premiers modes pour lisibilité
            n_modes = min(20, len(eigenvals))
            ax.plot(range(1, n_modes+1), eigenvals[:n_modes], 
                   color=color, label=sys_name, linewidth=linewidth, marker='o', markersize=3)
    
    ax.set_xlabel('Eigenvector Index')
    ax.set_ylabel('Eigenvalue (nm²)')
    ax.set_title('PCA Eigenvalues (First 20 modes)')
    ax.legend(loc='best', framealpha=0.8)
    ax.set_yscale('log')
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'PCA_Eigenvalues')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_cumulative_variance():
    """Plot de la variance cumulée pour l'analyse PCA"""
    print("Generating PCA cumulative variance plot...")
    
    fig, ax = plt.subplots(figsize=(8,4))
    
    for sys_name, (prefix, color) in systems.items():
        eigenvals = read_eigenval_file(f"{prefix}_eigenval.xvg")
        if len(eigenvals) > 0:
            # Calculer variance cumulée
            total_var = np.sum(eigenvals)
            cumvar = np.cumsum(eigenvals) / total_var * 100
            
            # Limiter à 20 premiers modes
            n_modes = min(20, len(cumvar))
            ax.plot(range(1, n_modes+1), cumvar[:n_modes], 
                   color=color, label=sys_name, linewidth=linewidth, marker='o', markersize=3)
    
    ax.set_xlabel('Number of Principal Components')
    ax.set_ylabel('Cumulative Variance Explained (%)')
    ax.set_title('PCA Cumulative Variance')
    ax.legend(loc='best', framealpha=0.8)
    ax.axhline(y=80, color='gray', linestyle='--', linewidth=0.8, alpha=0.5)
    ax.axhline(y=95, color='gray', linestyle='--', linewidth=0.8, alpha=0.5)
    ax.set_ylim([0, 105])
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'PCA_Cumulative_Variance')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

# ============================================================
# MAIN
# ============================================================
def main():
    print("=" * 60)
    print("🔬 SCRIPT D'ANALYSES SECONDAIRES MD")
    print("=" * 60)
    print(f"📊 Systèmes: {', '.join(systems.keys())}")
    print("📈 Utilisation de toutes les données disponibles")
    print("📊 Tests statistiques: Wilcoxon Rank-Sum (Mann-Whitney U)")
    print("=" * 60)
    
    # Boxplots avec tests statistiques
    print("\n📦 Génération des boxplots avec tests de Wilcoxon...")
    plot_rmsd_protein_boxplot()
    plot_rmsd_ligand_boxplot()
    
    # RMSD vs Energy
    print("\n🔗 Génération de la corrélation RMSD-Energy...")
    plot_rmsd_vs_energy_scatter()
    
    # PCA analyses avec temps extrait du RMSD
    print("\n🎨 Génération de l'analyse PCA temporelle...")
    plot_pca_colored_by_rmsd_time()
    
    # Eigenvalues analysis
    print("\n📈 Génération des analyses des valeurs propres...")
    plot_eigenvalues()
    plot_cumulative_variance()
    
    print("\n" + "=" * 60)
    print("✅ Toutes les analyses secondaires terminées!")
    print("=" * 60)
    print("\nFichiers générés:")
    print("  - RMSD_Protein_Boxplot.svg & .png (avec tests statistiques)")
    print("  - RMSD_Ligand_Boxplot.svg & .png (avec tests statistiques)")
    print("  - RMSD_vs_Energy_Scatter.svg & .png")
    print("  - PCA_TimeColored_CNP007.svg & .png")
    print("  - PCA_TimeColored_CNP018.svg & .png")
    print("  - PCA_TimeColored_CNP019.svg & .png")
    print("  - PCA_Eigenvalues.svg & .png")
    print("  - PCA_Cumulative_Variance.svg & .png")
    print("\nLégende des tests statistiques:")
    print("  *** : p < 0.001")
    print("  **  : p < 0.01")
    print("  *   : p < 0.05")
    print("  ns  : non significatif")

if __name__ == "__main__":
    main()