# fmt: off
import numpy as np
import pytest

from ase import Atoms, units
from ase.build import bulk, molecule
from ase.calculators.emt import EMT
from ase.optimize import QuasiNewton
from ase.phonons import Phonons
from ase.thermochemistry import (
    CrystalThermo,
    HarmonicThermo,
    HinderedThermo,
    IdealGasThermo,
    MSRRHOThermo,
    QuasiHarmonicThermo,
)
from ase.vibrations import Vibrations


def test_ideal_gas_thermo_monoatomic(testdir):
    "We do a basic test on a monoatomic system"
    atoms = Atoms("N", positions=[(0, 0, 0)])
    atoms.calc = EMT()
    energy = atoms.get_potential_energy()
    vib = Vibrations(atoms)
    vib.run()
    vib_energies = vib.get_energies()
    # should have 3 modes, all zero
    assert len(vib_energies) == 3
    assert vib_energies == pytest.approx(0.0, abs=1e-8)

    thermo = IdealGasThermo(
        vib_energies=vib_energies,
        geometry="monatomic",
        atoms=atoms,
        symmetrynumber=1,
        spin=0,
        potentialenergy=energy
    )
    assert len(thermo.vib_energies) == 0
    G = thermo.get_gibbs_energy(temperature=298.15, pressure=101325.0)
    assert G == pytest.approx(4.726469689824511, abs=1e-8)


def test_ideal_gas_thermo_n2(testdir):
    "We do a basic test on N2"
    atoms = Atoms("N2", positions=[(0, 0, 0), (0, 0, 1.1)])
    atoms.calc = EMT()
    QuasiNewton(atoms).run(fmax=0.01)
    energy = atoms.get_potential_energy()
    vib = Vibrations(atoms, name="igt-vib1")
    vib.run()
    vib_energies = vib.get_energies()
    assert len(vib_energies) == 6
    assert vib_energies[0] == pytest.approx(0.0, abs=1e-8)
    assert vib_energies[-1] == pytest.approx(1.52647479e-01)

    # No selection of vibrational energies
    # Entering 6 when 1 is expected should fail
    with pytest.raises(ValueError):
        IdealGasThermo(vib_energies=vib_energies,
                       geometry="linear",
                       vib_selection='exact',
                       atoms=atoms)

    # Now, let's select

    # ---------------------
    #   #    meV     cm^-1
    # ---------------------
    #   0    0.0       0.0 <--- remove!
    #   1    0.0       0.0 <--- remove!
    #   2    0.0       0.0 <--- remove!
    #   3    1.7      13.5 <--- remove!
    #   4    1.7      13.5 <--- remove!
    #   5  152.6    1231.2
    # ---------------------

    thermo = IdealGasThermo(
        vib_energies=vib_energies,
        geometry="linear",
        atoms=atoms,
        symmetrynumber=2,
        spin=0,
        potentialenergy=energy,
        vib_selection='highest'
    )
    assert len(thermo.vib_energies) == 1
    assert thermo.vib_energies[0] == vib_energies[-1]
    assert thermo.geometry == "linear"
    assert thermo.get_ZPE_correction() == pytest.approx(0.07632373926263808)
    assert thermo.get_enthalpy(1000) == pytest.approx(0.6719935644272014)
    assert thermo.get_entropy(1e3, 1e8) == pytest.approx(0.0017861226676818658)
    assert thermo.get_gibbs_energy(1000, 1e8) == pytest.approx(
        thermo.get_enthalpy(1000) - 1000 * thermo.get_entropy(1000, 1e8)
    )


def ideal_gas_thermo_ch3(
    vib_energies,
    geometry="nonlinear",
    atoms=None,
    symmetrynumber=6,
    potentialenergy=0.0,
    spin=0.5,
    vib_selection='exact',
    ignore_imag_modes=False,
):
    if atoms is None:
        atoms = molecule("CH3")
    return IdealGasThermo(
        vib_energies=vib_energies,
        geometry=geometry,
        atoms=atoms,
        symmetrynumber=symmetrynumber,
        potentialenergy=potentialenergy,
        spin=spin,
        vib_selection=vib_selection,
        ignore_imag_modes=ignore_imag_modes,
    )


CH3_THERMO = {
    "ZPE": 1.185,
    "enthalpy": 10.610695269124156,
    "entropy": 0.0019310086280219891,
    "gibbs": 8.678687641495167,
}


def test_ideal_gas_thermo_ch3(testdir):
    """
    Now we try something a bit harder. Let's consider a
    CH3 molecule, such that there should be 3*4-6 = 6 modes
    for calculating the thermochemistry. We will also provide
    the modes in an unsorted list to make sure the correct
    values are cut. Note that these vibrational energies
    are simply toy values.

    Input: [1.0, 0.05, 0.08, 0.1, 0.2, 0.3, 0.4, 0.35, 0.12]
    Expected: [0.12, 0.2, 0.3, 0.35, 0.4, 1.0]
    """
    thermo = ideal_gas_thermo_ch3(
        vib_energies=[1.0, 0.05, 0.08, 0.1, 0.2, 0.3, 0.4, 0.35, 0.12],
        potentialenergy=9,
        vib_selection='highest',
    )
    assert len(thermo.vib_energies) == 6
    assert list(thermo.vib_energies) == [0.12, 0.2, 0.3, 0.35, 0.4, 1.0]
    assert thermo.geometry == "nonlinear"
    assert thermo.get_ZPE_correction() == pytest.approx(CH3_THERMO["ZPE"])
    assert thermo.get_enthalpy(1000) == pytest.approx(CH3_THERMO["enthalpy"])
    assert thermo.get_entropy(1e3, 1e8) == pytest.approx(CH3_THERMO["entropy"])
    assert thermo.get_gibbs_energy(1000, 1e8) == pytest.approx(
        thermo.get_enthalpy(1000) - 1000 * thermo.get_entropy(1000, 1e8)
    )


def test_ideal_gas_thermo_ch3_v2(testdir):
    """
    Let's do the same as above but provide only
    the 6 modes to use.

    Input: [0.12, 0.2, 0.3, 0.35, 0.4, 1.0]
    Expected: [0.12, 0.2, 0.3, 0.35, 0.4, 1.0]
    """
    thermo = ideal_gas_thermo_ch3(
        vib_energies=[0.12, 0.2, 0.3, 0.35, 0.4, 1.0], potentialenergy=9
    )
    assert len(thermo.vib_energies) == 6
    assert list(thermo.vib_energies) == [0.12, 0.2, 0.3, 0.35, 0.4, 1.0]
    assert thermo.geometry == "nonlinear"
    assert thermo.get_ZPE_correction() == pytest.approx(CH3_THERMO["ZPE"])
    assert thermo.get_enthalpy(1000) == pytest.approx(CH3_THERMO["enthalpy"])
    assert thermo.get_entropy(1e3, 1e8) == pytest.approx(CH3_THERMO["entropy"])
    assert thermo.get_gibbs_energy(1000, 1e8) == pytest.approx(
        thermo.get_enthalpy(1000) - 1000 * thermo.get_entropy(1000, 1e8)
    )
    assert thermo.n_imag == 0


def test_ideal_gas_thermo_ch3_v3(testdir):
    """
    Now we give the module a more complicated set of
    vibrational frequencies to deal with to make sure
    the correct values are cut. This structure is not a
    minimum or TS and has several imaginary modes. However
    if we cut the first 6 modes, it'd look like all are
    real when they are not. We need to cut based on
    np.abs() of the vibrational energies.
    """

    # ---------------------
    #   #    meV     cm^-1
    # ---------------------
    #   0   63.8i    514.8i
    #   1   63.3i    510.7i
    #   2   42.4i    342.3i
    #   3    5.3i     43.1i <--- remove!
    #   4    0.0       0.0  <--- remove!
    #   5    0.0       0.0  <--- remove!
    #   6    0.0       0.0  <--- remove!
    #   7    5.6      45.5  <--- remove!
    #   8    6.0      48.1  <--- remove!
    #   9  507.9    4096.1
    #  10  547.2    4413.8
    #  11  547.7    4417.3
    # ---------------------
    vib_energies = [
        63.8j,
        63.3j,
        42.4j,
        5.3j,
        0.0,
        0.0,
        0.0,
        5.6,
        6.0,
        507.9,
        547.2,
        547.7,
    ]
    with pytest.raises(ValueError):
        # Imaginary frequencies present!!!
        ideal_gas_thermo_ch3(vib_energies=vib_energies)

    with pytest.raises(ValueError):
        # more than one imaginary mode present
        IdealGasThermo.from_transition_state(vib_energies)
    # this one should work with one imaginary mode
    tmp = [vib_energies[3]] + vib_energies[7:]
    thermo = IdealGasThermo.from_transition_state(tmp,
                                                  geometry="nonlinear",
                                                  atoms=None,
                                                  symmetrynumber=6,
                                                  potentialenergy=0.0,
                                                  spin=0.5,
                                                  vib_selection='all',
                                                  ignore_imag_modes=False)

    # but it should fail if we try to select only the highest frequencies
    with pytest.raises(ValueError):
        thermo = IdealGasThermo.from_transition_state(tmp,
                    geometry="nonlinear",
                    atoms=None,
                    symmetrynumber=6,
                    potentialenergy=0.0,
                    spin=0.5,
                    vib_selection='highest',
                    ignore_imag_modes=False)

    # and it should fail if we try to select only the highest frequencies
    # while setting natoms incorrectly
    with pytest.raises(ValueError):
        thermo = IdealGasThermo.from_transition_state(tmp,
                    geometry="nonlinear",
                    atoms=None,
                    natoms=1,
                    symmetrynumber=6,
                    potentialenergy=0.0,
                    spin=0.5,
                    vib_selection='highest',
                    ignore_imag_modes=False)

    # the same should not work when directly creating the class
    with pytest.raises(ValueError):
        ideal_gas_thermo_ch3(vib_energies=tmp)
        thermo = ideal_gas_thermo_ch3(vib_energies=vib_energies,
                                      vib_selection='abs_highest')

    # Same as above, but let's try ignoring the
    # imag modes. This should just use: 507.9, 547.2, 547.7
    with pytest.warns(UserWarning):
        thermo = ideal_gas_thermo_ch3(vib_energies=vib_energies,
                                      ignore_imag_modes=True,
                                      vib_selection='abs_highest')

    assert list(thermo.vib_energies) == [507.9, 547.2, 547.7]
    assert thermo.n_imag == 3

    # If we (in this case erroneously) cut the first 6 modes,
    # it should pass without crash or warning
    highest_freqs_thermo = ideal_gas_thermo_ch3(vib_energies=vib_energies,
                                                vib_selection='highest')
    assert list(highest_freqs_thermo.vib_energies) == [0.0, 5.6, 6.0,
                                                       507.9, 547.2, 547.7]


def test_ideal_gas_thermo_ch3_v4(testdir):
    """
    Let's do another test like above, just for thoroughness.
    Again, this is not a minimum or TS and has several
    imaginary modes.
    """
    atoms = molecule("CH3")
    atoms.calc = EMT()
    vib = Vibrations(atoms, name="igt-vib2")
    vib.run()
    vib_energies = vib.get_energies()
    assert len(vib_energies) == 12
    assert vib_energies[0] == pytest.approx(0.09599611291404943j)
    assert vib_energies[-1] == pytest.approx(0.39035516516367375)

    # ---------------------
    #   #    meV     cm^-1
    # ---------------------
    #   0   96.0i    774.3i
    #   1   89.4i    721.0i
    #   2   89.3i    720.4i
    #   3   85.5i    689.7i <-- remove!
    #   4   85.4i    689.1i <-- remove!
    #   5   85.4i    689.1i <-- remove!
    #   6    0.0       0.0 <-- remove!
    #   7    0.0       0.0 <-- remove!
    #   8    0.0       0.0 <-- remove!
    #   9  369.4    2979.1
    #  10  369.4    2979.3
    #  11  390.4    3148.4
    # ---------------------

    with pytest.raises(ValueError):
        ideal_gas_thermo_ch3(vib_energies=vib_energies,
                             vib_selection='abs_highest')


VIB_ENERGIES_HARMONIC = np.array(
    [0.00959394 + 0.0j, 0.00959394 + 0.0j, 0.01741657 + 0.0j]
)


def test_ideal_gas_thermo_ch3_with_small_real_frequency(testdir):
    """
    It is possible for a molecule to have smaller real frequencies than
    the largest imaginary. In that case, vib_selection='abs_highest' will
    fail unintentionally. Here are possible vibrations.
    """

    vib_energies = [
        65.0j,  # <-- remove!
        20.0j,  # <-- remove!
        14.0j,  # <-- remove!
        1.0,  # <-- remove!
        4.0,  # <-- remove!
        7.0,  # <-- remove!
        60.0,
        200.0,
        500.0,
        800.0,
        1200.0,
        3000.0,
    ]

    # Expected (but erroneous) failure
    with pytest.raises(ValueError):
        ideal_gas_thermo_ch3(vib_energies=vib_energies,
                             vib_selection='abs_highest')

    # Selecting the highest is correct in this case
    highest_freqs_thermo = ideal_gas_thermo_ch3(vib_energies=vib_energies,
                                                vib_selection='highest')
    assert list(highest_freqs_thermo.vib_energies) == [60.0, 200.0, 500.0,
                                                       800.0, 1200.0, 3000.0]


def test_ideal_gas_thermo_ch3_too_few_frequencies(testdir):
    """Too few frequencies should fail except vib_selection='all' that
    skips all checks. CH3 needs 3 * 4 - 6 = 6 frequencies"""

    vib_energies = [
        200.0,
        500.0,
        800.0,
        1200.0,
        3000.0,
    ]

    for selection in ('exact', 'highest', 'abs_highest'):
        with pytest.raises(ValueError):
            ideal_gas_thermo_ch3(vib_energies=vib_energies,
                                 vib_selection=selection)

    all_freqs_thermo = ideal_gas_thermo_ch3(vib_energies=vib_energies,
                                            vib_selection='all')
    assert list(all_freqs_thermo.vib_energies) == [200.0, 500.0, 800.0,
                                                   1200.0, 3000.0]


def harmonic_thermo(
    vib_energies=None,
    potentialenergy=4.120517148154894,
    ignore_imag_modes=False,
):
    return HarmonicThermo(
        vib_energies=vib_energies if vib_energies else VIB_ENERGIES_HARMONIC,
        potentialenergy=potentialenergy,
        ignore_imag_modes=ignore_imag_modes,
    )


HELMHOLTZ_HARMONIC = 4.060698673180732


def test_harmonic_thermo(testdir):
    "Basic test of harmonic thermochemistry"
    thermo = harmonic_thermo()
    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz == pytest.approx(HELMHOLTZ_HARMONIC)


def test_harmonic_thermo_v2(testdir):
    "Test with a reversed list giving the same results"
    vib_energies = list(VIB_ENERGIES_HARMONIC)
    vib_energies.sort(reverse=True)
    thermo = harmonic_thermo(vib_energies=vib_energies)
    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz == pytest.approx(HELMHOLTZ_HARMONIC)
    assert thermo.n_imag == 0


def test_harmonic_thermo_v3(testdir):
    "Test that a proper error is raised with imag modes"
    with pytest.raises(ValueError):
        harmonic_thermo(vib_energies=[10j])


def test_harmonic_thermo_v4(testdir):
    "Test that a proper warning is raised with non-crucial imag modes"
    with pytest.warns(UserWarning):
        thermo = harmonic_thermo(
            vib_energies=list(VIB_ENERGIES_HARMONIC) + [10j],
            ignore_imag_modes=True
        )
    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz == pytest.approx(HELMHOLTZ_HARMONIC)
    assert thermo.n_imag == 1


def test_crystal_thermo(asap3, testdir):
    atoms = bulk("Al", "fcc", a=4.05)
    calc = asap3.EMT()
    atoms.calc = calc
    energy = atoms.get_potential_energy()

    # Phonon calculator
    N = 7
    ph = Phonons(atoms, calc, supercell=(N, N, N), delta=0.05)
    ph.run()

    ph.read(acoustic=True)
    dos = ph.get_dos(kpts=(4, 4, 4)).sample_grid(npts=30, width=5e-4)
    phonon_energies = dos.get_energies()
    phonon_DOS = dos.get_weights()

    thermo = CrystalThermo(
        phonon_energies=phonon_energies,
        phonon_DOS=phonon_DOS,
        potentialenergy=energy,
        formula_units=4,
    )
    thermo.get_helmholtz_energy(temperature=298.15)


HELMHOLTZ_QUASI_HARMONIC = -0.04644196376152279


def quasi_harmonic_thermo(
    vib_energies=None,
    potentialenergy=0.0,
    raise_to=100 * units.invcm
):
    return QuasiHarmonicThermo(
        vib_energies=vib_energies if vib_energies else VIB_ENERGIES_HARMONIC,
        potentialenergy=potentialenergy,
        ignore_imag_modes=False,
        raise_to=raise_to
    )


def test_quasi_harmonic_thermo():
    "Basic test of quasi-harmonic thermochemistry"
    thermo = quasi_harmonic_thermo()
    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz == pytest.approx(HELMHOLTZ_QUASI_HARMONIC)


def test_quasi_harmonic_thermo_convergence():
    "Basic test of quasi-harmonic to harmonic convergence thermochemistry"
    thermo = harmonic_thermo(potentialenergy=0.0)
    helmholtz_harm = thermo.get_helmholtz_energy(temperature=298.15)
    thermo_quasi = quasi_harmonic_thermo(raise_to=0.0)
    helmholtz_quasi = thermo_quasi.get_helmholtz_energy(temperature=298.15)
    assert helmholtz_harm == pytest.approx(helmholtz_quasi)
    # now also test that it actually changes when a higher value is used
    thermo_quasi = quasi_harmonic_thermo(raise_to=1000 * units.invcm)
    helmholtz_quasi = thermo_quasi.get_helmholtz_energy(temperature=298.15)
    assert helmholtz_harm != pytest.approx(helmholtz_quasi)


def msRRHO_thermo(
    atoms=None,
    tau=35,
    vib_energies=None,
    potentialenergy=0.0,
    **kwargs
):
    return MSRRHOThermo(
        atoms=atoms,
        tau=tau,
        vib_energies=vib_energies if vib_energies else VIB_ENERGIES_HARMONIC,
        potentialenergy=potentialenergy,
        **kwargs
    )


HELMHOLTZ_msRRHO = -0.05665130354741103
HELMHOLTZ_msRRHO_internal_energy_treatment = -0.057760466601084096


def test_msRRHO():
    "Test proper functionality of msRRHO method"
    thermo = msRRHO_thermo(atoms=Atoms('H'), tau=35)
    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz == pytest.approx(HELMHOLTZ_msRRHO)


def test_msRRHO_enthalpy_treatment():
    "Test proper functionality of msRRHO enthalpy method"
    thermo = msRRHO_thermo(atoms=Atoms('H'), tau=35, treat_int_energy=True)
    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz == pytest.approx(
        HELMHOLTZ_msRRHO_internal_energy_treatment)


def test_msRRHO_converge_to_harmonic():
    "Test that msRRHO converges to harmonic limit when tau is 0"
    thermo = harmonic_thermo(potentialenergy=0.0)
    helmholtz_harm = thermo.get_helmholtz_energy(temperature=298.15)
    thermo_msrrho = msRRHO_thermo(atoms=Atoms('H'), tau=0)
    helmholtz_msrrho = thermo_msrrho.get_helmholtz_energy(temperature=298.15)
    assert helmholtz_harm == pytest.approx(helmholtz_msrrho)
    # now also test that it actually changes when a higher tau is used
    thermo_msrrho = msRRHO_thermo(atoms=Atoms('H'), tau=10)
    helmholtz_msrrho = thermo_msrrho.get_helmholtz_energy(temperature=298.15)
    assert helmholtz_harm != pytest.approx(helmholtz_msrrho)


def test_msRRHO_scaling():
    "Test proper functionality of the scaling factor in the msRRHO method"
    tmp = np.multiply(VIB_ENERGIES_HARMONIC, 0.25)
    thermo = harmonic_thermo(vib_energies=list(tmp), potentialenergy=0.0)
    helmholtz_harm = thermo.get_helmholtz_energy(temperature=298.15)
    thermo = msRRHO_thermo(atoms=Atoms('H'), tau=0, nu_scal=0.25)
    helmholtz_msrrho = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz_harm == pytest.approx(helmholtz_msrrho)


def test_msRRHO_imag():
    "Test that the result is the same if one of the modes is imaginary"
    tmp = list(VIB_ENERGIES_HARMONIC)
    tmp[0] *= 1.0j
    thermo = msRRHO_thermo(atoms=Atoms('H'), tau=35, vib_energies=tmp)
    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz == pytest.approx(HELMHOLTZ_msRRHO)


VIB_ENERGIES_HINDERED = (
    np.array(
        [
            3049.060670,
            3040.796863,
            3001.661338,
            2997.961647,
            2866.153162,
            2750.855460,
            1436.792655,
            1431.413595,
            1415.952186,
            1395.726300,
            1358.412432,
            1335.922737,
            1167.009954,
            1142.126116,
            1013.918680,
            803.400098,
            783.026031,
            310.448278,
            136.112935,
            112.939853,
            103.926392,
            77.262869,
            60.278004,
            25.825447,
        ]
    )
    / 8065.54429
)


def hindered_thermo(
    atoms=None,
    vib_energies=None,
    trans_barrier_energy=0.049313,
    rot_barrier_energy=0.017675,
    sitedensity=1.5e15,
    rotationalminima=6,
    symmetrynumber=1,
    mass=30.07,
    inertia=73.149,
    ignore_imag_modes=False,
):
    return HinderedThermo(
        atoms=atoms,
        vib_energies=vib_energies if vib_energies else VIB_ENERGIES_HINDERED,
        trans_barrier_energy=trans_barrier_energy,
        rot_barrier_energy=rot_barrier_energy,
        sitedensity=sitedensity,
        rotationalminima=rotationalminima,
        symmetrynumber=symmetrynumber,
        mass=mass,
        inertia=inertia,
        ignore_imag_modes=ignore_imag_modes,
    )


HELMHOLTZ_HINDERED = 1.5932242071261076


def test_hindered_thermo1():
    """
    Hindered translator / rotor.
    (Taken directly from the example given in the documentation.)
    """
    thermo = hindered_thermo()
    assert len(thermo.vib_energies) == 21
    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert helmholtz == pytest.approx(HELMHOLTZ_HINDERED)


def test_hindered_thermo2():
    """
    Now reverse the vib energies and make sure results are the same
    """
    vib_energies = list(VIB_ENERGIES_HINDERED)
    vib_energies.sort(reverse=True)
    thermo = hindered_thermo(vib_energies=vib_energies)

    helmholtz = thermo.get_helmholtz_energy(temperature=298.15)
    assert len(thermo.vib_energies) == 21
    assert helmholtz == pytest.approx(HELMHOLTZ_HINDERED)
    assert thermo.n_imag == 0


def test_hindered_thermo3():
    "Now add an imaginary mode and make sure it is removed"
    with pytest.warns(UserWarning):
        thermo = hindered_thermo(
            vib_energies=list(VIB_ENERGIES_HINDERED) + [10j],
            ignore_imag_modes=True
        )
    assert thermo.get_helmholtz_energy(temperature=298.15) == pytest.approx(
        HELMHOLTZ_HINDERED
    )
    assert thermo.n_imag == 1


def test_hindered_thermo4():
    "Make sure a ValueError is raised if imag modes are present"
    with pytest.raises(ValueError):
        hindered_thermo(vib_energies=[100 + 0.1j] * 24)


def test_hindered_thermo5():
    "Make sure appropriate amount are cut"
    atoms = bulk("Cu") * (2, 2, 2)
    thermo = hindered_thermo(atoms=atoms)
    assert len(thermo.vib_energies) == 3 * len(atoms) - 3
