import numpy as np
import pytest
from numpy.testing import assert_array_almost_equal

from ase import Atoms
from ase.calculators.calculator import CalculationFailed, CalculatorSetupError
from ase.calculators.emt import EMT
from ase.calculators.harmonic import HarmonicCalculator, HarmonicForceField
from ase.calculators.mixing import MixedCalculator
from ase.geometry.geometry import get_distances_derivatives
from ase.md.andersen import Andersen
from ase.md.velocitydistribution import (MaxwellBoltzmannDistribution,
                                         Stationary, ZeroRotation)
from ase.optimize import BFGS
from ase.units import fs
from ase.vibrations import Vibrations

ref_pos = np.asarray([[8.7161, 7.96276, 8.48206], [8.60594, 8.04985, 9.44464],
                      [8.0154, 8.52264, 8.10545]])
ref_atoms = Atoms('OH2', positions=ref_pos)  # relaxed water molecule
ref_energy = -14.222189  # y shift of the 'parabola' (harmonic potential)

# example Hessian matrix as obtained from DFT
hessian_x = np.asarray([[2.82630333e+01, -2.24763667e+01, 7.22478333e+00,
                         -2.96970000e+00, 2.34363333e+00, 2.72788333e+00,
                         -2.52159833e+01, 2.01307833e+01, -9.94651667e+00],
                        [-2.24763667e+01, 1.78621333e+01, -5.77378333e+00,
                         2.33703333e+00, -1.85276667e+00, -2.15118333e+00,
                         2.01258667e+01, -1.60350833e+01, 7.93248333e+00],
                        [7.22478333e+00, -5.77378333e+00, 5.72735000e+01,
                         7.25470000e+00, -5.75313333e+00, -4.69477333e+01,
                         -1.44613000e+01, 1.15504833e+01, -1.03523333e+01],
                        [-2.96970000e+00, 2.33703333e+00, 7.25470000e+00,
                         2.96170000e+00, -2.36901667e+00, -3.76841667e+00,
                         -2.83833333e-02, 3.06833333e-02, -3.49190000e+00],
                        [2.34363333e+00, -1.85276667e+00, -5.75313333e+00,
                         -2.36901667e+00, 1.89046667e+00, 2.95495000e+00,
                         2.90666667e-02, -1.80666667e-02, 2.79565000e+00],
                        [2.72788333e+00, -2.15118333e+00, -4.69477333e+01,
                         -3.76841667e+00, 2.95495000e+00, 4.89340000e+01,
                         1.03146667e+00, -8.18450000e-01, -1.96118333e+00],
                        [-2.52159833e+01, 2.01258667e+01, -1.44613000e+01,
                         -2.83833333e-02, 2.90666667e-02, 1.03146667e+00,
                         2.52034000e+01, -2.01516833e+01, 1.34293167e+01],
                        [2.01307833e+01, -1.60350833e+01, 1.15504833e+01,
                         3.06833333e-02, -1.80666667e-02, -8.18450000e-01,
                         -2.01516833e+01, 1.60592333e+01, -1.07369667e+01],
                        [-9.94651667e+00, 7.93248333e+00, -1.03523333e+01,
                         -3.49190000e+00, 2.79565000e+00, -1.96118333e+00,
                         1.34293167e+01, -1.07369667e+01, 1.23150000e+01]])


def assert_water_is_relaxed(atoms):
    forces = atoms.get_forces()
    assert np.allclose(np.zeros(forces.shape), forces)
    assert np.allclose(ref_energy, atoms.get_potential_energy())
    assert np.allclose(atoms.get_angle(1, 0, 2), ref_atoms.get_angle(1, 0, 2))
    assert np.allclose(atoms.get_distance(0, 1), ref_atoms.get_distance(0, 1))
    assert np.allclose(atoms.get_distance(0, 2), ref_atoms.get_distance(0, 2))


def run_optimize(atoms):
    opt = BFGS(atoms)
    opt.run(fmax=1e-9)


@pytest.mark.optimize()
def test_cartesians():
    """In Cartesian coordinates the first 6 trash eigenvalues (translations and
    rotations) can be slightly different from zero; hence set them to zero
    using an increased parameter zero_thresh.
    """
    zero_thresh = 0.06  # set eigvals to zero if abs(eigenvalue) < zero_thresh
    hff = HarmonicForceField(ref_atoms=ref_atoms, ref_energy=ref_energy,
                             hessian_x=hessian_x, zero_thresh=zero_thresh)
    assert np.allclose(hff.hessian_q, hff.hessian_x)
    atoms = ref_atoms.copy()
    atoms.calc = HarmonicCalculator(hff)
    assert_water_is_relaxed(atoms)  # atoms has not been distorted
    run_optimize(atoms)             # nothing should happen
    assert_water_is_relaxed(atoms)  # atoms should still be relaxed
    atoms.set_distance(0, 1, 3.5)   # now distort atoms along axis, no rotation
    run_optimize(atoms)             # optimization should recover original
    assert_water_is_relaxed(atoms)    # relaxed geometry

    with pytest.raises(AssertionError):
        atoms.rattle()                  # relaxation should fail to recover the
        atoms.rotate(90, 'x')           # original geometry of the atoms,
        run_optimize(atoms)             # because Cartesian coordinates are
        assert_water_is_relaxed(atoms)  # not rotationally invariant.


@pytest.mark.optimize()
def test_constraints_with_cartesians():
    """Project out forces along x-component of H-atom (index 0 in the q-vector
    with the Cartesian coordinates (here: x=q)). A change in the x-component of
    the H-atom should not result in restoring forces, when they were projected
    out from the Hessian matrix.
    """
    def test_forces(calc):
        atoms = ref_atoms.copy()
        atoms.calc = calc
        pos = ref_pos.copy()
        pos[0, 0] *= 2
        atoms.set_positions(pos)
        run_optimize(atoms)  # (no) restoring force along distorted x-component
        xdiff = atoms.get_positions() - ref_pos
        return all(xdiff[xdiff != 0] == pos[0, 0] / 2)

    zero_thresh = 0.06  # set eigvals to zero if abs(eigenvalue) < zero_thresh
    parameters = {'ref_atoms': ref_atoms, 'ref_energy': ref_energy,
                  'hessian_x': hessian_x, 'zero_thresh': zero_thresh}
    hff = HarmonicForceField(**parameters)
    calc = HarmonicCalculator(hff)
    assert not test_forces(calc)  # restoring force along distorted x-component

    parameters['constrained_q'] = [0]  # project out the coordinate with index 0
    hff = HarmonicForceField(**parameters)
    calc = HarmonicCalculator(hff)
    assert test_forces(calc)  # no restoring force along distorted x-component


def setup_water(calc):
    atoms = ref_atoms.copy()
    atoms.calc = calc
    assert_water_is_relaxed(atoms)
    atoms.rattle(0.3)
    atoms.rotate(160, 'x')
    assert not np.allclose(ref_energy, atoms.get_potential_energy())
    return atoms


# start doc example 3
dist_defs = [[0, 1], [1, 2], [2, 0]]  # define three distances by atom indices


def water_get_q_from_x(atoms):
    """Simple internal coordinates to describe water with three distances."""
    q_vec = [atoms.get_distance(i, j) for i, j in dist_defs]
    return np.asarray(q_vec)


def water_get_jacobian(atoms):
    """Function to return the Jacobian for the water molecule described by
    three distances."""
    pos = atoms.get_positions()
    dist_vecs = [pos[j] - pos[i] for i, j in dist_defs]
    derivs = get_distances_derivatives(dist_vecs)
    jac = []
    for i, defin in enumerate(dist_defs):
        dqi_dxj = np.zeros(ref_pos.shape)
        for j, deriv in enumerate(derivs[i]):
            dqi_dxj[defin[j]] = deriv
        jac.append(dqi_dxj.flatten())
    return np.asarray(jac)
# end doc example 3


def test_raise_Errors():
    with pytest.raises(CalculatorSetupError):
        HarmonicForceField(ref_atoms=ref_atoms, hessian_x=hessian_x,
                           get_q_from_x=lambda x: x)
    with pytest.raises(CalculatorSetupError):
        HarmonicForceField(ref_atoms=ref_atoms, hessian_x=hessian_x,
                           variable_orientation=True)
    with pytest.raises(CalculatorSetupError):
        HarmonicForceField(ref_atoms=ref_atoms, hessian_x=hessian_x,
                           cartesian=False)
    with pytest.raises(CalculationFailed):
        hff = HarmonicForceField(ref_atoms=ref_atoms, ref_energy=ref_energy,
                                 hessian_x=hessian_x,
                                 get_q_from_x=water_get_q_from_x,
                                 get_jacobian=lambda x: np.ones((3, 9)),
                                 cartesian=True, variable_orientation=True)
        calc = HarmonicCalculator(hff)
        setup_water(calc)


@pytest.mark.optimize()
def test_internals():
    parameters = {'ref_atoms': ref_atoms, 'ref_energy': ref_energy,
                  'hessian_x': hessian_x, 'get_q_from_x': water_get_q_from_x,
                  'get_jacobian': water_get_jacobian, 'cartesian': False}
    hff = HarmonicForceField(**parameters)  # calculation in internals
    calc = HarmonicCalculator(hff)
    atoms = setup_water(calc)  # distorted copy of ref_atoms
    run_optimize(atoms)        # recover original configuration
    assert_water_is_relaxed(atoms)

    parameters['cartesian'] = True  # calculation in Cartesian Coordinates
    hff = HarmonicForceField(**parameters)
    calc = HarmonicCalculator(hff)
    atoms = setup_water(calc)       # 'variable_orientation' not set to True!
    # but water has rotational degrees of freedom
    run_optimize(atoms)
    with pytest.raises(AssertionError):  # hence forces were incorrect
        assert_water_is_relaxed(atoms)   # original configuration not recovered

    parameters['variable_orientation'] = True
    hff = HarmonicForceField(**parameters)
    calc = HarmonicCalculator(hff)
    atoms = setup_water(calc)
    run_optimize(atoms)
    assert_water_is_relaxed(atoms)  # relaxation succeeded despite rotation


@pytest.mark.optimize()
def test_compatible_with_ase_vibrations():
    atoms = ref_atoms.copy()
    atoms.calc = EMT()
    run_optimize(atoms)
    opt_atoms = atoms.copy()
    opt_energy = atoms.get_potential_energy()
    vib = Vibrations(atoms, nfree=2)
    vib.run()
    energies = vib.get_energies()
    vib_data = vib.get_vibrations()
    hessian_2d = vib_data.get_hessian_2d()
    vib.clean()

    hff = HarmonicForceField(ref_atoms=opt_atoms, ref_energy=opt_energy,
                             hessian_x=hessian_2d)
    calc_harmonic = HarmonicCalculator(hff)
    atoms = ref_atoms.copy()
    atoms.calc = calc_harmonic
    vib = Vibrations(atoms, nfree=4, delta=1e-5)
    vib.run()
    assert np.allclose(energies, vib.get_energies())
    vib.clean()
    hff = HarmonicForceField(ref_atoms=ref_atoms, ref_energy=ref_energy,
                             hessian_x=hessian_2d,
                             get_q_from_x=water_get_q_from_x,
                             get_jacobian=water_get_jacobian, cartesian=True)
    calc_harmonic = HarmonicCalculator(hff)
    atoms = ref_atoms.copy()
    atoms.calc = calc_harmonic
    vib = Vibrations(atoms, nfree=4, delta=1e-5)
    vib.run()  # 3 transl and 3 rot are removed by internal coordinates
    assert_array_almost_equal(energies[-3:], vib.get_energies()[-3:], decimal=2)


def test_thermodynamic_integration():
    parameters = {'ref_atoms': ref_atoms, 'ref_energy': ref_energy,
                  'hessian_x': hessian_x, 'get_q_from_x': water_get_q_from_x,
                  'get_jacobian': water_get_jacobian, 'cartesian': True,
                  'variable_orientation': True}
    hff_1 = HarmonicForceField(**parameters)
    calc_harmonic_1 = HarmonicCalculator(hff_1)
    parameters['cartesian'] = False
    hff_0 = HarmonicForceField(**parameters)
    calc_harmonic_0 = HarmonicCalculator(hff_0)
    ediffs = {}  # collect energy difference for varying lambda coupling
    lambs = [0.00, 0.25, 0.50, 0.75, 1.00]  # integration grid
    for lamb in lambs:
        ediffs[lamb] = []
        calc_linearCombi = MixedCalculator(calc_harmonic_0, calc_harmonic_1,
                                           1 - lamb, lamb)
        atoms = ref_atoms.copy()
        atoms.calc = calc_linearCombi
        MaxwellBoltzmannDistribution(atoms, temperature_K=300, force_temp=True)
        Stationary(atoms)
        ZeroRotation(atoms)
        with Andersen(atoms, 0.5 * fs, temperature_K=300, andersen_prob=0.05,
                      fixcm=False) as dyn:
            for _ in dyn.irun(50):  # should be much longer for production runs
                e0, e1 = calc_linearCombi.get_energy_contributions(atoms)
                ediffs[lamb].append(float(e1) - float(e0))
            ediffs[lamb] = np.mean(ediffs[lamb])
    dA = np.trapz([ediffs[lamb] for lamb in lambs], x=lambs)  # anharm. corr.
    assert -0.005 < dA < 0.005  # the MD run is to short for convergence
    if dA == 0.0:
        raise ValueError('there is most likely something wrong, but it could '
                         'also be sheer coincidence')
