"""
Macer: Machine-learning accelerated Atomic Computational Environment for automated Research workflows
Copyright (c) 2025 The Macer Package Authors
Author: Soungmin Bae <soungminbae@gmail.com>
License: MIT
"""

import sys
import os
from pathlib import Path
import yaml
from monty.serialization import loadfn
from pymatgen.core import Structure, Composition
from pymatgen.analysis.phase_diagram import PhaseDiagram, PDEntry

from pydefect.defaults import defaults
from pydefect.analyzer.make_defect_structure_info import MakeDefectStructureInfo
from pydefect.cli.make_defect_vesta_file import MakeDefectVestaFile

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, _model_root, DEFAULT_DEVICE, resolve_model_path

def generate_composition_energies_direct(directories, yaml_filename="composition_energies.yaml"):
    """
    Directly generates composition_energies.yaml from calc_results.json files 
    in the specified directories, bypassing vise/OUTCAR parsing.
    """
    data = {}
    for d in directories:
        d_path = Path(d)
        calc_results_path = d_path / "calc_results.json"
        
        if not calc_results_path.exists():
            print(f"Warning: {calc_results_path} not found. Skipping {d_path.name}.")
            continue
        
        try:
            # Load calc_results.json generated by macer
            calc_results = loadfn(calc_results_path)
            
            # Check if calc_results is a dict or an object
            if isinstance(calc_results, dict):
                energy = calc_results.get("energy")
                structure_dict = calc_results.get("structure")
                if structure_dict:
                    structure = Structure.from_dict(structure_dict)
                else:
                    structure = None
            else:
                # Assume it's a CalcResults object
                energy = getattr(calc_results, "energy", None)
                structure = getattr(calc_results, "structure", None)
            
            if energy is None or structure is None:
                print(f"Warning: Incomplete data in {calc_results_path}. Skipping.")
                continue
                
            # Use formula with spaces removed (e.g. "Zn2 O2" -> "Zn2O2") as key
            formula = structure.composition.formula.replace(" ", "")
            
            # If multiple phases have the same formula, we should theoretically keep the most stable one.
            # However, pydefect usually handles competing phases by ensuring they are distinct or 
            # by processing them such that the most stable ones define the hull.
            # Here we simply store it. If duplicates exist, the last one visited overwrites.
            # Ideally, directories should correspond to unique stable phases.
            
            data[formula] = {
                "energy": float(energy),
                "source": str(d_path.absolute())
            }
        except Exception as e:
            print(f"Error processing {d_path.name}: {e}")
            
    if data:
        with open(yaml_filename, "w") as f:
            yaml.dump(data, f)
        print(f"Directly generated {yaml_filename} from {len(data)} entries.")
        return True
    else:
        print(f"No data collected for {yaml_filename}.")
        return False

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 stabilize_target(relative_energies_path, target_formula, manual_shift=0.0, buffer=1e-5):
    """
    Checks if the target is above convex hull. If so, shifts its energy in relative_energies.yaml
    to make it stable. Returns the total applied shift value (0.0 if stable and no manual shift).
    """
    if not Path(relative_energies_path).exists():
        print(f"Error: {relative_energies_path} does not exist.")
        return 0.0

    with open(relative_energies_path, 'r') as f:
        rel_energies = yaml.safe_load(f)

    # 1. Apply manual shift first (temporarily to check stability or permanently if stable)
    if target_formula not in rel_energies:
        print(f"Error: Target {target_formula} not found in {relative_energies_path}")
        return 0.0
    
    original_energy = rel_energies[target_formula]
    # We apply the manual shift (subtracted, as per convention: positive shift means lowering energy? 
    # The user said "--energy-shift-target 0.05 이면 target energy 를 0.05 eV 낮춤" -> new = old - 0.05)
    # So if manual_shift is 0.05, we subtract 0.05.
    current_energy_after_manual = original_energy - manual_shift
    rel_energies[target_formula] = current_energy_after_manual
    
    entries = []
    target_entry = None

    # Re-build entries with the manual shift applied
    for formula, energy_per_atom in rel_energies.items():
        try:
            comp = Composition(formula)
            total_energy = energy_per_atom * comp.num_atoms
            entry = PDEntry(comp, total_energy)
            entries.append(entry)
            if formula == target_formula:
                target_entry = entry
        except Exception as e:
            print(f"Warning: Could not process formula {formula}: {e}")

    if target_entry is None:
        return 0.0

    entries = []
    target_entry = None
    target_formula_norm = Composition(target_formula).reduced_formula

    # Re-build entries with the manual shift applied and normalized formulas
    for formula, energy_per_atom in rel_energies.items():
        try:
            comp = Composition(formula)
            norm_f = comp.reduced_formula
            total_energy = energy_per_atom * comp.num_atoms
            # Use normalized formula as the name for consistent matching
            entry = PDEntry(comp, total_energy, name=norm_f)
            entries.append(entry)
            
            if norm_f == target_formula_norm and formula == target_formula:
                target_entry = entry
        except Exception as e:
            print(f"Warning: Could not process formula {formula}: {e}")

    if target_entry is None:
        print(f"Error: Target entry for {target_formula} could not be created.")
        return 0.0

    pd = PhaseDiagram(entries)
    e_above_hull = pd.get_e_above_hull(target_entry)
    
    # --- Generate convex_hull_energy.yaml ---
    try:
        # Use a small threshold to determine stability
        is_stable = e_above_hull <= 1e-8
        
        # 1. Standard stability info (including target in PD)
        if is_stable:
            competing_phases_info = {target_formula: 1.0}
            convex_hull_energy_per_atom = float(target_entry.energy_per_atom)
        else:
            decomp = pd.get_decomposition(target_entry.composition)
            convex_hull_energy_per_atom = target_entry.energy_per_atom - e_above_hull
            competing_phases_info = {entry.name: float(amount) for entry, amount in decomp.items()}

        # 2. Benchmark stability info (excluding target from PD)
        decomposition_info = None
        try:
            entries_excl_target = [e for e in entries if e.composition.reduced_formula != target_formula_norm]
            if entries_excl_target:
                pd_benchmark = PhaseDiagram(entries_excl_target)
                hull_energy_total = pd_benchmark.get_hull_energy(target_entry.composition)
                hull_energy_per_atom = hull_energy_total / target_entry.composition.num_atoms
                e_above_hull_benchmark = target_entry.energy_per_atom - hull_energy_per_atom
                
                comp_phases_bench = {}
                try:
                    decomp_benchmark = pd_benchmark.get_decomposition(target_entry.composition)
                    comp_phases_bench = {entry.name: float(amount) for entry, amount in decomp_benchmark.items()}
                except Exception: pass
                
                decomposition_info = {
                    "hull_distance": float(e_above_hull_benchmark),
                    "competing_phases": comp_phases_bench
                }
        except Exception as e_bench:
            print(f"Warning: Benchmark stability calculation failed: {e_bench}")

        # Construct dictionary in specific order
        convex_hull_data = {
            "target": target_formula,
            "formation_energy_per_atom": float(target_entry.energy_per_atom),
            "convex_hull_energy_per_atom": float(convex_hull_energy_per_atom),
            "e_above_hull": float(e_above_hull)
        }
        if decomposition_info:
            convex_hull_data["decomposition"] = decomposition_info

        # Write to yaml in the same directory as relative_energies.yaml
        output_dir = Path(relative_energies_path).parent
        with open(output_dir / "convex_hull_energy.yaml", "w") as f:
            yaml.dump(convex_hull_data, f, default_flow_style=False)
        print(f"Generated convex_hull_energy.yaml")
        
    except Exception as e:
        print(f"Warning: Failed to generate convex_hull_energy.yaml details: {e} (e_h: {e_above_hull})")
    # ----------------------------------------
    # ----------------------------------------

    extra_stabilization_shift = 0.0

    if e_above_hull > 1e-8:  # Use a small epsilon
        # It's still unstable after manual shift (or manual shift was 0)
        required_shift = -1.0 * e_above_hull
        # Add buffer
        extra_stabilization_shift = required_shift - buffer
        
        print(f"Target {target_formula} is unstable (e_above_hull={e_above_hull:.4f} eV/atom).")
        print(f"Additional stabilization shift required: {extra_stabilization_shift:.6f} eV/atom.")

    # Total shift = manual_shift (already subtracted) + extra_stabilization_shift (negative value to lower energy further)
    # But wait, user said "shift 0.05 means lower by 0.05".
    # My stabilize logic was "return recommended_shift" where recommended_shift was negative.
    # So let's keep consistent: return value is the net change in energy.
    # manual_shift argument is positive -> net change is -manual_shift.
    # extra_stabilization_shift is negative.
    
    total_change = -manual_shift + extra_stabilization_shift
    
    if total_change != 0.0:
        final_energy = float(original_energy + total_change)
        rel_energies[target_formula] = final_energy
        
        with open(relative_energies_path, 'w') as f:
            yaml.dump(rel_energies, f)
            
        print(f"Total applied energy shift: {total_change:.6f} eV/atom (Manual: -{manual_shift}, Stabilization: {extra_stabilization_shift:.6f}).")
        return total_change

    return 0.0

def write_summary_at_vertices(summary_path: Path, applied_shift=0.0):
    """
    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-vertices.txt"

    with open(output_filename, "w") as f:
        f.write(f"title: {summary.title}\n")
        if applied_shift != 0.0:
            f.write(f"# Energy shift applied to make target stable: {applied_shift} eV/atom\n")
        f.write("\n")

        f.write("--- Standard Formation Energies (Chemical Potential = 0) ---\n")
        f.write(f"{'Defect Name':<20} {'Charge':<10} {'Formation Energy (eV)'}\n")
        f.write(f"{'-'*20} {'-'*10} {'-'*25}\n")
        
        for defect_name, defect_group in defect_energies_summary.items():
             if 0 in defect_group.charges:
                 idx = defect_group.charges.index(0)
                 energy = defect_group.defect_energies[idx].formation_energy
                 f.write(f"{defect_name:<20} {'0':<10} {energy:>25.4f}\n")
        f.write("\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:
                    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[0])
            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 run_macer_relax(target_dirs, isif=3, supercell_info=None, ff=None, model_path=None, device=None, fmax=0.03, verbose=False, modal=None):
    """
    Unified relaxation function for CPD and Defect workflows.
    Iterates over target_dirs and runs macer relaxation.
    """
    if not target_dirs:
        return []

    # Determine FF and Model if not provided
    if ff is None:
        available_ffs = get_available_ffs()
        if not available_ffs:
            print("Error: No MLFF packages installed.")
            return []
        ff = available_ffs[0]
    
    if device is None:
        device = DEFAULT_DEVICE
    
    if model_path is None:
        default_model_name = DEFAULT_MODELS.get(ff)
        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 = resolve_model_path(default_model_name)
    else:
        model_path = resolve_model_path(model_path)
    
    print(f"Initializing {ff.upper()} calculator on {device}...")
    calc_kwargs = {"model_path": model_path, "device": device, "modal": modal}
    if ff == "mace":
        calc_kwargs["model_paths"] = [calc_kwargs["model_path"]]
        del calc_kwargs["model_path"]

    model_info_str = ""
    FFS_USING_MODEL_NAME = {"fairchem", "orb", "chgnet", "m3gnet"}
    if calc_kwargs.get("model_path") or calc_kwargs.get("model_paths"):
        # This is simplified, usually resolve_model_path was already called before run_macer_relax 
        # in the main entry points, but let's ensure info is printed.
        pass

    try:
        calculator = get_calculator(ff_name=ff, **calc_kwargs)
        # Accessing private attribute to get resolved path for logging
        actual_path = None
        if hasattr(calculator, "model_path"): actual_path = calculator.model_path
        elif hasattr(calculator, "model_paths") and calculator.model_paths: actual_path = calculator.model_paths[0]
        
        if actual_path:
            print(f"  MLFF Model: {actual_path}")
        
        print("Calculator initialized.")
    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:
        if verbose:
            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=fmax,
                    isif=isif,
                    device=device,
                    calculator=calculator, 
                    ff=ff,
                    outcar_name="OUTCAR", 
                    contcar_name="CONTCAR",
                    xml_name="vasprun.xml",
                    make_pdf=True,
                    write_json=True,
                    modal=modal
                )
            
                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
                     # This logic comes from run_auto_defect[-full].py
                     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}")

                                # Generate VESTA files for defect visualization
                                try:
                                    # Reload from file to ensure clean state matching CLI usage
                                    dsi_loaded = loadfn("defect_structure_info.json")
                                    vesta_maker = MakeDefectVestaFile(dsi_loaded)
                                    vesta_maker.initial_vesta.write_file("defect_initial.vesta")
                                    vesta_maker.final_vesta.write_file("defect.vesta")
                                    print(f"Generated VESTA files for {d.name}")
                                except Exception as e_vesta:
                                    print(f"Warning: Failed to generate VESTA files: {e_vesta}")

                         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
