import numpy as np
import pytest
from ase.build import bulk
from ase.calculators.lj import LennardJones
from ase.calculators.emt import EMT
from ase.calculators.qmmm import ForceQMMM, RescaledCalculator
from ase.eos import EquationOfState
from ase.optimize import FIRE
from ase.neighborlist import neighbor_list


@pytest.mark.slow
def test_forceqmmm():
    # parameters
    N_cell = 2
    R_QMs = np.array([3, 7])

    # setup bulk and MM region
    bulk_at = bulk("Cu", cubic=True)
    sigma = (bulk_at*2).get_distance(0, 1)*(2.**(-1./6))
    mm = LennardJones(sigma=sigma, epsilon=0.05)
    qm = EMT()

    # compute MM and QM equations of state
    def strain(at, e, calc):
        at = at.copy()
        at.set_cell((1.0 + e)*at.cell, scale_atoms=True)
        at.calc = calc
        v = at.get_volume()
        e = at.get_potential_energy()
        return v, e

    eps = np.linspace(-0.01, 0.01, 13)
    v_qm, E_qm = zip(*[strain(bulk_at, e, qm) for e in eps])
    v_mm, E_mm = zip(*[strain(bulk_at, e, mm) for e in eps])

    eos_qm = EquationOfState(v_qm, E_qm)
    v0_qm, E0_qm, B_qm = eos_qm.fit()
    a0_qm = v0_qm**(1.0/3.0)

    eos_mm = EquationOfState(v_mm, E_mm)
    v0_mm, E0_mm, B_mm = eos_mm.fit()
    a0_mm = v0_mm**(1.0/3.0)

    mm_r = RescaledCalculator(mm, a0_qm, B_qm, a0_mm, B_mm)
    v_mm_r, E_mm_r = zip(*[strain(bulk_at, e, mm_r) for e in eps])

    eos_mm_r = EquationOfState(v_mm_r, E_mm_r)
    v0_mm_r, E0_mm_r, B_mm_r = eos_mm_r.fit()
    a0_mm_r = v0_mm_r**(1.0/3)

    # check match of a0 and B after rescaling is adequete
    assert abs((a0_mm_r - a0_qm)/a0_qm) < 1e-3 # 0.1% error in lattice constant
    assert abs((B_mm_r - B_qm)/B_qm) < 0.05 # 5% error in bulk modulus

    # plt.plot(v_mm, E_mm - np.min(E_mm), 'o-', label='MM')
    # plt.plot(v_qm, E_qm - np.min(E_qm), 'o-', label='QM')
    # plt.plot(v_mm_r, E_mm_r - np.min(E_mm_r), 'o-', label='MM rescaled')
    # plt.legend()

    at0 = bulk_at * N_cell
    r = at0.get_distances(0, np.arange(1, len(at0)), mic=True)
    print(len(r))
    del at0[0] # introduce a vacancy
    print("N_cell", N_cell, 'N_MM', len(at0))

    ref_at = at0.copy()
    ref_at.calc = qm
    opt = FIRE(ref_at)
    opt.run(fmax=1e-3)
    u_ref = ref_at.positions - at0.positions

    us = []
    for R_QM in R_QMs:
        at = at0.copy()
        mask = r < R_QM
        print('R_QM', R_QM, 'N_QM', mask.sum(), 'N_total', len(at))
        qmmm = ForceQMMM(at, mask, qm, mm, buffer_width=2*qm.rc)
        at.calc = qmmm
        opt = FIRE(at)
        opt.run(fmax=1e-3)
        us.append(at.positions - at0.positions)

    # compute error in energy norm |\nabla u - \nabla u_ref|
    def strain_error(at0, u_ref, u, cutoff, mask):
        I, J = neighbor_list('ij', at0, cutoff)
        I, J = np.array([(i,j) for i, j in zip(I, J) if mask[i]]).T
        v = u_ref - u
        dv = np.linalg.norm(v[I, :] - v[J, :], axis=1)
        return np.linalg.norm(dv)

    du_global = [strain_error(at0, u_ref, u, 1.5*sigma, np.ones(len(r))) for u in us]
    du_local = [strain_error(at0, u_ref, u, 1.5*sigma, r < 3.0) for u in us]

    print('du_local', du_local)
    print('du_global', du_global)

    # check local errors are monotonically decreasing
    assert np.all(np.diff(du_local) < 0)

    # check global errors are monotonically converging
    assert np.all(np.diff(du_global) < 0)

    # biggest QM/MM should match QM result
    assert du_local[-1] < 1e-10
    assert du_global[-1] < 1e-10
