#!/usr/bin/env python3
import sys
import os
import argparse
from argparse import Namespace
from pathlib import Path
import yaml
import warnings

# Add current directory to sys.path
sys.path.insert(0, str(Path.cwd()))

from pymatgen.core import IStructure
from monty.serialization import loadfn

# pydefect imports
from pydefect.cli.main_functions import (
    make_supercell,
    make_defect_energy_infos_main_func,
    make_defect_energy_summary_main_func
)
from pydefect.cli.vasp.main_vasp_functions import make_defect_entries, make_calc_results
from pydefect.input_maker.defect_set_maker import DefectSetMaker
from pydefect.chem_pot_diag.chem_pot_diag import StandardEnergies, TargetVertices
from pydefect.analyzer.unitcell import Unitcell
from pydefect.analyzer.band_edge_states import PerfectBandEdgeState, EdgeInfo, OrbitalInfo
from pydefect.analyzer.defect_energy import DefectEnergyInfo
from pydefect.analyzer.make_defect_structure_info import MakeDefectStructureInfo
from pydefect.defaults import defaults
from pymatgen.core import Structure

# macer imports
from macer.calculator.factory import get_available_ffs, get_calculator
from macer.relaxation.optimizer import relax_structure
from macer.utils.logger import Logger
from macer.io.writers import write_pydefect_dummy_files
from macer.defaults import DEFAULT_MODELS, _macer_root, DEFAULT_DEVICE

# warnings
warnings.filterwarnings("ignore")

def write_summary_at_vertices(summary_path: Path):
    """
    Calculates and prints defect formation energies for each chemical potential
    vertex from a defect_energy_summary.json file.
    """
    if not summary_path.exists():
        print(f"Error: {summary_path} not found.")
        return

    try:
        # loadfn will deserialize into DefectEnergySummary object
        summary = loadfn(summary_path)
    except Exception as e:
        print(f"Error loading {summary_path}: {e}")
        return

    rel_chem_pots = summary.rel_chem_pots
    defect_energies_summary = summary.defect_energies
    output_filename = "defect_energy_summary-at-verticies.txt"

    with open(output_filename, "w") as f:
        f.write(f"title: {summary.title}\n\n")

        for vertex_label, chem_pots in rel_chem_pots.items():
            f.write(f"--- Vertex: {vertex_label} ---\n")
            chem_pot_str = ", ".join([f"{elem}: {pot:.4f}" for elem, pot in chem_pots.items()])
            f.write(f"Relative chemical potentials: {chem_pot_str}\n")
            f.write("-" * 40 + "\n")
            f.write(f"{'Defect Name':<20} {'Formation Energy (eV)'}\n")
            f.write(f"{'-'*20} {'-'*25}\n")

            formation_energies = {}
            for defect_name, defect_group in defect_energies_summary.items():
                charge_0_defect_energy = None
                try:
                    charge_0_index = defect_group.charges.index(0)
                    charge_0_defect_energy = defect_group.defect_energies[charge_0_index]
                except (ValueError, AttributeError):
                    print(f"Warning: Defect {defect_name} does not have a charge 0 state in the summary. Skipping.")
                    continue

                if not charge_0_defect_energy:
                    continue

                atom_io = defect_group.atom_io
                base_energy = charge_0_defect_energy.formation_energy
                
                reservoir_contribution = 0.0
                for element, num_diff in atom_io.items():
                    if element in chem_pots:
                        reservoir_contribution -= num_diff * chem_pots[element]

                formation_energy = base_energy + reservoir_contribution
                formation_energies[defect_name] = formation_energy
            
            sorted_defects = sorted(formation_energies.items(), key=lambda item: item[1])
            for name, energy in sorted_defects:
                f.write(f"{name:<20} {energy:>25.4f}\n")
            
            f.write("\n")

    print(f"Output written to {output_filename}")


def main():
    parser = argparse.ArgumentParser(description="Auto Defect Workflow with Macer")
    parser.add_argument("-p", "--poscar", type=str, required=True, help="Unitcell POSCAR file")
    parser.add_argument("-s", "--std_energies", type=str, required=True, help="Path to standard_energies.yaml")
    parser.add_argument("-t", "--target_vertices", type=str, required=True, help="Path to target_vertices.yaml")
    args = parser.parse_args()

    poscar_path = Path(args.poscar)
    if not poscar_path.exists():
        print(f"Error: {poscar_path} not found.")
        sys.exit(1)
    
    std_energies_path = Path(args.std_energies)
    if not std_energies_path.exists():
        print(f"Error: {std_energies_path} not found.")
        sys.exit(1)

    target_vertices_path = Path(args.target_vertices)
    if not target_vertices_path.exists():
        print(f"Error: {target_vertices_path} not found.")
        sys.exit(1)

    print("--- Step 1: Make Supercell ---")
    unitcell = IStructure.from_file(str(poscar_path))
    args_sc = Namespace(unitcell=unitcell, matrix=None, min_num_atoms=50, max_num_atoms=300, 
                        analyze_symmetry=True, sites_yaml_filename=None)
    make_supercell(args_sc)
    
    if not Path("supercell_info.json").exists():
        print("Error: supercell_info.json not created.")
        sys.exit(1)

    print("\n--- Step 2: Make Defect Set (Charge 0 only) ---")
    supercell_info = loadfn("supercell_info.json")
    maker = DefectSetMaker(supercell_info)
    
    defect_dict = {defect.name: [0] for defect in maker.defect_set}
    
    with open("defect_in.yaml", "w") as f:
        yaml.dump(defect_dict, f)
    print("defect_in.yaml created with charge states [0].")

    print("\n--- Step 3: Make Defect Entries ---")
    make_defect_entries(Namespace())
    
    dirs_to_calc = ["perfect"] + [f"{name}_0" for name in defect_dict if Path(f"{name}_0").exists()]
    print(f"Directories to calculate: {dirs_to_calc}")

    print("\n--- Step 4: Run Macer Relax ---")
    
    available_ffs = get_available_ffs()
    if not available_ffs:
        print("Error: No MLFF packages installed.")
        return
    
    ff = available_ffs[0]
    device = DEFAULT_DEVICE
    
    default_model_name = DEFAULT_MODELS.get(ff, "")
    model_path = os.path.join(_macer_root, "mlff-model", default_model_name) if default_model_name and ff not in {"fairchem", "orb", "chgnet", "m3gnet"} else default_model_name

    print(f"Initializing {ff.upper()} calculator on {device}...")
    calc_kwargs = {"device": device}
    if ff == "mace":
        calc_kwargs["model_paths"] = [model_path]
    else:
        calc_kwargs["model_path"] = model_path

    try:
        calculator = get_calculator(ff_name=ff, **calc_kwargs)
    except Exception as e:
        print(f"Failed to initialize calculator: {e}")
        return

    cwd = Path.cwd()
    successful_dirs = []
    
    for d_name in dirs_to_calc:
        d = cwd / d_name
        if not d.exists():
            print(f"Skipping {d_name}, directory not found.")
            continue
            
        print(f"Relaxing in {d_name}...")
        os.chdir(d)
        try:
            write_pydefect_dummy_files(".")
            relax_structure(input_file="POSCAR", fmax=0.03, isif=2, device=device, calculator=calculator, ff=ff,
                            outcar_name="OUTCAR", contcar_name="CONTCAR", xml_name="vasprun.xml", make_pdf=False, write_json=True)
            if Path("vasprun.xml").exists() and Path("OUTCAR").exists():
                successful_dirs.append(d)

                # Generate defect_structure_info.json directly using CONTCAR
                try:
                    if Path("CONTCAR").exists() and Path("defect_entry.json").exists():
                        final_structure = Structure.from_file("CONTCAR")
                        defect_entry = loadfn("defect_entry.json")
                        
                        dsi_maker = MakeDefectStructureInfo(
                            perfect=supercell_info.structure,
                            initial=defect_entry.structure,
                            final=final_structure,
                            symprec=defaults.symmetry_length_tolerance,
                            dist_tol=defaults.dist_tol
                        )
                        dsi_maker.defect_structure_info.to_json_file("defect_structure_info.json")
                        print(f"Generated defect_structure_info.json for {d_name}")
                except Exception as e:
                    print(f"Warning: Failed to generate defect_structure_info.json for {d_name}: {e}")

        except Exception as e:
            print(f"Relaxation failed in {d_name}: {e}")
        finally:
            os.chdir(cwd)

    print("\n--- Step 4.5: Generate calc_results.json ---")
    if successful_dirs:
        make_calc_results(Namespace(dirs=[d for d in successful_dirs], verbose=False, check_calc_results=False))
    else:
        print("No successful calculations found.")
        return

    print("\n--- Step 5: Generate Defect Energy Info ---")
    
    dummy_unitcell_str = """system: dummy
vbm: 0.0
cbm: 5.0
ele_dielectric_const: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
ion_dielectric_const: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
"""
    Path("unitcell.yaml").write_text(dummy_unitcell_str)
    
    perfect_calc_results_path = Path("perfect/calc_results.json")
    if not perfect_calc_results_path.exists():
        print("Error: perfect/calc_results.json not found.")
        return

    perfect_calc_results = loadfn(str(perfect_calc_results_path))
    unitcell_obj = Unitcell.from_yaml("unitcell.yaml")
    std_energies_obj = StandardEnergies.from_yaml(str(std_energies_path))
    
    defect_dirs = [d for d in successful_dirs if d.name != "perfect"]
    
    args_dei = Namespace(dirs=defect_dirs, check_calc_results=False, perfect_calc_results=perfect_calc_results,
                         std_energies=std_energies_obj, unitcell=unitcell_obj, verbose=False)
    
    make_defect_energy_infos_main_func(args_dei)
    print("Done. Defect energy infos generated.")

    print("\n--- Step 6: Generate Defect Energy Summary ---")
    
    vbm_info = EdgeInfo(band_idx=0, kpt_coord=(0.0, 0.0, 0.0),
                        orbital_info=OrbitalInfo(energy=0.0, orbitals={}, occupation=1.0))
    cbm_info = EdgeInfo(band_idx=0, kpt_coord=(0.0, 0.0, 0.0),
                        orbital_info=OrbitalInfo(energy=5.0, orbitals={}, occupation=0.0))
    p_state = PerfectBandEdgeState(vbm_info=vbm_info, cbm_info=cbm_info)
    
    args_des = Namespace(
        dirs=defect_dirs,
        verbose=False,
        target_vertices_yaml=str(target_vertices_path),
        unitcell=unitcell_obj,
        p_state=p_state
    )
    
    make_defect_energy_summary_main_func(args_des)
    print("defect_energy_summary.json created.")

    print("\n--- Step 7: Write summary at each vertex ---")
    summary_path = Path("defect_energy_summary.json")
    if summary_path.exists():
        write_summary_at_vertices(summary_path)
    else:
        print("defect_energy_summary.json not found, skipping final summary generation.")


if __name__ == "__main__":
    main()
