from pathlib import Path
from subprocess import check_call

import numpy as np
from numpy.testing import assert_allclose

from pymatgen.core import Lattice, Structure
from pymatgen.io.aims.inputs import AimsControlIn, AimsCube, AimsGeometryIn
from pymatgen.io.aims.outputs import AimsOutput

# AIMS_CMD should be modified to match your system
AIMS_CMD = "aims.x"
AIMS_OUTPUT = "aims.out"
AIMS_SD = "species_dir"
AIMS_TEST_DIR = "../../tests/io/aims/species_directory/light/"

# Create test structure
structure = Structure(
    lattice=Lattice(np.array([[0, 2.715, 2.715], [2.715, 0, 2.715], [2.715, 2.715, 0]])),
    species=["Si", "Si"],
    coords=np.array([np.zeros(3), np.ones(3) * 0.25]),
)

# Create the geometry file from the structure
geo_in = AimsGeometryIn.from_structure(structure)

# Create the control.in file
cont_in = AimsControlIn(
    {
        "xc": "pw-lda",
        "relax_geometry": "trm 0.01",
        "relax_unit_cell": "full",
        "species_dir": AIMS_SD,
    }
)

# Add new parameters as if AimsControl
cont_in["k_grid"] = [1, 1, 1]

# Output options to control in automatically append the list
cont_in["output"] = "hirshfeld"
cont_in["output"] = ["eigenvectors"]

# Cube file output controlled by the AimsCube class
cont_in["cubes"] = [
    AimsCube("total_density", origin=[0, 0, 0], points=[11, 11, 11]),
    AimsCube("eigenstate_density 1", origin=[0, 0, 0], points=[11, 11, 11]),
]

# Write the input files
work_dir = Path.cwd() / "workdir/"
work_dir.mkdir(exist_ok=True)

geo_in.write_file(work_dir, overwrite=True)
cont_in.write_file(structure, work_dir, overwrite=True)

# Run the calculation
with open(f"{work_dir}/{AIMS_OUTPUT}", "w") as outfile:
    aims_run = check_call([AIMS_CMD], cwd=work_dir, stdout=outfile)

# Read the aims output file and the final relaxed geometry
outputs = AimsOutput.from_outfile(f"{work_dir}/{AIMS_OUTPUT}")
relaxed_structure = AimsGeometryIn.from_file(f"{work_dir}/geometry.in.next_step")

# Check the results
assert outputs.get_results_for_image(-1).lattice == relaxed_structure.structure.lattice
assert_allclose(outputs.get_results_for_image(-1).frac_coords, relaxed_structure.structure.frac_coords)
assert_allclose(outputs.get_results_for_image(-1).properties["stress"], outputs.stress)
assert_allclose(outputs.get_results_for_image(-1).site_properties["force"], outputs.forces)