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

labels = [
    "Norm (Normalization)",
    "nv (Vesicle concencentration x10^6 (nm^-3))",
    "Rm (Radius of the vesicle (nm))",
    "Z (Radius polydispersity)",
    "n_TR (Tris fraction )",
    "d_TR (Tris width (nm))",
    "s_TR (Tris position (nm))",
    "d_Chol (CholCH3 position (nm))",
    "s_Chol (CholCH3 width (nm))",
    "d_PCN (PCN position (nm))",
    "s_PCN (PCN width (nm))",
    "d_CG (CG position (nm))",
    "s_CG (CG width (nm))",
    "A_L (Area per lipid (nm²))",
    "s_D_C (HC excess std (nm))",
    "s_CH2 (HC smearing (nm))",
    "d_CH (CH position (nm))",
    "s_CH (CH width (nm))",
    "s_CH3 (CH3 width (nm))",
    "r_PCN (V_PCN/V_HL)",
    "r_CG (V_GG/V_HL)",
    "r12 (V_CH/V_CH2)",
    "r32 (V_CH3/V_CH2)",
    "T (Temperature (°C) )",
    "V_BW (Bound-water volume (nm³))",
    "Con (Constant (nm⁻¹))"
]


def print_x2_stats(dir, threshold=1.0):
    stats = get_x2_stats(dir, threshold=threshold)
    thres_str = str(int(100*threshold))

    if threshold < 1:
        print("  x2 (best %s%%)" % thres_str)
    else:
        print("  x2 ")

    print(" min: %f" % stats["min"])
    print(" max: %f" % stats["max"])
    print("mean: %f" % stats["mean"])
    print(" std: %f" % stats["std"])

def std_vs_threshold(dir, obs):
    x2 = np.genfromtxt(os.path.join(dir, "Results_collection-X2.dat"))
    if obs == "x2":
        d = x2
        suptitle= "x2"
    else:
        d = np.genfromtxt(os.path.join(dir, "Results_collection.dat"))[:,obs]
        suptitle = labels[obs]

    x2_min = x2.min()
    x2_p2p = x2.max() - x2.min()

    d_first_decile = d[x2 < x2_min + 0.1*x2_p2p]

    thres = np.linspace(0.01, 1, 100)

    std = np.zeros_like(thres)
    card = np.zeros_like(thres)

    for (i,t) in enumerate(thres):
        d_filtered = d[x2 < x2_min + t*x2_p2p]
        std[i] = d_filtered.std()
        card[i] = len(d_filtered)

    fig, axes = plt.subplots(1, 2)

    plt.suptitle(suptitle)

    ax1 = axes[0]
    ax3 = axes[1]

    ax1.set_xlabel("threshold")
    ax1.set_ylabel("stddev", color="blue")
    ax1.tick_params(axis='y', labelcolor="blue")

    ax2 = ax1.twinx()
    ax2.set_ylabel("cardinality", color="red")
    ax2.tick_params(axis='y', labelcolor="red")

    ax3.set_ylabel("all samples", color="blue")
    ax3.tick_params(axis='y', labelcolor="blue")

    ax4 = ax3.twinx()
    ax4.set_ylabel("1st decile", color="red")
    ax4.tick_params(axis='y', labelcolor="red")

    ax1.plot(thres, std, color="blue")
    ax2.semilogy(thres, card, color="red")
    ax3.hist(d, color="blue")
    ax4.hist(d_first_decile, color="red")
    plt.show()

def get_x2_stats(dir, threshold=1.0):
    x2 = np.genfromtxt(os.path.join(dir, "Results_collection-X2.dat"))

    p2p = x2.max() - x2.min()
    x2_filtered = x2[x2 < (x2.min() + threshold * p2p)]

    stats = {
        "min" : x2_filtered.min(),
        "max" : x2_filtered.max(),
        "mean" : x2_filtered.mean(),
        "std" : x2_filtered.std()
    }

    return stats

if __name__ == "__main__":
    if len(sys.argv) > 1 and os.path.isdir(sys.argv[1]):
        dir = sys.argv[1]
    else:
        dir = "POPC-test"

    print_x2_stats(dir)
    std_vs_threshold(dir, 13)