#!/usr/bin/env python3
import sys
import os
import shutil
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 Structure, IStructure, Composition
from monty.serialization import loadfn

# pydefect imports
from pydefect.cli.vasp.main_vasp_functions import (
    make_competing_phase_dirs, 
    make_composition_energies,
    make_defect_entries,
    make_calc_results
)
from pydefect.cli.main_functions import (
    make_standard_and_relative_energies, 
    make_cpd_and_vertices,
    make_supercell,
    make_defect_energy_infos_main_func,
    make_defect_energy_summary_main_func
)
from pydefect.defaults import defaults as pydefect_defaults
from pydefect.defaults import defaults
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.make_defect_structure_info import MakeDefectStructureInfo

# 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:
        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}")
            f.write("\n")

            formation_energies = {}
            for defect_name, defect_group in defect_energies_summary.items():
                charge_0_defect_energy = None
                try:
                    if 0 in defect_group.charges:
                        charge_0_index = defect_group.charges.index(0)
                        charge_0_defect_energy = defect_group.defect_energies[charge_0_index]
                except (ValueError, AttributeError):
                    pass

                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 get_unique_dir_name(base_name):
    if not Path(base_name).exists():
        return Path(base_name)
    i = 1
    while True:
        new_name = f"{base_name}-NEW{i:02d}"
        if not Path(new_name).exists():
            return Path(new_name)
        i += 1

def run_macer_relax(target_dirs, isif=3, supercell_info=None):
    if not target_dirs:
        return []

    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 = None
    if default_model_name:
        FFS_USING_MODEL_NAME = {"fairchem", "orb", "chgnet", "m3gnet"}
        if ff in FFS_USING_MODEL_NAME:
            model_path = default_model_name
        else:
            model_path = os.path.join(_macer_root, "mlff-model", default_model_name)
    
    print(f"Initializing {ff.upper()} calculator on {device}...")
    calc_kwargs = {"model_path": model_path, "device": device}
    if ff == "mace":
        calc_kwargs["model_paths"] = [calc_kwargs["model_path"]]
        del calc_kwargs["model_path"]

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

    successful_dirs = []
    cwd = Path.cwd()
    original_stdout = sys.stdout

    for d in target_dirs:
        print(f"Relaxing in {d.name}...")
        os.chdir(d) 
        try:
            log_name = f"relax-POSCAR_log.txt"
            with Logger(log_name) as lg:
                sys.stdout = lg
                write_pydefect_dummy_files(".") 
                relax_structure(
                    input_file="POSCAR",
                    fmax=0.03,
                    isif=isif,
                    device=device,
                    calculator=calculator, 
                    ff=ff,
                    outcar_name="OUTCAR", 
                    contcar_name="CONTCAR",
                    xml_name="vasprun.xml",
                    make_pdf=True,
                    write_json=True
                )
            
            if Path("vasprun.xml").exists() or Path("OUTCAR").exists():
                 successful_dirs.append(d)
                 
                 # Generate defect_structure_info.json directly using CONTCAR if supercell_info is provided
                 if supercell_info:
                     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:
            sys.stdout = original_stdout
            print(f"Relaxation failed in {d.name}: {e}")
        finally:
            sys.stdout = original_stdout
            os.chdir(cwd)
            
    return successful_dirs

def main():
    parser = argparse.ArgumentParser(description="Full Auto Defect Workflow (CPD + Defects) using Macer")
    parser.add_argument("-p", "--poscar", type=str, required=True, help="Input POSCAR file (Perfect Unitcell)")
    parser.add_argument("--matrix", nargs="+", type=int, help="Supercell matrix applied to the conventional cell. 1, 3 or 9 components are accepted.")
    parser.add_argument("--min_atoms", type=int, default=50, help="Minimum number of atoms (default: 50)")
    parser.add_argument("--max_atoms", type=int, default=300, help="Maximum number of atoms (default: 300)")
    parser.add_argument("--no_symmetry_analysis", dest="analyze_symmetry", action="store_false", help="Set if symmetry is not analyzed. If set, sites.yaml file is required.")
    parser.set_defaults(analyze_symmetry=True)
    parser.add_argument("-s", "--sites_yaml", type=str, dest="sites_yaml_filename", help="Path to sites.yaml file.")
    args = parser.parse_args()

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

    sites_yaml_path = None
    if args.sites_yaml_filename:
        sites_yaml_path = Path(args.sites_yaml_filename).absolute()
        if not sites_yaml_path.exists():
            print(f"Error: {sites_yaml_path} not found.")
            sys.exit(1)

    # 1. Setup Directory
    try:
        structure = Structure.from_file(str(input_poscar_path))
        formula = structure.composition.reduced_formula
        input_poscar_name = input_poscar_path.name
        
        base_dir_name = f"DEFECT-{input_poscar_name}-formula={formula}"
        work_dir = get_unique_dir_name(base_dir_name)
        
        print(f"Creating workspace: {work_dir}")
        work_dir.mkdir()
        os.chdir(work_dir)
        
        # Copy input POSCAR to POSCAR-unitcell
        shutil.copy(str(input_poscar_path), "POSCAR-unitcell")
        
        # Reload structure to ensure formula matches the file we work with
        structure = Structure.from_file("POSCAR-unitcell")
        formula = structure.composition.reduced_formula
        host_elements = [str(e) for e in structure.composition.elements]
        
        print(f"Target Formula: {formula}")
        print(f"Host Elements: {host_elements}")

    except Exception as e:
        print(f"Initialization failed: {e}")
        sys.exit(1)

    # ------------------------------------------------------------------
    # Step 2: CPD Workflow
    # ------------------------------------------------------------------
    print("\n--- Starting CPD Workflow ---")
    
    # Create cpd directory and work inside it
    root_dir = Path.cwd()
    cpd_dir = Path("cpd")
    cpd_dir.mkdir(exist_ok=True)
    shutil.copy("POSCAR-unitcell", cpd_dir / "POSCAR-unitcell")
    
    os.chdir(cpd_dir)
    print(f"Changed working directory to: {os.getcwd()}")
    
    args_mp = Namespace(elements=host_elements, e_above_hull=pydefect_defaults.e_above_hull)
    try:
        make_competing_phase_dirs(args_mp)
    except Exception as e:
        print(f"Failed to generate competing phase directories: {e}")
    
    target_dirs = []
    # Identify competing phase directories.
    for d in Path.cwd().iterdir():
        if d.is_dir() and d.name not in ["cpd", "defect"] and (d / "POSCAR").exists():
             target_dirs.append(d)
    
    # Ensure host material is also calculated for CPD
    target_comp_dir = Path(f"{formula}_host")
    if not target_comp_dir.exists():
        target_comp_dir.mkdir()
        shutil.copy("POSCAR-unitcell", target_comp_dir / "POSCAR")
    
    if target_comp_dir not in target_dirs:
        target_dirs.append(target_comp_dir)

    successful_cpd_dirs = run_macer_relax(target_dirs, isif=3)
    
    yaml_filename = "composition_energies.yaml"
    yaml_input = yaml_filename if Path(yaml_filename).exists() else None
    args_mce = Namespace(yaml_file=yaml_input, dirs=successful_cpd_dirs, verbose=False)
    make_composition_energies(args_mce)
    
    if Path(yaml_filename).exists():
        try:
            with open(yaml_filename, 'r') as f:
                comp_energies = yaml.safe_load(f)
            
            filtered_comp_energies = {}
            removed_compounds = []
            for formula_key, data in comp_energies.items():
                comp = Composition(formula_key)
                elements_in_comp = [str(e) for e in comp.elements]
                impurities_in_comp = [e for e in elements_in_comp if e not in host_elements]
                if len(impurities_in_comp) <= 1:
                    filtered_comp_energies[formula_key] = data
                else:
                    removed_compounds.append(formula_key)
            
            if removed_compounds:
                print(f"Filtering out compounds with multiple impurities: {removed_compounds}")
            
            with open(yaml_filename, 'w') as f:
                yaml.dump(filtered_comp_energies, f)
        except Exception as e:
            print(f"Failed to filter composition energies: {e}")

    args_sre = Namespace(composition_energies_yaml=yaml_filename)
    make_standard_and_relative_energies(args_sre)
    
    # Verify target exists in relative energies
    if Path("relative_energies.yaml").exists():
        try:
            with open("relative_energies.yaml") as f:
                rel_energies_keys = yaml.safe_load(f).keys()
            if formula not in rel_energies_keys:
                print(f"WARNING: Target '{formula}' not found in relative_energies.yaml keys: {list(rel_energies_keys)}")
        except Exception as e:
            print(f"Error checking relative_energies: {e}")

    args_cv = Namespace(rel_energy_yaml="relative_energies.yaml", target=formula, elements=host_elements)
    make_cpd_and_vertices(args_cv)
    
    print("CPD Workflow Finished.")
    
    # Return to root directory
    os.chdir(root_dir)

    # ------------------------------------------------------------------
    # Step 3: Defect Workflow
    # ------------------------------------------------------------------
    print("\n--- Starting Defect Workflow ---")
    defect_dir = Path("defect")
    defect_dir.mkdir(exist_ok=True)
    
    shutil.copy("POSCAR-unitcell", defect_dir / "POSCAR-unitcell")
    
    # Paths relative to the defect dir (which we will enter)
    # std_energies_path is in ../cpd/standard_energies.yaml
    std_energies_path = (cpd_dir / "standard_energies.yaml").absolute()
    target_vertices_path = (cpd_dir / "target_vertices.yaml").absolute()
    
    os.chdir(defect_dir)
    
    unitcell_structure = IStructure.from_file("POSCAR-unitcell")
    args_sc = Namespace(
        unitcell=unitcell_structure, 
        matrix=args.matrix, 
        min_num_atoms=args.min_atoms, 
        max_num_atoms=args.max_atoms, 
        analyze_symmetry=args.analyze_symmetry, 
        sites_yaml_filename=str(sites_yaml_path) if sites_yaml_path else None
    )
    make_supercell(args_sc)
    
    if not Path("supercell_info.json").exists():
        print("Error: supercell_info.json not created.")
        return

    supercell_info = loadfn("supercell_info.json")
    maker = DefectSetMaker(supercell_info)
    
    # Restrict to charge 0 only, as per run_auto_defect.py
    defect_dict = {defect.name: [0] for defect in maker.defect_set}
    with open("defect_in.yaml", "w") as f:
        yaml.dump(defect_dict, f)
    
    make_defect_entries(Namespace())
    
    defect_calc_dirs = []
    for d in Path.cwd().iterdir():
        if d.is_dir() and (d / "POSCAR").exists() and (d / "defect_entry.json").exists():
            defect_calc_dirs.append(d)
    if Path("perfect").exists() and Path("perfect/POSCAR").exists():
        defect_calc_dirs.append(Path("perfect"))
    
    defect_calc_dirs = list(set(defect_calc_dirs))
    
    successful_defect_dirs = run_macer_relax(defect_calc_dirs, isif=2, supercell_info=supercell_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)
    unitcell_obj = Unitcell.from_yaml("unitcell.yaml")
    
    if successful_defect_dirs:
        make_calc_results(Namespace(dirs=successful_defect_dirs, verbose=False, check_calc_results=False))
    
    perfect_res_path = Path("perfect/calc_results.json")
    if not perfect_res_path.exists():
        print("Error: perfect/calc_results.json not found. Cannot proceed to analysis.")
        return

    perfect_calc_results = loadfn(str(perfect_res_path))
    
    try:
        std_energies_obj = StandardEnergies.from_yaml(str(std_energies_path))
    except FileNotFoundError:
        print(f"Standard energies not found at {std_energies_path}")
        return

    analysis_dirs = [d for d in successful_defect_dirs if d.name != "perfect"]
    
    args_dei = Namespace(dirs=analysis_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)
    
    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=analysis_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)
    
    summary_path = Path("defect_energy_summary.json")
    if summary_path.exists():
        write_summary_at_vertices(summary_path)
    
    print("\nFull Auto Defect Workflow Completed.")

if __name__ == "__main__":
    main()
