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 VCSGCEnsemble
from mchammer.observers import BinaryShortRangeOrderObserver


def run_simulation(
    supercell: Atoms,
    temperature: float,
    phi: float,
    mc_cycles: int,
    dc_filename: str,
    kappa: float = 200,  # a conservative default value
):
    """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.
    phi
        Average concentration constraint.
    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.
    kappa
        Variance of concentration constraint.
    """
    # 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 = VCSGCEnsemble(
        supercell,
        calculator=calculator,
        temperature=temperature,
        phis={'Pd': phi},
        kappa=kappa,
        dc_filename=dc_filename,
    )

    # 3. Set up short-range order observer
    sro = BinaryShortRangeOrderObserver(
        ce.get_cluster_space_copy(), supercell, radius=3.2)
    # ... and attach it to the simulation.
    mc.attach_observer(sro)

    # 4. 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(300, 800 + 1, 50)
# - Range of average constraint values to sample.
phis = np.arange(-2.9, 1.5 + 1e-3, 0.025)
# - Variance constraint.
kappa = 50

# 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 phi in phis:
        phi = np.round(phi, 5)
        print(f'temperature: {temperature}  phi: {phi:.3f}  kappa: {kappa}  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/vcsgc-T{temperature:.0f}-phi{phi:.3f}-kappa{kappa}-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,
            phi,
            mc_cycles=100,
            dc_filename=dc_filename,
            kappa=kappa)