import os
import zipfile
from chgnet.model.model import CHGNet
from pymatgen.core import Structure

# Define the path to the ZIP archive
zip_file_path = '/teamspace/studios/this_studio/mattergen/results/chemical_system_energy_above_hull/generated_crystals_cif.zip'

# Define a directory to extract the CIF files
extract_dir = 'extracted_cif_files'
os.makedirs(extract_dir, exist_ok=True)

# Extract all files from the ZIP archive
with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
    zip_ref.extractall(extract_dir)

# Load the CHGNet model
chgnet = CHGNet.load()

results = []

# Loop over the 16 expected CIF files (named gen_0.cif to gen_15.cif)
for i in range(16):
    cif_filename = f'gen_{i}.cif'
    cif_path = os.path.join(extract_dir, cif_filename)
    
    if not os.path.exists(cif_path):
        print(f"File {cif_filename} not found. Skipping.")
        continue
    
    # Load the structure from the CIF file
    structure = Structure.from_file(cif_path)
    
    # Get the CHGNet predictions for the structure
    prediction = chgnet.predict_structure(structure)
    
    # Extract magnetic moments and energy from CHGNet's output
    magmom = prediction.get('m', None)
    magmom_list = magmom.tolist() if magmom is not None else None
    energy = prediction.get('e', None)
    
    # Calculate the total magnetic moment (sum over all atoms)
    total_mag = sum(magmom) if magmom is not None else None
    
    # Calculate magnetic density = (total magnetic moment) / (unit cell volume)
    # Unit cell volume is given in ų, so the resulting unit is μB/ų.
    mag_density = total_mag / structure.volume if total_mag is not None else None

    # Get the structure's mass density (from pymatgen) for reference (in g/cm³)
    structure_density = structure.density
    
    results.append({
        "filename": cif_filename,
        "magnetic_moments (μB)": magmom_list,
        "energy (eV)": energy,
        "structure_density (g/cm³)": structure_density,
        "magnetic_density (μB/ų)": mag_density,
    })

# Pretty print the results in a formatted layout
print("\nCHGNet Predictions:")
print("=" * 40)
for result in results:
    print(f"File: {result['filename']}")
    print("-" * 40)
    print(f"Magnetic Moments (μB): {result['magnetic_moments (μB)']}")
    print(f"Energy (eV): {result['energy (eV)']}")
    print(f"Structure Density (g/cm³): {result['structure_density (g/cm³)']:.4f}")
    if result["magnetic_density (μB/ų)"] is not None:
        print(f"Magnetic Density (μB/ų): {result['magnetic_density (μB/ų)']:.4e}\n")
    else:
        print("Magnetic Density (μB/ų): None\n")

File: gen_6.cif
----------------------------------------
Magnetic Moments (μB): [3.226381778717041, 3.2163245677948, 3.2346415519714355, 3.238424777984619, 3.242114543914795, 3.246438503265381, 0.04040634632110596, 0.040341854095458984, 0.026384830474853516, 0.026986122131347656, 0.02667105197906494, 0.026878714561462402, 0.02722954750061035, 0.026971817016601562]
Energy (eV): -5.461234092712402
Structure Density (g/cm³): 5.8438
Magnetic Density (μB/ų): 1.4821e-01

File: gen_12.cif
----------------------------------------
Magnetic Moments (μB): [2.4848427772521973, 3.066071033477783, 3.071133852005005, 3.0688695907592773, 3.066573143005371, 2.483609437942505, 0.03734135627746582, 0.03735232353210449, 0.055411696434020996, 0.05534994602203369, 0.05536293983459473, 0.05529987812042236]
Energy (eV): -5.8867878913879395
Structure Density (g/cm³): 6.4693
Magnetic Density (μB/ų): 1.4710e-01

File: gen_13.cif
----------------------------------------
Magnetic Moments (μB): [2.7075822353363037, 2.737920045852661, 2.7561793327331543, 2.947281837463379, 2.884899139404297, 2.6326050758361816, 2.7517340183258057, 0.029694557189941406, 0.025991439819335938, 0.040144920349121094, 0.030662059783935547, 0.038489460945129395, 0.04551875591278076, 0.039252400398254395, 0.038976430892944336, 0.03884077072143555]
Energy (eV): -5.592731952667236
Structure Density (g/cm³): 6.2167
Magnetic Density (μB/ų): 1.4150e-01