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 SemiGrandCanonicalEnsemble


def run_simulation(
    supercell: Atoms,
    temperature: float,
    dmu: float,
    mc_cycles: int,
    dc_filename: str,
):
    """Convenience function for running a Monte Carlo simulation.
    Simulations are carried out at constant temperature and chemical potential
    in the semi-grandcanonical ensemble.

    Parameters
    ----------
    supercell
        Atomic structure to sample.
    temperature
        Temperature in Kelvin at which to sample.
    dmu
        Chemical potential difference in eV/atom.
    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.
    """
    # 1. Set up a cluster expansion calculator.
    # Here, we are using a cluster expansion that has been constructed to
    # describe the AgPd system (doi: 10.1002/adts.201900015).
    calculator = ClusterExpansionCalculator(supercell, ce)

    # 2. Set up the canonical ensemble object.
    mc = SemiGrandCanonicalEnsemble(
        supercell,
        calculator=calculator,
        temperature=temperature,
        chemical_potentials={'Ag': 0, 'Pd': dmu},
        dc_filename=dc_filename,
    )

    # 3. Run the simulation.
    steps = len(supercell) * mc_cycles
    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 = 6
# - Temperature range to sample.
temperatures = range(100, 800 + 1, 50)
# - Chemical potential range to sample.
dmus = np.arange(-1.04, 1.04 + 1e-3, 0.02)

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

# 2. 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.
conventional_structure = make_supercell(
    ce.primitive_structure, [[-1, 1, 1], [1, -1, 1], [1, 1, -1]])
supercell = conventional_structure.repeat(size)

# 3. Run Monte Carlo simulations.
makedirs('runs', exist_ok=True)
for temperature in temperatures:
    for dmu in dmus:
        dmu = np.round(dmu, 5)
        print(f'temperature: {temperature}  dmu: {dmu:.3f}  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/sgc-T{temperature:.0f}-dmu{dmu:.3f}-size{size}.dc'
        if exists(dc_filename):
            dc = DataContainer.read(dc_filename)
            if len(dc.data) > 100:
                continue
            supercell = dc.get('trajectory')[-1]
    
        supercell = run_simulation(
            supercell,
            temperature,
            dmu,
            mc_cycles=100,
            dc_filename=dc_filename)