import os

import matplotlib.pyplot as plt
from scm.plams import CRSJob, JobRunner, Settings, 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"),
]
DATABASE_PATH = CRSJob.database()
IMAGE_NAME = "TernaryStabilityLLE.png"


def ensure_compound_files(compounds):
    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}")


def run_demo(df_stability, df_lle):
    outputs = []
    for _, row in df_stability.iterrows():
        settings = Settings()
        settings.input.compound = []
        for component_index, compound in enumerate(SOLVENTS, start=1):
            block = Settings()
            block._h = os.path.join(DATABASE_PATH, compound.coskf)
            block.frac1 = float(row[f"x{component_index}"])
            block.name = compound.name.lower()
            settings.input.compound.append(block)

        settings.input.method = row["Method"]
        settings.input.property._h = "STABILITY"
        settings.input.temperature = row["Temperature"]

        job = CRSJob(settings=settings, name=f"STABILITY_{row['mix_idx']}")
        outputs.append(job.run())

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

    outputs = []
    for _, row in df_lle.iterrows():
        settings = Settings()
        settings.input.compound = []
        for component_index, compound in enumerate(SOLVENTS, start=1):
            block = Settings()
            block._h = os.path.join(DATABASE_PATH, compound.coskf)
            block.frac1 = float(row[f"x{component_index}"])
            block.name = compound.name.lower()
            settings.input.compound.append(block)

        settings.input.method = row["Method"]
        settings.input.property._h = "LLE"
        settings.input.temperature = row["Temperature"]

        job = CRSJob(settings=settings, name=f"LLE_{row['mix_idx']}")
        outputs.append(job.run())

    for index, out in zip(df_lle.index, 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


ensure_compound_files(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)
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()