import argparse
import math
import sys
import os
import glob
import subprocess
import shutil
import numpy as np

from phonopy import Phonopy
from phonopy.interface.vasp import read_vasp, write_vasp
from phonopy.interface.calculator import write_supercells_with_displacements

def run_macer_workflow(
    input_filename='POSCAR',
    min_length=20.0,
    displacement_distance=0.01,
    is_plusminus='auto',
    is_diagonal=True,
    macer_model_path=None,
    macer_device='cpu',
    tolerance=1e-3,
):
    """
    Automates the process for a single input file in its own directory.
    """
    output_prefix = os.path.basename(input_filename)
    print(f"\n{'='*20} Processing: {input_filename} in {os.getcwd()} {'='*20}")

    # Use default filenames that phonopy tools expect
    DEFAULT_DISP_YAML = "phonopy_disp.yaml"
    DEFAULT_FORCE_SETS = "FORCE_SETS"
    DEFAULT_FORCE_CONSTANTS = "FORCE_CONSTANTS"
    DEFAULT_BAND_CONF = "band.conf"
    DEFAULT_BAND_PDF = "band.pdf"
    DEFAULT_BAND_YAML = "band.yaml"
    DEFAULT_SPOSCAR = "SPOSCAR"
    DEFAULT_DISP_DIR = "displacements"

    # === Step 0: Relax and Symmetrize Unit Cell ===
    print("\n--- Step 0: Relaxing and symmetrizing unit cell ---")
    symmetrized_poscar_filename = f"{output_prefix}-symmetrized"
    macer_ru_command = [
        'macer_phonopy', 'sr',
        '-p', input_filename,
        '--output-prefix', output_prefix,
        '--tolerance', str(tolerance)
    ]
    print(f"Running command: {' '.join(macer_ru_command)}")
    try:
        subprocess.run(macer_ru_command, check=True, capture_output=True, text=True)
        print(f"Symmetrized unit cell created: {symmetrized_poscar_filename}")
    except subprocess.CalledProcessError as e:
        print("Error during macer_phonopy sr execution:")
        print(e.stderr)
        return

    # === Step 1: Create displaced supercells using phonopy ===
    print("\n--- Step 1: Creating supercells and displacements ---")
    try:
        unitcell = read_vasp(symmetrized_poscar_filename)
    except Exception as e:
        print(f"Error reading file {symmetrized_poscar_filename}: {e}")
        return

    cell = unitcell.cell
    vector_lengths = [np.linalg.norm(v) for v in cell]

    if any(v == 0 for v in vector_lengths):
        print("Error: One of the lattice vectors has a length of zero.")
        return

    scaling_factors = [math.ceil(min_length / v) if v > 0 else 1 for v in vector_lengths]
    supercell_matrix = np.diag(scaling_factors)

    print(f"Supercell matrix determined to be: {scaling_factors}")

    phonon = Phonopy(unitcell, supercell_matrix, primitive_matrix='auto')
    phonon.generate_displacements(
        distance=displacement_distance,
        is_plusminus=is_plusminus,
        is_diagonal=is_diagonal,
    )

    supercell = phonon.supercell
    displaced_cells = phonon.supercells_with_displacements

    write_vasp(DEFAULT_SPOSCAR, supercell)
    phonon.save(filename=DEFAULT_DISP_YAML)

    if os.path.exists(DEFAULT_DISP_DIR):
        shutil.rmtree(DEFAULT_DISP_DIR)
    os.makedirs(DEFAULT_DISP_DIR, exist_ok=True)
    
    disp_filenames_for_macer = [os.path.join(DEFAULT_DISP_DIR, f'POSCAR-{i+1:03d}') for i in range(len(displaced_cells))]
    
    for i, cell in enumerate(displaced_cells):
        write_vasp(disp_filenames_for_macer[i], cell)

    print(f"{DEFAULT_SPOSCAR}, {DEFAULT_DISP_YAML}, and {len(displaced_cells)} displaced POSCARs created.")

    # === Step 2: Relax structures with macer ===
    print("\n--- Step 2: Relaxing structures with macer ---")
    macer_command = [
        'macer', 'relax',
        '--poscar', DEFAULT_SPOSCAR, *disp_filenames_for_macer,
        '--isif', '0',
        '--device', macer_device,
    ]
    if macer_model_path:
        macer_command.extend(['--model', macer_model_path])

    print(f"Running command: {' '.join(macer_command)}")
    try:
        subprocess.run(macer_command, check=True, capture_output=True, text=True)
        print("macer relaxation completed successfully.")
    except subprocess.CalledProcessError as e:
        print("Error during macer execution:")
        print(e.stderr)
        return

    # === Step 3: Create FORCE_SETS with phonopy ===
    print("\n--- Step 3: Creating FORCE_SETS ---")
    vasprun_files = sorted(glob.glob(f'{DEFAULT_DISP_DIR}/vasprun-*.xml'))
    if not vasprun_files:
        print("Error: No vasprun-*.xml files found after macer relaxation.")
        return

    phonopy_f_command = ['phonopy', '-f'] + vasprun_files
    print(f"Running command: {' '.join(phonopy_f_command)}")
    try:
        result = subprocess.run(phonopy_f_command, check=True, capture_output=True, text=True)
        print(result.stdout)
        print(f"{DEFAULT_FORCE_SETS} created successfully.")
    except subprocess.CalledProcessError as e:
        print("Error during phonopy -f execution:")
        print(e.stderr)
        return

    # === Step 3.5: Create FORCE_CONSTANTS from FORCE_SETS ===
    print("\n--- Step 3.5: Creating FORCE_CONSTANTS ---")
    dim_string = " ".join(map(str, scaling_factors))
    phonopy_fc_command = [
        'phonopy',
        '-c', symmetrized_poscar_filename,
        '--dim', dim_string,
        '--writefc'
    ]
    print(f"Running command: {' '.join(phonopy_fc_command)}")
    try:
        subprocess.run(phonopy_fc_command, check=True, capture_output=True, text=True)
        print(f"{DEFAULT_FORCE_CONSTANTS} created successfully.")
    except subprocess.CalledProcessError as e:
        print("Error during phonopy --writefc execution:")
        print(e.stderr)
        return

    # === Step 4: Create band.conf using macer_phonopy ===
    print("\n--- Step 4: Creating band.conf ---")
    macer_bp_command = ['macer_phonopy', 'bp', '--poscar', symmetrized_poscar_filename, '--out', DEFAULT_BAND_CONF]
    print(f"Running command: {' '.join(macer_bp_command)}")
    try:
        subprocess.run(macer_bp_command, check=True, capture_output=True, text=True)
        print(f"{DEFAULT_BAND_CONF} created successfully.")
    except subprocess.CalledProcessError as e:
        print("Error during macer_phonopy bp execution:")
        print(e.stderr)
        return

    # === Step 5: Plot band structure using phonopy ===
    print("\n--- Step 5: Plotting band structure ---")
    phonopy_plot_command = ['phonopy', '-p', DEFAULT_BAND_CONF, '-s', '-c', symmetrized_poscar_filename]
    print(f"Running command: {' '.join(phonopy_plot_command)}")
    try:
        result = subprocess.run(phonopy_plot_command, check=True, capture_output=True, text=True)
        print(result.stdout)
        print(f"{DEFAULT_BAND_PDF} and {DEFAULT_BAND_YAML} created successfully.")
    except subprocess.CalledProcessError as e:
        print("Error during phonopy plot execution:")
        print(e.stderr)
        return

    # === Step 6: Rename and Clean up ===
    print("\n--- Step 6: Renaming outputs and cleaning up ---")
    rename_map = {
        DEFAULT_DISP_YAML: f"phonopy_disp-{output_prefix}.yaml",
        DEFAULT_FORCE_SETS: f"FORCE_SETS_{output_prefix}",
        DEFAULT_FORCE_CONSTANTS: f"FORCE_CONSTANTS_{output_prefix}",
        DEFAULT_BAND_CONF: f"band-{output_prefix}.conf",
        DEFAULT_BAND_PDF: f"band-{output_prefix}.pdf",
        DEFAULT_BAND_YAML: f"band-{output_prefix}.yaml",
    }
    for old, new in rename_map.items():
        if os.path.exists(old):
            os.rename(old, new)
            print(f"Renamed {old} -> {new}")

    # Cleanup intermediate files
    cleanup_items = glob.glob('vasprun-*.xml') + glob.glob('OUTCAR-*') + glob.glob('CONTCAR-*') + glob.glob('relax-*.log.*') + [DEFAULT_SPOSCAR, 'phonopy.yaml']
    if os.path.exists(DEFAULT_DISP_DIR):
        cleanup_items.append(DEFAULT_DISP_DIR)
    
    for item in cleanup_items:
        try:
            if os.path.isdir(item):
                shutil.rmtree(item)
            else:
                os.remove(item)
        except OSError as e:
            print(f"Error cleaning up {item}: {e}")

    print(f"\nWorkflow for {input_filename} completed.")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Full phonopy workflow using macer for force calculations."
    )
    parser.add_argument(
        "-p", "--poscar",
        dest="input_files",
        required=True,
        nargs='+',
        help="One or more input cell files in VASP POSCAR format."
    )
    parser.add_argument(
        "-l", "--length",
        type=float,
        default=20.0,
        help="Minimum length of supercell lattice vectors in Angstroms (default: 20.0)"
    )
    parser.add_argument(
        "--amplitude",
        type=float,
        default=0.01,
        help="Displacement amplitude in Angstroms (default: 0.01)"
    )
    parser.add_argument(
        '--tolerance',
        type=float,
        default=1e-3,
        help='Symmetry tolerance for spglib in macer_phonopy sr (default: 1e-3)'
    )
    parser.add_argument(
        '--pm',
        dest='is_plusminus',
        action="store_true",
        help='Generate plus and minus displacements for each direction.',
    )
    parser.add_argument(
        '--nodiag',
        dest='is_diagonal',
        action="store_false",
        help='Do not generate diagonal displacements.',
    )
    parser.add_argument(
        '--model', 
        type=str, 
        default=None, 
        help='Path to the MACE model file for macer.'
    )
    parser.add_argument(
        '--device', 
        type=str, 
        default='cpu', 
        choices=['cpu', 'mps', 'cuda'],
        help='Device for macer computation.'
    )

    args = parser.parse_args()

    is_plusminus = 'auto'
    if args.is_plusminus:
        is_plusminus = True

    original_cwd = os.getcwd()
    for filepath in args.input_files:
        abs_filepath = os.path.abspath(filepath)
        output_dir = os.path.dirname(abs_filepath)
        base_filename = os.path.basename(abs_filepath)

        os.makedirs(output_dir, exist_ok=True)

        try:
            os.chdir(output_dir)
            run_macer_workflow(
                input_filename=base_filename,
                min_length=args.length,
                displacement_distance=args.amplitude,
                is_plusminus=is_plusminus,
                is_diagonal=args.is_diagonal,
                macer_model_path=args.model,
                macer_device=args.device,
                tolerance=args.tolerance,
            )
        finally:
            os.chdir(original_cwd)
