# This script validates the endpoint correction protocol on
# relative hydration free energies by comparing two absolute differences
# against the relative difference.

import argparse
import multiprocessing
from importlib import resources

import numpy as np
from rdkit import Chem

from timemachine.fe import model_rabfe
from timemachine.fe.free_energy import get_romol_conf
from timemachine.fe.lambda_schedule import (
    construct_pre_optimized_absolute_lambda_schedule_solvent,
    validate_lambda_schedule,
)
from timemachine.fe.restraints import setup_relative_restraints_by_distance
from timemachine.ff import Forcefield
from timemachine.md import builders, minimizer
from timemachine.parallel.client import CUDAPoolClient, GRPCClient
from timemachine.parallel.utils import get_gpu_count
from timemachine.training.dataset import Dataset

def construct_relative_lambda_schedule(num_windows, nonbonded_cutoff=1.2):
    """Generate a length-num_windows list of lambda values from 0.0 up to 1.0
    Notes
    -----
    * manually optimized by YTZ
    * assumes nonbonded cutoff = 1.2 nm
        (since decoupling_distance = lambda * nonbonded_cutoff,
        this schedule will not be appropriate for nonbonded_cutoff != 1.2!)
    """
    assert nonbonded_cutoff == 1.2

    A = int(0.15 * num_windows)
    B = int(0.60 * num_windows)
    C = num_windows - A - B

    # optimizing the overlap based on eyeballing absolute hydration free energies
    # there's probably some better way to deal with this by inspecting the curvature
    lambda_schedule = np.concatenate(
        [
            np.linspace(0.00, 0.08, A, endpoint=False),
            np.linspace(0.08, 0.27, B, endpoint=False),
            np.linspace(0.27, 1.00, C, endpoint=True),
        ]
    )

    validate_lambda_schedule(lambda_schedule, num_windows)

    return lambda_schedule

if __name__ == "__main__":

    multiprocessing.set_start_method("spawn")

    parser = argparse.ArgumentParser(
        description="Absolute Hydration Free Energy Testing",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument("--hosts", nargs="*", default=None, help="Hosts running GRPC worker to use for compute")

    parser.add_argument("--num_gpus", type=int, help="number of gpus", default=get_gpu_count())

    parser.add_argument("--num_windows", type=int, help="number of solvent lambda windows", required=True)

    parser.add_argument(
        "--num_equil_steps",
        type=int,
        help="number of equilibration steps for each solvent lambda window",
        required=True,
    )

    parser.add_argument(
        "--num_prod_steps", type=int, help="number of production steps for each solvent lambda window", required=True
    )

    cmd_args = parser.parse_args()

    if not cmd_args.hosts:
        num_gpus = cmd_args.num_gpus
        # set up multi-GPU client
        client = CUDAPoolClient(max_workers=num_gpus)
    else:
        # Setup GRPC client
        print("Connecting to GRPC workers...")
        client = GRPCClient(hosts=cmd_args.hosts)
    client.verify()

    with resources.path("timemachine.testsystems.data", "ligands_40.sdf") as path_to_ligand:
        suppl = Chem.SDMolSupplier(str(path_to_ligand), removeHs=False)

    forcefield = Forcefield.load_from_file("smirnoff_1_1_0_ccc.py")
    mols = [x for x in suppl]

    dataset = Dataset(mols)

    absolute_solvent_schedule = construct_pre_optimized_absolute_lambda_schedule_solvent(cmd_args.num_windows)
    relative_solvent_schedule = construct_relative_lambda_schedule(cmd_args.num_windows - 1)
    solvent_system, solvent_coords, solvent_box, solvent_topology = builders.build_water_system(4.0)

    # pick the largest mol as the blocker
    largest_size = 0
    ref_mol = None
    for mol in mols:
        if mol.GetNumAtoms() > largest_size:
            largest_size = mol.GetNumAtoms()
            ref_mol = mol

    print("Reference Molecule:", ref_mol.GetProp("_Name"), Chem.MolToSmiles(ref_mol))

    temperature = 300.0
    pressure = 1.0
    dt = 2.5e-3

    model_relative = model_rabfe.RelativeHydrationModel(
        client,
        forcefield,
        solvent_system,
        relative_solvent_schedule,
        solvent_topology,
        temperature,
        pressure,
        dt,
        cmd_args.num_equil_steps,
        cmd_args.num_prod_steps,
    )

    model_absolute = model_rabfe.AbsoluteHydrationModel(
        client,
        forcefield,
        solvent_system,
        absolute_solvent_schedule,
        solvent_topology,
        temperature,
        pressure,
        dt,
        cmd_args.num_equil_steps,
        cmd_args.num_prod_steps,
    )

    ordered_params = forcefield.get_ordered_params()
    ordered_handles = forcefield.get_ordered_handles()

    M = len(dataset.data)

    # generate initial coordinates
    def minimize_absolute(mol):
        hc = minimizer.minimize_host_4d([mol], solvent_system, solvent_coords, forcefield, solvent_box)
        return np.concatenate([hc, get_romol_conf(mol)])

    def minimize_relative(mol_a, mol_b):
        hc = minimizer.minimize_host_4d([mol_a, mol_b], solvent_system, solvent_coords, forcefield, solvent_box)
        return np.concatenate([hc, get_romol_conf(mol_a), get_romol_conf(mol_b)])

    for epoch in range(100):

        for i in range(M):

            for j in range(i + 1, M):

                mol_a = dataset.data[i]
                mol_b = dataset.data[j]

                # relative calculation
                core_idxs = setup_relative_restraints_by_distance(mol_a, mol_b)

                xab = minimize_relative(mol_a, mol_b)
                ddG_ab, ddG_ab_err = model_relative.predict(
                    ordered_params,
                    mol_a,
                    mol_b,
                    core_idxs,
                    x0=xab,
                    box0=solvent_box,
                    prefix="epoch_"
                    + str(epoch)
                    + "_solvent_relative_"
                    + mol_a.GetProp("_Name")
                    + "_"
                    + mol_b.GetProp("_Name"),
                )

                # absolute calculation
                xa = minimize_absolute(mol_a)
                xb = minimize_absolute(mol_b)

                dG_a, dG_a_err = model_absolute.predict(
                    ordered_params, mol_a, x0=xa, box0=solvent_box, prefix="solvent_absolute_" + mol_a.GetProp("_Name")
                )
                dG_b, dG_b_err = model_absolute.predict(
                    ordered_params, mol_b, x0=xb, box0=solvent_box, prefix="solvent_absolute_" + mol_b.GetProp("_Name")
                )
                dG_ab_err = np.sqrt(dG_a_err ** 2 + dG_b_err ** 2)

                print(
                    f"mol_i {i} {mol_a.GetProp('_Name')} mol_j {j} {mol_b.GetProp('_Name')} ddG_ab {ddG_ab:.3f} +- {ddG_ab_err:.3f} dG_a-dG_b {dG_a-dG_b:.3f} +- {dG_ab_err:.3f}"
                )
