import argparse
import os
import sys
import subprocess
import shutil
import numpy as np
from pathlib import Path
import re
import math

from ase.io import read, write
from ase import Atoms

# Import necessary components from macer
from macer.defaults import DEFAULT_MODELS, _macer_root
from macer.calculator.factory import get_available_ffs


def run_qha_workflow(args):
    # --- Configuration ---
    input_poscar_path = Path(args.poscar).resolve()

    # Create a dedicated output directory for QHA results
    base_qha_output_dir = input_poscar_path.parent / f"qha_results_{input_poscar_path.stem}"

    # If a folder with the same name already exists, keep appending "-NEW" repeatedly
    qha_output_dir = base_qha_output_dir
    while qha_output_dir.exists():
        qha_output_dir = Path(str(qha_output_dir) + "-NEW")

    qha_output_dir.mkdir(parents=True, exist_ok=False)

    # Force field and model settings
    ff = args.ff
    model_path = args.model
    modal = args.modal
    device = args.device

    # Volume/Length variation settings
    length_scale_range = args.length_scale
    num_volumes = args.num_volumes

    # Mesh settings for phonopy --mesh
    mesh_x, mesh_y, mesh_z = args.mesh

    # Phonopy displacement settings
    min_length = args.min_length
    displacement_distance = args.amplitude
    is_plusminus = args.is_plusminus
    is_diagonal = args.is_diagonal
    tolerance_sr = args.tolerance_sr  # Tolerance for macer_phonopy sr

    # QHA settings
    tmax = args.tmax
    atom_fix = args.atom_fix

    print(f"--- Starting QHA workflow for {input_poscar_path.name} ---")
    print(f"All QHA results will be saved in: {qha_output_dir}")

    # --- Step 1: Initial Unit Cell Relaxation ---
    print("\n--- Step 1: Initial Unit Cell Relaxation ---")

    copied_input_poscar_path = qha_output_dir / input_poscar_path.name
    shutil.copy(input_poscar_path, copied_input_poscar_path)
    print(f"Copied input POSCAR to {copied_input_poscar_path}")

    relaxed_unit_cell_path = None
    original_cwd = os.getcwd()

    try:
        if args.initial_relax_mode == 'sr':
            print(f"Mode: 'sr'. Performing iterative relaxation and symmetrization with fmax={args.initial_fmax}.")
            
            symmetrized_poscar_name = f"{input_poscar_path.stem}-symmetrized"
            relaxed_unit_cell_path = qha_output_dir / symmetrized_poscar_name

            macer_ru_command = [
                'macer_phonopy', 'sr',
                '-p', str(copied_input_poscar_path),
                '--output-prefix', input_poscar_path.stem,
                '--tolerance', str(tolerance_ru),
                '--fmax', str(args.initial_fmax),
                '--ff', ff,
                '--device', device,
            ]
            if model_path:
                macer_ru_command.extend(['--model', str(model_path)])
            if modal:
                macer_ru_command.extend(['--modal', modal])

            print(f"Running command: {' '.join(macer_ru_command)}")
            result = subprocess.run(macer_ru_command, check=True, capture_output=True, text=True)
            print(result.stdout)
            print(f"Symmetrized unit cell created: {relaxed_unit_cell_path.name}")

        elif args.initial_relax_mode == 'relax':
            print(f"Mode: 'relax'. Performing standard relaxation (ISIF=3) with fmax={args.initial_fmax}.")
            
            relaxed_poscar_name = f"CONTCAR-{copied_input_poscar_path.name}"
            relaxed_unit_cell_path = qha_output_dir / relaxed_poscar_name

            macer_relax_command = [
                'macer', 'relax',
                '-p', copied_input_poscar_path.name,
                '--isif', '3',
                '--fmax', str(args.initial_fmax),
                '--ff', ff,
                '--device', device,
                '--contcar', relaxed_poscar_name,
                '--no-pdf',
            ]
            if model_path:
                macer_relax_command.extend(['--model', str(model_path)])
            if modal:
                macer_relax_command.extend(['--modal', modal])

            print(f"Running command (in {qha_output_dir}): {' '.join(macer_relax_command)}")
            os.chdir(qha_output_dir)
            result = subprocess.run(macer_relax_command, check=True, capture_output=True, text=True)
            print(result.stdout)
            print(f"Relaxed unit cell created: {relaxed_poscar_name}")

    except subprocess.CalledProcessError as e:
        print(f"Error during initial relaxation (mode: {args.initial_relax_mode}):")
        print(e.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"An unexpected error occurred during initial relaxation: {e}")
        sys.exit(1)
    finally:
        os.chdir(original_cwd)

    if not relaxed_unit_cell_path or not relaxed_unit_cell_path.exists():
        print(f"Error: Relaxed/symmetrized POSCAR not found at '{relaxed_unit_cell_path}'. Exiting.")
        sys.exit(1)

    # --- Step 1.5: Prepare QHA Input Structure ---
    qha_input_poscar_path = qha_output_dir / f"{input_poscar_path.stem}-input"
    shutil.move(relaxed_unit_cell_path, qha_input_poscar_path)
    print(f"Initial relaxed structure prepared for QHA: {qha_input_poscar_path.name}")

    base_structure_atoms = read(str(qha_input_poscar_path), format="vasp")
    initial_cell = base_structure_atoms.get_cell()

    # Determine supercell dimensions
    if args.dim:
        scaling_factors = args.dim
        print(f"Using user-provided supercell DIM: {' '.join(map(str, scaling_factors))}")
    else:
        base_vector_lengths = [np.linalg.norm(v) for v in initial_cell]
        scaling_factors = [math.ceil(min_length / v) if v > 0 else 1 for v in base_vector_lengths]
        print(f"Supercell DIM for phonopy -d (from --min-length): {' '.join(map(str, scaling_factors))}")

    dim_str = " ".join(map(str, scaling_factors))

    length_scales = np.linspace(1 - length_scale_range, 1 + length_scale_range, num_volumes)
    print(f"Length scales to be used: {length_scales}")

    qha_energies = []
    qha_volumes = []
    thermal_properties_with_volumes = []

    # --- Step 2: Volume Variation Loop ---
    print("\n--- Step 2: Performing calculations for each volume variation ---")
    for i, l_scale in enumerate(length_scales):
        volume_scale = l_scale**3
        volume_dir_name = f"length-scale={l_scale:.5f}"
        volume_dir_path = qha_output_dir / volume_dir_name
        volume_dir_path.mkdir(exist_ok=True)

        print(f"\n--- Processing length scale: {l_scale:.5f} (volume scale: {volume_scale:.5f}) ---")

        # Use the QHA input structure as the base
        scaled_atoms = read(str(qha_input_poscar_path), format="vasp")
        scaled_atoms.set_cell(initial_cell * l_scale, scale_atoms=True)
        scaled_poscar_path = volume_dir_path / "POSCAR"
        write(str(scaled_poscar_path), scaled_atoms, format="vasp")
        print(f"  Scaled POSCAR written to {scaled_poscar_path}")

        original_cwd = os.getcwd()
        try:
            os.chdir(volume_dir_path)
            if atom_fix:
                print("  Performing macer relax --isif 0 for the scaled unit cell (due to --atom-fix)...")
                relax_scaled_poscar_command = [
                    'macer', 'relax',
                    '-p', "POSCAR",
                    '--isif', '0',
                    '--no-pdf',
                    '--ff', ff,
                    '--device', device,
                    '--outcar', "OUTCAR-POSCAR",
                    '--vasprun', "vasprun-POSCAR.xml",
                    '--contcar', "CONTCAR-POSCAR",
                ]
                if model_path:
                    relax_scaled_poscar_command.extend(['--model', str(model_path)])
                if modal:
                    relax_scaled_poscar_command.extend(['--modal', modal])

                subprocess.run(relax_scaled_poscar_command, check=True, capture_output=True, text=True)
                print("  Single-point calculation for scaled unit cell completed.")

            elif not atom_fix:
                print("  Performing macer relax --isif 2 (atoms only relaxation)...")
                relax_isif2_command = [
                    'macer', 'relax',
                    '-p', "POSCAR",
                    '--isif', '2',
                    '--ff', ff,
                    '--device', device,
                    '--contcar', 'CONTCAR-POSCAR',
                    '--outcar', 'OUTCAR-POSCAR',
                    '--vasprun', 'vasprun-POSCAR.xml',
                ]
                if model_path:
                    relax_isif2_command.extend(['--model', str(model_path)])
                if modal:
                    relax_isif2_command.extend(['--modal', modal])

                subprocess.run(relax_isif2_command, check=True, capture_output=True, text=True)
                shutil.move("CONTCAR-POSCAR", "POSCAR")
                print(f"  Relaxed POSCAR (ISIF=2) updated at {scaled_poscar_path}")

            print("  Generating displacements with phonopy -d...")
            phonopy_d_command = ['phonopy', '-d', '--dim', dim_str, '-c', "POSCAR"]
            subprocess.run(phonopy_d_command, check=True, capture_output=True, text=True)

            print("  Performing macer relax --isif 0 for displaced supercells...")
            disp_poscar_files = sorted(list(Path(".").glob("POSCAR-*")))
            if not disp_poscar_files:
                raise FileNotFoundError(f"No displaced POSCAR-* files found in {volume_dir_path}")

            relax_isif0_command_base = [
                'macer', 'relax',
                '--isif', '0',
                '--no-pdf',
                '--ff', ff,
                '--device', device,
            ]
            if model_path:
                relax_isif0_command_base.extend(['--model', str(model_path)])
            if modal:
                relax_isif0_command_base.extend(['--modal', modal])

            for disp_file in disp_poscar_files:
                disp_prefix = disp_file.name
                relax_isif0_command = relax_isif0_command_base + [
                    '-p', str(disp_file),
                    '--outcar', f"OUTCAR-{disp_prefix}",
                    '--vasprun', f"vasprun-{disp_prefix}.xml",
                    '--contcar', f"CONTCAR-{disp_prefix}",
                ]
                subprocess.run(relax_isif0_command, check=True, capture_output=True, text=True)
            print("  macer relax --isif 0 completed for all displaced supercells.")

            print("  Creating FORCE_SETS with phonopy -f...")
            vasprun_xml_files = sorted(list(Path(".").glob("vasprun-POSCAR-*.xml")))
            vasprun_xml_filenames = [f.name for f in vasprun_xml_files]
            phonopy_f_command = ['phonopy', '-f'] + vasprun_xml_filenames
            subprocess.run(phonopy_f_command, check=True, capture_output=True, text=True)

            print(f"  Generating thermal_properties.yaml with phonopy --mesh {mesh_x} {mesh_y} {mesh_z}...")
            thermal_properties_path = volume_dir_path / "thermal_properties.yaml"
            phonopy_mesh_command = [
                'phonopy', '--mesh', f"{mesh_x} {mesh_y} {mesh_z}", '-t',
                '-c', "POSCAR",
                '--dim', dim_str,
            ]
            subprocess.run(phonopy_mesh_command, check=True, capture_output=True, text=True)
            
            # We need the volume of the relaxed cell inside the volume directory
            final_vol_atoms = read("POSCAR", format="vasp")
            thermal_properties_with_volumes.append((final_vol_atoms.get_volume(), str(thermal_properties_path)))

            outcar_path = "OUTCAR-POSCAR"
            with open(outcar_path, 'r') as f:
                outcar_content = f.read()
            energy_match = re.search(r"energy\s+without\s+entropy=\s*([-\d.]+)", outcar_content)
            volume_match = re.search(r"volume\s+of\s+cell\s*:\s*([-\d.]+)", outcar_content)
            if energy_match and volume_match:
                qha_energies.append(float(energy_match.group(1)))
                qha_volumes.append(float(volume_match.group(1)))
            else:
                raise ValueError(f"Could not parse energy/volume from {outcar_path}")

        except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
            print(f"  Error during processing for scale {l_scale:.5f}:")
            if isinstance(e, subprocess.CalledProcessError):
                print(e.stderr)
            else:
                print(e)
            sys.exit(1)
        finally:
            os.chdir(original_cwd)


    # --- Step 3: Run phonopy-qha ---
    print("\n--- Step 3: Running phonopy-qha ---")
    e_v_dat_path = qha_output_dir / "e-v.dat"
    with open(e_v_dat_path, 'w') as f:
        f.write("# Volume(A^3)  Energy(eV)\n")
        sorted_qha_data = sorted(zip(qha_volumes, qha_energies))
        for vol, energy in sorted_qha_data:
            f.write(f"{vol:.4f} {energy:.8f}\n")
    print(f"  e-v.dat created at {e_v_dat_path}")

    sorted_thermal_properties_paths = [path for vol, path in sorted(thermal_properties_with_volumes)]
    phonopy_qha_command = [
        'phonopy-qha',
        str(e_v_dat_path),
    ] + sorted_thermal_properties_paths + [
        '--tmax', str(tmax),
        '-s',
    ]
    print(f"Running command: {' '.join(phonopy_qha_command)}")
    try:
        os.chdir(qha_output_dir)
        subprocess.run(phonopy_qha_command, check=True, capture_output=True, text=True)
        os.chdir(original_cwd)
        print("  phonopy-qha completed successfully.")
    except subprocess.CalledProcessError as e:
        print("  Error during phonopy-qha execution:")
        print(e.stderr)
        sys.exit(1)

    print(f"\n--- QHA workflow for {input_poscar_path.name} completed successfully! ---")
    print(f"All results are in: {qha_output_dir}")


def main():
    parser = argparse.ArgumentParser(
        description="Perform Quasiharmonic Approximation (QHA) workflow using macer and phonopy.",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    # Group for initial relaxation
    init_relax_group = parser.add_argument_group('Initial Relaxation Settings')
    init_relax_group.add_argument("--initial-relax-mode", type=str, default="relax", choices=["relax", "sr"],
                                  help="Initial relaxation mode. 'relax' for standard 'macer relax --isif 3', 'sr' for 'macer_phonopy sr'. Default: relax.")
    init_relax_group.add_argument("--initial-fmax", type=float, default=1e-4,
                                  help="Force convergence threshold for the initial relaxation step. Default: 1e-4.")
    init_relax_group.add_argument("--tolerance-ru", type=float, default=0.01,
                                  help="Symmetry tolerance (Å) for 'sr' mode. Default: 0.01.")

    # Group for general settings
    general_group = parser.add_argument_group('General Settings')
    general_group.add_argument("--poscar", "-p", type=str, required=True,
                               help="Input POSCAR file for QHA workflow.")
    general_group.add_argument("--ff", type=str, default="mace",
                               choices=get_available_ffs(),
                               help="Force field to use for macer calculations.")
    general_group.add_argument("--model", type=str, default=None,
                               help="Path to the MLFF model file.")
    general_group.add_argument("--modal", type=str, default=None,
                               help="Modal for SevenNet model.")
    general_group.add_argument("--device", type=str, default="cpu",
                               choices=["cpu", "mps", "cuda"],
                               help="Compute device.")

    # Group for QHA volume and displacement settings
    qha_group = parser.add_argument_group('QHA and Phonopy Settings')
    qha_group.add_argument("--length-scale", "-ls", type=float, default=0.05,
                           help="Fractional range for lattice parameter scaling. Default: 0.05 (i.e., +/- 5%%).")
    qha_group.add_argument("--num-volumes", type=int, default=5,
                           help="Number of volume points to sample (min 4).")
    qha_group.add_argument("--atom-fix", action="store_true",
                           help="Skip 'macer relax --isif 2' for scaled structures.")
    qha_group.add_argument("--min-length", "-l", type=float, default=20.0,
                           help="Minimum supercell length (Å). This is ignored if --dim is provided.")
    qha_group.add_argument("--dim", type=int, nargs=3, default=None,
                           help="Override supercell dimensions (e.g., 2 2 2). Overrides --min-length.")
    qha_group.add_argument("--amplitude", type=float, default=0.01,
                           help="Displacement amplitude (Å).")
    qha_group.add_argument('--pm', dest='is_plusminus', action="store_true",
                           help='Generate plus and minus displacements.')
    qha_group.add_argument('--nodiag', dest='is_diagonal', action="store_false",
                           help='Do not generate diagonal displacements.')
    qha_group.add_argument("--mesh", type=int, nargs=3, default=[7, 7, 7],
                           help="Phonopy mesh grid (X Y Z).")
    qha_group.add_argument("--tmax", type=int, default=1300,
                           help="Maximum temperature for phonopy-qha (K).")


    args = parser.parse_args()

    if args.num_volumes < 4:
        parser.error("Number of volume points (--num-volumes) must be at least 4 for QHA.")

    run_qha_workflow(args)


if __name__ == "__main__":
    main()