from chgnet.model import CHGNetCalculator
from ase.calculators.singlepoint import SinglePointCalculator
from raffle.generator import raffle_generator
from ase import build, Atoms
from ase.optimize import FIRE
from ase.io import write, read
import numpy as np
import os
from joblib import Parallel, delayed
from pathlib import Path

script_dir = Path(__file__).resolve().parent


# Function to relax a structure
def process_structure(i, atoms, num_structures_old, calc_params, optimise_structure, iteration):
    # Check if the structure has already been processed
    if i < num_structures_old:
        return None, None, None
    
    # calc = Vasp(**calc_params, label=f"struct{i}", directory=f"iteration{iteration}/struct{i}/", txt=f"stdout{i}.o")
    inew = i - num_structures_old
    atoms.calc = calc

    # Calculate and save the initial energy per atom
    energy_unrlxd = atoms.get_potential_energy() / len(atoms)
    print(f"Initial energy per atom: {energy_unrlxd}")

    # Optimise the structure if requested
    if optimise_structure:
        optimizer = FIRE(atoms, trajectory = f"traje{inew}.traj", logfile=f"optimisation{inew}.log")
        try:
            optimizer.run(fmax=0.05, steps=100)
        except Exception as e:
            print(f"Optimisation failed: {e}")
            return None, None, None
    
    # Save the optimised structure and its energy per atom
    energy_rlxd = atoms.get_potential_energy() / len(atoms)

    # Get the distance matrix
    distances = atoms.get_all_distances(mic=True)

    # Set self-distances (diagonal entries) to infinity to ignore them
    np.fill_diagonal(distances, np.inf)

    # Check if the minimum non-self distance is below 1.5
    if distances.min() < 1.0:
        print(f"Distance too small: {atoms.get_all_distances(mic=True).min()}")
        return None, None, None
    
    if abs(energy_rlxd - energy_unrlxd) > 10.0:
        print(f"Energy difference too large: {energy_rlxd} vs {energy_unrlxd}")
        return None, None, None
    
    return atoms, energy_unrlxd, energy_rlxd


if __name__ == "__main__":

    # set up the calculator
    calc_params = {}
    calc = CHGNetCalculator()

    hosts = []
    # get all hosts in ../../data/C-MgO_hosts/ directory
    for filename in os.listdir(script_dir / ".." / ".." / ".." / "data" / "C-MgO_hosts"):
        if filename.startswith("POSCAR"):
            atoms = read(script_dir / ".." / ".." / ".." / "data" / "C-MgO_hosts" / filename)
            if not all(element == 'C' for element in set(atoms.get_chemical_symbols())):
                print(f"Skipping {filename} as it has non-carbon atoms")
                continue
            elif len(atoms) > 20:
                print(f"Skipping {filename} as it has more than 20 atoms (want to do smaller cells for now)")
                continue
            hosts.append(read(script_dir / ".." / ".." / ".." / "data" / "C-MgO_hosts" / filename))
            hosts[-1].calc = calc

    # set up element reference energies
    C_reference = Atoms('C4', positions=[
        [0.000000,    0.000000,    1.950768],
        [0.000000,    0.000000,    5.852305],
        [0.000000,    1.424491,    1.950768],
        [1.233646,    0.712246,    5.852305],
        ], cell=[[1.2336, -2.1367, 0.0], [1.2336, 2.1367, 0.0], [0.0, 0.0, 7.8031]], pbc=True)
    C_reference.calc = calc
    C_reference_energy = C_reference.get_potential_energy() / len(C_reference)
    Mg_reference = Atoms("Mg2", positions=[
            [0.000000,    1.831369,    1.285301],
            [1.586012,    0.915684,    3.855904],
        ], cell=[
            [1.58601,   -2.7471,    0.0000],
            [1.58601,    2.7471,    0.0000],
            [0.00000,    0.0001,    5.1412],
        ], pbc=True)
    Mg_reference.calc = calc
    Mg_reference_energy = Mg_reference.get_potential_energy() / len(Mg_reference)
    O_reference = Atoms("O2", positions=[
        [0.0, 0.0, 0.0],
        [1.23, 0.0, 0.0]], cell=[10, 10, 10], pbc=False)
    O_reference.calc = calc
    O_reference_energy = O_reference.get_potential_energy() / len(O_reference)

    generator = raffle_generator()
    generator.distributions.set_element_energies(
        {
            'C': C_reference_energy,
            'Mg': Mg_reference_energy,
            'O': O_reference_energy
        }
    )

    # set the covalent radii
    generator . distributions . set_bond_radii (
        {
            ('C', 'Mg'): 1.4,
            ('Mg', 'Mg'): 1.4,
            ('C', 'O'): 1.4,
            # ('C', 'Mg'): 1.15, # 2.3 = what Joe used https://github.com/ExeQuantCode/RAFFLE/blob/1cfdf2d220dce41f20ace4cce5a4135a37d1f1b1/chem.in
            # ('C', 'O'): 0.587,
            # ('Mg', 'Mg'): 1.592,
            # INITIALISE O2 BOND TO LARGER VALUE
        }
    )

    # set energy scale
    generator.distributions.set_kBT(0.4)

    # set the distribution function widths (2-body, 3-body, 4-body)
    generator.distributions.set_width([0.025, np.pi/200.0, np.pi/200.0])

    # generate MgO bulk
    MgO_atom = Atoms("Mg4O4", positions=[
            [0.0, 0.0, 0.0],
            [0.0, 2.097, 2.097],
            [2.097, 0.0, 2.097],
            [2.097, 2.097, 0.0],
            [0.0, 0.0, 2.097],
            [0.0, 2.097, 0.0],
            [2.097, 0.0, 0.0],
            [2.097, 2.097, 2.097],
        ], cell=[4.1940, 4.1940, 4.1940], pbc=True)
    MgO_atom.calc = calc

    # read C-Mg-O database, obtained from Materials Project
    print("Reading database")
    # database = read("../database.xyz", index=":")
    database = [ C_reference, Mg_reference, O_reference, MgO_atom ]
    generator.distributions.create(database, deallocate_systems=False)

    # # set up the initial database
    # initial_database = []
    # for i, atoms in enumerate(database):
    #     if atoms.calc is None:
    #         database.remove(atoms)
    #         continue
    #     # check if any other species in the system other than C, Mg, O
    #     if not all(element in ['C', 'Mg', 'O'] for element in set(atoms.get_chemical_symbols())):
    #         print(f"Skipping structure {i} as it has non-C-Mg-O atoms")
    #         continue
    #     atoms.calc = calc
    #     initial_database.append(atoms)
    # initial_database.append(C_reference)
    # initial_database.append(Mg_reference)
    # initial_database.append(O_reference)
    # initial_database.append(MgO_atom)
    # generator.distributions.create(initial_database, deallocate_systems=False)

    # set parameters
    seed = 0

    # check if the energies file exists, if not create it
    energies_rlxd_filename = f"energies_rlxd_seed{seed}.txt"
    energies_unrlxd_filename = f"energies_unrlxd_seed{seed}.txt"
    if os.path.exists(energies_rlxd_filename):
        with open(energies_rlxd_filename, "w") as energy_file:
            pass
    else:
        open(energies_rlxd_filename, "w").close()

    if os.path.exists(energies_unrlxd_filename):
        with open(energies_unrlxd_filename, "w") as energy_file:
            pass
    else:
        open(energies_unrlxd_filename, "w").close()
    
    # initialise the number of structures generated
    iter = -1
    num_structures_old = 0
    unrlxd_structures = []
    rlxd_structures = []
    optimise_structure = True
    mass_C = 12.011
    mass_Mg = 24.305
    mass_O = 15.999
    generator.init_seed(put=seed)
    # loop over all hosts
    for host in hosts:
        print("setting host")
        generator.set_host(host)
        volume = host.get_volume()
        num_atoms_C = len(host)
    
        host_positions = host.get_scaled_positions(wrap=True)
        host_height = np.linalg.norm(host.get_cell()[2, :])
        print(f"host_height: {host_height}")

        # get all unique z values
        z_values = np.unique(host_positions[:, 2])
        print(f"z_values: {z_values}")
        # return the two closes to 0.5 in the cell
        z_values = z_values[np.argsort(np.abs(z_values - 0.5))[:2]]
        print(f"z_values: {z_values}")
        # set the bounds for the generator
        z_min = min(z_values) + 1.0 / host_height
        z_max = max(z_values) - 1.0 / host_height
        generator.set_bounds([[0, 0, z_min], [1, 1, z_max]])

        # loop over range of MgO densities to intercalate
        for num_atoms in range(1, 50):

            num_atoms_Mg = num_atoms
            num_atoms_O = num_atoms
            density = ( mass_C * num_atoms_C + mass_Mg * num_atoms_Mg + mass_O * num_atoms_O ) / volume
            if density > 2.2:
                # print("Density too high:", density, "u/A^3")
                continue
            elif density < 1.0:
                # print("Density too low", density, "u/A^3")
                continue
            iter += 1

            # generate structures
            generator.generate(
                num_structures = 5,
                stoichiometry = { 'Mg': num_atoms_Mg, 'O': num_atoms_O },
                method_ratio = {"void": 0.1, "rand": 0.01, "walk": 0.25, "grow": 0.25, "min": 1.0},
                verbose = 0,
            )

            # print the number of structures generated
            print("Total number of structures generated: ", generator.num_structures)
            generated_structures = generator.get_structures(calc)
            num_structures_new = len(generated_structures)

            # check if directory iteration[iter] exists, if not create it
            iterdir = f"iteration{iter}/"
            if not os.path.exists(iterdir):
                os.makedirs(iterdir)
            generator.print_settings(iterdir+"generator_settings.txt")

            # set up list of energies
            energy_unrlxd = np.zeros(num_structures_new - num_structures_old)
            energy_rlxd = np.zeros(num_structures_new - num_structures_old)
            for i in range(num_structures_new - num_structures_old):
                write(iterdir+f"POSCAR_unrlxd_{i}", generated_structures[num_structures_old + i])
                print(f"Structure {i} energy per atom: {generated_structures[num_structures_old + i].get_potential_energy() / len(generated_structures[num_structures_old + i])}")
                unrlxd_structures.append(generated_structures[num_structures_old + i].copy())
                unrlxd_structures[-1].calc = SinglePointCalculator(
                    generated_structures[num_structures_old + i],
                    energy=generated_structures[num_structures_old + i].get_potential_energy(),
                    forces=generated_structures[num_structures_old + i].get_forces()
                )
            
            # Start parallel execution
            print("Starting parallel execution")
            results = Parallel(n_jobs=5)(
                delayed(process_structure)(i, generated_structures[i].copy(), num_structures_old, calc_params, optimise_structure, iteration=seed)
                for i in range(num_structures_old, num_structures_new)
            )

            # Wait for all futures to complete
            for j, result in enumerate(results):
                generated_structures[num_structures_old + j], energy_unrlxd[j], energy_rlxd[j] = result
                if generated_structures[num_structures_old + j] is None:
                    print("Structure failed the checks")
                    continue
                rlxd_structures.append(generated_structures[num_structures_old + j].copy())
                rlxd_structures[-1].calc = SinglePointCalculator(
                    generated_structures[j+num_structures_old],
                    energy=generated_structures[num_structures_old + j].get_potential_energy(),
                    forces=generated_structures[num_structures_old + j].get_forces()
                )
            print("All futures completed")

            # Remove structures that failed the checks
            for j, atoms in reversed(list(enumerate(generated_structures))):
                if j < num_structures_old:
                    continue
                if atoms is None:
                    energy_unrlxd = np.delete(energy_unrlxd, j-num_structures_old)
                    energy_rlxd = np.delete(energy_rlxd, j-num_structures_old)
                    del generated_structures[j]
                    del unrlxd_structures[j]
                    del rlxd_structures[j]
                    generator.remove_structure(j)
            num_structures_new = len(generated_structures) 

            # write the structures to files
            for i in range(num_structures_new - num_structures_old):
                write(iterdir+f"POSCAR_{i}", generated_structures[num_structures_old + i])
                print(f"Structure {i} energy per atom: {energy_rlxd[i]}")
                # append energy per atom to the 'energies_unrlxd_filename' file
                with open(energies_unrlxd_filename, "a") as energy_file:
                    energy_file.write(f"{i+num_structures_old} {energy_unrlxd[i]}\n")
                # append energy per atom to the 'energies_rlxd_filename' file
                with open(energies_rlxd_filename, "a") as energy_file:
                    energy_file.write(f"{i+num_structures_old} {energy_rlxd[i]}\n")

            # update the distribution functions
            print("Updating distributions")
            generator.distributions.update(generated_structures[num_structures_old:], from_host=False, deallocate_systems=False)

            # print the new distribution functions to a file
            print("Printing distributions")
            generator.distributions.write_dfs(iterdir+"distributions.txt")
            generator.distributions.write_2body(iterdir+"df2.txt")
            generator.distributions.write_3body(iterdir+"df3.txt")
            generator.distributions.write_4body(iterdir+"df4.txt")
            generator.distributions.deallocate_systems()

            # update the number of structures generated
            num_structures_old = num_structures_new

    # Write the final distribution functions to a file
    generator.distributions.write_gdfs(f"gdfs_seed{seed}.txt")

    # Read energies from the file
    with open(energies_rlxd_filename, "r") as energy_file:
        energies = energy_file.readlines()

    # Parse and sort the energies
    energies = [line.strip().split() for line in energies]
    energies = sorted(energies, key=lambda x: float(x[1]))

    # Write the sorted energies back to the file
    with open(f"sorted_{energies_rlxd_filename}", "w") as energy_file:
        for entry in energies:
            energy_file.write(f"{int(entry[0])} {float(entry[1])}\n")

    # Write the structures to files
    write(f"unrlxd_structures_seed{seed}.traj", unrlxd_structures)
    write(f"rlxd_structures_seed{seed}.traj", rlxd_structures)
    print("All generated and relaxed structures written")

    print("Learning complete")