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")

Provided below are the results of the script for 3 notable systems and their .cifs

I passed these cifs along to Matt from Newfoundmaterials and he mentioned the following paper as a potentially interesting resource to investigate that centers around high pressure synthesis of Manganese Hydrides: https://journals.aps.org/prb/abstract/10.1103/PhysRevB.100.224102, as well as: https://academic.oup.com/nsr/article/11/7/nwad307/7462326, that investigates ternary hydrides and their applicability as superconductors.

In the hand-off I did mention that I was skeptical of the 'stability' of these systems despite having included an 'energy_above_hull' boundary in the property sampling job. Matt echoed this sentiment but had his own MLIP up and ready:

I just ran gen_6.cif through a relaxation with the eqV2_dens_31M model, and actually, the e_hull seems to be around +0.107 eV/atom, which is far lower than I would have expected.

Next up is looping in more property evals around magnetic anisotropy and Curie temperature prediction: