from __future__ import annotations
import argparse
import os
import sys
import subprocess
import shutil
from concurrent.futures import ProcessPoolExecutor
import numpy as np
from pathlib import Path
import math
import yaml
from tqdm import tqdm



# Determine default force field based on installed extras
from types import SimpleNamespace

# Import necessary components from macer
from macer.defaults import DEFAULT_MODELS, _macer_root, _model_root, DEFAULT_DEVICE, DEFAULT_FF, resolve_model_path
from macer.calculator.factory import get_available_ffs, get_calculator, ALL_SUPPORTED_FFS
from macer.utils.logger import Logger
from macer import __version__


# Determine default force field based on installed extras
available_ffs = get_available_ffs()
_dynamic_default_ff = available_ffs[0] if available_ffs else None

_QHA_WORKER_CALC = None


def _init_qha_force_worker(ff_name, device, modal, model_path):
    global _QHA_WORKER_CALC
    calc_kwargs = {"device": device, "modal": modal}
    if ff_name == "mace":
        calc_kwargs["model_paths"] = [model_path]
    else:
        calc_kwargs["model_path"] = model_path
    _QHA_WORKER_CALC = get_calculator(ff_name=ff_name, **calc_kwargs)


def _qha_force_energy_worker(atoms):
    atoms.calc = _QHA_WORKER_CALC
    return atoms.get_forces(), atoms.get_potential_energy()


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


def _resolve_model_path(ff: str, model_path: str | None) -> str | None:
    """Resolve model path: prefer user-provided, else DEFAULT_MODELS."""
    if model_path:
        return resolve_model_path(str(model_path))
    
    FFS_USING_MODEL_NAME = {"fairchem", "orb", "chgnet", "m3gnet"}
    default_model_name = DEFAULT_MODELS.get(ff)

    if default_model_name:
        if ff in FFS_USING_MODEL_NAME:
            return default_model_name
        else:
            return resolve_model_path(default_model_name)
            
    return None


def _get_strain_scaling_factor(B_GPa: float) -> float:
    """
    Calculates a scaling factor for strain based on the bulk modulus.
    The function is an exponential decay of the form f(B) = 2**(1.4 - B/50).
    """
    if B_GPa <= 0:
        return 1.0
    return 2.5**(2.5 - B_GPa / 40.0)


def _plot_qha_results(output_dir, results):
    """Plots all QHA results and saves them to PDF files."""
    
    import matplotlib.pyplot as plt
    
    def _plot_single(filename_pdf, xlabel, ylabel, key_x, key_y):
        fig, ax = plt.subplots()
        ax.xaxis.set_ticks_position("both")
        ax.yaxis.set_ticks_position("both")
        ax.xaxis.set_tick_params(which="both", direction="in")
        ax.yaxis.set_tick_params(which="both", direction="in")
        
        # Filter out non-finite values to avoid plotting errors
        finite_mask = np.isfinite(results[key_x]) & np.isfinite(results[key_y])
        x_data = np.array(results[key_x])[finite_mask]
        y_data = np.array(results[key_y])[finite_mask]

        if len(x_data) > 0:
            ax.plot(x_data, y_data, 'r-')
            ax.set_xlim(left=0, right=max(x_data))
        
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.grid(True)
        plt.savefig(output_dir / filename_pdf)
        plt.close(fig)
        print(f"  {filename_pdf} generated.")

    _plot_single("gibbs-temperature.pdf", "Temperature (K)", "Gibbs free energy (eV)", 'temperature', 'gibbs')
    _plot_single("volume-temperature.pdf", "Temperature (K)", r"Volume $(\AA^3)$", 'temperature', 'volume')
    _plot_single("bulk_modulus-temperature.pdf", "Temperature (K)", "Bulk modulus (GPa)", 'temperature', 'bulk_modulus')
    _plot_single("thermal_expansion.pdf", "Temperature (K)", r"Thermal expansion ($10^{-6} \mathrm{K}^{-1}$)", 'temperature', 'thermal_expansion')
    _plot_single("Cp-temperature.pdf", "Temperature (K)", r"$C_P$ $\mathrm{(J/mol\cdot K)}$", 'temperature', 'cp')
    _plot_single("gruneisen-temperature.pdf", "Temperature (K)", "Gruneisen parameter", 'temperature', 'gruneisen')


def _plot_helmholtz_volume_polyfit(
    output_dir,
    temperatures,
    volumes,
    all_helmholtz_energies,
    all_fit_polynomials,
    all_equilibrium_points,
    thin_number=10
):
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots()
    ax.xaxis.set_ticks_position("both")
    ax.yaxis.set_ticks_position("both")
    ax.xaxis.set_tick_params(which="both", direction="in")
    ax.yaxis.set_tick_params(which="both", direction="in")

    volume_range = np.linspace(min(volumes), max(volumes), 100)
    
    e_min = all_equilibrium_points[0]['energy'] if all_equilibrium_points else 0

    for i, T in enumerate(temperatures):
        if i % thin_number == 0 and all_fit_polynomials[i] is not None:
            ax.plot(volumes, all_helmholtz_energies[i] - e_min, 'bo', markersize=3)
            ax.plot(volume_range, all_fit_polynomials[i](volume_range) - e_min, 'b-')

    eq_volumes = [p['volume'] for p in all_equilibrium_points if p]
    eq_energies = [p['energy'] for p in all_equilibrium_points if p]
    ax.plot(eq_volumes, np.array(eq_energies) - e_min, 'ro-', markersize=3)

    ax.set_xlabel(r"Volume $(\AA^3)$")
    ax.set_ylabel("Free energy (eV)")
    ax.grid(True)
    plt.savefig(output_dir / "helmholtz-volume.pdf")
    plt.close(fig)
    print(f"  helmholtz-volume.pdf generated.")


def run_local_poly_qha_analysis(
    volumes,
    electronic_energies,
    thermal_properties_with_volumes,
    output_dir,
    poly_degree,
    poly_points,
    tmax,
    smooth_deg,
):
    """
    Performs QHA analysis using a local polynomial fit for the F-V curve.
    """
    from scipy.interpolate import interp1d
    from scipy.optimize import minimize_scalar
    from phonopy.physical_units import get_physical_units

    AVOGADRO = 6.02214076e23
    EV_TO_KJ_MOL = get_physical_units().EvTokJmol

    # 1. Load thermal properties
    thermal_props = {}
    for vol, path in thermal_properties_with_volumes:
        with open(path) as f:
            data = yaml.safe_load(f)
        
        temps_list = [item['temperature'] for item in data['thermal_properties']]
        free_energy_list = [item['free_energy'] for item in data['thermal_properties']]
        entropy_list = [item['entropy'] for item in data['thermal_properties']]
        heat_capacity_list = [item['heat_capacity'] for item in data['thermal_properties']]

        thermal_props[vol] = {
            'temperatures': np.array(temps_list),
            'free_energy': np.array(free_energy_list),
            'entropy': np.array(entropy_list),
            'heat_capacity': np.array(heat_capacity_list),
        }

    # 2. Prepare data structures
    volumes = np.array(volumes)
    electronic_energies = np.array(electronic_energies)
    sort_idx = np.argsort(volumes)
    volumes = volumes[sort_idx]
    electronic_energies = electronic_energies[sort_idx]
    sorted_thermal_props = [thermal_props[vol] for vol in volumes]

    temps = np.linspace(0, tmax, 101)
    raw_results = { 'temperature': [], 'volume': [] }
    all_helmholtz_energies, all_fit_polynomials = [], []

    # 3. First pass: Calculate F(V,T) and find raw equilibrium V_eq(T)
    for T in temps:
        helmholtz_energies = []
        for i, V in enumerate(volumes):
            tp = sorted_thermal_props[i]
            f_interp = interp1d(tp['temperatures'], tp['free_energy'], bounds_error=False, fill_value='extrapolate')
            phonon_F_kJ_mol = f_interp(T)
            phonon_F_eV = phonon_F_kJ_mol / EV_TO_KJ_MOL
            helmholtz_energies.append(electronic_energies[i] + phonon_F_eV)
        
        helmholtz_energies = np.array(helmholtz_energies)
        all_helmholtz_energies.append(helmholtz_energies)

        min_F_idx = np.argmin(helmholtz_energies)
        num_fit_points = poly_points if poly_points is not None else len(volumes)
        start = max(0, min_F_idx - num_fit_points // 2)
        end = start + num_fit_points
        if end > len(volumes):
            end = len(volumes)
            start = max(0, end - num_fit_points)
        fit_volumes, fit_energies = volumes[start:end], helmholtz_energies[start:end]

        if len(fit_volumes) < poly_degree + 1:
            p_fit = None
        else:
            coeffs = np.polyfit(fit_volumes, fit_energies, poly_degree)
            p_fit = np.poly1d(coeffs)
        all_fit_polynomials.append(p_fit)
        
        if p_fit is None:
            V_eq = np.nan
        else:
            res = minimize_scalar(p_fit, bounds=(volumes[0], volumes[-1]), method='bounded')
            V_eq = res.x if res.success else np.nan
        
        # Sanity check: V_eq must be within the sampled volume range
        if not (volumes[0] <= V_eq <= volumes[-1]):
            print(f"Warning: Found equilibrium volume ({V_eq:.2f}) is outside the sample range "
                  f"[{volumes[0]:.2f}, {volumes[-1]:.2f}] at T={T:.1f}K. Fit is unreliable. Skipping.")
            V_eq = np.nan
        
        # Stability check: curvature must be positive for a valid minimum
        if np.isfinite(V_eq) and p_fit is not None and p_fit.deriv(2)(V_eq) <= 0:
            print(f"Warning: Non-positive curvature at T={T:.1f}K. Fit is unreliable. Skipping.")
            V_eq = np.nan
        
        # Sanity check: V_eq must be within the sampled volume range
        if not (volumes[0] <= V_eq <= volumes[-1]):
            V_eq = np.nan
        
        # Stability check: curvature must be positive for a valid minimum
        if np.isfinite(V_eq) and p_fit.deriv(2)(V_eq) <= 0:
            V_eq = np.nan

        raw_results['temperature'].append(T)
        raw_results['volume'].append(V_eq)

    # 4. Smooth V(T) curve and re-calculate all properties
    smooth_results = { 'temperature': temps, 'gibbs': [], 'volume': [], 'bulk_modulus': [], 'thermal_expansion': [], 'cp': [], 'entropy': [], 'gruneisen': [] }
    all_eos_parameters = []

    valid_mask = np.isfinite(raw_results['volume'])
    if np.sum(valid_mask) > smooth_deg:
        v_smooth_coeffs = np.polyfit(np.array(raw_results['temperature'])[valid_mask], np.array(raw_results['volume'])[valid_mask], smooth_deg)
        v_smooth_fit = np.poly1d(v_smooth_coeffs)
        
        for i, T in enumerate(temps):
            V_smooth = v_smooth_fit(T)
            p_fit = all_fit_polynomials[i]

            if p_fit is None:
                for key in smooth_results:
                    if key != 'temperature': smooth_results[key].append(np.nan)
                all_eos_parameters.append(None)
                continue

            F_smooth = p_fit(V_smooth)
            B_GPa = V_smooth * p_fit.deriv(2)(V_smooth) * 160.21766208
            B_prime = -1 - V_smooth * p_fit.deriv(3)(V_smooth) / p_fit.deriv(2)(V_smooth) if abs(p_fit.deriv(2)(V_smooth)) > 1e-9 else np.nan
            
            alpha = (1 / V_smooth) * v_smooth_fit.deriv(1)(T) if V_smooth != 0 else 0

            cv_at_T_points = [interp1d(tp['temperatures'], tp['heat_capacity'], bounds_error=False, fill_value='extrapolate')(T) for tp in sorted_thermal_props]
            Cv_at_Veq = interp1d(volumes, cv_at_T_points, bounds_error=False, fill_value='extrapolate')(V_smooth)
            
            # Add entropy interpolation
            entropy_at_T_points = [interp1d(tp['temperatures'], tp['entropy'], bounds_error=False, fill_value='extrapolate')(T) for tp in sorted_thermal_props]
            Entropy_at_Veq = interp1d(volumes, entropy_at_T_points, bounds_error=False, fill_value='extrapolate')(V_smooth)

            V_m3, B_Pa = V_smooth * 1e-30, B_GPa * 1e9
            Cp_corr = T * V_m3 * B_Pa * alpha**2 * AVOGADRO
            Cp = Cv_at_Veq + Cp_corr
            
            Cv_per_cell = Cv_at_Veq / AVOGADRO
            gamma = (alpha * B_Pa * V_m3 / Cv_per_cell) if Cv_per_cell > 1e-12 else 0

            smooth_results['volume'].append(V_smooth); smooth_results['gibbs'].append(F_smooth)
            smooth_results['bulk_modulus'].append(B_GPa); smooth_results['thermal_expansion'].append(alpha * 1e6)
            smooth_results['cp'].append(Cp); smooth_results['entropy'].append(Entropy_at_Veq); smooth_results['gruneisen'].append(gamma)
            all_eos_parameters.append((F_smooth, B_GPa, B_prime, V_smooth))
    else:
        print("Warning: Not enough valid points to create a smooth V(T) curve. Skipping smoothing.")
        smooth_results = raw_results # Fallback to raw results

    # 5. Write and plot results
    print("  Writing result files and generating plots...")

    dat_headers = {
        'gibbs': "# Temperature (K), Gibbs free energy (eV)\n",
        'volume': "# Temperature (K), Equilibrium volume (A^3)\n",
        'bulk_modulus': "# Temperature (K), Bulk modulus (GPa)\n",
        'thermal_expansion': "# Temperature (K), Thermal expansion (10^-6 K^-1)\n",
        'cp': "# Temperature (K), Heat capacity Cp (J/mol-K)\n",
        'entropy': "# Temperature (K), Entropy (J/mol-K)\n",
        'gruneisen': "# Temperature (K), Gruneisen parameter\n"
    }

    def write_dat(filename, header, key1, key2):
        with open(output_dir / filename, 'w') as f:
            f.write(header)
            if key1 in smooth_results and key2 in smooth_results and len(smooth_results[key1]) > 0:
                for val1, val2 in zip(smooth_results[key1], smooth_results[key2]):
                    if np.isfinite(val1) and np.isfinite(val2):
                        f.write(f"{val1:.3f} {val2:.8f}\n")
    
    for key in ['gibbs', 'volume', 'bulk_modulus', 'thermal_expansion', 'cp', 'entropy', 'gruneisen']:
        filename_dat = f'{key}-temperature.dat'
        header = dat_headers.get(key, f"# Temperature (K)  {key}\n")
        write_dat(filename_dat, header, 'temperature', key)
        print(f"  {filename_dat} generated.")

    # Write thermal_properties.yaml
    tp_data = {
        "unit-temperature": "K",
        "unit-free-energy": "kJ/mol",
        "unit-entropy": "J/mol-K",
        "unit-heat-capacity": "J/mol-K",
        "thermal_properties": []
    }
    for i in range(len(smooth_results['temperature'])):
        if i >= len(smooth_results['gibbs']) or np.isnan(smooth_results['gibbs'][i]): continue
        tp_data["thermal_properties"].append({
            "temperature": float(smooth_results['temperature'][i]),
            "free_energy": float(smooth_results['gibbs'][i] * EV_TO_KJ_MOL),
            "entropy": float(smooth_results['entropy'][i]),
            "heat_capacity": float(smooth_results['cp'][i])
        })
    with open(output_dir / "thermal_properties.yaml", "w") as f:
        yaml.dump(tp_data, f, default_flow_style=False, sort_keys=False)
    print("  thermal_properties.yaml generated.")

    with open(output_dir / 'helmholtz-volume.dat', 'w') as f:
        f.write("# Volume (A^3), Free energy (eV)\n")
        for i, T in enumerate(temps):
            if i < len(all_eos_parameters) and all_eos_parameters[i] is not None:
                f.write(f"# Temperature: {T:.5f}\n")
                f.write("# Parameters: %f %f %f %f\n" % all_eos_parameters[i])
                for j, V in enumerate(volumes):
                    f.write(f"{V:20.15f} {all_helmholtz_energies[i][j]:25.15f}\n")
                f.write("\n\n")
    print("  helmholtz-volume.dat generated.")

    valid_points = [{'volume': v, 'energy': p(v)} for v, p in zip(smooth_results.get('volume', []), all_fit_polynomials) if p is not None and np.isfinite(v)]
    if valid_points:
        _plot_helmholtz_volume_polyfit(output_dir, temps, volumes, all_helmholtz_energies, all_fit_polynomials, valid_points)
    
    _plot_qha_results(output_dir, smooth_results)
    print("  Local polynomial QHA analysis finished.")


def _write_standard_qha_extra_outputs(output_dir, e_v_dat, thermal_paths, eos_name):
    """
    Post-processes standard phonopy-qha results to add entropy-temperature.dat
    and thermal_properties.yaml (similar to local_poly output).
    """
    try:
        from phonopy import PhonopyQHA
        from phonopy.file_IO import read_v_e, read_thermal_properties_yaml
        import numpy as np
        import yaml
    except ImportError:
        print("  Warning: Could not import phonopy for post-processing extra QHA outputs.")
        return

    EV_TO_KJ_MOL = 96.485391 # Approximate conversion used by phonopy

    try:
        # Read data
        volumes, electronic_energies = read_v_e(str(e_v_dat))
        (temperatures,
         cv,
         entropy,
         fe_phonon,
         num_modes,
         num_integrated_modes) = read_thermal_properties_yaml([str(p) for p in thermal_paths])

        # Initialize QHA
        phonopy_qha = PhonopyQHA(
            volumes=volumes,
            electronic_energies=electronic_energies,
            temperatures=temperatures,
            free_energy=fe_phonon,
            cv=cv,
            entropy=entropy,
            eos=eos_name,
            verbose=False
        )
        
        # Get results from internal QHA object
        qha = phonopy_qha._qha
        t_qha = qha._temperatures[:qha._len]
        v_eq = qha._equiv_volumes[:qha._len]
        g_qha = qha._equiv_energies[:qha._len]
        cp_qha = qha._cp_numerical[:qha._len]
        
        # Interpolate entropy and Cv at equilibrium volumes
        s_eq = []
        cv_eq = []
        for i, t in enumerate(t_qha):
            s_val = np.interp(v_eq[i], volumes, entropy[i])
            cv_val = np.interp(v_eq[i], volumes, cv[i])
            s_eq.append(s_val)
            cv_eq.append(cv_val)

        s_eq = np.array(s_eq)
        cv_eq = np.array(cv_eq)

        # Write entropy-temperature.dat
        with open(output_dir / "entropy-temperature.dat", "w") as f:
            f.write("# Temperature (K), Entropy (J/mol-K)\n")
            for t, s in zip(t_qha, s_eq):
                f.write(f"{t:20.10f} {s:20.10f}\n")
        print("  entropy-temperature.dat generated.")

        # Write thermal_properties.yaml
        tp_data = {
            "unit-temperature": "K",
            "unit-free-energy": "kJ/mol",
            "unit-entropy": "J/mol-K",
            "unit-heat-capacity": "J/mol-K",
            "thermal_properties": []
        }
        for i in range(len(t_qha)):
            tp_data["thermal_properties"].append({
                "temperature": float(t_qha[i]),
                "free_energy": float(g_qha[i] * EV_TO_KJ_MOL),
                "entropy": float(s_eq[i]),
                "heat_capacity": float(cv_eq[i])
            })
        
        with open(output_dir / "thermal_properties.yaml", "w") as f:
            yaml.dump(tp_data, f, default_flow_style=False, sort_keys=False)
        print("  thermal_properties.yaml generated.")

    except Exception as e:
        print(f"  Warning: Error during standard QHA post-processing: {e}")


def _add_headers_to_qha_dat_files(output_dir):
    """Adds headers with units to .dat files generated by phonopy-qha or macer."""
    headers = {
        "bulk_modulus-temperature.dat": "# Temperature (K), Bulk modulus (GPa)\n",
        "Cv-volume.dat": "# Volume (A^3), Heat capacity Cv (J/mol-K)\n",
        "entropy-temperature.dat": "# Temperature (K), Entropy (J/mol-K)\n",
        "gruneisen-temperature.dat": "# Temperature (K), Gruneisen parameter\n",
        "thermal_expansion.dat": "# Temperature (K), Thermal expansion (K^-1)\n",
        "Cp-temperature_polyfit.dat": "# Temperature (K), Heat capacity Cp (J/mol-K)\n",
        "Cp-temperature.dat": "# Temperature (K), Heat capacity Cp (J/mol-K)\n",
        "entropy-volume.dat": "# Volume (A^3), Entropy (J/mol-K)\n",
        "gibbs-temperature.dat": "# Temperature (K), Gibbs free energy (eV)\n",
        "volume-temperature.dat": "# Temperature (K), Equilibrium volume (A^3)\n",
        "helmholtz-volume.dat": "# Volume (A^3), Free energy (eV)\n",
        "helmholtz-volume_fitted.dat": "# Volume (A^3), Free energy (eV)\n",
        "dsdv-temperature.dat": "# Temperature (K), dS/dV (GPa/K)\n",
        "e-v.dat": "# Volume (A^3), Energy (eV)\n"
    }

    for filename, header in headers.items():
        file_path = output_dir / filename
        if file_path.exists():
            with open(file_path, 'r') as f:
                content = f.read()
            if not content: continue
            
            # If our header is already there, skip
            if header in content: continue
            
            # If there's an existing simple header, we might want to replace it
            # But the simplest is just to prepend if header not in content.
            # phonopy-qha often writes no header.
            with open(file_path, 'w') as f:
                f.write(header + content)


def generate_displacements_for_qha(
    poscar_path: Path,
    supercell_matrix,
    displacement_distance: float,
    is_plusminus: bool,
    is_diagonal: bool,
    tolerance_phonopy: float,
):
    """Only generates displacement files, does not calculate forces."""
    from phonopy import Phonopy
    from phonopy.interface.vasp import read_vasp, write_vasp

    unitcell = read_vasp(str(poscar_path))

    phonon = Phonopy(
        unitcell,
        supercell_matrix=supercell_matrix,
        primitive_matrix="auto",
        symprec=tolerance_phonopy,
    )

    phonon.generate_displacements(
        distance=displacement_distance,
        is_plusminus=is_plusminus,
        is_diagonal=is_diagonal,
    )

    phonon.save("phonopy_disp.yaml")
    write_vasp("SPOSCAR", phonon.supercell)
    displaced_cells = phonon.supercells_with_displacements
    for i, cell in enumerate(displaced_cells):
        write_vasp(f"POSCAR-{(i + 1):03d}", cell)
    return phonon.dataset, len(displaced_cells)


def create_force_constants_from_force_sets_api(
    poscar_path,
    supercell_matrix,
    tolerance_phonopy,
    fc_calculator='symfc'
):
    """
    Creates FORCE_CONSTANTS from FORCE_SETS using phonopy API.
    This is an API-based replacement for the CLI command:
    `phonopy -c POSCAR --dim ... --writefc --symfc`
    """
    from phonopy import load as load_phonopy
    from phonopy.file_IO import write_FORCE_CONSTANTS

    phonon = load_phonopy(
        unitcell_filename=str(poscar_path),
        supercell_matrix=supercell_matrix,
        primitive_matrix="auto",
        force_sets_filename="FORCE_SETS",
        fc_calculator=fc_calculator,
        produce_fc=True,
        symprec=tolerance_phonopy,
        log_level=0
    )
    
    if phonon.force_constants is None:
        raise RuntimeError("Force constants could not be calculated via API.")

    write_FORCE_CONSTANTS(phonon.force_constants, filename="FORCE_CONSTANTS")


def run_qha_workflow(args):
    # Lazy imports
    from ase.io import read, write
    from macer.relaxation.bulk_modulus import get_bulk_modulus_and_volume
    from phonopy.file_IO import write_FORCE_SETS
    from phonopy import load as load_phonopy
    # --- Parse selected axes ---
    selected_axes_str = args.select_axis
    if selected_axes_str:
        selected_axes = {axis.strip().lower() for axis in selected_axes_str.split(',')}
        valid_axes = {'a', 'b', 'c'}
        if not selected_axes.issubset(valid_axes):
            raise ValueError(f"Error: Invalid axis in --select-axis. Allowed values are 'a', 'b', 'c'. You provided: {selected_axes_str}")
    else:
        selected_axes = {'a', 'b', 'c'}  # Default is all axes
    if getattr(args, "n_jobs", 1) < 1:
        raise ValueError("Error: --n-jobs must be >= 1.")

    # --- Configuration ---
    from macer.utils.struct_tools import resolve_structure_inputs
    input_list = resolve_structure_inputs(
        poscar_paths=[args.poscar] if args.poscar else None,
        cif_paths=[args.cif] if args.cif else None,
        formula=getattr(args, "formula", None),
        mpid=getattr(args, "mpid", None),
    )
    if not input_list:
        raise ValueError("Error: Please provide structure input via -p/--poscar, -c/--cif, -f/--formula, or -m/--mpid.")
    input_poscar_path = Path(input_list[0]).resolve()

    if args.output_dir:
        base_output_dir = Path(args.output_dir).resolve()
    else:
        base_output_dir = input_poscar_path.parent / f"qha_{input_poscar_path.stem}-mlff={args.ff}"
    
    qha_output_dir = base_output_dir
    i = 1
    while qha_output_dir.exists():
        qha_output_dir = Path(f"{base_output_dir}-NEW{i:03d}")
        i += 1
    qha_output_dir.mkdir(parents=True, exist_ok=False)

    log_file = qha_output_dir / "macer_phonopy_qha.log"
    orig_stdout = sys.stdout
    
    with Logger(str(log_file)) as lg:
        sys.stdout = lg
        try:
            from macer import __version__
            print(f"--- Macer QHA Workflow (v{__version__}) ---")
            print(f"Command: {' '.join(sys.argv)}")
            ff = args.ff
            model_path = _resolve_model_path(ff, args.model)
            if model_path and os.path.exists(model_path):
                model_path = os.path.abspath(model_path)

            modal = args.modal
            device = args.device

            length_scale_range = args.length_scale
            lf_min_arg = args.length_factor_min
            lf_max_arg = args.length_factor_max
            num_volumes = args.num_volumes
            mesh_x, mesh_y, mesh_z = args.mesh
            min_length = args.min_length
            displacement_distance = args.amplitude
            is_plusminus = True
            is_diagonal = args.is_diagonal
            tolerance_phonopy = args.tolerance_phonopy
            tmax = args.tmax
            relax_atom = args.relax_atom

            # --- Parse mass override ---
            mass_map = {}
            if args.mass:
                if len(args.mass) % 2 != 0:
                    raise ValueError("Error: --mass option requires pairs of Symbol Mass (e.g. H 2.014).")
                for i in range(0, len(args.mass), 2):
                    sym = args.mass[i]
                    try:
                        m = float(args.mass[i+1])
                        mass_map[sym] = m
                    except ValueError:
                        raise ValueError(f"Error: Invalid mass value for {sym}: {args.mass[i+1]}")
                print(f"--- Atomic mass override ---")
                for sym, m in mass_map.items():
                    print(f"  {sym:3s} : {m:8.4f} amu")
                print("----------------------------")

            # --- Input validation for length scales ---
            if length_scale_range is not None and (lf_min_arg is not None or lf_max_arg is not None):
                raise ValueError("Error: --length-scale cannot be used with --length-factor-min/max.")
            if (lf_min_arg is not None and lf_max_arg is None) or (lf_min_arg is None and lf_max_arg is not None):
                raise ValueError("Error: --length-factor-min and --length-factor-max must be used together.")

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

            model_info_str = ""
            FFS_USING_MODEL_NAME = {"fairchem", "orb", "chgnet", "m3gnet"}
            if args.model:
                model_info_str = f" (from --model option)"
            else:
                default_model_name = DEFAULT_MODELS.get(ff)
                if default_model_name:
                    model_info_str = f" (default for {ff.upper()}: {default_model_name})"
            
            if model_path:
                print(f"  MLFF Model: {model_path}{model_info_str}")
            else:
                print(f"  MLFF Model:{model_info_str}")

            # --- Step 1: Initial Unit Cell Relaxation & Preparation ---
            print("\n--- Step 1: Initial Unit Cell Relaxation & Preparation ---")
            original_cwd = os.getcwd()
            try:
                # Copy initial structure
                copied_input_poscar_path = qha_output_dir / input_poscar_path.name
                shutil.copy(input_poscar_path, copied_input_poscar_path)
        
                os.chdir(qha_output_dir)

                if args.isif == 0:
                    print("ISIF=0 specified, skipping initial relaxation.")
                    qha_input_poscar_path = copied_input_poscar_path
                else:
                    print(f"Performing initial cell relaxation (ISIF={args.isif}).")
                    relaxed_poscar_name = f"CONTCAR-{copied_input_poscar_path.name}"
                    macer_relax_command = [
                        sys.executable, '-m', 'macer.cli.main', 'relax', '-p', str(copied_input_poscar_path), '--isif', str(args.isif),
                        '--fmax', str(args.initial_fmax), '--device', device,
                        '--contcar', relaxed_poscar_name, '--no-pdf',
                        '--symprec', str(args.initial_symprec)
                    ]
                    if not args.initial_use_symmetry:
                        macer_relax_command.append('--symmetry-off')
                    if ff: macer_relax_command.extend(['--ff', ff])
                    if model_path: macer_relax_command.extend(['--model', str(model_path)])
                    if modal: macer_relax_command.extend(['--modal', modal])
            
                    try:
                        # Run without capturing output to allow seeing segfault details or C-level errors directly
                        subprocess.run(macer_relax_command, check=True, text=True)
                    except subprocess.CalledProcessError as e:
                        print("Error during initial 'macer relax' execution.")
                        # e.stderr is None when not captured, so we don't print it.
                        raise e

                    qha_input_poscar_path = qha_output_dir / relaxed_poscar_name

                if not qha_input_poscar_path.exists():
                    print("Initial relaxation finished successfully but output file was not found.")
                    if 'result' in locals():
                         print("STDOUT:", result.stdout)
                         print("STDERR:", result.stderr)
                    raise FileNotFoundError("Relaxed/symmetrized POSCAR not found.")
        
                base_structure_atoms = read(str(qha_input_poscar_path), format="vasp")
                initial_cell = base_structure_atoms.get_cell()

                # --- Determine supercell matrix ---
                supercell_matrix = None
                if args.dim:
                    if len(args.dim) == 3:
                        supercell_matrix = np.diag(args.dim)
                        print(f"Using user-provided diagonal supercell DIM = {args.dim}")
                    elif len(args.dim) == 9:
                        supercell_matrix = np.array(args.dim).reshape(3, 3)
                        print(f"Using user-provided supercell matrix:\n{supercell_matrix}")
                    else:
                        raise ValueError("Error: --dim must be followed by 3 or 9 integers.")
                else:
                    # Auto-determine from min_length
                    scaling_factors_list = [max(1, math.ceil(min_length / np.linalg.norm(v))) for v in initial_cell]
                    supercell_matrix = np.diag(scaling_factors_list)
                    print(f"Auto-determining supercell from --min-length. DIM = {scaling_factors_list}")

                # For phonopy CLI calls, we need a string representation
                dim_str = " ".join(map(str, supercell_matrix.flatten().astype(int)))
        
                # Determine strain range
                if lf_min_arg is not None:  # Absolute user-provided range
                    length_scales = np.linspace(lf_min_arg, lf_max_arg, num_volumes)
                    print(f"Using user-provided absolute length factor range: {lf_min_arg} to {lf_max_arg}")
                else:
                    if length_scale_range is not None:  # Symmetric user-provided strain range
                        ls_min = ls_max = length_scale_range
                        print(f"Using user-provided symmetric strain range: \u00b1{length_scale_range*100:.1f}%")
                    else:  # Estimate from bulk modulus
                        print("\n--- Estimating optimal length scale from bulk modulus ---")
                        calc_args = SimpleNamespace(ff=ff, model=model_path, device=device, modal=modal)
                        B_GPa, V0_per_atom = get_bulk_modulus_and_volume(base_structure_atoms, calc_args)
                        if B_GPa and V0_per_atom:
                            E_target_eV = args.target_energy / 1000.0
                            B_eV_per_A3 = B_GPa / 160.21766208
                            epsilon_V = math.sqrt(2 * E_target_eV / (B_eV_per_A3 * V0_per_atom))
                    
                            scaling_factor = _get_strain_scaling_factor(B_GPa)
                            estimated_range = epsilon_V * scaling_factor
                            ls_min = ls_max = estimated_range
                    
                            print(f"  - Calculated Bulk Modulus: {B_GPa:.2f} GPa")
                            print(f"  - Base estimated strain: \u00b1{epsilon_V:.4f}")
                            print(f"  - Exponential scaling factor for strain: {scaling_factor:.3f}")
                            print(f"  Final estimated length-scale range: \u00b1{estimated_range*100:.1f}%")
                        else:
                            print("  Bulk modulus estimate failed; fallback to \u00b15%.")
                            ls_min = ls_max = 0.05
                    length_scales = np.linspace(1 - ls_min, 1 + ls_max, num_volumes)

                print(f"Length scales = {length_scales}")

            finally:
                os.chdir(original_cwd)

            # --- Step 2: Generate all displacements for all volumes first ---
            print("\n--- Step 2: Generating all displacements for all volumes ---")
            all_disp_atoms = []
            volume_dirs = []
            datasets = {}
            structure_info = []
            for l_scale in length_scales:
                volume_dir = qha_output_dir / f"length-scale={l_scale:.5f}"
                volume_dirs.append(volume_dir)
                volume_dir.mkdir(exist_ok=True)
                print(f"\n--- Preparing volume for length scale {l_scale:.5f} ---")

                scaled_atoms = read(str(qha_input_poscar_path), format="vasp")
        
                # Apply selective scaling based on --select-axis
                new_cell = initial_cell.copy()
                scale_a = l_scale if 'a' in selected_axes else 1.0
                scale_b = l_scale if 'b' in selected_axes else 1.0
                scale_c = l_scale if 'c' in selected_axes else 1.0
                new_cell[0] *= scale_a
                new_cell[1] *= scale_b
                new_cell[2] *= scale_c
                scaled_atoms.set_cell(new_cell, scale_atoms=True)

                poscar_path_vol = volume_dir / "POSCAR"
                write(str(poscar_path_vol), scaled_atoms, format="vasp")

                os.chdir(volume_dir)
                try:
                    if relax_atom:
                        print("  Relaxing scaled structure (atoms only)...")
                        relax_cmd = ['macer', 'relax', '-p', 'POSCAR', '--isif', '2', '--device', device, '--contcar', 'CONTCAR-POSCAR', '--no-pdf', '--symprec', str(args.initial_symprec)]
                        if not args.initial_use_symmetry:
                            relax_cmd.append('--symmetry-off')
                        if ff: relax_cmd.extend(['--ff', ff])
                        if model_path: relax_cmd.extend(['--model', str(model_path)])
                        if modal: relax_cmd.extend(['--modal', modal])
                        subprocess.run(relax_cmd, check=True, capture_output=True, text=True)
                        shutil.move("CONTCAR-POSCAR", "POSCAR")
            
                    dataset, num_disps = generate_displacements_for_qha(
                        poscar_path=Path("POSCAR"),
                        supercell_matrix=supercell_matrix,
                        displacement_distance=displacement_distance,
                        is_plusminus=is_plusminus,
                        is_diagonal=is_diagonal,
                        tolerance_phonopy=tolerance_phonopy,
                    )
                    datasets[l_scale] = dataset
            
                    # Collect all atoms from POSCAR and POSCAR-xxx for batch calculation
                    all_disp_atoms.append(read("POSCAR", format="vasp")) # Relaxed, undisplaced
                    structure_info.append({'l_scale': l_scale, 'disp_num': 0})
                    for i in range(1, num_disps + 1):
                        all_disp_atoms.append(read(f"POSCAR-{(i):03d}", format="vasp"))
                        structure_info.append({'l_scale': l_scale, 'disp_num': i})
                finally:
                    os.chdir(original_cwd)

                shutil.copy(volume_dir / "SPOSCAR", qha_output_dir / f"SPOSCAR_scale_{l_scale:.5f}")
                print(f"  Copied SPOSCAR to {qha_output_dir.name}/SPOSCAR_scale_{l_scale:.5f}")

            # --- Step 3: Batch Force & Energy Calculation ---
            print("\n--- Step 3: Batch force and energy calculation for all structures ---")
            try:
                all_forces = []
                all_energies = []

                n_jobs = max(1, int(getattr(args, "n_jobs", 1)))
                use_parallel = n_jobs > 1 and device == "cpu"
                if n_jobs > 1 and device != "cpu":
                    print(f"  WARNING: --n-jobs>1 is supported only with --device cpu (current: {device}). Falling back to serial evaluation.")

                if use_parallel:
                    chunksize = max(1, len(all_disp_atoms) // (n_jobs * 4))
                    with ProcessPoolExecutor(
                        max_workers=n_jobs,
                        initializer=_init_qha_force_worker,
                        initargs=(ff, device, modal, model_path),
                    ) as pool:
                        results_iter = pool.map(_qha_force_energy_worker, all_disp_atoms, chunksize=chunksize)
                        for forces_i, energy_i in tqdm(results_iter, total=len(all_disp_atoms), desc="Batch force/energy calculation"):
                            all_forces.append(forces_i)
                            all_energies.append(energy_i)
                else:
                    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"]
                    calculator = get_calculator(ff_name=ff, **calc_kwargs)
                    for atoms in tqdm(all_disp_atoms, desc="Batch force/energy calculation"):
                        atoms.calc = calculator
                        all_forces.append(atoms.get_forces())
                        all_energies.append(atoms.get_potential_energy())
                print("  All forces and energies calculated in a single batch.")

            except Exception as e:
                print(f"Error during batch calculation: {e}")
                # Use raise instead of sys.exit
                raise e

            # --- Step 4: Process results for each volume ---
            print("\n--- Step 4: Processing results for each volume ---")
            qha_energies, qha_volumes, thermal_properties_with_volumes = [], [], []
            atom_idx_counter = 0
            for l_scale, volume_dir in zip(length_scales, volume_dirs):
                print(f"\n--- Post-processing for length scale {l_scale:.5f} ---")
                os.chdir(volume_dir)
                try:
                    # Get energy and volume for the relaxed, undisplaced structure
                    energy_ev = all_energies[atom_idx_counter]
                    vol_atoms = read("POSCAR", format="vasp")
                    volume_a3 = vol_atoms.get_volume()
                    qha_energies.append(energy_ev)
                    qha_volumes.append(volume_a3)
                    atom_idx_counter += 1

                    # Create FORCE_SETS
                    dataset = datasets[l_scale]
                    num_disps = len(dataset.get('first_atoms', []))
                    force_sets = all_forces[atom_idx_counter : atom_idx_counter + num_disps]
                    atom_idx_counter += num_disps

                    if "first_atoms" in dataset:
                        for i, forces in enumerate(force_sets):
                            dataset["first_atoms"][i]["forces"] = forces
                    else:
                        dataset["forces"] = np.array(force_sets)
            
                    force_sets_path = volume_dir / "FORCE_SETS"
                    write_FORCE_SETS(dataset, filename=str(force_sets_path))
                    print(f"  FORCE_SETS created successfully.")

                    if args.use_force_constants:
                        # Create FORCE_CONSTANTS from FORCE_SETS using API
                        print("  Creating FORCE_CONSTANTS from FORCE_SETS (using API)...")
                        create_force_constants_from_force_sets_api(
                            poscar_path=Path("POSCAR"),
                            supercell_matrix=supercell_matrix,
                            tolerance_phonopy=tolerance_phonopy,
                            fc_calculator='symfc'
                        )
                        print("  FORCE_CONSTANTS created successfully.")
            
                    # Compute thermal properties
                    print(f"  Calculating thermal properties (mesh: {mesh_x} {mesh_y} {mesh_z})...")
            
                    # Load phonopy via API
                    # This automatically picks up FORCE_SETS or FORCE_CONSTANTS in the current directory
                    phonon = load_phonopy(
                        unitcell_filename="POSCAR",
                        supercell_matrix=supercell_matrix,
                        primitive_matrix="auto",
                        log_level=0
                    )

                    # Apply mass override
                    if mass_map:
                        current_masses = phonon.masses
                        symbols = phonon.primitive.symbols # Symbols of the primitive cell atoms
                        new_masses = []
                        for s, m in zip(symbols, current_masses):
                            if s in mass_map:
                                new_masses.append(mass_map[s])
                            else:
                                new_masses.append(m)
                        phonon.masses = new_masses
                        print(f"  Applied mass override for thermal properties.")

                    phonon.run_mesh(
                        mesh=[mesh_x, mesh_y, mesh_z],
                        is_gamma_center=True
                    )
                    phonon.run_thermal_properties(
                        t_step=10,
                        t_max=tmax,
                        t_min=0
                    )
                    phonon.write_yaml_thermal_properties(filename="thermal_properties.yaml")
            
                    thermal_properties_with_volumes.append(
                        (volume_a3, str(volume_dir / "thermal_properties.yaml"))
                    )
                    print(f"  Collected E={energy_ev:.8f} eV, V={volume_a3:.4f} Å³")

                except (subprocess.CalledProcessError, RuntimeError, Exception) as e:
                    print("  Error during post-processing:")
                    if hasattr(e, 'stderr') and e.stderr:
                        print(e.stderr)
                    else:
                        print(e)
                        import traceback
                        traceback.print_exc()
                    raise e
                finally:
                    os.chdir(original_cwd)

            # --- Step 5: Final QHA analysis ---
            print("\n--- Step 5: Running phonopy-qha analysis ---")
            e_v_dat = qha_output_dir / "e-v.dat"
            with open(e_v_dat, "w") as f:
                f.write("# Volume (A^3), Energy (eV)\n")
                for V, E in sorted(zip(qha_volumes, qha_energies)):
                    f.write(f"{V:20.15f} {E:25.15f}\n")
            print(f"  e-v.dat written: {e_v_dat}")

            if args.eos == 'local_poly':
                print("  Using local polynomial fit for EOS.")
                run_local_poly_qha_analysis(
                    volumes=qha_volumes,
                    electronic_energies=qha_energies,
                    thermal_properties_with_volumes=thermal_properties_with_volumes,
                    output_dir=qha_output_dir,
                    poly_degree=args.poly_degree,
                    poly_points=args.poly_points,
                    tmax=args.tmax,
                    smooth_deg=args.smooth_deg
                )
            else:
                print(f"  Using '{args.eos}' EOS with phonopy-qha.")
                thermal_paths = [p for _, p in sorted(thermal_properties_with_volumes)]
                qha_cmd = ["phonopy-qha", "--eos", args.eos, str(e_v_dat)] + thermal_paths + ["--tmax", str(tmax), "-s"]
                os.chdir(qha_output_dir)
                subprocess.run(qha_cmd, check=True, capture_output=True, text=True)
                
                # Post-process for extra outputs (entropy-temperature.dat and thermal_properties.yaml)
                _write_standard_qha_extra_outputs(qha_output_dir, e_v_dat, thermal_paths, args.eos)
                
                os.chdir(original_cwd)

            # Final post-processing to ensure all .dat files have headers
            _add_headers_to_qha_dat_files(qha_output_dir)

        except Exception as e:
            # Catch any other exceptions and re-raise so interactive shell can catch them
            raise e
        finally:
            sys.stdout = orig_stdout

    print("\nQHA workflow completed successfully!")


def add_qha_parser(subparsers):
    qha_parser = subparsers.add_parser(
        "run-qha",
        help="Quasi-Harmonic Approximation (Thermal expansion, Free energy)",
        description=MACER_LOGO + f"\nmacer_phonopy qha (v{__version__}): Perform Quasiharmonic Approximation (QHA) workflow using MLFFs and phonopy.",
        epilog="""
Examples:
  # 1. Standard QHA workflow: Auto-detect strain range, compute up to 1300K
  macer phonopy qha -p POSCAR --dim 2 2 2 --tmax 1300 --ff mattersim

  # 2. Manual strain range: +/- 5% volume change with 11 points
  macer phonopy qha -p POSCAR --dim 2 2 2 --num-volumes 11 --length-scale 0.05

  # 3. Anisotropic QHA: Apply strain only to the c-axis (e.g., 2D materials)
  macer phonopy qha -p POSCAR --dim 3 3 1 --select-axis c --eos vinet

  # 4. High-precision: Relax atomic positions at each volume step
  macer phonopy qha -p POSCAR --dim 2 2 2 --relax-atom --initial-fmax 0.005
""",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        aliases=["qha"],
    )

    # General & Input
    general_group = qha_parser.add_argument_group("General & Input")
    general_group.add_argument("--poscar", "-p", required=False, default=None, help="Input crystal structure file (POSCAR or CIF).")
    general_group.add_argument("--cif", "-c", required=False, default=None, help="Input CIF file.")
    general_group.add_argument("-f", "--formula", type=str, help="Formula to retrieve from Materials Project (e.g., MgAl2O4)")
    general_group.add_argument("-m", "--mpid", type=str, help="Materials Project ID (e.g., mp-3536)")
    general_group.add_argument("--output-dir", help="Directory to save all output files.")

    # MLFF Settings
    mlff_group = qha_parser.add_argument_group("MLFF Model Settings")
    mlff_group.add_argument(
        "--ff",
        choices=ALL_SUPPORTED_FFS,
        default=DEFAULT_FF,
        help=f"Force field to use. (default: {DEFAULT_FF})"
    )
    mlff_group.add_argument("--model", type=str, default=None, help="Path to the force field model file.")
    mlff_group.add_argument("--modal", type=str, default=None, help="Modal for certain force fields like SevenNet.")
    mlff_group.add_argument(
        "--device",
        choices=["cpu", "mps", "cuda"],
        default=DEFAULT_DEVICE,
        help=f"Compute device. (default: {DEFAULT_DEVICE})"
    )
    mlff_group.add_argument("--n-jobs", type=int, default=1, help="Parallel workers for force calculations (CPU only, default: 1).")

    # Relaxation Settings
    relax_group = qha_parser.add_argument_group("Structure Relaxation Settings")
    relax_group.add_argument("--isif", type=int, default=3,
                       help="ISIF flag for initial cell relaxation, passed to 'macer relax'. "
                            "ISIF=3 (default) relaxes atoms, shape, and volume. ISIF=0 skips relaxation.")
    relax_group.add_argument("--initial-fmax", type=float, default=1e-2,
                        help="Force convergence threshold for initial relaxation in eV/Å. (default: 0.001)")
    relax_group.add_argument("--initial-symprec", type=float, default=1e-3,
                        help="Symmetry tolerance for FixSymmetry during initial relaxation (default: 1e-3 Å).")
    relax_group.add_argument(
        "--initial-symmetry-off",
        dest="initial_use_symmetry",
        action="store_false",
        help="Disable FixSymmetry constraint during initial relaxation."
    )
    relax_group.add_argument("--relax-atom", action="store_true",
                       help="Relax atomic positions at each volume scaling. If not set, atoms are fixed. (default: False)")

    # QHA Workflow Settings
    qha_wf_group = qha_parser.add_argument_group("QHA Workflow Settings")
    qha_wf_group.add_argument("--eos", type=str, default="vinet",
                       choices=["vinet", "birch_murnaghan", "murnaghan", "local_poly"],
                       help="Equation of state for fitting Helmholtz free energy. (default: vinet)")
    qha_wf_group.add_argument("--num-volumes", type=int, default=5,
                       help="Number of volume points to sample for the E-V curve. (default: 5)")
    qha_wf_group.add_argument("--length-scale", "-ls", type=float, default=None,
                       help="Symmetric strain range for E-V curve, e.g., 0.05 for \u00b15%%. "
                            "If not set, estimated automatically. Mutually exclusive with --length-factor-*.")
    qha_wf_group.add_argument("--length-factor-min", type=float, default=None,
                       help="Minimum absolute length scaling factor. Requires --length-factor-max. "
                            "Mutually exclusive with --length-scale.")
    qha_wf_group.add_argument("--length-factor-max", type=float, default=None,
                       help="Maximum absolute length scaling factor. Requires --length-factor-min. "
                            "Mutually exclusive with --length-scale.")
    qha_wf_group.add_argument("--target-energy", type=float, default=10.0,
                       help="Target energy in meV for bulk modulus-based strain estimation. (default: 10.0)")
    qha_wf_group.add_argument("--select-axis", type=str, default=None,
                       help="Comma-separated list of axes (a,b,c) to apply length scaling to. "
                            "e.g., 'c' or 'a,b'. Default is all axes ('a,b,c').")

    # Supercell & Displacement Settings
    supercell_group = qha_parser.add_argument_group("Supercell & Displacement Settings")
    supercell_group.add_argument("--dim", type=int, nargs='+', default=None,
                       help='Set supercell dimension. Accepts 3 integers for a diagonal matrix (e.g., "2 2 2") '
                            'or 9 integers for a full matrix (e.g., "1 0 1 -1 1 1 0 -1 1"). '
                            'Overrides --min-length.')
    supercell_group.add_argument("--min-length", "-l", type=float, default=15.0,
                       help="Minimum supercell lattice vector length in Å to determine DIM automatically if --dim is not set. (default: 15.0)")
    supercell_group.add_argument("--amplitude", type=float, default=0.01,
                       help="Displacement amplitude for phonopy in Å. (default: 0.01)")
    pm_group = supercell_group.add_mutually_exclusive_group()
    pm_group.add_argument("--pm", dest="is_plusminus", action="store_true",
                        help="Set plus-minus displacements for all directions (always enabled).")
    pm_group.add_argument("--no-pm", dest="is_plusminus", action="store_false",
                        help="Deprecated: ignored. Conservative mode always uses plus-minus displacements.")
    qha_parser.set_defaults(is_plusminus=True)
    supercell_group.add_argument("--nodiag", dest="is_diagonal", action="store_false",
                       help="Do not generate diagonal displacements. (default: enabled)")
    supercell_group.add_argument("--mass", nargs='+', help="Specify atomic masses. Format: Symbol Mass Symbol Mass ... (e.g. --mass H 2.014 D 2.014)")

    # Phonopy Settings
    phonopy_group = qha_parser.add_argument_group("Phonopy Settings")
    phonopy_group.add_argument("--tolerance-phonopy", type=float, default=1e-3,
                         help="Symmetry tolerance for phonopy in Å. (default: 1e-3)")
    phonopy_group.add_argument("--mesh", type=int, nargs=3, default=[7, 7, 7],
                       help="Reciprocal space mesh for thermal property calculation. (default: 7 7 7)")
    phonopy_group.add_argument("--tmax", type=int, default=1300,
                       help="Maximum temperature for thermal property calculation in Kelvin. (default: 1300)")
    phonopy_group.add_argument("--use-force-constants", action="store_true",
                       help="Use force constants for phonopy calculation. If not set, FORCE_SETS will be used. (default: False)")

    # Local Polynomial EOS Settings
    poly_group = qha_parser.add_argument_group("Local Polynomial EOS Settings")
    poly_group.add_argument("--poly-degree", type=int, default=2, help="Degree of polynomial for 'local_poly' EOS fit. (default: 2)")
    poly_group.add_argument("--poly-points", type=int, default=3, help="Number of points around the minimum to use for 'local_poly' EOS fit. (default: 3)")
    poly_group.add_argument("--smooth-deg", type=int, default=10, help="Degree of polynomial for smoothing the final V-T curve. (default: 3)")

    qha_parser.set_defaults(func=run_qha_workflow)
    return qha_parser
