import os
import glob
import shutil
from pathlib import Path
import openmc

# --- inputs ---
LIB_ROOT = Path("/home/stafie_c/my_afs/openmc/lib_mine")

LIB_DIRS = [
    "endfb71_official",  # ENDF/B-VII.1
    "endfb80_official",  # ENDF/B-VIII.0
    "endfb81_official",  # ENDF/B-VIII.1
    "jeff32_other",  # JEFF3.2
    "jeff33_official",  # JEFF3.3
    "jeff40/300",  # JEFF4
    "jendl4",  # JENDL4
    "jendl5",  # JENDL5
]

LIB_TAGS = {
    "endfb71_official": "e71",
    "endfb80_official": "e80",
    "endfb81_official": "e81",
    "jeff32_other": "j32",
    "jeff33_official": "j33",
    "jeff40/300": "j40",
    "jendl4": "jn4",
    "jendl5": "jn5",
}

# Natural carbon split by atom fraction (approx.)
C12_FRAC = 0.9889
C13_FRAC = 0.0111

CARBON_MODE = {
    "endfb71_official": "TO_C0", 
    "endfb80_official": "TO_C12C13",  
    "endfb81_official": "TO_C12C13",  
    "jeff32_other": "TO_C0",
    "jeff33_official": "TO_C0",
    "jeff40/300": "TO_C12C13", 
    "jendl4": "TO_C0", 
    "jendl5": "TO_C12C13",  
}


def cleanup_previous_run_artifacts():
    # Remove artifacts from previous run in the current folder (do NOT delete base model.xml)
    for pat in [
        "statepoint.*.h5",
        "summary.h5",
        "tallies.out.h5",
        "openmc_simulation_n*.h5",
    ]:
        for f in glob.glob(pat):
            try:
                os.remove(f)
            except OSError:
                pass


def rename_outputs(tag: str):
    # Rename any file ending with .h5 to .h5_{tag}
    for f in glob.glob("*.h5"):
        new_name = f"{f}_{tag}"
        os.rename(f, new_name)


def _get_af(mat: openmc.Material, nuc: str) -> float:
    try:
        return float(mat.get_nuclide_atom_densities(nuc).get(nuc, 0.0))
    except Exception:
        return 0.0


def _remove_if_present(mat, nuc: str) -> None:
    try:
        mat.remove_nuclide(nuc)
    except Exception:
        pass


def _ensure_to_c0(mat):
    
    # print("from C12/C13 to C0")
    # breakpoint()

    c0 = _get_af(mat, "C0")
    c12 = _get_af(mat, "C12")
    c13 = _get_af(mat, "C13")

    total = c0 + c12 + c13
    if total > 0.0:
        _remove_if_present(mat, "C0")
        _remove_if_present(mat, "C12")
        _remove_if_present(mat, "C13")
        mat.add_nuclide("C0", total)


def _ensure_to_c12c13(mat):

    # print("from C0 to C12/C13")
    # breakpoint()

    c0 = _get_af(mat, "C0")
    if c0 > 0.0:
        _remove_if_present(mat, "C0")
        mat.add_nuclide("C12", c0 * C12_FRAC)
        mat.add_nuclide("C13", c0 * C13_FRAC)
        

def apply_carbon_mode(materials, mode: str):
    for mat in materials:
        if mode == "TO_C0":
            _ensure_to_c0(mat)
        elif mode == "TO_C12C13":
            _ensure_to_c12c13(mat)
        else:
            raise ValueError(f"Unknown carbon mode: {mode}")
    return materials


def main() -> None:
    if not Path("model.xml").exists():
        raise FileNotFoundError(f"Missing model.xml in current directory")

    for lib_dir in LIB_DIRS:
        tag = LIB_TAGS[lib_dir]
        print(f"--- Running with library: {lib_dir} (tag: {tag}) ---")
        xs = LIB_ROOT / lib_dir / "cross_sections.xml"
        if not xs.exists():
            raise FileNotFoundError(f"Missing cross_sections.xml: {xs}")

        # Make a temporary copy of model.xml -> original_model.xml
        original_model_xml = Path("original_model.xml")
        with open("model.xml", "rb") as src, open(original_model_xml, "wb") as dst:
            dst.write(src.read())

        # Load, apply library + carbon rules, save back
        model = openmc.Model.from_model_xml("original_model.xml")
        model.materials.cross_sections = str(xs)
        mode = CARBON_MODE[lib_dir]
        model.materials = apply_carbon_mode(model.materials, mode)
        model.export_to_model_xml()

        # breakpoint()

        openmc.run()

        # Rename outputs with tag and model.xml to model_{tag}.xml
        rename_outputs(tag)
        staged = Path("model.xml")
        if staged.exists():
            staged.rename(f"model_{tag}.xml")

        # Restore original model.xml from original_model.xml
        if original_model_xml.exists():
            with open(original_model_xml, "rb") as src, open("model.xml", "wb") as dst:
                dst.write(src.read())
            original_model_xml.unlink()


if __name__ == "__main__":
    main()