import torchani
import torch
from simtk import unit
import numpy as np
import matplotlib.pyplot as plt

hartree_to_kJ_mol = 2625.499638
nm_to_angstroms = (1.0 * unit.nanometer) / (1.0 * unit.angstrom)

device = torch.device('cpu')
model = torchani.models.ANI1ccx()
species = model.species_to_tensor('CHHHH').to(device).unsqueeze(0)

def ANI1ccx_energy(x):
    """
    Parameters
    ----------
    x : numpy array, unit'd
        coordinates

    Returns
    -------
    energy : float, unit'd
        energy, with unit of kJ/mol attached
    """

    assert(type(x) == unit.Quantity)

    x_in_nm = x.value_in_unit(unit.nanometer)

    coordinates = torch.tensor([x_in_nm],
                           requires_grad=False, device=device, dtype=torch.float32)

    # convert from nm to angstroms
    coordinates_in_angstroms = coordinates * nm_to_angstroms
    _, energy_in_hartree = model((species, coordinates_in_angstroms))

    # convert energy from hartrees to kJ/mol
    energy_in_kJ_mol = energy_in_hartree * hartree_to_kJ_mol
    return float(energy_in_kJ_mol.detach()) * unit.kilojoule_per_mole


def ANI1ccx_energy_ensemble(x):
    """
    Parameters
    ----------
    x : numpy array, unit'd
        coordinates

    Returns
    -------
    energies : numpy array, unit'd
        8 energies from the ensemble, with units of kJ/mol attached
    """

    assert(type(x) == unit.Quantity)

    x_in_nm = x.value_in_unit(unit.nanometer)

    coordinates = torch.tensor([x_in_nm],
                           requires_grad=False, device=device, dtype=torch.float32)

    # convert from nm to angstroms
    coordinates_in_angstroms = coordinates * nm_to_angstroms

    energies_in_kJ_mol = [float(hartree_to_kJ_mol * m((species, coordinates_in_angstroms))[1].detach()) for m in model]

    return np.array(energies_in_kJ_mol) * unit.kilojoule_per_mole

scale_factors = np.linspace(0,2,1000)

tetrahedral_geometry = np.array([
    (0, 0, 0),
    (1, 1, 1),
    (-1, -1, 1),
    (-1, 1, -1),
    (1, -1, -1),
], dtype=float)
tetrahedral_geometry /= np.linalg.norm(tetrahedral_geometry[0] - tetrahedral_geometry[1])

# lowest energy found
bent_optimum = np.array([
    [ 0.        ,  0.        ,  0.        ],
    [ 0.64702535,  0.64647233,  0.8811238 ],
    [-0.64734113, -0.6475463 ,  0.8808765 ],
    [-0.7168874 ,  0.7172707 , -0.33484012],
    [ 0.7182246 , -0.71852016, -0.33497408]
]) # in angstroms
lowest_energy = ANI1ccx_energy(bent_optimum * unit.angstroms)

# scaled so that the average bond length is 1 angstrom
normalized_bent_optimum = bent_optimum / np.mean(np.linalg.norm(bent_optimum[1:], axis=1))


# energy scan from global optimum (bent geometry), with ensemble uncertainty
bent_scan = np.array([normalized_bent_optimum * scale_factor for scale_factor in scale_factors])
bent_energies = [ANI1ccx_energy_ensemble(g * unit.angstrom) for g in bent_scan]
bent_energy_array = np.array([(e - lowest_energy) / unit.kilojoule_per_mole for e in bent_energies])
bent_means = np.mean(bent_energy_array, 1)
bent_stdevs = np.std(bent_energy_array, 1)


# energy scan from tetrahedral geometry, with ensemble uncertainty
tetrahedral_scan = np.array([tetrahedral_geometry * scale_factor for scale_factor in scale_factors])
tetra_energies = [ANI1ccx_energy_ensemble(g * unit.angstrom) for g in tetrahedral_scan]
tetra_energy_array = np.array([(e - lowest_energy) / unit.kilojoule_per_mole for e in tetra_energies])
tetra_means = np.mean(tetra_energy_array, 1)
tetra_stdevs = np.std(tetra_energy_array, 1)


def average_bond_length_in_snapshot(xyz):
    return np.mean(np.linalg.norm(xyz[1:], axis=1))


# plot
plt.figure()
ax = plt.subplot(1,1,1)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

plt.ylim(-20,400)
plt.xlim(0.85,1.35)

plt.xlabel('C-H bond length (Å)')
plt.ylabel('$\Delta E \pm$ stddev (kJ/mol)\nrelative to global minimum')
plt.title('methane energy scans: varying C-H bond length')

# tetrahedral
# x: average C-H bond length
tetrahedral_x = [average_bond_length_in_snapshot(xyz) for xyz in tetrahedral_scan]
# y: mean energy +/- standard deviation
plt.plot(tetrahedral_x, tetra_means, label='scan from tetrahedron')
plt.fill_between(tetrahedral_x, tetra_means - tetra_stdevs, tetra_means + tetra_stdevs, alpha=0.5)

# bent
# x: average C-H bond length
bent_x = [average_bond_length_in_snapshot(xyz) for xyz in bent_scan]
# y: mean energy +/- standard deviation
plt.plot(bent_x, bent_means, label='scan from global minimum (bent)')
plt.fill_between(bent_x, bent_means - bent_stdevs, bent_means + bent_stdevs, alpha=0.5)

# horizontal line indicating global minimum
plt.hlines(0, 0,2,
           linestyles='dotted', linewidth=1, label='global minimum (bent)')


# clip x and y limits
plt.ylim(-150,400)
plt.xlim(0.85,1.35)


plt.legend(loc='lower left')
plt.tight_layout()

plt.savefig('methane_energy_scans_with_uncertainties.png', dpi=300, bbox_inches='tight')
plt.close()