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

# Add current directory to sys.path to ensure local vise/pydefect/macer are used
sys.path.insert(0, str(Path.cwd()))

from pymatgen.core import Structure
from vise.cli.main_functions import get_poscar_from_mp
from pydefect.cli.vasp.main_vasp_functions import make_competing_phase_dirs, make_composition_energies
from pydefect.cli.main_functions import make_standard_and_relative_energies, make_cpd_and_vertices
from pydefect.defaults import defaults as pydefect_defaults

# 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

def main():
    parser = argparse.ArgumentParser(
        description="Auto CPD workflow using Macer, Vise, and Pydefect",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("-f", "--formula", type=str, help="Formula to retrieve from Materials Project (e.g., MgAl2O4)")
    parser.add_argument("-m", "--mpid", type=str, help="Materials Project ID (e.g., mp-3536)")
    parser.add_argument("-d", "--doping", type=str, nargs='+', help="Dopant element(s) (e.g., Cl)")
    
    args = parser.parse_args()

    # Custom validation for mutually exclusive arguments
    if args.formula and args.mpid:
        print("\nError: Ambiguous structure specification.")
        print("Please provide the structure using either the -m option (Material Project ID) or the -f option (Chemical Formula), but not both.\n")
        sys.exit(1)

    if not args.formula and not args.mpid:
        parser.print_help()
        print("\nError: No structure specified.")
        print("Please provide the structure using either the -m option (Material Project ID) or the -f option (Chemical Formula).\n")
        sys.exit(1)

    # ------------------------------------------------------------------
    # Step 1: Get POSCAR
    # ------------------------------------------------------------------
    target_str = args.formula if args.formula else args.mpid
    print(f"--- Step 1: Retrieve POSCAR for {target_str} ---")
    
    # Prepare arguments for vise's get_poscar_from_mp
    args_gp = Namespace(mpid=args.mpid, formula=args.formula)
    
    try:
        get_poscar_from_mp(args_gp)
        if Path("POSCAR").exists():
            print(f"Successfully retrieved POSCAR for {target_str}.")
        else:
            print("Error: POSCAR file was not created.")
            return
    except Exception as e:
        print(f"Failed to retrieve POSCAR: {e}")
        return

    # Determine elements and formula from the retrieved POSCAR
    target_formula = None
    host_elements = []
    competing_elements = []
    
    try:
        structure = Structure.from_file("POSCAR")
        host_elements = [str(e) for e in structure.composition.elements]
        target_formula = structure.composition.reduced_formula
        print(f"Host Elements: {host_elements}, Target Formula: {target_formula}")
        
        competing_elements = host_elements.copy()
        if args.doping:
            for dopant in args.doping:
                if dopant not in competing_elements:
                    competing_elements.append(dopant)
            print(f"Dopant(s) added: {args.doping}. Full competing element list: {competing_elements}")
            
    except Exception as e:
        print(f"Failed to read POSCAR or determine elements: {e}")
        return

    # ------------------------------------------------------------------
    # Create CPD Output Directory and Move Files
    # ------------------------------------------------------------------
    dopants_part = ""
    if args.doping:
        dopants_part = f"-dopant={'-'.join(args.doping)}"
    
    base_dir_name = f"CPD-target={target_formula}{dopants_part}"
    cpd_output_dir = Path.cwd() / base_dir_name
    
    i = 1
    original_base = cpd_output_dir
    while cpd_output_dir.exists():
        cpd_output_dir = Path(f"{original_base}-NEW{i:03d}")
        i += 1
    
    print(f"\nCreating output directory: {cpd_output_dir.name}")
    cpd_output_dir.mkdir()
    
    # Move initial files to the new directory
    for filename in ["POSCAR", "prior_info.yaml"]:
        if Path(filename).exists():
            shutil.move(filename, cpd_output_dir / filename)
            
    # Change working directory
    os.chdir(cpd_output_dir)
    print(f"Changed working directory to: {os.getcwd()}")

    # ------------------------------------------------------------------
    # Step 2: Generate competing phase directories
    # Command: pydefect_vasp mp -e [competing_elements]
    # ------------------------------------------------------------------
    print(f"\n--- Step 2: Generate competing phase directories ({'-'.join(competing_elements)}) ---")
    args_mp = Namespace(elements=competing_elements, 
                        e_above_hull=pydefect_defaults.e_above_hull)
    try:
        make_competing_phase_dirs(args_mp)
        print("Competing phase directories created.")
    except Exception as e:
        print(f"Failed to generate competing phase directories: {e}")
        return

    # ------------------------------------------------------------------
    # Step 3: Run macer relax in each directory (Single MLFF Load)
    # ------------------------------------------------------------------
    print("\n--- Step 3: Run macer relax (Single MLFF Load) ---")
    
    cwd = Path.cwd()
    target_dirs = []
    for d in cwd.iterdir():
        if d.is_dir() and (d / "POSCAR").exists() and (d / "prior_info.yaml").exists():
             target_dirs.append(d)
    
    target_dirs.sort(key=lambda x: x.name)
    
    if not target_dirs:
        print("No defect directories found to relax.")
        return

    # Initialize Calculator (once)
    available_ffs = get_available_ffs()
    if not available_ffs:
        print("Error: No MLFF packages installed.")
        return
    
    ff = available_ffs[0] # Default to first available
    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,
    }
    # Special handling for MACE
    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)
        print("Calculator initialized.")
    except Exception as e:
        print(f"Failed to initialize calculator: {e}")
        return

    successful_dirs = []
    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=3,
                    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)
            
        except Exception as e:
            sys.stdout = original_stdout
            print(f"Relaxation failed in {d.name}: {e}")
        finally:
            sys.stdout = original_stdout
            os.chdir(cwd)

    # ------------------------------------------------------------------
    # Step 4: Generate composition_energies.yaml
    # ------------------------------------------------------------------
    print("\n--- Step 4: Generate composition_energies.yaml ---")
    if successful_dirs:
        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_dirs,
                             verbose=False)
        try:
            make_composition_energies(args_mce)
            print(f"Created/Updated {yaml_filename} from {len(successful_dirs)} calculations.")
        except Exception as e:
            print(f"Failed to create composition_energies.yaml: {e}")
    else:
        print("No successful calculations to process.")

    # ------------------------------------------------------------------
    # Step 5: Standard and Relative Energies
    # Command: pydefect sre
    # ------------------------------------------------------------------
    print("\n--- Step 5: Generate Standard and Relative Energies ---")
    if Path("composition_energies.yaml").exists():
        # Filter composition_energies to exclude cross-impurity compounds
        # (compounds containing more than one non-host element) to prevent pydefect errors.
        import yaml
        from pymatgen.core import Composition
        
        try:
            with open("composition_energies.yaml", 'r') as f:
                comp_energies = yaml.safe_load(f)
            
            filtered_comp_energies = {}
            removed_compounds = []
            
            for formula, data in comp_energies.items():
                comp = Composition(formula)
                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] = data
                else:
                    removed_compounds.append(formula)
            
            if removed_compounds:
                print(f"Filtering out compounds with multiple impurities: {removed_compounds}")
            
            filtered_yaml_filename = "composition_energies_filtered.yaml"
            with open(filtered_yaml_filename, 'w') as f:
                yaml.dump(filtered_comp_energies, f)
                
            args_sre = Namespace(composition_energies_yaml=filtered_yaml_filename)
            make_standard_and_relative_energies(args_sre)
            print("Generated standard_energies.yaml and relative_energies.yaml (from filtered data).")
            
        except Exception as e:
            print(f"Failed to generate standard/relative energies: {e}")
    else:
        print("composition_energies.yaml not found. Skipping Step 5.")

    # ------------------------------------------------------------------
    # Step 6: Chemical Potential Diagram
    # Command: pydefect cv -t {target_formula}
    # ------------------------------------------------------------------
    print(f"\n--- Step 6: Generate Chemical Potential Diagram for {target_formula} ---")
    if Path("relative_energies.yaml").exists() and target_formula:
        # Note: We use host_elements for the vertices of the CPD
        args_cv = Namespace(rel_energy_yaml="relative_energies.yaml", 
                            target=target_formula, 
                            elements=host_elements) 
        try:
            make_cpd_and_vertices(args_cv)
            print("Generated chem_pot_diag.json and target_vertices.yaml.")
        except Exception as e:
            print(f"Failed to generate chemical potential diagram: {e}")
    else:
        print("relative_energies.yaml not found or target formula missing. Skipping Step 6.")

if __name__ == "__main__":
    main()
