import matplotlib.pyplot as plt
import numpy as np
import os

# ============================================================
# 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_rmsf_file(filename):
    """Lecteur RMSF (residue, rmsf)"""
    residues = []
    rmsf = []
    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()
            if len(parts) >= 2:
                try:
                    residues.append(float(parts[0]))
                    rmsf.append(float(parts[1]))
                except ValueError:
                    continue
    return np.array(residues), np.array(rmsf)

def read_pca_file(filename):
    """Lecture PCA simple (pc1, pc2) - sans temps"""
    if not os.path.exists(filename):
        return np.array([]), np.array([])
    pc1 = []
    pc2 = []
    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) >= 3:
                    a = float(parts[1])
                    b = float(parts[2])
                elif len(parts) >= 2:
                    a = float(parts[0])
                    b = float(parts[1])
                else:
                    continue
                pc1.append(a)
                pc2.append(b)
            except ValueError:
                continue
    return np.array(pc1), np.array(pc2)

# ============================================================
# 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)

# ============================================================
# PLOTS PRIMAIRES
# ============================================================
def plot_rmsd_protein_grouped():
    print("Generating Protein RMSD grouped plot...")
    fig, ax = plt.subplots(figsize=(8,4))
    for sys_name, (prefix, color) in systems.items():
        t, v = read_xvg_file(f"{prefix}_rmsd_prot.xvg")
        if len(t)>0:
            ax.plot(t, v, color=color, label=sys_name, linewidth=linewidth)
    ax.set_xlabel('Time (ns)')
    ax.set_ylabel(r'RMSD ($\AA$)')
    ax.set_title('Protein RMSD')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'RMSD_Protein_Grouped')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_rmsd_ligand_grouped():
    print("Generating Ligand RMSD grouped plot...")
    fig, ax = plt.subplots(figsize=(8,4))
    for sys_name, (prefix, color) in systems.items():
        t, v = read_xvg_file(f"{prefix}_rmsd_lig.xvg")
        if len(t)>0:
            ax.plot(t, v, color=color, label=sys_name, linewidth=linewidth)
    ax.set_xlabel('Time (ns)')
    ax.set_ylabel(r'RMSD ($\AA$)')
    ax.set_title('Ligand RMSD')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'RMSD_Ligand_Grouped')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_rmsf_grouped():
    print("Generating RMSF grouped plot...")
    fig, ax = plt.subplots(figsize=(8,4))
    for sys_name, (prefix, color) in systems.items():
        res, v = read_rmsf_file(f"{prefix}_rmsf.xvg")
        if len(res)>0:
            ax.plot(res, v, color=color, label=sys_name, linewidth=linewidth)
    ax.set_xlabel('Residue Number')
    ax.set_ylabel(r'RMSF ($\AA$)')
    ax.set_title('Root Mean Square Fluctuation')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'RMSF_Grouped')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_rg_grouped():
    """Create grouped Radius of Gyration plot for 3 systems"""
    print("Generating Radius of Gyration grouped plot...")

    fig, ax = plt.subplots(figsize=(8, 4))

    for system_name, (prefix, color) in systems.items():
        time, rg = read_xvg_file(f"{prefix}_rg.xvg")
        if len(time) > 0:
            ax.plot(time, rg, color=color, label=system_name, linewidth=linewidth)

    ax.set_xlabel('Time (ns)')
    ax.set_ylabel(r'Radius of Gyration ($\AA$)')
    ax.set_title('Radius of Gyration')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'Rg_Grouped')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_hbonds_grouped():
    """Create grouped Hydrogen Bonds plot for 3 systems"""
    print("Generating Hydrogen Bonds grouped plot...")
    
    fig, ax = plt.subplots(figsize=(8, 4))
    
    for system_name, (prefix, color) in systems.items():
        time, hb = read_xvg_file(f"{prefix}_hbonds.xvg")
        if len(time) > 0:
            ax.plot(time, hb, color=color, label=system_name, linewidth=linewidth)
    
    ax.set_xlabel('Time (ns)')
    ax.set_ylabel('Number of H-bonds')
    ax.set_title('Protein-Ligand Hydrogen Bonds')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'HBonds_Grouped')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_energy_grouped():
    print("Generating Energy grouped plot...")
    fig, ax = plt.subplots(figsize=(8, 4))
    for system_name, (prefix, color) in systems.items():
        time, energy = read_xvg_file(f"{prefix}_energy.xvg")
        if len(time) > 0:
            ax.plot(time, energy, color=color, label=system_name, linewidth=linewidth)
    ax.set_xlabel('Time (ns)')
    ax.set_ylabel('Energy (kJ/mol)')
    ax.set_title('Potential Energy')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'Energy_Grouped')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_pca_grouped():
    print("Generating PCA grouped plot...")
    fig, ax = plt.subplots(figsize=(6,5))
    for sys_name, (prefix, color) in systems.items():
        pc1, pc2 = read_pca_file(f"{prefix}_pca.xvg")
        if len(pc1)>0:
            ax.scatter(pc1, pc2, color=color, label=sys_name, alpha=0.6, s=10,
                      edgecolors='black', linewidth=0.2)
    ax.set_xlabel('PC1')
    ax.set_ylabel('PC2')
    ax.set_title('Principal Component Analysis')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'PCA_Grouped')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

def plot_sasa_grouped():
    """Create grouped SASA plot for 3 systems"""
    print("Generating SASA grouped plot...")
    
    fig, ax = plt.subplots(figsize=(8, 4))
    
    for system_name, (prefix, color) in systems.items():
        time, sasa = read_xvg_file(f"{prefix}_sasa.xvg")
        if len(time) > 0:
            ax.plot(time, sasa, color=color, label=system_name, linewidth=linewidth)
    
    ax.set_xlabel('Time (ns)')
    ax.set_ylabel(r'SASA ($nm^2$)')
    ax.set_title('Solvent Accessible Surface Area')
    ax.legend(loc='best', framealpha=0.8)
    style_axes(ax)
    plt.tight_layout()
    
    svg_file, png_file = save_plot_multiformat(fig, 'SASA_Grouped')
    plt.close()
    print(f"✓ Saved: {svg_file}")
    print(f"✓ Saved: {png_file}")

# ============================================================
# MAIN
# ============================================================
def main():
    print("=" * 60)
    print("🚀 SCRIPT D'ANALYSES PRIMAIRES MD")
    print("=" * 60)
    print(f"📊 Systèmes: {', '.join(systems.keys())}")
    print("📈 Utilisation de toutes les données disponibles")
    print("=" * 60)
    
    plot_rmsd_protein_grouped()
    plot_rmsd_ligand_grouped()
    plot_rmsf_grouped()
    plot_rg_grouped()
    plot_sasa_grouped()
    plot_hbonds_grouped()
    plot_energy_grouped()
    plot_pca_grouped()
    
    print("=" * 60)
    print("✅ Toutes les analyses primaires terminées!")
    print("=" * 60)
    print("\nFichiers générés:")
    print("  - RMSD_Protein_Grouped.svg & .png")
    print("  - RMSD_Ligand_Grouped.svg & .png")
    print("  - RMSF_Grouped.svg & .png")
    print("  - Rg_Grouped.svg & .png")
    print("  - SASA_Grouped.svg & .png")
    print("  - HBonds_Grouped.svg & .png")
    print("  - Energy_Grouped.svg & .png")
    print("  - PCA_Grouped.svg & .png")

if __name__ == "__main__":
    main()