import os
import numpy as np
import itertools
import shutil
from collections import Counter
from pymatgen.core import Structure
from smact import element_dictionary, metals, neutral_ratios
from smact.screening import pauling_test
from pymatgen.analysis.structure_matcher import StructureMatcher
from tqdm import tqdm

# --- FIX: Use relative paths for Hugging Face ---
input_folder = "generation_cifs"
output_folder = "valid_generation_cifs"
unique_folder = "unique_structures"
novel_folder = "novel_structures"
# ------------------------------------------------

def get_safe_space_group_number(structure):
    try:
        spg_info = structure.get_space_group_info()
        return spg_info[1] if spg_info else None
    except Exception:
        return None
    
def structure_validity(structure, cutoff=0.5):
    dist_mat = structure.distance_matrix
    dist_mat = dist_mat + np.diag(np.ones(dist_mat.shape[0]) * (cutoff + 10.))
    if dist_mat.min() < cutoff or structure.volume < 0.1:
        return False
    spg_number = get_safe_space_group_number(structure)
    if spg_number is None:
        return False
    return True

def smact_validity(atom_types, use_pauling_test=True, include_alloys=True):
    elem_counter = Counter(atom_types)
    elems = [(elem, elem_counter[elem]) for elem in sorted(elem_counter.keys())]
    comp, elem_counts = list(zip(*elems))
    elem_counts = np.array(elem_counts)
    elem_counts = elem_counts / np.gcd.reduce(elem_counts)
    count = tuple(elem_counts.astype("int").tolist())

    elem_symbols = tuple(comp)
    try:
        space = element_dictionary(elem_symbols)
    except:
        return False # Handle elements not in SMACT

    smact_elems = [e[1] for e in space.items()]
    electronegs = [e.pauling_eneg for e in smact_elems]
    ox_combos = [e.oxidation_states for e in smact_elems]

    if len(set(elem_symbols)) == 1:
        return True
    if include_alloys:
        is_metal_list = [elem_s in metals for elem_s in elem_symbols]
        if all(is_metal_list):
            return True

    threshold = np.max(count)
    oxn = 1
    for oxc in ox_combos:
        oxn *= len(oxc)
    if oxn > 1e7:
        return False

    for ox_states in itertools.product(*ox_combos):
        stoichs = [(c,) for c in count]
        cn_e, cn_r = neutral_ratios(ox_states, stoichs=stoichs, threshold=threshold)
        if cn_e:
            if use_pauling_test:
                electroneg_OK = pauling_test(ox_states, electronegs)
            else:
                electroneg_OK = True
            if electroneg_OK:
                return True
    return False

def get_atom_types_from_structure(structure):
    return [str(element) for element in structure.species]

# Create directories
os.makedirs(output_folder, exist_ok=True)
os.makedirs(unique_folder, exist_ok=True)
os.makedirs(novel_folder, exist_ok=True)

total = 0
structure_valid_count = 0
composition_valid_count = 0
both_valid_count = 0

if os.path.exists(input_folder):
    for filename in os.listdir(input_folder):
        if filename.endswith(".cif"):
            total += 1
            full_path = os.path.join(input_folder, filename)
            try:
                structure = Structure.from_file(full_path)
                atom_types = get_atom_types_from_structure(structure)

                struct_valid = structure_validity(structure)
                comp_valid = smact_validity(atom_types)

                if struct_valid:
                    structure_valid_count += 1
                if comp_valid:
                    composition_valid_count += 1
                if struct_valid and comp_valid:
                    both_valid_count += 1
                    shutil.copy(full_path, os.path.join(output_folder, filename))
            except Exception as e:
                print(f"Error processing {filename}: {e}")

print(f"\n=== Validity Summary ===")
print(f"Total CIFs checked:              {total}")
print(f"Structural validity passed:     {structure_valid_count}")
print(f"Compositional validity passed:  {composition_valid_count}")
print(f"Both valid (saved):             {both_valid_count}")