"""
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 argparse
import sys
import glob
import copy
import os
from pathlib import Path

from macer.calculator.factory import get_available_ffs, ALL_SUPPORTED_FFS
from macer.defaults import DEFAULT_DEVICE, DEFAULT_FF
# from macer import __version__

MACER_LOGO = r"""
███╗   ███╗  █████╗   ██████╗ ███████╗ ██████╗
████╗ ████║ ██╔══██╗ ██╔════╝ ██╔════╝ ██╔══██╗
██╔████╔██║ ███████║ ██║      █████╗   ██████╔╝
██║╚██╔╝██║ ██╔══██║ ██║      ██╔══╝   ██╔══██╗
██║ ╚═╝ ██║ ██║  ██║ ╚██████╗ ███████╗ ██║  ██║
╚═╝     ╚═╝ ╚═╝  ╚═╝  ╚═════╝ ╚══════╝ ╚═╝  ╚═╝
macer: ML-accelerated Atomic Computational Environment for Research
"""

# Determine default force field based on installed extras
available_ffs = get_available_ffs()

def add_pydefect_parsers(subparsers):
    """Add pydefect sub-commands (cpd, defect, full) directly to the provided subparsers."""
    
    # CPD command
    cpd_epilog = """
Examples:
  # 1. Basic CPD generation using Materials Project ID (Recommended)
  macer pydefect cpd -m mp-3536

  # 2. CPD using formula (retrieves most stable from MP)
  macer pydefect cpd -f MgAl2O4

  # 3. CPD with a dopant element
  macer pydefect cpd -m mp-3536 -d Ca
"""
    cpd_parser = subparsers.add_parser(
        "cpd",
        description=MACER_LOGO + "\nRun Chemical Potential Diagram (CPD) workflow.",
        help="Generate CPD and target vertices.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=cpd_epilog
    )
    cpd_general = cpd_parser.add_argument_group('General & Input')
    cpd_general.add_argument("-f", "--formula", type=str, help="Formula to retrieve from Materials Project (e.g., MgAl2O4)")
    cpd_general.add_argument("-m", "--mpid", type=str, help="Materials Project ID (e.g., mp-3536)")
    cpd_general.add_argument("-d", "--doping", type=str, nargs='+', help="Dopant element(s) (e.g., Cl)")
    cpd_general.add_argument("-p", "--poscar", type=str, nargs='+', help="Input POSCAR file(s) (Optional)")
    _add_mlff_group(cpd_parser)
    cpd_calc = cpd_parser.add_argument_group('Calculation Settings')
    cpd_calc.add_argument("--fmax", type=float, default=0.03, help="Force convergence threshold (eV/Å). Default: 0.03")
    cpd_calc.add_argument("--energy-shift-target", type=float, default=0.0, help="Manually shift target energy (eV/atom). Default: 0.0")
    cpd_out = cpd_parser.add_argument_group('Output Settings')
    cpd_out.add_argument("--output-dir", type=str, help="Directory to save CPD output files.")
    cpd_parser.set_defaults(func=run_batch_workflow, workflow_name="cpd")

    # Defect command
    defect_epilog = """
Examples:
  # 1. Basic defect analysis with required inputs
  macer pydefect defect -p POSCAR --std_energies standard_energies.yaml --target_vertices target_vertices.yaml

  # 2. Defect analysis with a specific supercell matrix
  macer pydefect defect -p POSCAR -s standard_energies.yaml -t target_vertices.yaml --matrix 2 2 2

  # 3. Analysis with dopant and atom count limits
  macer pydefect defect -p POSCAR -s standard_energies.yaml -t target_vertices.yaml -d Mg --min_atoms 100 --max_atoms 400

  # 4. Tighter force convergence for defect relaxation
  macer pydefect defect -p POSCAR -s standard_energies.yaml -t target_vertices.yaml --fmax-defect 0.01

  # 5. Automatically find and add interstitial sites using Voronoi tessellation
  macer pydefect defect -p POSCAR -s standard_energies.yaml -t target_vertices.yaml --interstitial

  # 6. Disable symmetry analysis (if spglib fails or manual setup is preferred)
  macer pydefect defect -p POSCAR -s standard_energies.yaml -t target_vertices.yaml --no_symmetry_analysis
"""
    defect_parser = subparsers.add_parser(
        "defect",
        description=MACER_LOGO + "\nRun Defect Analysis workflow (Supercell generation, Relaxation, Analysis).",
        help="Calculate defect formation energies.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=defect_epilog
    )
    defect_general = defect_parser.add_argument_group('General & Input')
    defect_general.add_argument("-p", "--poscar", type=str, nargs='+', help="Input POSCAR file(s) (Perfect Unitcell)")
    defect_general.add_argument("-f", "--formula", type=str, help="Formula to retrieve from Materials Project (e.g., MgAl2O4)")
    defect_general.add_argument("-m", "--mpid", type=str, help="Materials Project ID (e.g., mp-3536)")
    defect_general.add_argument("-d", "--doping", type=str, nargs='+', help="Dopant element(s) (e.g., Cl)")
    defect_general.add_argument("-s", "--std_energies", type=str, help="Path to standard_energies.yaml")
    defect_general.add_argument("-t", "--target_vertices", type=str, help="Path to target_vertices.yaml")
    _add_mlff_group(defect_parser)
    defect_sc = defect_parser.add_argument_group('Supercell & Symmetry')
    defect_sc.add_argument("--matrix", nargs="+", type=int, help="Supercell matrix applied to the conventional cell.")
    defect_sc.add_argument("--dim", nargs="+", type=int, help="Supercell matrix applied directly to the input POSCAR (3 or 9 integers).")
    defect_sc.add_argument("--min_atoms", type=int, default=50, help="Minimum number of atoms (default: 50)")
    defect_sc.add_argument("--max_atoms", type=int, default=300, help="Maximum number of atoms (default: 300)")
    defect_sc.add_argument("--no_symmetry_analysis", dest="analyze_symmetry", action="store_false", help="Disable symmetry analysis.")
    defect_parser.set_defaults(analyze_symmetry=True)
    defect_sc.add_argument("--sites_yaml", type=str, dest="sites_yaml_filename", help="Path to sites.yaml file.")
    defect_sc.add_argument("--interstitial", action="store_true", help="Automatically find and add interstitial sites using Voronoi tessellation.")
    defect_relax = defect_parser.add_argument_group('Relaxation Settings')
    defect_relax.add_argument("--fmax", type=float, default=0.03, help="Force threshold (default: 0.03)")
    defect_relax.add_argument("--fmax-defect", type=float, help="Force threshold specifically for defect relaxation. Default: same as --fmax")
    defect_out = defect_parser.add_argument_group('Output Settings')
    defect_out.add_argument("--output-dir", type=str, help="Directory to save defect analysis output files.")
    defect_parser.set_defaults(func=run_batch_workflow, workflow_name="defect")

    # Full command
    full_epilog = """
Examples:
  # 1. Full defect workflow from scratch using MPID (Recommended)
  macer pydefect full -m mp-2605

  # 2. Full workflow using formula
  macer pydefect full -f MgO

  # 3. Full workflow with explicit input POSCAR and dopant
  macer pydefect full -p POSCAR_MgO -d Al

  # 4. Custom force thresholds for different stages
  macer pydefect full -m mp-2133 --fmax-cpd 0.02 --fmax-defect 0.05

  # 5. Full workflow with explicit supercell matrix
  macer pydefect full -m mp-149 --matrix 3 3 3

  # 6. Full workflow with automatic interstitial site finding
  macer pydefect full -m mp-2605 --interstitial
"""
    full_parser = subparsers.add_parser(
        "full",
        description=MACER_LOGO + "\nRun Full Defect Analysis workflow (CPD + Defect Analysis).",
        help="Run both CPD and Defect Analysis workflows.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=full_epilog
    )
    full_general = full_parser.add_argument_group('General & Input')
    full_general.add_argument("-f", "--formula", type=str, help="Formula (e.g., MgAl2O4)")
    full_general.add_argument("-m", "--mpid", type=str, help="Materials Project ID (e.g., mp-3536)")
    full_general.add_argument("-p", "--poscar", type=str, nargs='+', help="Input POSCAR file(s)")
    full_general.add_argument("-d", "--doping", type=str, nargs='+', help="Dopant element(s)")
    _add_mlff_group(full_parser)
    full_sc = full_parser.add_argument_group('Supercell & Symmetry')
    full_sc.add_argument("--matrix", nargs="+", type=int, help="Supercell matrix.")
    full_sc.add_argument("--dim", nargs="+", type=int, help="Supercell matrix applied directly to the input POSCAR (3 or 9 integers).")
    full_sc.add_argument("--min_atoms", type=int, default=50, help="Min atoms (default: 50)")
    full_sc.add_argument("--max_atoms", type=int, default=300, help="Max atoms (default: 300)")
    full_sc.add_argument("--no_symmetry_analysis", dest="analyze_symmetry", action="store_false", help="Disable symmetry analysis.")
    full_parser.set_defaults(analyze_symmetry=True)
    full_sc.add_argument("--sites_yaml", type=str, dest="sites_yaml_filename", help="Path to sites.yaml file.")
    full_sc.add_argument("--interstitial", action="store_true", help="Automatically find and add interstitial sites using Voronoi tessellation.")
    full_relax = full_parser.add_argument_group('Relaxation & Correction')
    full_relax.add_argument("--fmax", type=float, default=0.03, help="Global force threshold (default: 0.03)")
    full_relax.add_argument("--fmax-cpd", type=float, help="CPD force threshold.")
    full_relax.add_argument("--fmax-defect", type=float, help="Defect force threshold.")
    full_relax.add_argument("--energy-shift-target", type=float, default=0.0, help="Shift target energy. Default: 0.0")
    full_relax.add_argument("--without-cpd", action="store_true",
                            help="Skip CPD/target vertices; compute only elemental standard energies.")
    full_out = full_parser.add_argument_group('Output Settings')
    full_out.add_argument("--output-dir", type=str, help="Directory to save full defect analysis output files.")
    full_parser.set_defaults(func=run_batch_workflow, workflow_name="full")

    return {
        "cpd": cpd_parser,
        "defect": defect_parser,
        "full": full_parser
    }

def _add_mlff_group(parser):
    """Helper to add MLFF arguments to a parser group."""
    mlff_group = parser.add_argument_group('MLFF Model Settings')
    mlff_group.add_argument("--ff", type=str, default=DEFAULT_FF, choices=ALL_SUPPORTED_FFS, help=f"Force field to use. (default: {DEFAULT_FF})")
    mlff_group.add_argument("--model", type=str, default=None, help="Path to the MLFF model file.")
    mlff_group.add_argument("--device", type=str, default=DEFAULT_DEVICE, choices=["cpu", "mps", "cuda"], help="Compute device.")
    mlff_group.add_argument("--modal", type=str, default=None, help="Modal for SevenNet model.")
    mlff_group.add_argument("--n-jobs", type=int, default=1, help="Parallel workers for relaxations (CPU only, default: 1).")

def run_batch_workflow(args):
    workflow_name = args.workflow_name
    launch_cwd = Path.cwd()

    # Keep relative --model stable even if workflows chdir into work directories.
    if hasattr(args, "model") and args.model:
        model_path = Path(args.model)
        if not model_path.is_absolute():
            candidate = (launch_cwd / model_path).resolve()
            if candidate.exists():
                args.model = str(candidate)
    
    if workflow_name == "cpd":
        from macer.pydefect.cpd import run_cpd_workflow as workflow_func
    elif workflow_name == "defect":
        from macer.pydefect.defect import run_defect_workflow as workflow_func
    elif workflow_name == "full":
        from macer.pydefect.full import run_full_workflow as workflow_func
    else:
        raise ValueError(f"Unknown workflow: {workflow_name}")

    if workflow_name == "defect" and not (hasattr(args, 'poscar') and args.poscar):
        if getattr(args, "formula", None) or getattr(args, "mpid", None):
            from macer.utils.struct_tools import resolve_structure_inputs
            resolved = resolve_structure_inputs(
                poscar_paths=None,
                cif_paths=None,
                formula=getattr(args, "formula", None),
                mpid=getattr(args, "mpid", None), show_mp_progress=True)
            if resolved:
                args.poscar = resolved

    if hasattr(args, 'poscar') and args.poscar:
        input_patterns = args.poscar
        input_files = []
        for pat in input_patterns:
            if os.path.exists(pat): input_files.append(pat)
            else: input_files.extend(glob.glob(pat))
        input_files = sorted(list(set(input_files)))
        if not input_files:
            print(f"No input files found matching: {input_patterns}")
            sys.exit(1)
        original_cwd = os.getcwd()
        for f in input_files:
            os.chdir(original_cwd)
            print(f"\n{'='*60}\nProcessing Input File: {f}\n{'='*60}\n")
            single_args = copy.copy(args)
            single_args.poscar = f
            try: workflow_func(single_args)
            except Exception as e:
                print(f"Error processing {f}: {e}")
                import traceback
                traceback.print_exc()
                continue
    else: workflow_func(args)

def main():
    from macer import __version__
    parser = argparse.ArgumentParser(
        description=MACER_LOGO + f"\nmacer_pydefect (v{__version__}): Automated Point Defect Calculations with MLFFs",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--version", "-v", action="version", version=f"macer_pydefect {__version__}")
    subparsers = parser.add_subparsers(dest="subcommand", help="Available commands")
    p_map = add_pydefect_parsers(subparsers)
    args = parser.parse_args()
    launch_cwd = Path.cwd()
    if hasattr(args, "model") and args.model:
        model_path = Path(args.model)
        if not model_path.is_absolute():
            candidate = (launch_cwd / model_path).resolve()
            if candidate.exists():
                args.model = str(candidate)

    if not args.subcommand:
        parser.print_help()
        return

    # Manual checks
    if args.subcommand == "cpd" and not (args.formula or args.mpid or args.poscar):
        p_map["cpd"].print_help(); return
    if args.subcommand == "defect" and not args.poscar:
        p_map["defect"].print_help(); return
    if args.subcommand == "full" and not (args.formula or args.mpid or args.poscar):
        p_map["full"].print_help(); return

    if hasattr(args, "func"): args.func(args)
    else: parser.print_help(); sys.exit(1)

if __name__ == "__main__":
    main()
