from ase.atoms import Cell
from ase.io import read
import subprocess
from tempfile import TemporaryDirectory

COMP = {}
ELEMENTS = {}
CELL = {}

def load_compositions(comp_list):

    for comp in comp_list:

        name = comp["name"]
        COMP[name] = comp["elements"]
        ELEMENTS[name] = list(comp["elements"].keys())

        # define the cell or the volume
        if "cell" in comp.keys():
            CELL[name] = comp["cell"]
        else:
            volume = comp["volume"]
            CELL[name] = [[2.0, 0, 0],
                        [0, 2.0, 0],
                        [0, 0, volume/4.0]]
                              

    # returns random structure generated by AIRSS
    # with volume defined within CELL
    # cell_fixed = True the cell will be exactly the same as in CELL[comp]
    # cell_fixed = False the cell will be randomised but the volume will be almost the same
def random_airss_structure(comp, cell_fixed=True):

    cell = CELL[comp]
    comp = COMP[comp]

    l = list(comp.values())
    num_atoms = sum([i[0] for i in l])
    volume = Cell(cell).volume / num_atoms * 3.11

    cell_fixed = "#FIX\n" if cell_fixed else ""

    # create input file in .cell format for the buildcell tool
    input_str = f"%BLOCK LATTICE_CART\n"
    for i in range(3):
        input_str += f"{cell[i][0]} {cell[i][1]} {cell[i][2]}\n"
    input_str += f"{cell_fixed}%%ENDBLOCK LATTICE_CART\n\n#VARVOL={volume}\n# #FIX=true\n%BLOCK POSITIONS_FRAC"

    for a in comp.keys():
        input_str += f"\n{a} 0.0 0.0 0.0 # {a}1 %% NUM={comp[a][0]}"

    input_str += "\n%%ENDBLOCK POSITIONS_FRAC\n\n#MINSEP=1"

    result = subprocess.run(["buildcell"], input = bytes(input_str, "utf-8"), stdout=subprocess.PIPE)

    # Create a temporary directory
    with TemporaryDirectory() as temp_dir:

        # save output .cell text to a file for the following reading by the ASE library
        f = open(f"{temp_dir}/tmp.cell", "w")
        f.write(result.stdout.decode("utf-8"))
        f.close()

        atoms = read(f"{temp_dir}/tmp.cell")
        # print(f"Random structure created. Volume: {atoms.cell.volume}")

    return atoms