#!/usr/bin/env python
# coding: utf-8

# ## ADF and COSMO-RS workflow
#
# This example uses ADF to generate the .coskf file for benzene. You can also modify it to instead use the Benzene.coskf from the ADFCRS-2018 database. Note that you first need to install the ADFCRS-2018 database.
#
# The script then plots the sigma profile. This is not a necessary step but is done for illustration purposes only.
#
# Then a solubility calculation is performed for benzene in water between 0 and 10 degrees C. The melting point and enthalpy of fusion can either be estimated using the property prediction tool, or the experimental numbers can be given (recommended).
#
# A second section shows the same workflow using pyCRS.

from __future__ import annotations

from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
from pyCRS.CRSManager import CRSSystem
from pyCRS.Database import COSKFDatabase
from scm.plams import CRSJob, CRSResults, Settings, Units, from_smiles, view
from scm.plams.recipes.adfcosmorscompound import ADFCOSMORSCompoundJob


# ### Helper functions


def generate_benzene_coskf(smiles: str, job_name: str, mol_info: dict[str, Any] | None = None) -> Path:
    molecule = from_smiles(smiles, nconfs=100, forcefield="uff")[0]
    view(molecule)
    job = ADFCOSMORSCompoundJob(name=job_name, molecule=molecule, mol_info=mol_info)
    job.run()
    return Path(job.results.coskfpath())


def get_plams_sigma_profile(coskf_file: Path) -> dict[str, list[float]]:
    settings = Settings()
    settings.input.property._h = "PURESIGMAPROFILE"
    settings.input.compound._h = str(coskf_file)
    job = CRSJob(name="sigma_profile_plams", settings=settings)
    job.run()
    return job.results.get_sigma_profile()


def run_plams_solubility(
    solvent_coskf: Path,
    solute_coskf: Path,
    temperature_spec: str = "273.15 283.15 10",
    pressure_spec: str = "1.01325 1.01325 10",
    solvent_density: float = 1.0,
    melting_point: float = 278.7,
    hfusion_kjmol: float = 9.91,
) -> CRSJob:
    settings = Settings()
    settings.input.property._h = "solubility"
    settings.input.temperature = temperature_spec
    settings.input.pressure = pressure_spec
    settings.input.property.DensitySolvent = solvent_density

    settings.input.compound = [Settings(), Settings()]
    settings.input.compound[0]._h = str(solvent_coskf)
    settings.input.compound[0].frac1 = 1.0

    settings.input.compound[1]._h = str(solute_coskf)
    settings.input.compound[1].nring = 6
    settings.input.compound[1].meltingpoint = melting_point
    settings.input.compound[1].hfusion = hfusion_kjmol * Units.convert(1.0, "kJ/mol", "kcal/mol")

    job = CRSJob(name="benzene_in_water_plams", settings=settings)
    job.run()
    return job


def run_pycrs_solubility(db_path: Path, temperature_spec: str = "273.15 283.15 10") -> CRSResults:
    crs = CRSSystem()
    additional_settings = Settings()
    additional_settings.input.property.DensitySolvent = 1.0

    crs.add_Mixture(
        mixture={"water": 1.0, "benzene": 0.0},
        temperature=temperature_spec,
        database=str(db_path),
        problem_type="solubility",
        additional_sett=additional_settings,
    )
    crs.runCRSJob()
    return crs.outputs[0]


def get_pycrs_sigma_profile(compound_name: str, db: COSKFDatabase) -> dict[str, list[float]]:
    crs = CRSSystem()
    crs.add_Mixture(mixture={compound_name: 1.0}, problem_type="PURESIGMAPROFILE", database=db.path)
    crs.runCRSJob()
    return crs.outputs[0].get_sigma_profile()


def extract_solubility_series(results_obj: CRSResults | Any) -> tuple[list[float], list[float]]:
    if isinstance(results_obj, CRSResults):
        results = results_obj.get_results("SOLUBILITY")
    else:
        results = results_obj.get_results("SOLUBILITY")

    temperatures = list(results["temperature"])
    solubility_g_l = list(results["solubility g_per_L_solvent"][1])
    return temperatures, solubility_g_l


def print_solubility_table(temperatures: list[float], solubility_g_l: list[float]) -> None:
    for temperature, value in zip(temperatures, solubility_g_l):
        print(f"{temperature:.2f} K -> {value:.4f} g/L solvent")


def plot_sigma_profile(sigma_profile: dict[str, list[float]], title: str) -> Any:
    x_label = "σ (e/A**2)"
    fig, ax = plt.subplots()
    for key, values in sigma_profile.items():
        if key == x_label:
            continue
        ax.plot(sigma_profile[x_label], values, label=key.split(".")[0])

    ax.set_xlabel("σ (e/A**2)")
    ax.set_ylabel("p(σ)")
    ax.set_title(title)
    ax.legend()
    fig.tight_layout()
    return ax


def plot_solubility(temperatures: list[float], solubility_g_l: list[float], title: str) -> Any:
    fig, ax = plt.subplots()
    ax.plot(temperatures, solubility_g_l, marker="o")
    ax.set_xlabel("Temperature (K)")
    ax.set_ylabel("Solubility (g/L solvent)")
    ax.set_title(title)
    fig.tight_layout()
    return ax


# ### PLAMS workflow

database_dir = Path(CRSJob.database())
if not database_dir.exists():
    raise FileNotFoundError(f"COSMO-RS database was not found: {database_dir}")

benzene_coskf = generate_benzene_coskf("c1ccccc1", "adf_benzene")
sigma_profile_plams = get_plams_sigma_profile(benzene_coskf)
ax = plot_sigma_profile(sigma_profile_plams, "Benzene sigma profile (PLAMS)")
ax
ax.figure.savefig("picture1.png")


plams_job = run_plams_solubility(
    solvent_coskf=database_dir / "Water.coskf",
    solute_coskf=benzene_coskf,
    temperature_spec="273.15 283.15 10",
    pressure_spec="1.01325 1.01325 10",
)
temps_plams, sol_plams = extract_solubility_series(plams_job.results)
print_solubility_table(temps_plams, sol_plams)
ax = plot_solubility(temps_plams, sol_plams, "Benzene solubility in water (PLAMS)")
ax
ax.figure.savefig("picture2.png")


# ### pyCRS workflow
#
# The pyCRS section reuses the generated benzene compound and stores compounds in a local COSKF database file.

db_path = Path("my_coskf_db.db")
db = COSKFDatabase(str(db_path))

if db.get_compounds_id("water")[0] is None:
    db.add_compound(str(database_dir / "Water.coskf"))

if db.get_compounds_id("benzene")[0] is None:
    mol_info = {
        "IUPAC": "Benzene",
        "CAS": "71-43-2",
        "SMILES": "c1ccccc1",
    }
    benzene_coskf_pycrs = generate_benzene_coskf("c1ccccc1", "adf_benzene_pycrs", mol_info=mol_info)
    db.add_compound(str(benzene_coskf_pycrs))
    db.add_physical_property("benzene", "meltingpoint", 278.7)
    db.add_physical_property("benzene", "hfusion", 9.91, unit="kJ/mol")

sigma_profile_pycrs = get_pycrs_sigma_profile("benzene", db)
ax = plot_sigma_profile(sigma_profile_pycrs, "Benzene sigma profile (pyCRS)")
ax
ax.figure.savefig("picture3.png")


pycrs_results = run_pycrs_solubility(db_path)
temps_pycrs, sol_pycrs = extract_solubility_series(pycrs_results)
print_solubility_table(temps_pycrs, sol_pycrs)
ax = plot_solubility(temps_pycrs, sol_pycrs, "Benzene solubility in water (pyCRS)")
ax
ax.figure.savefig("picture4.png")