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

# ## Initial Imports

import sys
import multiprocessing
from scm.plams import (
    JobRunner,
    config,
    from_smiles,
    Settings,
    AMSJob,
    JobAnalysis,
    jobs_in_directory,
    view,
    plot_image_grid,
)
import numpy as np


# ## Set Up Job Runner
# Set up job runner, running as many jobs as possible in parallel.

config.default_jobrunner = JobRunner(parallel=True, maxjobs=multiprocessing.cpu_count())
config.job.runscript.nproc = 1


# ## Set Up Molecules
# Create the molecules we want to use in our benchmark from SMILES.

# The molecules we want to use in our benchmark:
mol_smiles = {"Methane": "C", "Ethane": "C-C", "Ethylene": "C=C", "Acetylene": "C#C"}
molecules = {}
for name, smiles in mol_smiles.items():
    # Compute 10 conformers, optimize with UFF and pick the lowest in energy.
    molecules[name] = from_smiles(smiles, nconfs=10, forcefield="uff")[0]

plot_image_grid({k: view(v) for k, v in molecules.items()}, rows=1, save_path="picture1.png")
# ## Initialize Calculation Settings
# Set up the settings which are common across jobs. The basis type is added later for each job.

common_settings = Settings()
common_settings.input.ams.Task = "SinglePoint"
common_settings.input.ams.System.Symmetrize = "Yes"
common_settings.input.adf.Basis.Core = "None"


basis = ["QZ4P", "TZ2P", "TZP", "DZP", "DZ", "SZ"]
reference_basis = "QZ4P"


# ## Run Calculations

results = {}
jobs = []
for bas in basis:
    for name, molecule in molecules.items():
        with jobs_in_directory(name):
            settings = common_settings.copy()
            settings.input.adf.Basis.Type = bas
            job = AMSJob(name=f"{name}_{bas}", molecule=molecule, settings=settings)
            jobs.append(job)
            results[(name, bas)] = job.run()

for job in jobs:
    job.ok()


# ## Results
# Extract the energy from each calculation. Calculate the average absolute error in bond energy per atom for each basis set.

import pandas as pd

results = []
for job in jobs:
    mol = job.results.get_main_molecule()
    results.append(
        {
            "Formula": mol.get_formula(),
            "Basis": job.settings.input.adf.basis.type,
            "NAtoms": len(mol),
            "Energy": job.results.get_energy(unit="kcal/mol"),
        }
    )

results_df = pd.DataFrame(results)
print(results_df)


ref_results_df = results_df[results_df["Basis"] == "QZ4P"]
comp_results_df = results_df[results_df["Basis"] != "QZ4P"]
comp_results_df = comp_results_df.merge(ref_results_df, on="Formula", suffixes=["", " (Ref)"])
comp_results_df["Average Error"] = (
    abs(comp_results_df["Energy"] - comp_results_df["Energy (Ref)"]) / comp_results_df["NAtoms"]
)
print(comp_results_df)


benchmark_results_df = comp_results_df.groupby("Basis", as_index=False)["Average Error"].mean()
benchmark_results_df.rename(columns={"Average Error": "MAE [kcal/mol]"}, inplace=True)
benchmark_results_df.sort_values(by="MAE [kcal/mol]", inplace=True, ignore_index=True)
print(benchmark_results_df)