from os import makedirs
from os.path import exists
import numpy as np

from ase import Atoms
from ase.build import make_supercell

from icet import ClusterExpansion
from mchammer import DataContainer
from mchammer.calculators import ClusterExpansionCalculator
from mchammer.ensembles import CanonicalEnsemble
from mchammer.observers import StructureFactorObserver


def run_simulation(
    supercell: Atoms,
    temperature: float,
    mc_cycles: int,
    dc_filename: str,
    lattice_parameter: float,
):
    """Convenience function for running a Monte Carlo simulation.
    Simulations are carried out at constant temperature and composition
    in the canonical ensemble. The long-range order in the system is observed
    by monitoring the structure factor at a set of q-points.

    Parameters
    ----------
    supercell
        Atomic structure to sample.
    temperature
        Temperature in Kelvin at which to sample.
    mc_cycles
        Number of MC cycles to sample. 1 MC cycle equals n_sites MC trials,
        where n_sites is the number of sites in the structure.
    dc_filename
        Name of output file.
    lattice_parameter
        Lattice parameter of the underlying lattice.
        This value is used to specify the q-points at which the
        structure factor is sampled during the simulation.
    """
    # 1. Set up a cluster expansion calculator.
    # Here, we are using a cluster expansion that has been constructed to
    # describe the AuPd:H system (doi: 10.1016/j.actamat.2021.116893).
    # We are, however, only interested in the Au-Pd lattice.
    # To sample the energy per Au/Pd atom (rather than per both Au/Pd and
    # H/vacancy sites), we therefore overwrite the default scaling parameter.
    # In the vast majority of cases, such consideration is not needed but the
    # present situation serves as a useful demonstration for the (occasional)
    # utility of the `scaling` keyword.
    calculator = ClusterExpansionCalculator(supercell, ce, scaling=len(supercell) / 2)

    # 2. Set up the canonical ensemble object.
    n_sites = int(len(supercell) / 2)  # number of sites on Au/Pd lattice
    interval = n_sites
    mc = CanonicalEnsemble(
        supercell,
        calculator=calculator,
        temperature=temperature,
        dc_filename=dc_filename,
        ensemble_data_write_interval=interval,
    )

    # 3. Set up the long range order (LRO) observer.
    qpoints = 2 * np.pi / lattice_parameter * np.eye(3)
    LRO = StructureFactorObserver(
        supercell,
        q_points=qpoints,
        interval=interval,
        symbol_pairs=[['Pd', 'Pd'], ['Au', 'Au']])
    # ... and attach it to the simulation.
    mc.attach_observer(LRO)

    # 4. Run the simulation.
    steps = int(len(supercell) * mc_cycles / 2)
    mc.run(steps)

    # We return the last configuration.
    return mc.data_container.get('trajectory')[-1]


# Simulation parameters.
# - Supercell size in number of unit cells in each direction.
size = 4
# - Temperature range to sample.
#   To cut down the equilibration time, (as you will see below), we are
#   reusing the configuration from the previous temperature.
#   We therefore sample from high to low temperatures.
temperatures = range(300, 120 - 1, -10)
# - Pd concentration at which to sample.
concentration = 0.75

# 1. Load the cluster expansion.
ce = ClusterExpansion.read('AuPd.ce')

# 2a. Prepare primitive structure
# Here, we are using a cluster expansion that has been constructed to
# describe the AuPd:H system (doi: 10.1016/j.actamat.2021.116893).
# We are, however, only interested in the H-free case. We therefore
# set the species on the H-lattice to "X", which implies a vacancy.
primitive_structure = ce.primitive_structure
primitive_structure.symbols = ['Au', 'X']

# 2b. Set up supercell structure for simulation.
# The primitive structure is based on an FCC lattice. To maximize
# the spacing between equivalent sites (and also for ease of
# visualization) it is preferable to use a conventional (simple
# cubic) cell. We therefore first convert the FCC lattice to the
# conventional cell.
supercell = make_supercell(
    primitive_structure, [[-1, 1, 1], [1, -1, 1], [1, 1, -1]])

# We remember the lattice parameter of the conventional cell, which we
# will use later to set up the q-points at which to sample the structure
# factor.
lattice_parameter = np.linalg.norm(supercell.cell, axis=1)[0]

# Then we repeat the cell.
supercell = supercell.repeat(size)

# And finally set the occupation of the Au/Pd lattice to the
# desired Pd concentration.
available_sites = [a.index for a in supercell if a.symbol == 'Au']
n_Pd = int(concentration * len(available_sites))
selection = np.random.choice(available_sites, size=n_Pd, replace=False)
supercell.symbols[selection] = 'Pd'

# 3. Run Monte Carlo simulations.
makedirs('runs-canonical', exist_ok=True)
for temperature in temperatures:
    print(f'temperature: {temperature}  size: {size}')

    # If we can read the data container from a previous run, we will
    # restart from the last configuration and add to the existing data.
    dc_filename = f'runs/T{temperature:.0f}-size{size}.dc'
    if exists(dc_filename):
        dc = DataContainer.read(dc_filename)
        supercell = dc.get('trajectory')[-1]

    supercell = run_simulation(
        supercell,
        temperature,
        mc_cycles=2000,
        dc_filename=dc_filename,
        lattice_parameter=lattice_parameter)