"""
Thermodynamic integration
=========================
This example shows how to perform thermodynamic integration (TI).

We here demonstrate the TI from an effective harmonic potential to either EMT
or MTP.

Note that the ASE EMT is very slow; the timing is:

- With MTP: takes around 16 s on Apple M1 Pro single core.
- With ASAP EMT: takes around 12 s on Apple M1 Pro single core.
- With ASE EMT: takes around 110 s on Apple M1 Pro single core.

Therefore, if ASAP is available, this exmaple uses its EMT implementation.

You can run the script either on a single core as ``python ./ti_example.py``
or in parallel as ``mpirun -n 9 python ./ti_example.py``, where 3 * 3 = 9
(volume, temperature) grid points are parallelized over.
"""

from itertools import product
from pathlib import Path

import numpy as np
from ase.build import bulk
from ase.parallel import parprint

import directupsampling
from directupsampling.calculators.emt import EMT
from directupsampling.effqh.effqh import EffQH
from directupsampling.plot_tools import plot
from directupsampling.ti import ThermoInt

# from directupsampling.calculators.mtp import MTP

# %%
# Get pretrained EffQH and the target calculator.
resource_path = Path(directupsampling.__file__).parents[1] / "tests/resources"
calc = EMT()

# calc = MTP.read_potential(resource_path / "full_test/twostep_emt_4g_with_hs.mtp")
# calc.species = (13,)  # motep
# calc.species = ("Al",)  # mlippy

# %%
# Make grid.
vmin = 13.0
vmax = 17.0
nvolumes = 2
tmin = 298.0
tmax = 500.0
ntemperatures = 2
grid_fah = product(
    np.linspace(vmin, vmax, nvolumes),
    np.linspace(tmin, tmax, ntemperatures),
)

# %%
# Make an initial ASE Atoms object.
atoms = bulk("Al", "fcc", cubic=True) * 3

# %%
# Make EffQH
effqh = EffQH.read_fc_fit(resource_path / "full_test/fc_fit.hdf5", atoms)

# %%
# Run thermodynamic integration.
nlambdas = 5
ti = ThermoInt(
    atoms,
    effqh.calc,
    calc,
    grid_fah,
    lambdas=np.linspace(0.0, 1.0, nlambdas),
    sampling_setting={"sampling_offset": 0},
)
ti.run_sampling(nsnapshots=5, file="test.db")
fdiff_qh_to_mtp = ti.free_energy_differences

# %%
# Write the TI snapshots to an ASE db format and print and save *F*\ :sup:`ah`.
ti.snapshot_container.write("thermodynamic_integration.db")
parprint("Results in meV/atom:")
parprint(fdiff_qh_to_mtp * 1e3)
fdiff_qh_to_mtp.to_csv("fah.csv")
# %%
# Plot Δ\ *U* as a function of λ.
plot(ti).show()