import os

from jax.config import config

config.update("jax_enable_x64", True)
import functools
import multiprocessing
from importlib import resources

import jax
import jax.numpy as jnp
import numpy as np
from rdkit import Chem

from timemachine.ff import handlers
from timemachine.ff.handlers.deserialize import deserialize_handlers
from timemachine.integrator import langevin_coefficients
from timemachine.potentials import bonded, shape


def pmi_restraints_new(conf, params, box, lamb, a_idxs, b_idxs, masses, angle_force, com_force):

    a_com, a_tensor = inertia_tensor(conf[a_idxs], masses[a_idxs])
    b_com, b_tensor = inertia_tensor(conf[b_idxs], masses[b_idxs])

    a_eval, a_evec = jnp.linalg.eigh(a_tensor)  # already sorted
    b_eval, b_evec = jnp.linalg.eigh(b_tensor)  # already sorted

    # convert from column to row eigenvectors
    a_rvec = jnp.transpose(a_evec)
    b_rvec = jnp.transpose(b_evec)

    loss = []
    for a, b in zip(a_rvec, b_rvec):
        delta = 1 - jnp.abs(jnp.dot(a, b))
        loss.append(delta * delta)

    return angle_force * jnp.sum(loss) + com_force * jnp.linalg.norm(b_com - a_com)


def recenter(conf):
    return conf - jnp.mean(conf, axis=0)


def inertia_tensor(conf, masses):
    com = jnp.average(conf, axis=0, weights=masses)
    conf = conf - com

    xs = conf[:, 0]
    ys = conf[:, 1]
    zs = conf[:, 2]
    xx = jnp.average(ys * ys + zs * zs, weights=masses)
    yy = jnp.average(xs * xs + zs * zs, weights=masses)
    zz = jnp.average(xs * xs + ys * ys, weights=masses)
    xy = jnp.average(-xs * ys, weights=masses)
    xz = jnp.average(-xs * zs, weights=masses)
    yz = jnp.average(-ys * zs, weights=masses)
    tensor = jnp.array([[xx, xy, xz], [xy, yy, yz], [xz, yz, zz]])

    return com, tensor


def get_conf(romol, idx):
    conformer = romol.GetConformer(idx)
    guest_conf = jnp.array(conformer.GetPositions(), dtype=jnp.float64)
    guest_conf /= 10
    return recenter(guest_conf)


def make_conformer(mol, conf_a, conf_b):
    mol.RemoveAllConformers()
    mol = Chem.CombineMols(mol, mol)
    cc = Chem.Conformer(mol.GetNumAtoms())
    conf = jnp.concatenate([conf_a, conf_b])
    conf *= 10
    for idx, pos in enumerate(np.asarray(conf)):
        cc.SetAtomPosition(idx, (float(pos[0]), float(pos[1]), float(pos[2])))
    mol.AddConformer(cc)

    return mol


def get_heavy_atom_idxs(mol):

    idxs = []
    for a_idx, a in enumerate(mol.GetAtoms()):
        if a.GetAtomicNum() > 1:
            idxs.append(a_idx)
    return jnp.array(idxs, dtype=jnp.int32)


def convergence(args):
    epoch, lamb, lamb_idx = args

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

    ligands = []
    for mol in suppl:
        ligands.append(mol)

    ligand_a = ligands[0]
    ligand_b = ligands[1]

    # print(ligand_a.GetNumAtoms())
    # print(ligand_b.GetNumAtoms())

    # ligand_a = Chem.AddHs(Chem.MolFromSmiles("CCCC1=NN(C2=C1N=C(NC2=O)C3=C(C=CC(=C3)S(=O)(=O)N4CCN(CC4)C)OCC)C"))
    # ligand_b = Chem.AddHs(Chem.MolFromSmiles("CCCC1=NN(C2=C1N=C(NC2=O)C3=C(C=CC(=C3)S(=O)(=O)N4CCN(CC4)C)OCC)C"))
    # ligand_a = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1CC"))
    # ligand_b = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1CC"))
    # AllChem.EmbedMolecule(ligand_a, randomSeed=2020)
    # AllChem.EmbedMolecule(ligand_b, randomSeed=2020)

    coords_a = get_conf(ligand_a, idx=0)
    coords_b = get_conf(ligand_b, idx=0)
    # coords_b = np.matmul(coords_b, special_ortho_group.rvs(3))

    coords_a = recenter(coords_a)
    coords_b = recenter(coords_b)

    coords = jnp.concatenate([coords_a, coords_b])

    a_idxs = get_heavy_atom_idxs(ligand_a)
    b_idxs = get_heavy_atom_idxs(ligand_b)

    # NOTE: used in commented code below
    # a_full_idxs = np.arange(0, ligand_a.GetNumAtoms())
    # b_full_idxs = np.arange(0, ligand_b.GetNumAtoms())

    b_idxs += ligand_a.GetNumAtoms()
    # b_full_idxs += ligand_a.GetNumAtoms()

    nrg_fns = []

    forcefield = "timemachine/ff/params/smirnoff_1_1_0_ccc.py"
    ff_raw = open(forcefield, "r").read()
    ff_handlers = deserialize_handlers(ff_raw)

    combined_mol = Chem.CombineMols(ligand_a, ligand_b)

    for handler in ff_handlers:
        if isinstance(handler, handlers.HarmonicBondHandler):
            bond_idxs, (bond_params, _) = handler.parameterize(combined_mol)
            nrg_fns.append(functools.partial(bonded.harmonic_bond, params=bond_params, box=None, bond_idxs=bond_idxs))
        elif isinstance(handler, handlers.HarmonicAngleHandler):
            angle_idxs, (angle_params, _) = handler.parameterize(combined_mol)
            nrg_fns.append(
                functools.partial(bonded.harmonic_angle, params=angle_params, box=None, angle_idxs=angle_idxs)
            )
        # elif isinstance(handler, handlers.ImproperTorsionHandler):
        #     torsion_idxs, (torsion_params, _) = handler.parameterize(combined_mol)
        #     print(torsion_idxs)
        #     assert 0
        #     nrg_fns.append(
        #         functools.partial(bonded.periodic_torsion,
        #             params=torsion_params,
        #             box=None,
        #             lamb=None,
        #             torsion_idxs=torsion_idxs
        #         )
        #     )
        # elif isinstance(handler, handlers.ProperTorsionHandler):
        #     torsion_idxs, (torsion_params, _) = handler.parameterize(combined_mol)
        #     # print(torsion_idxs)
        #     nrg_fns.append(
        #         functools.partial(bonded.periodic_torsion,
        #             params=torsion_params,
        #             box=None,
        #             lamb=None,
        #             torsion_idxs=torsion_idxs
        #         )
        #     )

    masses_a = np.array([a.GetMass() for a in ligand_a.GetAtoms()]) * 10000
    masses_b = np.array([a.GetMass() for a in ligand_b.GetAtoms()])

    combined_masses = jnp.concatenate([masses_a, masses_b])

    # com_restraint_fn = functools.partial(bonded.centroid_restraint,
    #     params=None,
    #     box=None,
    #     lamb=None,
    #     # masses=combined_masses, # try making this ones-like
    #     masses=np.ones_like(combined_masses),
    #     group_a_idxs=a_idxs,
    #     group_b_idxs=b_idxs,
    #     kb=50.0,
    #     b0=0.0)

    pmi_restraint_fn = functools.partial(
        pmi_restraints_new,
        params=None,
        box=None,
        lamb=None,
        # masses=np.ones_like(combined_masses),
        masses=combined_masses,
        # a_idxs=a_full_idxs,
        # b_idxs=b_full_idxs,
        a_idxs=a_idxs,
        b_idxs=b_idxs,
        angle_force=100.0,
        com_force=100.0,
    )

    prefactor = 2.7  # unitless
    shape_lamb = (4 * jnp.pi) / (3 * prefactor)  # unitless
    kappa = jnp.pi / (jnp.power(shape_lamb, 2 / 3))  # unitless
    sigma = 0.15  # 1 angstrom std, 95% coverage by 2 angstroms
    alpha = kappa / (sigma * sigma)

    alphas = jnp.zeros(combined_mol.GetNumAtoms()) + alpha
    weights = jnp.zeros(combined_mol.GetNumAtoms()) + prefactor

    shape_restraint_fn = functools.partial(
        shape.harmonic_overlap,
        box=None,
        lamb=None,
        params=None,
        a_idxs=a_idxs,
        b_idxs=b_idxs,
        alphas=alphas,
        weights=weights,
        k=150.0,
    )

    # shape_restraint_4d_fn = functools.partial(
    #     shape.harmonic_4d_overlap,
    #     box=None,
    #     params=None,
    #     a_idxs=a_idxs,
    #     b_idxs=b_idxs,
    #     alphas=alphas,
    #     weights=weights,
    #     k=200.0
    # )

    def restraint_fn(conf, lamb):

        return pmi_restraint_fn(conf) + lamb * shape_restraint_fn(conf)
        # return (1-lamb)*pmi_restraint_fn(conf) + lamb*shape_restraint_fn(conf)

    nrg_fns.append(restraint_fn)

    def nrg_fn(conf, lamb):
        s = []
        for u in nrg_fns:
            s.append(u(conf, lamb=lamb))
        return jnp.sum(s)

    grad_fn = jax.grad(nrg_fn, argnums=(0, 1))
    grad_fn = jax.jit(grad_fn)

    du_dx_fn = jax.grad(nrg_fn, argnums=(0))
    du_dx_fn = jax.jit(du_dx_fn)

    x_t = coords
    v_t = jnp.zeros_like(x_t)

    Chem.SDWriter("frames_heavy_" + str(epoch) + "_" + str(lamb_idx) + ".sdf")

    dt = 1.5e-3
    ca, cb, cc = langevin_coefficients(300.0, dt, 1.0, combined_masses)
    cb = -1 * np.expand_dims(cb, axis=-1)
    cc = np.expand_dims(cc, axis=-1)

    du_dls = []

    # re-seed since forking
    np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))

    # for step in range(100000):
    for step in range(100000):

        # if step % 1000 == 0:
        #     u = nrg_fn(x_t, lamb)
        #     print("step", step, "nrg", np.asarray(u), "avg_du_dl",  np.mean(du_dls))
        #     mol = make_conformer(combined_mol, x_t[:ligand_a.GetNumAtoms()], x_t[ligand_a.GetNumAtoms():])
        #     w.write(mol)
        #     w.flush()

        if step % 5 == 0 and step > 10000:
            du_dx, du_dl = grad_fn(x_t, lamb)
            du_dls.append(du_dl)
        else:
            du_dx = du_dx_fn(x_t, lamb)

        v_t = ca * v_t + cb * du_dx + cc * np.random.normal(size=x_t.shape)
        x_t = x_t + v_t * dt

    return jnp.mean(np.mean(du_dls))


if __name__ == "__main__":

    pool = multiprocessing.Pool()  # defaults to # of cpus

    # lambda_schedule = np.linspace(0, 1.0, os.cpu_count())
    lambda_schedule = np.linspace(0, 1.0, 24)
    # lambda_schedule = np.array([0.0])
    # lambda_schedule = [0.81, 0.81, 0.81, 0.81, 0.81, 0.81, 0.81, 0.81, 0.81, 0.81, 0.81]
    # lambda_schedule = np.array([1e-4, 5e-4, 1e-3, 5e-3, 0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.11, 0.12, 0.13, 0.15,0.175, 0.2, 0.225, 0.25, 0.275, 0.3, 0.35, 0.4, 0.5])
    # lambda_schedule = np.array([0.0])
    # lambda_schedule = np.linspace(0.2, 0.6, 24)

    print("cpu count:", os.cpu_count())

    for epoch in range(100):
        args = []
        for l_idx, lamb in enumerate(lambda_schedule):
            args.append((epoch, lamb, l_idx))
        avg_du_dls = pool.map(convergence, args)
        avg_du_dls = jnp.asarray(avg_du_dls)

        for lamb, ddl in zip(lambda_schedule, avg_du_dls):
            print("lambda", lamb, "du_dl", ddl)
        print(epoch, "epoch", "deltaG", np.trapz(avg_du_dls, lambda_schedule))
