import os
from pathlib import Path

import matplotlib.pyplot as plt
from pyCRS.CRSManager import CRSSystem
from pyCRS.Database import COSKFDatabase
from scm.plams import CRSJob, JobRunner, config
from TernaryStabilityLLE_common import Compound, build_ternary_grid_df, build_ternary_lle_df, plot_demo_figure

TEMPERATURE = 298.15
METHOD = "COSMORS"
MAX_JOBS = 8
SOLVENTS = [
    Compound("Water", "Water.coskf"),
    Compound("Ethanol", "Ethanol.coskf"),
    Compound("Benzene", "Benzene.coskf"),
]
DB_NAME = "ternary_stability_lle_demo.db"
DATABASE_PATH = CRSJob.database()
IMAGE_NAME = "TernaryStabilityLLE.png"


def prepare_database(db_name, compounds):
    db_path = Path(db_name)
    if db_path.exists():
        db_path.unlink()

    db = COSKFDatabase(db_name)
    for compound in compounds:
        coskf_path = os.path.join(DATABASE_PATH, compound.coskf)
        if not os.path.exists(coskf_path):
            raise FileNotFoundError(f"COSKF file not found: {coskf_path}")
        db.add_compound(coskf_file=compound.coskf, name=compound.name, coskf_path=DATABASE_PATH)

    return db


def run_demo(df_stability, df_lle, db):
    cal_stability = CRSSystem()
    for _, row in df_stability.iterrows():
        mixture = {row["s1"]: row["x1"], row["s2"]: row["x2"], row["s3"]: row["x3"]}
        cal_stability.add_Mixture(
            mixture,
            temperature=row["Temperature"],
            problem_type="STABILITY",
            database=db,
            method=row["Method"],
            jobname=f"STABILITY_{row['mix_idx']}",
        )
    cal_stability.runCRSJob()

    for index, out in zip(df_stability.index, cal_stability.outputs):
        results = out.get_results()
        df_stability.loc[index, "unstable"] = bool(results.get("unstable", False))

    cal_lle = CRSSystem()
    for _, row in df_lle.iterrows():
        mixture = {row["s1"]: row["x1"], row["s2"]: row["x2"], row["s3"]: row["x3"]}
        cal_lle.add_Mixture(
            mixture,
            temperature=row["Temperature"],
            problem_type="LLE",
            database=db,
            method=row["Method"],
            jobname=f"LLE_{row['mix_idx']}",
        )
    cal_lle.runCRSJob()

    for index, out in zip(df_lle.index, cal_lle.outputs):
        results = out.get_results()
        converged = bool(results.get("converged", False))
        llle_detected = bool(results.get("llle_detected", False))
        df_lle.loc[index, "converged"] = converged
        df_lle.loc[index, "llle_detected"] = llle_detected
        if converged:
            x_i = results.get("xI")
            x_ii = results.get("xII")
            for component_index in range(3):
                component_number = component_index + 1
                if x_i is not None:
                    df_lle.loc[index, f"xI{component_number}"] = x_i[component_index][0]
                if x_ii is not None:
                    df_lle.loc[index, f"xII{component_number}"] = x_ii[component_index][0]

    return df_stability, df_lle


db = prepare_database(DB_NAME, SOLVENTS)

config.default_jobrunner = JobRunner(parallel=True, maxjobs=MAX_JOBS)
config.job.runscript.nproc = 1
config.log.stdout = 1

df_stability = build_ternary_grid_df(SOLVENTS, min_frac=0.1, step=0.1, temperature=TEMPERATURE, method=METHOD)
df_lle = build_ternary_lle_df(
    SOLVENTS,
    s1_ratio=1.0,
    s2_ratio=1.0,
    s3_min=0.1,
    s3_max=0.9,
    num_points=9,
    temperature=TEMPERATURE,
    method=METHOD,
)

df_stability, df_lle = run_demo(df_stability, df_lle, db)
fig = plot_demo_figure(df_stability, df_lle)

print(df_stability[["mix_idx", "x1", "x2", "x3", "unstable"]])
print(df_lle[["mix_idx", "x1", "x2", "x3", "converged", "llle_detected", "xI1", "xI2", "xI3", "xII1", "xII2", "xII3"]])
fig.savefig(IMAGE_NAME, dpi=200, bbox_inches="tight")
plt.show()