import os
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker


BASE = "/home/databank/ydn/DRAK2/MD-DRAK2/DCA-16C/Mg/system/"

FILES = {
    "md1": f"{BASE}/md1/analysis/RMSD/rmsd_analysis_refit/rmsd-combined.dat",
    "md2": f"{BASE}/md2/analysis/RMSD/rmsd_analysis_refit/rmsd-combined.dat",
    "md3": f"{BASE}/md3/analysis/RMSD/rmsd_analysis_refit/rmsd-combined.dat",
    "md4": f"{BASE}/md4/analysis/RMSD/rmsd_analysis_refit/rmsd-combined.dat",
}


OUT1 = "Mg_md1_md2_md3_md4_complex_RMSD.png"
OUT2 = "Mg_md1_md2_md3_md4_ligand_RMSD.png"


NS_PER_FRAME = 0.002
WINDOW = 500


COLORS = {
    "md1":"blue",
    "md2":"green",
    "md3":"red",
    "md4":"orange"
}


def smooth(x):
    return pd.Series(x).rolling(
        WINDOW,
        center=True,
        min_periods=1
    ).mean()


data={}


for name,path in FILES.items():

    if not os.path.exists(path):
        raise FileNotFoundError(path)

    df=pd.read_csv(
        path,
        sep=r"\s+",
        comment="#",
        names=[
            "frame",
            "protein",
            "ligand",
            "complex"
        ]
    )

    data[name]=df


# ======================
# complex RMSD
# ======================

plt.figure(figsize=(12,5),dpi=300)


for name,df in data.items():

    t=df.frame*NS_PER_FRAME

    plt.plot(
        t,
        df["complex"],
        alpha=0.2,
        linewidth=0.5,
        color=COLORS[name]
    )

    plt.plot(
        t,
        smooth(df["complex"]),
        linewidth=1.5,
        color=COLORS[name],
        label=name
    )


plt.xlabel("Time (ns)",fontsize=16)
plt.ylabel("Complex RMSD (Å)",fontsize=16)

plt.xlim(0,max(data["md1"].frame)*NS_PER_FRAME)
plt.ylim(0,20)
plt.gca().yaxis.set_major_locator(ticker.MultipleLocator(5))

plt.legend(fontsize=12)

plt.grid(
    linestyle="--",
    alpha=0.3
)

plt.tight_layout()

plt.savefig(
    OUT1,
    bbox_inches="tight"
)

plt.close()



# ======================
# ligand RMSD
# ======================


plt.figure(figsize=(12,5),dpi=300)


for name,df in data.items():

    t=df.frame*NS_PER_FRAME

    plt.plot(
        t,
        df["ligand"],
        alpha=0.2,
        linewidth=0.5,
        color=COLORS[name]
    )


    plt.plot(
        t,
        smooth(df["ligand"]),
        linewidth=1.5,
        color=COLORS[name],
        label=name
    )


plt.xlabel("Time (ns)",fontsize=16)
plt.ylabel("Ligand RMSD (Å)",fontsize=16)


plt.xlim(0,max(data["md1"].frame)*NS_PER_FRAME)
plt.ylim(0,20)
plt.gca().yaxis.set_major_locator(ticker.MultipleLocator(5))


plt.legend(fontsize=12)

plt.grid(
    linestyle="--",
    alpha=0.3
)

plt.tight_layout()


plt.savefig(
    OUT2,
    bbox_inches="tight"
)

plt.close()


print("Finished")
print(OUT1)
print(OUT2)