"""
qSCAILD command implementation.

This module owns the `macer phonopy qscaild` execution path.
"""

from __future__ import annotations

import argparse
import itertools
import os
import random
import shutil
import subprocess
import sys
import time
import traceback
from collections import defaultdict
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import scipy.constants as codata
import yaml
from tqdm import tqdm

from macer.calculator.factory import get_calculator, ALL_SUPPORTED_FFS
from macer.defaults import DEFAULT_MODELS, DEFAULT_DEVICE, DEFAULT_FF, resolve_model_path
from macer.utils.logger import Logger
from macer.utils.struct_tools import resolve_structure_inputs

EV_A3_TO_GPA = 160.21766208


def _resolve_model_path(ff: str, model_path: str | None) -> str | None:
    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
        return resolve_model_path(default_model_name)
    return None


def _fit_fc2_and_fc3(U_scaled, F_scaled, ph):
    from symfc.utils.utils import SymfcAtoms
    from symfc.basis_sets import FCBasisSetO2, FCBasisSetO3
    from symfc.solvers import FCSolverO2O3

    print("    Fitting 2nd and 3rd order force constants simultaneously using symfc...")
    supercell = ph.supercell
    symfc_supercell = SymfcAtoms(
        cell=supercell.cell,
        scaled_positions=supercell.scaled_positions,
        numbers=supercell.numbers,
    )

    print("    Computing FC2 basis set...")
    basis_set_o2 = FCBasisSetO2(symfc_supercell).run()
    if basis_set_o2._blocked_basis_set is None:
        raise RuntimeError("FC2 basis set computation failed, _blocked_basis_set is None.")

    print("    Computing FC3 basis set...")
    basis_set_o3 = FCBasisSetO3(symfc_supercell).run()
    if basis_set_o3._blocked_basis_set is None:
        raise RuntimeError("FC3 basis set computation failed, _blocked_basis_set is None.")

    print("    Solving for FC2 and FC3...")
    solver = FCSolverO2O3([basis_set_o2, basis_set_o3])
    solver.solve(U_scaled, F_scaled)

    fc2, fc3 = solver.full_fc
    if fc2 is None or fc3 is None:
        raise RuntimeError("symfc failed to compute FC2/FC3.")

    print("    symfc fitting complete.")
    return fc2, fc3


def fit_weighted_fc(U, F, w, ph, include_third_order=False):
    sqrtw = np.sqrt(w)[:, None, None]
    U_scaled = sqrtw * U
    F_scaled = sqrtw * F

    if not include_third_order:
        ph.dataset = {"displacements": U_scaled, "forces": F_scaled}
        ph.produce_force_constants(fc_calculator="symfc")
        return ph.force_constants.copy(), None

    try:
        fc2, fc3 = _fit_fc2_and_fc3(U_scaled, F_scaled, ph)
        return fc2, fc3
    except RuntimeError as exc:
        print(f"  WARNING: FC3 fitting failed with error: {exc}")
        print("           Falling back to FC2-only fitting for this iteration.")
        ph.dataset = {"displacements": U_scaled, "forces": F_scaled}
        ph.produce_force_constants(fc_calculator="symfc")
        return ph.force_constants.copy(), None


def _compute_fc3_free_energy(fc3, U, w):
    if fc3 is None:
        return 0.0
    cubic = np.einsum("ijklmn,sil,sjm,skn->s", fc3, U, U, U)
    return -np.sum(w * cubic) / 6.0


def compute_weighted_free_energy(U, E, FC, FC3, w, ph, T, mesh, return_components=False):
    from phonopy.physical_units import get_physical_units

    ph.force_constants = FC
    ph.run_mesh(mesh)
    ph.run_thermal_properties(t_min=T, t_max=T, t_step=T)
    hfe_kJmol = ph.get_thermal_properties_dict()["free_energy"][0]
    hfe = hfe_kJmol / get_physical_units().EvTokJmol

    V_mean = np.sum(w * E)
    V_harm = 0.5 * np.sum(w * np.einsum("ijab,nia,njb->n", FC, U, U))
    n_prim = len(ph.supercell) / len(ph.primitive)
    F3 = _compute_fc3_free_energy(FC3, U, w)
    F_total = hfe + (V_mean - V_harm) / n_prim + F3 / n_prim

    if return_components:
        return F_total, hfe, V_mean / n_prim, V_harm / n_prim
    return F_total


def _read_band_yaml_data(band_yaml_path):
    with open(band_yaml_path, "r") as file:
        data = yaml.safe_load(file)

    phonons = data["phonon"]
    distances = []
    frequencies = []
    qpoints = []
    for p in phonons:
        distances.append(p["distance"])
        frequencies.append([b["frequency"] for b in p["band"]])
        qpoints.append(p["q-position"])

    distances = np.array(distances)
    frequencies = np.array(frequencies)
    qpoints = np.array(qpoints)
    segment_nqpoint = data["segment_nqpoint"]

    bands = []
    q_idx = 0
    for nq in segment_nqpoint:
        bands.append(qpoints[q_idx:q_idx + nq])
        q_idx += nq

    return distances, frequencies, bands, segment_nqpoint


def _write_band_dat(band_yaml_path, dat_filename):
    try:
        distances, frequencies, _, segment_nqpoint = _read_band_yaml_data(band_yaml_path)

        with open(dat_filename, "w") as f:
            f.write("# q-distance, frequency\n")
            for i, freqs_band in enumerate(frequencies.T):
                f.write(f"# mode {i + 1}\n")
                q_idx = 0
                for nq in segment_nqpoint:
                    for d, freq in zip(distances[q_idx:q_idx + nq], freqs_band[q_idx:q_idx + nq]):
                        f.write(f"{d:12.8f} {freq:15.8f}\n")
                    q_idx += nq
                    f.write("\n")
                f.write("\n")
    except Exception as exc:
        print(f"    Could not write .dat file: {exc}")


def _plot_iterative_band_structure(ph, band_conf_path, iter_num, input_poscar_stem):
    from phonopy.cui.settings import PhonopyConfParser
    from phonopy.phonon.band_structure import get_band_qpoints

    print(f"  Plotting band structure for iteration {iter_num}...")
    t0 = time.time()
    try:
        conf_parser = PhonopyConfParser(filename=str(band_conf_path))
        settings = conf_parser.settings

        npoints = settings.band_points if settings.band_points else 51
        qpoints = get_band_qpoints(settings.band_paths, npoints=npoints)
        labels = settings.band_labels
        path_connections = []
        for paths in settings.band_paths:
            path_connections += [True] * (len(paths) - 2)
            path_connections.append(False)

        ph.run_band_structure(qpoints, path_connections=path_connections, labels=labels)

        pdf_name = f"band-{input_poscar_stem}_{iter_num}.pdf"
        yaml_name = f"band-{input_poscar_stem}_{iter_num}.yaml"
        dat_name = f"band-{input_poscar_stem}_{iter_num}.dat"

        ph.plot_band_structure().savefig(pdf_name)
        ph.write_yaml_band_structure(filename=yaml_name)
        _write_band_dat(yaml_name, dat_name)
        print(f"    Band export for iteration {iter_num} finished in {time.time() - t0:.2f} s.")
    except Exception as exc:
        print(f"    Could not plot band structure for iteration {iter_num}: {exc}")
        traceback.print_exc()


def _append_text(paths, text: str):
    for p in paths:
        with p.open("a") as f:
            f.write(text)


def _describe_top_file(name: str) -> str:
    if name.startswith("FC3_") and name.endswith(".hdf5"):
        return "3rd-order force constants (fc3) snapshot in HDF5."
    if name.startswith("adp_iter_") and name.endswith(".cif"):
        return "Primitive ADP CIF with anisotropic displacement tensor."
    if name.startswith("adp_iter_") and name.endswith(".dat"):
        return "Primitive ADP table: U11/U22/U33/U23/U13/U12/Uiso/Biso."
    if name.startswith("distribution_proj_iter_") and name.endswith(".dat"):
        return "Element-averaged displacement distribution table (x/y/z and optional projected axis)."
    if name.startswith("distribution_proj_iter_") and name.endswith(".pdf"):
        return "Element-averaged displacement distribution plot."
    if name.startswith("band-") and name.endswith(".pdf"):
        return "Band structure plot for a specific saved iteration (or harmonic)."
    if name.startswith("band-") and name.endswith(".yaml"):
        return "Band structure raw YAML for a specific saved iteration (or harmonic)."
    if name.startswith("band-") and name.endswith(".dat"):
        return "Band structure text table for a specific saved iteration (or harmonic)."
    if name == "band.pdf":
        return "Latest promoted band plot."
    if name == "band.yaml":
        return "Latest promoted band YAML."
    if name == "band.dat":
        return "Latest promoted band text table."
    if name == "band.conf":
        return "Band-path configuration used for plotting."
    if name.startswith("cycle_"):
        return "Cycle folder with ensemble snapshots and reweighting logs."
    if name.startswith("FORCE_CONSTANTS_"):
        return "Force-constant tensor snapshot."
    if name == "FORCE_CONSTANTS_init":
        return "Initial harmonic force constants."
    if name == "FORCE_CONSTANTS_CURRENT":
        return "Current force constants after latest update."
    if name == "FORCE_CONSTANTS_PREVIOUS":
        return "Previous-cycle force constants."
    if name == "FORCE_CONSTANTS_QSCAILD_final":
        return "Final force constants at termination."
    if name.startswith("POSCAR_average_primitive_"):
        return "Weighted average primitive structure at saved iteration."
    if name.startswith("POSCAR_average_"):
        return "Backward-compatible weighted average POSCAR file."
    if name == "POSCAR_qscaild_input":
        return "Input POSCAR copied into run directory."
    if name == "CONTCAR-relax":
        return "Initial relaxed structure (if initial-isif != 0)."
    if name == "qscaild_fit_r2_history.log":
        return "qSCAILD fit R^2 convergence table by cycle."
    if name == "qscaild_fit_r2_history.pdf":
        return "qSCAILD fit R^2 convergence plot."
    if name == "qscaild_score_history.log":
        return "Compatibility alias of fit-R^2 history log."
    if name == "qscaild_score_history.pdf":
        return "Compatibility alias of fit-R^2 history plot."
    if name == "qscaild_free_energy_history.log":
        return "Compatibility alias of fit-R^2 history."
    if name == "qscaild_free_energy_history.pdf":
        return "Compatibility alias of fit-R^2 history plot."
    if name == "macer_qscaild.log":
        return "Full stdout/stderr log for this run."
    if name.startswith("weights_iter_") and name.endswith(".dat"):
        return "Normalized weights at a saved iteration (text)."
    if name.startswith("weights_iter_") and name.endswith(".npy"):
        return "Normalized weights at a saved iteration (binary)."
    if name == "weights_final.dat":
        return "Final normalized weights (text)."
    if name == "weights_final.npy":
        return "Final normalized weights (binary)."
    if name in ("out_fit", "out_fit.txt"):
        return "Fit summary per iteration."
    if name in ("out_convergence", "out_convergence.txt"):
        return "Convergence summary per iteration."
    if name in ("out_energy", "out_energy.txt"):
        return "Per-cycle mean energy and raw reweighting factors."
    if name in ("out_volume", "out_volume.txt"):
        return "Pressure/volume convergence history (pressure mode)."
    if name == "iteration":
        return "Last completed iteration index."
    if name.startswith("relax-") and name.endswith(".pdf"):
        return "Initial relaxation diagnostics plot."
    if name.startswith("relax-") and name.endswith(".traj"):
        return "Initial relaxation trajectory."
    if name == "README_QSCAILD_OUTPUTS.md":
        return "This guide file."
    return "Run artifact."


def _build_generated_checklist(output_dir: Path) -> str:
    lines = []
    lines.append("## Generated File Checklist (Current Snapshot)")
    lines.append("This section is refreshed during the run and at run end.")
    entries = sorted(output_dir.iterdir(), key=lambda p: p.name)
    if not entries:
        lines.append("- [ ] No files yet.")
        return "\n".join(lines) + "\n"

    def _category(name: str) -> str:
        if name.startswith("FORCE_CONSTANTS") or name.startswith("FC3_"):
            return "FC"
        if name.startswith("band") or name == "band.conf":
            return "Band"
        if name.startswith("adp_iter_") or name.startswith("POSCAR_average"):
            return "ADP/Structure"
        if name.startswith("distribution_proj_iter_"):
            return "Distribution"
        if name.startswith("weights_"):
            return "Weights"
        if name.startswith("out_") or name.startswith("qscaild_free_energy_history") or name.startswith("qscaild_score_history") or name.startswith("qscaild_fit_r2_history") or name == "iteration":
            return "Convergence/History"
        if name.startswith("cycle_"):
            return "Cycle"
        if name.startswith("relax-") or name in ("CONTCAR-relax", "POSCAR_qscaild_input"):
            return "Relax/Input"
        if name in ("macer_qscaild.log", "README_QSCAILD_OUTPUTS.md"):
            return "Logs/Docs"
        return "Other"

    category_order = [
        "Logs/Docs",
        "Relax/Input",
        "FC",
        "Band",
        "ADP/Structure",
        "Distribution",
        "Weights",
        "Convergence/History",
        "Cycle",
        "Other",
    ]
    grouped = {k: [] for k in category_order}
    for p in entries:
        grouped[_category(p.name)].append(p)

    for cat in category_order:
        if not grouped[cat]:
            continue
        lines.append("")
        lines.append(f"### {cat}")
        for p in grouped[cat]:
            mark = "[x]"
            suffix = "/" if p.is_dir() else ""
            desc = _describe_top_file(p.name)
            lines.append(f"- {mark} `{p.name}{suffix}`: {desc}")

    # Add one cycle-level detail block if cycle_001 exists.
    c1 = output_dir / "cycle_001"
    if c1.is_dir():
        lines.append("")
        lines.append("## `cycle_001` File Details (Example)")
        for q in sorted(c1.iterdir(), key=lambda x: x.name):
            qdesc = {
                "U.dat": "Displacements (Å), usually shape `(nconf, natom_supercell, 3)`.",
                "U.npy": "Binary version of `U.dat`.",
                "F.dat": "Forces (eV/Å), usually shape `(nconf, natom_supercell, 3)`.",
                "F.npy": "Binary version of `F.dat`.",
                "E.dat": "Energies (eV), usually shape `(nconf,)`.",
                "E.npy": "Binary version of `E.dat`.",
                "FC_origin.dat": "FC used to generate this cycle ensemble.",
                "FC_origin.npy": "Binary version of `FC_origin.dat`.",
                "logp_origin.dat": "Harmonic log-probability under `FC_origin`.",
                "logp_origin.npy": "Binary version of `logp_origin.dat`.",
            }.get(q.name, "Cycle artifact.")
            if q.name.startswith("logp_current_cycle_") and q.suffix == ".dat":
                qdesc = "Re-evaluated log-probability for this cycle under a later FC state."
            if q.name.startswith("logp_current_cycle_") and q.suffix == ".npy":
                qdesc = "Binary version of re-evaluated log-probability."
            lines.append(f"- [x] `cycle_001/{q.name}`: {qdesc}")
    return "\n".join(lines) + "\n"


def _write_output_guide(output_dir: Path, args, input_poscar_path: Path):
    guide_path = output_dir / "README_QSCAILD_OUTPUTS.md"
    projection_axis = getattr(args, "distribution_axis", None)
    proj_text = "not set"
    if projection_axis is not None:
        proj_text = " ".join(f"{x:g}" for x in projection_axis)
    txt = f"""# qSCAILD Output Guide

This file is auto-generated at run start so users can identify each output file quickly.

## Run Summary
- Input structure: `{input_poscar_path}`
- Output directory: `{output_dir}`
- Temperature (K): `{getattr(args, 'temperature', 'N/A')}`
- nconf: `{getattr(args, 'nconf', 'N/A')}`
- nsteps: `{getattr(args, 'nsteps', 'N/A')}`
- save-every: `{getattr(args, 'save_every', 'N/A')}`
- save-npy: `{getattr(args, 'save_npy', False)}`
- tolerance: `{getattr(args, 'tolerance', 'N/A')}`
- pressure mode: `{getattr(args, 'use_pressure', 'False')}`
- pressure tolerance pdiff (kbar): `{getattr(args, 'pdiff', 'N/A')}`
- distribution-axis (fractional): `{proj_text}`

## Top-Level Files
- `macer_qscaild.log`: full run log.
- `FORCE_CONSTANTS_init`: initial harmonic FC.
- `FORCE_CONSTANTS_CURRENT`: FC after the latest update.
- `FORCE_CONSTANTS_PREVIOUS`: FC from the previous cycle.
- `FORCE_CONSTANTS_FIT_<iter>`: fitted FC at saved cycles.
- `FORCE_CONSTANTS_QSCAILD_final`: final FC at termination.
- `FC3_*.hdf5` (if `--include-third-order`): 3rd-order FC snapshots.
- `iteration`: last completed iteration index.

- `out_fit.txt` (`out_fit`): fit summary table per iteration.
- `out_convergence.txt` (`out_convergence`): convergence metrics (`fc_diff_max`, ESS, entropy, pressure error if enabled).
- `out_energy.txt` (`out_energy`): per-cycle mean energy and raw reweighting factors.
- `out_volume.txt` (`out_volume`): pressure/volume history (pressure mode only).

- `qscaild_fit_r2_history.log`: fit-R^2 history table.
- `qscaild_fit_r2_history.pdf`: fit-R^2 convergence plot.
- `qscaild_score_history.log/.pdf`: compatibility aliases.
- `qscaild_free_energy_history.log/.pdf`: compatibility aliases.

- `weights_iter_<iter>.dat` (and optionally `.npy`): normalized weights at saved cycles.
- `weights_final.dat` (and optionally `.npy`): final normalized weights.

- `band-<stem>_harmonic.pdf/.yaml/.dat`: harmonic band output (initial FC).
- `band-<stem>_<iter>.pdf/.yaml/.dat`: band output at saved cycles.
- `band.pdf`, `band.yaml`, `band.dat`: latest promoted band files.
- `band.conf`: band-path configuration used for band exports.

- `POSCAR_qscaild_input`: input POSCAR copied into the run directory.
- `CONTCAR-relax`: initial relaxed structure if `initial-isif != 0`.
- `POSCAR_average_primitive_i<iter>`: weighted average primitive structure at saved cycles.
- `POSCAR_average_i<iter>`: compatibility alias (currently primitive-based output).

- `adp_iter_<iter>.dat`: primitive ADP table with columns
  `U11 U22 U33 U23 U13 U12 Uiso Biso(=8*pi^2*Uiso)`.
- `adp_iter_<iter>.cif`: primitive CIF including `_atom_site_U_iso_or_equiv` and anisotropic U tensor block.

- `distribution_proj_iter_<iter>_element_<El>_avg.dat/.pdf`:
  element-averaged displacement distributions on primitive atoms.
  - `dat` columns: `disp_A pdf_x pdf_y pdf_z` (+ `pdf_proj` if `--distribution-axis` is set).
  - `pdf` curves: `x, y, z` (+ `proj` if enabled), centered around displacement 0.

## `cycle_<NNN>/` Folder Structure
Each cycle folder stores ensemble data used for reweighting and fitting.

- `U.dat` (and optionally `U.npy`):
  displacements, shape typically `(nconf, natom_supercell, 3)`, unit Å.
- `F.dat` (and optionally `F.npy`):
  forces, shape `(nconf, natom_supercell, 3)`, unit eV/Å.
- `E.dat` (and optionally `E.npy`):
  energies, shape `(nconf,)`, unit eV.
- `FC_origin.dat` (and optionally `FC_origin.npy`):
  FC used when this cycle ensemble was generated.
- `logp_origin.dat` (and optionally `logp_origin.npy`):
  harmonic log-probability (up to normalization) under `FC_origin`.
- `logp_current_cycle_<iter>.dat` (and optionally `.npy`):
  the same ensemble re-evaluated with the FC of iteration `<iter>`.
  Reweighting uses `exp(logp_current - logp_origin)`.

## Notes
- `.txt` and no-extension compatibility files are both written.
- When pressure mode is off, volume/pressure files may be empty or minimal.
"""
    txt += "\n" + _build_generated_checklist(output_dir)
    guide_path.write_text(txt)


def _parse_mass_map(args) -> dict[str, float]:
    mass_map = {}
    if not getattr(args, "mass", None):
        return mass_map
    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 = str(args.mass[i])
        try:
            mass_map[sym] = float(args.mass[i + 1])
        except ValueError as exc:
            raise ValueError(f"Error: Invalid mass value for {sym}: {args.mass[i + 1]}") from exc
    return mass_map


def _apply_mass_override_to_phonopy_atoms(unitcell, mass_map: dict[str, float]):
    if not mass_map:
        return unitcell
    symbols = [str(s) for s in unitcell.symbols]
    if getattr(unitcell, "masses", None) is None:
        masses = [None] * len(symbols)
    else:
        masses = list(unitcell.masses)
    for i, sym in enumerate(symbols):
        if sym in mass_map:
            masses[i] = mass_map[sym]
    unitcell.masses = masses
    return unitcell


def _log_mass_override_application(unitcell, mass_map: dict[str, float]):
    if not mass_map:
        return
    symbols = [str(s) for s in unitcell.symbols]
    present = set(symbols)
    print("Applied mass override to unit cell:")
    for sym, m in mass_map.items():
        n_updated = sum(1 for s in symbols if s == sym)
        print(f"  {sym:3s} -> {m:8.4f} amu (atoms updated: {n_updated})")
    missing = [sym for sym in mass_map if sym not in present]
    if missing:
        print(f"WARNING: mass override symbol(s) not found in structure: {', '.join(missing)}")


def _build_output_dir(args, input_poscar_path: Path, create: bool = True) -> Path:
    cmd = str(
        getattr(args, "phonopy_command", None)
        or getattr(args, "command", "")
        or ""
    ).lower()

    if args.output_dir:
        base_output_dir = Path(args.output_dir).resolve()
    else:
        prefix = "sscha" if cmd in ("sscha", "run-sscha") else "qscaild"
        base_output_dir = input_poscar_path.parent / f"{prefix}-{input_poscar_path.stem}-ff={args.ff}"

    output_dir = base_output_dir
    i = 1
    while output_dir.exists():
        output_dir = Path(f"{base_output_dir}-NEW{i:02d}")
        i += 1
    if create:
        output_dir.mkdir(parents=True, exist_ok=True)
    return output_dir


def _prepare_input_poscar(args, output_dir: Path, convert_cif: bool = True) -> Path:
    from ase.io import read as ase_read, write as ase_write

    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/-c/-f/-m.")

    input_path = Path(input_list[0]).resolve()
    if not input_path.is_file():
        raise FileNotFoundError(f"Input file not found at '{input_path}'.")

    # Convert CIF-like inputs to POSCAR in output dir for a stable run path.
    if convert_cif and input_path.suffix.lower() == ".cif":
        converted_poscar = output_dir / "POSCAR_input"
        atoms_in = ase_read(str(input_path))
        ase_write(str(converted_poscar), atoms_in, format="vasp")
        return converted_poscar

    return input_path


def _compute_harmonic_logpdf(U: np.ndarray, fc: np.ndarray, beta: float) -> np.ndarray:
    quad = np.einsum("ijab,nia,njb->n", fc, U, U)
    return -0.5 * beta * quad


def _generate_cycle_ensemble(ph, fc_current, temperature: float, nconf: int):
    from phonopy.phonon.random_displacements import RandomDisplacements

    rd = RandomDisplacements(ph.supercell, ph.primitive, fc_current)
    rd.run(temperature, number_of_snapshots=nconf)
    return np.asarray(rd.u)


def _fbe(freq_thz: float, temperature: float) -> float:
    if temperature == 0.0:
        return 0.0
    x = codata.h * abs(freq_thz) * 1e12 / (codata.k * temperature)
    if x < 2.0:
        return 1.0 / np.expm1(x)
    return -np.exp(-x) / np.expm1(-x)


def _create_mesh(nq: int) -> np.ndarray:
    pts = []
    for i, j, k in itertools.product(range(nq), range(nq), range(nq)):
        pts.append([float(i) / nq, float(j) / nq, float(k) / nq])
    q = np.asarray(pts, dtype=float)
    center = np.mean(q, axis=0)
    q -= center
    return q


def _build_displacement_covariance(ph, fc_current, temperature: float, grid: int, imaginary_freq: float):
    """
    Build displacement covariance in Angstrom^2 using a q-mesh integration.
    This follows the qSCAILD thermal-displacement-matrix idea in a serial form.
    """
    ph.force_constants = fc_current
    supercell = ph.supercell
    primitive = ph.primitive
    nat_prim = len(primitive)
    nat_super = len(supercell)
    nmodes = 3 * nat_prim
    dim = np.array(ph.supercell_matrix, dtype=float)
    if dim.shape != (3, 3):
        raise ValueError("Unexpected supercell_matrix shape.")
    if not np.allclose(dim, np.diag(np.diag(dim))):
        raise ValueError("Only diagonal supercell matrices are supported for covariance mesh integration.")
    na, nb, nc = [int(round(v)) for v in np.diag(dim)]
    if na * nb * nc != (nat_super // nat_prim):
        raise ValueError("Inconsistent primitive/supercell relation for covariance assembly.")

    # positions in reduced coordinates (shape: 3 x nat_prim)
    positions = np.asarray(primitive.scaled_positions).T
    masses = np.asarray(supercell.masses, dtype=float) * codata.physical_constants["atomic mass constant"][0]
    nq = max(1, int(grid))
    qpoints = _create_mesh(nq) if nq > 1 else np.array([[0.0, 0.0, 0.0]])
    nqpoints = len(qpoints)

    ncells = na * nb * nc
    size = ncells * nmodes
    matrix = np.zeros((size, size), dtype=np.complex128)

    factors = np.empty(ncells, dtype=np.complex128)
    cell_triplets = list(itertools.product(range(nc), range(nb), range(na)))

    for q in qpoints:
        freqs, eigvecs = ph.get_frequencies_with_eigenvectors(q)
        freqs = np.array(freqs, dtype=float)
        for i in range(len(freqs)):
            if freqs[i] < -1.0e-4:
                if imaginary_freq == -1:
                    freqs[i] = abs(freqs[i])
                else:
                    freqs[i] = float(imaginary_freq)

        n_B = np.array([_fbe(fq, temperature) for fq in freqs], dtype=float)
        for pos, (k, j, i) in enumerate(cell_triplets):
            factors[pos] = np.exp(2j * np.pi * np.dot(q, [i, j, k]))

        for im in range(nmodes):
            if freqs[im] <= 1.0e-4:
                continue
            lpsi = np.empty(size, dtype=np.complex128)
            for ia in range(nat_prim):
                for pos, (k, j, i) in enumerate(cell_triplets):
                    for ic in range(3):
                        jm = np.ravel_multi_index((ia, ic), (nat_prim, 3))
                        li = np.ravel_multi_index((ia, k, j, i, ic), (nat_prim, nc, nb, na, 3))
                        lpsi[li] = eigvecs[jm, im] * factors[pos] * np.exp(
                            2j * np.pi * np.dot(q, positions[:, ia])
                        )
            matrix += np.outer(lpsi, lpsi.conj()) * (1.0 + 2.0 * n_B[im]) / freqs[im]

    matrix *= codata.hbar / 2.0 / nqpoints / 2.0 / np.pi / 1.0e12  # m^2
    for i, j in itertools.product(range(size), range(size)):
        matrix[i, j] /= np.sqrt(masses[i // 3] * masses[j // 3])
    matrix = matrix.real * 1.0e20  # Angstrom^2
    # Numerical stabilization for sampling.
    matrix = 0.5 * (matrix + matrix.T)
    return matrix


def _sample_displacements_from_covariance(cov: np.ndarray, nconf: int, n_atoms: int) -> np.ndarray:
    cov = np.asarray(cov, dtype=float)
    jitter = 1.0e-10
    cov = cov + np.eye(cov.shape[0]) * jitter
    disp = np.random.multivariate_normal(mean=np.zeros(cov.shape[0]), cov=cov, size=nconf, check_valid="ignore")
    disp = np.asarray(disp, dtype=float).reshape(nconf, n_atoms, 3)
    return disp


def _generate_cycle_ensemble_staged(
    ph,
    fc_current,
    temperature: float,
    nconf: int,
    use_smalldisp: bool,
    imaginary_freq: float,
    grid: int,
):
    # Phase-2 implementation:
    # - use_smalldisp: fixed tiny Gaussian displacement
    # - otherwise: covariance-based sampling with grid and imaginary handling
    if use_smalldisp:
        n_atoms = len(ph.supercell)
        return np.random.normal(loc=0.0, scale=1e-3, size=(nconf, n_atoms, 3))
    cov = _build_displacement_covariance(
        ph,
        fc_current,
        temperature=temperature,
        grid=grid,
        imaginary_freq=imaginary_freq,
    )
    return _sample_displacements_from_covariance(cov, nconf=nconf, n_atoms=len(ph.supercell))


def _evaluate_ensemble(
    U: np.ndarray,
    base_supercell,
    calculator,
    need_stress: bool = False,
    progress_desc: str | None = None,
):
    from ase import Atoms as AseAtoms

    base_positions = np.asarray(base_supercell.positions)
    forces_list = []
    energy_list = []
    stress_list = []
    iterator = U
    if progress_desc:
        iterator = tqdm(U, total=len(U), desc=progress_desc, leave=False)
    for disp in iterator:
        snapshot = AseAtoms(
            symbols=base_supercell.symbols,
            positions=base_positions + np.asarray(disp),
            cell=base_supercell.cell,
            pbc=True,
        )
        snapshot.calc = calculator
        forces_list.append(snapshot.get_forces())
        energy_list.append(snapshot.get_potential_energy())
        if need_stress:
            stress_list.append(snapshot.get_stress(voigt=True))
    stress_arr = np.asarray(stress_list) if need_stress else None
    return np.asarray(forces_list), np.asarray(energy_list), stress_arr


def _save_cycle_data(
    cycle_dir: Path,
    U: np.ndarray,
    F: np.ndarray,
    E: np.ndarray,
    fc_origin: np.ndarray,
    logp_origin: np.ndarray,
    save_npy: bool = False,
):
    _save_array_with_dat(cycle_dir / "U", U, save_npy=save_npy)
    _save_array_with_dat(cycle_dir / "F", F, save_npy=save_npy)
    _save_array_with_dat(cycle_dir / "E", E, save_npy=save_npy)
    _save_array_with_dat(cycle_dir / "FC_origin", fc_origin, save_npy=save_npy)
    _save_array_with_dat(cycle_dir / "logp_origin", logp_origin, save_npy=save_npy)


def _save_array_with_dat(stem: Path, arr: np.ndarray, save_npy: bool = False):
    """
    Save an array as both `.npy` and `.dat`.
    `.dat` format is a flattened text representation with shape metadata.
    """
    a = np.asarray(arr)
    if save_npy:
        np.save(f"{stem}.npy", a)

    dat_path = f"{stem}.dat"
    shape_str = " ".join(str(x) for x in a.shape)
    if a.ndim == 1:
        idx = np.arange(1, len(a) + 1, dtype=int)
        table = np.column_stack((idx, a))
        np.savetxt(
            dat_path,
            table,
            fmt=["%d", "%.16e"],
            header=f"shape {shape_str}\nindex value",
        )
    else:
        flat2d = a.reshape(a.shape[0], -1)
        np.savetxt(
            dat_path,
            flat2d,
            fmt="%.16e",
            header=f"shape {shape_str}",
        )


def _load_saved_array(stem: Path) -> np.ndarray:
    npy = Path(f"{stem}.npy")
    if npy.exists():
        return np.load(npy)
    dat = Path(f"{stem}.dat")
    if not dat.exists():
        raise FileNotFoundError(f"Missing both {npy.name} and {dat.name}")
    with open(dat, "r") as f:
        first = f.readline().strip()
    if not first.startswith("# shape "):
        raise ValueError(f"Missing shape header in {dat}")
    shape = tuple(int(x) for x in first.replace("#", "").split()[1:])
    raw = np.loadtxt(dat, comments="#")
    if len(shape) == 1:
        if raw.ndim == 1:
            return np.array([float(raw[-1])], dtype=float)
        return np.asarray(raw[:, -1], dtype=float)
    raw2d = np.atleast_2d(raw)
    return raw2d.reshape(shape)


def _save_weights_qscaild(weight_file_stem: Path, w: np.ndarray, save_npy: bool = False):
    w = np.asarray(w, dtype=float).reshape(-1)
    if save_npy:
        np.save(f"{weight_file_stem}.npy", w)
    idx = np.arange(1, len(w) + 1, dtype=int)
    table = np.column_stack((idx, w))
    np.savetxt(
        f"{weight_file_stem}.dat",
        table,
        fmt=["%d", "%.16e"],
        header="snapshot_index weight",
    )


def _load_history_for_fit(output_dir: Path, iteration: int, memory: float):
    iteration_min = max(1, int(np.floor(iteration * (1.0 - memory))))
    cycle_ids = list(range(iteration_min, iteration + 1))
    cycle_dirs = [output_dir / f"cycle_{i:03d}" for i in cycle_ids]
    return cycle_ids, cycle_dirs


def _assemble_fit_data(
    cycle_ids,
    cycle_dirs,
    fc_current,
    beta: float,
    iteration: int,
    use_pressure_mode: str,
    save_npy: bool = False,
):
    U_all = []
    F_all = []
    W_all = []
    E_all = []
    Stress_all = []
    energy_lines = [
        "\n"
        f"[Iteration {iteration:04d}] "
        f"history_cycles={cycle_ids} total_snapshots={0}\n"
        "cycle_id    mean_energy_eV        mean_raw_weight\n"
    ]

    for cid, cdir in zip(cycle_ids, cycle_dirs):
        U_hist = _load_saved_array(cdir / "U")
        F_hist = _load_saved_array(cdir / "F")
        E_hist = _load_saved_array(cdir / "E")
        logp_origin_hist = _load_saved_array(cdir / "logp_origin")
        logp_current_hist = _compute_harmonic_logpdf(U_hist, fc_current, beta)
        _save_array_with_dat(
            cdir / f"logp_current_cycle_{iteration:03d}",
            logp_current_hist,
            save_npy=save_npy,
        )

        w_hist = np.exp(logp_current_hist - logp_origin_hist)
        w_hist = np.clip(w_hist, 1e-300, None)
        U_all.append(U_hist)
        F_all.append(F_hist)
        W_all.append(w_hist)
        E_all.append(E_hist)
        if use_pressure_mode != "False":
            stress_dat = cdir / "Stress.dat"
            stress_npy = cdir / "Stress.npy"
            if stress_dat.exists() or stress_npy.exists():
                Stress_all.append(_load_saved_array(cdir / "Stress"))
        energy_lines.append(
            f"{cid:8d}    {np.mean(E_hist):16.8f}    {np.mean(w_hist):16.6e}\n"
        )

    U_fit = np.concatenate(U_all, axis=0)
    F_fit = np.concatenate(F_all, axis=0)
    E_fit = np.concatenate(E_all, axis=0)
    w_fit = np.concatenate(W_all, axis=0)
    w_fit /= np.sum(w_fit)
    energy_lines[1] = (
        f"[Iteration {iteration:04d}] "
        f"history_cycles={cycle_ids} total_snapshots={len(U_fit)}\n"
    )
    return U_fit, F_fit, E_fit, w_fit, Stress_all, energy_lines


def _weight_stats(w: np.ndarray):
    w = np.asarray(w, dtype=float).reshape(-1)
    w = np.clip(w, 1e-300, None)
    w /= np.sum(w)
    ess = 1.0 / np.sum(w**2)
    entropy = -np.sum(w * np.log(w))
    return ess, entropy


def _compute_kinetic_pressure_kbar(
    U_fit: np.ndarray,
    F_fit: np.ndarray,
    w_fit: np.ndarray,
    unitcell,
) -> np.ndarray:
    """
    Kinetic/virial-like pressure contribution in kbar, axis-resolved.
    Mirrors qSCAILD's weighted <u_alpha * f_alpha> / V correction.
    """
    U = np.asarray(U_fit, dtype=float)
    F = np.asarray(F_fit, dtype=float)
    w = np.asarray(w_fit, dtype=float).reshape(-1)
    if U.shape != F.shape or U.ndim != 3 or U.shape[0] != w.shape[0]:
        raise ValueError("Invalid shapes for kinetic pressure evaluation.")
    wsum = np.sum(w)
    if wsum <= 0:
        return np.zeros(3, dtype=float)
    w = w / wsum

    # axis-wise virial term: sum_n w_n sum_i u_{ni,alpha} f_{ni,alpha}
    virial_axis_eV = np.sum(U * F * w[:, None, None], axis=(0, 1))

    vol_a3 = abs(float(np.linalg.det(np.asarray(unitcell.cell, dtype=float))))
    if vol_a3 <= 1e-18:
        return np.zeros(3, dtype=float)

    # eV/Ang^3 -> kbar
    return -virial_axis_eV / vol_a3 * EV_A3_TO_GPA * 10.0


def _update_cell_from_pressure(
    unitcell,
    pressure_mode: str,
    pressure_diag_kbar,
    pdiff_kbar: float,
    stress_weighted_voigt,
    kinetic_pressure_kbar: np.ndarray | None = None,
):
    # stress in eV/A^3 -> pressure in kbar (sign: P = -sigma)
    pressure_kbar = -np.asarray(stress_weighted_voigt[:3]) * EV_A3_TO_GPA * 10.0
    if kinetic_pressure_kbar is not None:
        pressure_kbar = pressure_kbar + np.asarray(kinetic_pressure_kbar, dtype=float).reshape(3)
    target = np.asarray(pressure_diag_kbar, dtype=float)
    factor = np.array([1.0, 1.0, 1.0], dtype=float)
    for i in range(3):
        if pressure_kbar[i] > target[i] + pdiff_kbar / 2.0:
            factor[i] = 1.0 + 0.001 * min((pressure_kbar[i] - target[i]) / 10.0, 4.0)
        elif pressure_kbar[i] < target[i] - pdiff_kbar / 2.0:
            factor[i] = 1.0 - 0.001 * min((target[i] - pressure_kbar[i]) / 10.0, 4.0)

    mode = str(pressure_mode)
    if mode == "cubic":
        m = float(np.mean(factor))
        factor[:] = m
    elif mode == "tetragonal":
        m = 0.5 * (factor[0] + factor[1])
        factor[0] = m
        factor[1] = m
    elif mode == "orthorhombic":
        pass
    else:
        return unitcell, pressure_kbar, np.max(np.abs(pressure_kbar - target))

    new_unitcell = unitcell.copy()
    new_cell = np.array(new_unitcell.cell, dtype=float)
    for i in range(3):
        new_cell[i, :] *= factor[i]
    new_unitcell.cell = new_cell
    pressure_err = float(np.max(np.abs(pressure_kbar - target)))
    return new_unitcell, pressure_kbar, pressure_err


def _primitive_groups_from_phonopy(ph):
    p2p = ph.primitive.p2p_map
    s2p = ph.primitive.s2p_map
    nat_prim = len(ph.primitive)
    groups = [[] for _ in range(nat_prim)]
    for s_idx, p_rep in enumerate(s2p):
        p_idx = p2p[int(p_rep)]
        groups[p_idx].append(s_idx)
    return [np.array(g, dtype=int) for g in groups]


def _projection_unit_vector_from_fractional(cell: np.ndarray, frac_axis):
    if frac_axis is None:
        return None
    vfrac = np.asarray(frac_axis, dtype=float).reshape(3)
    if np.linalg.norm(vfrac) < 1e-12:
        raise ValueError("Error: --distribution-axis cannot be a zero vector.")
    vcart = vfrac @ np.asarray(cell, dtype=float)
    nv = np.linalg.norm(vcart)
    if nv < 1e-12:
        raise ValueError("Error: --distribution-axis becomes near-zero in Cartesian basis.")
    return vcart / nv


def _hist_pdf(values: np.ndarray, weights: np.ndarray, nbins: int = 100):
    max_abs = float(np.max(np.abs(values)))
    xlim = max(1e-4, 1.05 * max_abs)
    bins = np.linspace(-xlim, xlim, nbins + 1)
    centers = 0.5 * (bins[:-1] + bins[1:])
    hist, _ = np.histogram(values, bins=bins, weights=weights, density=True)
    return centers, hist


def _save_primitive_distributions(
    output_dir: Path,
    iteration: int,
    prim_symbols,
    primitive_cell: np.ndarray,
    du_prim_samples: list[np.ndarray],
    w_prim_samples: list[np.ndarray],
    projection_axis_frac=None,
):
    import matplotlib.pyplot as plt

    proj_unit = _projection_unit_vector_from_fractional(primitive_cell, projection_axis_frac)
    by_element_dxyz = defaultdict(list)
    by_element_w = defaultdict(list)
    for sym, dxyz, wxyz in zip(prim_symbols, du_prim_samples, w_prim_samples):
        if dxyz.size == 0:
            continue
        by_element_dxyz[sym].append(dxyz)
        by_element_w[sym].append(wxyz)

    for sym in sorted(by_element_dxyz.keys()):
        dxyz = np.concatenate(by_element_dxyz[sym], axis=0)
        wxyz = np.concatenate(by_element_w[sym], axis=0)
        wxyz /= np.sum(wxyz)

        vals_x = dxyz[:, 0]
        vals_y = dxyz[:, 1]
        vals_z = dxyz[:, 2]
        vals_proj = dxyz @ proj_unit if proj_unit is not None else None

        max_abs = max(
            float(np.max(np.abs(vals_x))),
            float(np.max(np.abs(vals_y))),
            float(np.max(np.abs(vals_z))),
            float(np.max(np.abs(vals_proj))) if vals_proj is not None else 0.0,
        )
        xlim = max(1e-4, 1.05 * max_abs)
        bins = np.linspace(-xlim, xlim, 101)  # centered at 0
        centers = 0.5 * (bins[:-1] + bins[1:])
        hx, _ = np.histogram(vals_x, bins=bins, weights=wxyz, density=True)
        hy, _ = np.histogram(vals_y, bins=bins, weights=wxyz, density=True)
        hz, _ = np.histogram(vals_z, bins=bins, weights=wxyz, density=True)

        avg_dat = output_dir / f"distribution_proj_iter_{iteration:04d}_element_{sym}_avg.dat"
        avg_pdf = output_dir / f"distribution_proj_iter_{iteration:04d}_element_{sym}_avg.pdf"
        if vals_proj is not None:
            hp, _ = np.histogram(vals_proj, bins=bins, weights=wxyz, density=True)
            data = np.column_stack((centers, hx, hy, hz, hp))
            header = "disp_A  pdf_x  pdf_y  pdf_z  pdf_proj"
        else:
            data = np.column_stack((centers, hx, hy, hz))
            header = "disp_A  pdf_x  pdf_y  pdf_z"
        np.savetxt(avg_dat, data, fmt="%.8e", header=header)

        fig, ax = plt.subplots(1, 1, figsize=(7.2, 4.2))
        ax.plot(centers, hx, lw=1.8, label="x")
        ax.plot(centers, hy, lw=1.8, label="y")
        ax.plot(centers, hz, lw=1.8, label="z")
        if vals_proj is not None:
            ax.plot(centers, hp, lw=2.0, label="proj")
        ax.axvline(0.0, lw=1.0, color="#777777", alpha=0.8)
        ax.set_xlabel("Displacement (A)")
        ax.set_ylabel("Probability density")
        ax.set_title(f"Element-average distribution: {sym} (iter {iteration})")
        ax.grid(True, alpha=0.3)
        ax.legend(frameon=False)
        fig.tight_layout()
        fig.savefig(avg_pdf)
        plt.close(fig)


def _save_weighted_average_and_adp(output_dir: Path, iteration: int, ph, U_fit: np.ndarray, w_fit: np.ndarray, projection_axis_frac=None):
    """
    Save weighted-average structure and per-atom ADP tensor from weighted ensemble.
    """
    from phonopy.interface.vasp import write_vasp

    w = np.asarray(w_fit, dtype=float).reshape(-1)
    w /= np.sum(w)
    U_avg = np.sum(w[:, None, None] * U_fit, axis=0)  # (natom, 3)

    # Primitive mapping: evaluate ADP/distributions on primitive atoms, not supercell atoms.
    prim = ph.primitive
    prim_symbols = list(prim.symbols)
    groups = _primitive_groups_from_phonopy(ph)

    # Weighted primitive mean displacement from supercell representatives.
    U_avg_prim = np.zeros((len(prim), 3), dtype=float)
    for p_idx, g in enumerate(groups):
        U_avg_prim[p_idx] = np.mean(U_avg[g], axis=0)

    primitive_average = prim.copy()
    primitive_average.positions = np.asarray(prim.positions) + U_avg_prim
    poscar_avg = output_dir / f"POSCAR_average_primitive_i{iteration:04d}"
    write_vasp(str(poscar_avg), primitive_average, direct=True)

    # Backward-compatible filename kept, now primitive-based.
    legacy_avg = output_dir / f"POSCAR_average_i{iteration:04d}"
    write_vasp(str(legacy_avg), primitive_average, direct=True)

    # ADP tensor from weighted covariance around weighted mean (primitive basis).
    du = U_fit - U_avg[None, :, :]
    adp_path = output_dir / f"adp_iter_{iteration:04d}.dat"
    adp_cif_path = output_dir / f"adp_iter_{iteration:04d}.cif"
    adp_tensors = []
    du_prim_samples = []
    w_prim_samples = []
    with open(adp_path, "w") as f:
        f.write("# atom_index symbol U11 U22 U33 U23 U13 U12 Uiso Biso(=8*pi^2*Uiso)\n")
        for p_idx, g in enumerate(groups):
            x = du[:, g, :].reshape(-1, 3)  # all translationally-equivalent primitive sites
            wp = np.repeat(w / len(g), len(g))
            wp /= np.sum(wp)
            cov = np.einsum("n,ni,nj->ij", wp, x, x)  # (3,3)
            u11, u22, u33 = cov[0, 0], cov[1, 1], cov[2, 2]
            u23, u13, u12 = cov[1, 2], cov[0, 2], cov[0, 1]
            uiso = (u11 + u22 + u33) / 3.0
            biso = 8.0 * np.pi * np.pi * uiso
            f.write(
                f"{p_idx+1:6d} {prim_symbols[p_idx]:>3s} "
                f"{u11:15.8e} {u22:15.8e} {u33:15.8e} "
                f"{u23:15.8e} {u13:15.8e} {u12:15.8e} {uiso:15.8e} {biso:15.8e}\n"
            )
            adp_tensors.append(cov)
            du_prim_samples.append(x)
            w_prim_samples.append(wp)

    _write_adp_cif(adp_cif_path, primitive_average, adp_tensors)
    _save_primitive_distributions(
        output_dir=output_dir,
        iteration=iteration,
        prim_symbols=prim_symbols,
        primitive_cell=np.asarray(prim.cell, dtype=float),
        du_prim_samples=du_prim_samples,
        w_prim_samples=w_prim_samples,
        projection_axis_frac=projection_axis_frac,
    )

    return poscar_avg, adp_path, adp_cif_path


def _write_adp_cif(cif_path: Path, atoms, adp_tensors):
    cell = np.array(atoms.cell, dtype=float)
    a = float(np.linalg.norm(cell[0]))
    b = float(np.linalg.norm(cell[1]))
    c = float(np.linalg.norm(cell[2]))

    def _angle(v1, v2):
        x = float(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
        x = min(1.0, max(-1.0, x))
        return float(np.degrees(np.arccos(x)))

    alpha = _angle(cell[1], cell[2])
    beta = _angle(cell[0], cell[2])
    gamma = _angle(cell[0], cell[1])

    symbols = list(atoms.symbols)
    scaled_positions = np.asarray(atoms.scaled_positions, dtype=float)
    with open(cif_path, "w") as f:
        f.write("data_qscaild_adp\n")
        f.write("_symmetry_space_group_name_H-M 'P 1'\n")
        f.write("_symmetry_Int_Tables_number 1\n")
        f.write(f"_cell_length_a {a:.10f}\n")
        f.write(f"_cell_length_b {b:.10f}\n")
        f.write(f"_cell_length_c {c:.10f}\n")
        f.write(f"_cell_angle_alpha {alpha:.10f}\n")
        f.write(f"_cell_angle_beta {beta:.10f}\n")
        f.write(f"_cell_angle_gamma {gamma:.10f}\n")
        f.write("\n")
        f.write("loop_\n")
        f.write("_symmetry_equiv_pos_as_xyz\n")
        f.write("'x, y, z'\n")
        f.write("\n")
        f.write("loop_\n")
        f.write("_atom_site_label\n")
        f.write("_atom_site_type_symbol\n")
        f.write("_atom_site_fract_x\n")
        f.write("_atom_site_fract_y\n")
        f.write("_atom_site_fract_z\n")
        f.write("_atom_site_U_iso_or_equiv\n")
        for i, (sym, spos) in enumerate(zip(symbols, scaled_positions), start=1):
            uiso = (adp_tensors[i - 1][0, 0] + adp_tensors[i - 1][1, 1] + adp_tensors[i - 1][2, 2]) / 3.0
            label = f"{sym}{i}"
            f.write(
                f"{label:<6s} {sym:<3s} "
                f"{spos[0]:12.8f} {spos[1]:12.8f} {spos[2]:12.8f} {uiso:12.8f}\n"
            )
        f.write("\n")
        f.write("loop_\n")
        f.write("_atom_site_aniso_label\n")
        f.write("_atom_site_aniso_U_11\n")
        f.write("_atom_site_aniso_U_22\n")
        f.write("_atom_site_aniso_U_33\n")
        f.write("_atom_site_aniso_U_23\n")
        f.write("_atom_site_aniso_U_13\n")
        f.write("_atom_site_aniso_U_12\n")
        for i, (sym, u) in enumerate(zip(symbols, adp_tensors), start=1):
            label = f"{sym}{i}"
            f.write(
                f"{label:<6s} "
                f"{u[0,0]:12.8f} {u[1,1]:12.8f} {u[2,2]:12.8f} "
                f"{u[1,2]:12.8f} {u[0,2]:12.8f} {u[0,1]:12.8f}\n"
            )


def _plot_free_energy_convergence(log_path: Path, output_dir: Path):
    if not log_path.exists():
        return
    x = []
    y = []
    with open(log_path, "r") as f:
        for line in f:
            if line.startswith("#") or not line.strip():
                continue
            parts = line.split()
            if len(parts) < 3:
                continue
            x.append(int(parts[0]))
            y.append(float(parts[2]))
    if not x:
        return
    plt.figure(figsize=(7, 4.5))
    plt.plot(x, y, marker="o", lw=1.8)
    plt.xlabel("Cycle")
    plt.ylabel("Free energy (meV/atom)")
    plt.title("qSCAILD Free-Energy Convergence")
    plt.grid(True, alpha=0.35)
    plt.tight_layout()
    plt.savefig(output_dir / "qscaild_free_energy_history.pdf")
    plt.close()


def _plot_sscha_history(output_dir: Path, input_stem: str):
    """
    Plot a compact SSCHA/qSCAILD-style convergence dashboard from out_fit.txt.
    """
    fit_path = output_dir / "out_fit.txt"
    if not fit_path.exists():
        return

    rows = []
    with open(fit_path, "r") as f:
        for ln in f:
            s = ln.strip()
            if (not s) or s.startswith("#"):
                continue
            parts = s.split()
            # iteration n_cycles n_snapshots fc_diff_max ess entropy free_energy_eV free_energy_meV_per_atom
            if len(parts) < 8:
                continue
            try:
                it = int(parts[0])
                fc_diff = float(parts[3])
                ess = float(parts[4])
                fe_mev = float(parts[7])
                rows.append((it, fc_diff, ess, fe_mev))
            except Exception:
                continue
    if not rows:
        return

    arr = np.asarray(rows, dtype=float)
    x = arr[:, 0]
    fc_diff = np.maximum(arr[:, 1], 1e-16)
    ess = arr[:, 2]
    fe_mev = arr[:, 3]

    fig, axes = plt.subplots(3, 1, figsize=(7.0, 8.0), sharex=True)

    axes[0].plot(x, fc_diff, marker="o", lw=1.5, color="tab:blue")
    axes[0].set_yscale("log")
    axes[0].set_ylabel("FC diff max")
    axes[0].grid(True, alpha=0.3)

    axes[1].plot(x, ess, marker="s", lw=1.5, color="tab:orange")
    axes[1].set_ylabel("ESS")
    axes[1].grid(True, alpha=0.3)

    axes[2].plot(x, fe_mev, marker="^", lw=1.5, color="tab:green")
    axes[2].set_ylabel("Free energy\n(meV/atom)")
    axes[2].set_xlabel("Cycle")
    axes[2].grid(True, alpha=0.3)

    fig.suptitle(f"SSCHA Convergence ({input_stem})")
    fig.tight_layout()
    out_pdf = output_dir / f"sscha_convergence-{input_stem}.pdf"
    fig.savefig(out_pdf)
    plt.close(fig)


def _promote_latest_band(output_dir: Path, input_stem: str, iteration: int):
    src_pdf = output_dir / f"band-{input_stem}_{iteration}.pdf"
    src_dat = output_dir / f"band-{input_stem}_{iteration}.dat"
    src_yaml = output_dir / f"band-{input_stem}_{iteration}.yaml"
    if src_pdf.exists():
        shutil.copyfile(src_pdf, output_dir / "band.pdf")
    if src_dat.exists():
        shutil.copyfile(src_dat, output_dir / "band.dat")
    if src_yaml.exists():
        shutil.copyfile(src_yaml, output_dir / "band.yaml")


def _run_qscaild_impl(args, input_poscar_path: Path, output_dir: Path):
    from ase import Atoms as AseAtoms
    from phonopy import Phonopy
    try:
        from phonopy.file_IO import parse_FORCE_CONSTANTS, write_FORCE_CONSTANTS, write_fc3_to_hdf5
        _HAS_WRITE_FC3 = True
    except ImportError:
        from phonopy.file_IO import parse_FORCE_CONSTANTS, write_FORCE_CONSTANTS
        write_fc3_to_hdf5 = None
        _HAS_WRITE_FC3 = False
    from phonopy.interface.vasp import read_vasp
    from scipy.constants import Boltzmann

    from macer.phonopy.band_path import generate_band_conf
    from macer.relaxation.optimizer import relax_structure

    if args.nconf < 1:
        raise ValueError("Error: --nconf must be >= 1.")
    if args.nsteps < 1:
        raise ValueError("Error: --nsteps must be >= 1.")
    if getattr(args, "save_every", 5) < 1:
        raise ValueError("Error: --save-every must be >= 1.")
    if not (0.0 <= args.memory <= 1.0):
        raise ValueError("Error: --memory must be in [0, 1].")
    if not (0.0 <= args.mixing <= 1.0):
        raise ValueError("Error: --mixing must be in [0, 1].")
    ess_ratio = getattr(args, "ess_collapse_ratio", None)
    if ess_ratio is not None and not (0.0 < float(ess_ratio) <= 1.0):
        raise ValueError("Error: --ess-collapse-ratio must be in (0, 1] when set.")
    if getattr(args, "max_regen", 0) < 0:
        raise ValueError("Error: --max-regen must be >= 0.")
    if getattr(args, "temperature", None) is None:
        raise ValueError("Error: --temperature is required.")
    user_set_use_pressure = "--use-pressure" in sys.argv
    use_pressure_mode = str(args.use_pressure)
    if (not args.optimize_volume) and (not user_set_use_pressure):
        use_pressure_mode = "False"
        args.use_pressure = "False"
    if args.optimize_volume:
        if use_pressure_mode == "False":
            use_pressure_mode = "orthorhombic"
            args.use_pressure = "orthorhombic"
        print(
            f"INFO: --optimize-volume enabled; pressure-coupled volume update is active "
            f"(mode={use_pressure_mode}, target={args.pressure_diag} kbar)."
        )
    elif use_pressure_mode != "False":
        print(
            "INFO: --use-pressure is set without --optimize-volume. "
            "Proceeding with pressure-coupled volume update for backward compatibility."
        )
    if args.reference_method != "random":
        print("WARNING: qscaild phase currently supports random ensemble generation only. Forcing --reference-method random.")
        args.reference_method = "random"
    if args.grid < 0:
        raise ValueError("Error: --grid must be >= 0.")
    if args.imaginary_freq <= -1 and args.imaginary_freq != -1:
        raise ValueError("Error: --imaginary-freq must be -1 or a positive value.")
    mass_map = _parse_mass_map(args)
    if mass_map:
        print("--- Atomic mass override ---")
        for sym, m in mass_map.items():
            print(f"  {sym:3s} : {m:8.4f} amu")
        print("----------------------------")
    if use_pressure_mode != "False":
        print(f"INFO: pressure mode enabled: {use_pressure_mode} (target={args.pressure_diag} kbar).")

    print("--- Starting phonopy qSCAILD workflow ---")
    print(f"All results will be saved in: {output_dir}")
    _write_output_guide(output_dir, args, input_poscar_path)
    if getattr(args, "dry_run", False):
        print("Dry-run enabled. Parsed settings:")
        print(
            f"  T={args.temperature}K nconf={args.nconf} nsteps={args.nsteps} "
            f"memory={args.memory} mixing={args.mixing} tol={args.tolerance}"
        )
        print(
            f"  pressure_mode={args.use_pressure} pressure_diag={args.pressure_diag} "
            f"pdiff={args.pdiff}"
        )
        print("Dry-run finished without executing force evaluations.")
        return

    model_path = _resolve_model_path(args.ff, args.model)
    if model_path and os.path.exists(model_path):
        model_path = os.path.abspath(model_path)

    calc_kwargs = {"device": args.device, "modal": args.modal}
    if args.ff == "mace":
        calc_kwargs["model_paths"] = [model_path]
    else:
        calc_kwargs["model_path"] = model_path
    calculator = get_calculator(ff_name=args.ff, **calc_kwargs)

    run_poscar = output_dir / "POSCAR_qscaild_input"
    shutil.copy(input_poscar_path, run_poscar)

    relaxed_poscar_path = run_poscar
    if args.initial_isif != 0:
        relaxed_poscar = output_dir / "CONTCAR-relax"
        relax_structure(
            input_file=str(run_poscar),
            isif=args.initial_isif,
            fmax=args.initial_fmax,
            device=args.device,
            contcar_name=str(relaxed_poscar),
            outcar_name=str(output_dir / "OUTCAR-relax"),
            xml_name=str(output_dir / "vasprun-relax.xml"),
            ff=args.ff,
            model_path=model_path,
            modal=args.modal,
            quiet=True,
            make_pdf=False,
            symprec=args.symprec,
            use_symmetry=getattr(args, "initial_use_symmetry", True),
        )
        unitcell = read_vasp(str(relaxed_poscar))
        relaxed_poscar_path = relaxed_poscar
    else:
        unitcell = read_vasp(str(run_poscar))
        relaxed_poscar_path = run_poscar
    unitcell = _apply_mass_override_to_phonopy_atoms(unitcell, mass_map)
    _log_mass_override_application(unitcell, mass_map)

    if args.dim:
        if len(args.dim) == 3:
            supercell_matrix = np.diag(args.dim)
        elif len(args.dim) == 9:
            supercell_matrix = np.array(args.dim).reshape(3, 3)
        else:
            raise ValueError("Error: --dim must be 3 or 9 integers.")
    else:
        vector_lengths = [np.linalg.norm(v) for v in unitcell.cell]
        scaling_factors = [int(np.ceil(args.min_length / v)) if v > 0 else 1 for v in vector_lengths]
        supercell_matrix = np.diag(scaling_factors)

    tolerance_phonopy = float(getattr(args, "tolerance_phonopy", 5e-3))
    ph = Phonopy(unitcell, supercell_matrix=supercell_matrix, primitive_matrix="auto", symprec=tolerance_phonopy)
    base_supercell = ph.supercell
    band_conf_path = output_dir / "band.conf" if getattr(args, "plot_bands", True) else None

    if args.read_initial_fc:
        print(f"Using provided initial FC2: {args.read_initial_fc}")
        fc_current = parse_FORCE_CONSTANTS(filename=args.read_initial_fc)
    else:
        t_fc2 = time.time()
        ph.generate_displacements(distance=args.amplitude, is_plusminus=True, is_diagonal=not args.nodiag)
        forces = []
        disp_cells = ph.supercells_with_displacements
        print(f"Initial FC2 setup: {len(disp_cells)} displaced supercells to evaluate.")
        for scell in tqdm(disp_cells, total=len(disp_cells), desc="Initial FC2 force evaluations"):
            atoms = AseAtoms(
                symbols=scell.symbols,
                cell=scell.cell,
                scaled_positions=scell.scaled_positions,
                pbc=True,
            )
            atoms.calc = calculator
            forces.append(atoms.get_forces())
        ph.forces = forces
        ph.produce_force_constants()
        fc_current = ph.force_constants.copy()
        print(f"Initial FC2 calculation finished in {time.time() - t_fc2:.2f} s.")

    write_FORCE_CONSTANTS(fc_current, filename=str(output_dir / "FORCE_CONSTANTS_init"))
    write_FORCE_CONSTANTS(fc_current, filename=str(output_dir / "FORCE_CONSTANTS_CURRENT"))
    if band_conf_path is not None:
        try:
            dim_override = " ".join(map(str, supercell_matrix.flatten().astype(int)))
            generate_band_conf(
                poscar_path=relaxed_poscar_path,
                out_path=band_conf_path,
                symprec=args.symprec,
                gamma_label=args.gamma_label,
                print_summary_flag=False,
                dim_override=dim_override,
            )
            old_cwd = Path.cwd()
            os.chdir(output_dir)
            try:
                ph.force_constants = fc_current
                _plot_iterative_band_structure(ph, band_conf_path, "harmonic", input_poscar_path.stem)
            finally:
                os.chdir(old_cwd)
        except Exception as e:
            print(f"WARNING: failed to export harmonic band outputs: {e}")
    if args.temperature <= 0:
        beta = 1.0e12
    else:
        beta = 1.0 / (Boltzmann / 1.602176634e-19 * args.temperature)

    out_fit_txt = output_dir / "out_fit.txt"
    out_conv_txt = output_dir / "out_convergence.txt"
    out_energy_txt = output_dir / "out_energy.txt"
    out_volume_txt = output_dir / "out_volume.txt"
    # Keep legacy names for compatibility.
    out_fit_legacy = output_dir / "out_fit"
    out_conv_legacy = output_dir / "out_convergence"
    out_energy_legacy = output_dir / "out_energy"
    out_volume_legacy = output_dir / "out_volume"
    out_fit_paths = [out_fit_txt, out_fit_legacy]
    out_conv_paths = [out_conv_txt, out_conv_legacy]
    out_energy_paths = [out_energy_txt, out_energy_legacy]
    out_volume_paths = [out_volume_txt, out_volume_legacy]
    fe_hist_path = output_dir / "qscaild_free_energy_history.log"
    for p in out_fit_paths + out_conv_paths + out_energy_paths + out_volume_paths:
        p.write_text("")
    _append_text(
        out_fit_paths,
        "# qSCAILD fit history\n"
        "# columns: iteration n_cycles n_snapshots fc_diff_max ess entropy free_energy_eV free_energy_meV_per_atom\n",
    )
    _append_text(
        out_conv_paths,
        "# qSCAILD convergence history\n"
        "# columns: iteration fc_diff_max ess entropy pressure_err_kbar(optional)\n",
    )
    _append_text(
        out_energy_paths,
        "# qSCAILD energy/reweight history\n"
        "# per iteration: cycle_id mean_energy_eV mean_raw_weight\n",
    )
    _append_text(
        out_volume_paths,
        "# qSCAILD pressure/volume history\n"
        "# columns: iteration Pxx_kbar Pyy_kbar Pzz_kbar pressure_err_kbar\n",
    )
    with open(fe_hist_path, "w") as f:
        f.write("# cycle  free_energy_eV  free_energy_meV_per_atom\n")

    converged = False
    final_reason = "max_cycles"
    pressure_converged = (use_pressure_mode == "False")
    pressure_err = np.nan
    pressure_kbar = np.array([np.nan, np.nan, np.nan])
    w_fit = np.array([1.0], dtype=float)
    fc3_current = None
    last_free_energy = np.nan
    regen_count = 0
    for iteration in range(1, args.nsteps + 1):
        cycle_dir = output_dir / f"cycle_{iteration:03d}"
        cycle_dir.mkdir(parents=True, exist_ok=True)

        try:
            U_cycle = _generate_cycle_ensemble_staged(
                ph,
                fc_current,
                args.temperature,
                args.nconf,
                use_smalldisp=args.use_smalldisp,
                imaginary_freq=args.imaginary_freq,
                grid=args.grid,
            )
        except Exception as e:
            print(f"WARNING: covariance-based sampling failed ({e}); falling back to phonopy RandomDisplacements.")
            U_cycle = _generate_cycle_ensemble(ph, fc_current, args.temperature, args.nconf)
        F_cycle, E_cycle, Stress_cycle = _evaluate_ensemble(
            U_cycle,
            base_supercell,
            calculator,
            need_stress=(use_pressure_mode != "False"),
            progress_desc=f"Cycle {iteration:03d} force/stress evaluations",
        )
        logp_origin = _compute_harmonic_logpdf(U_cycle, fc_current, beta)
        _save_cycle_data(
            cycle_dir,
            U_cycle,
            F_cycle,
            E_cycle,
            fc_current,
            logp_origin,
            save_npy=getattr(args, "save_npy", False),
        )
        if Stress_cycle is not None:
            _save_array_with_dat(cycle_dir / "Stress", Stress_cycle, save_npy=getattr(args, "save_npy", False))

        cycle_ids, cycle_dirs = _load_history_for_fit(output_dir, iteration, args.memory)
        U_fit, F_fit, E_fit, w_fit, Stress_all, energy_lines = _assemble_fit_data(
            cycle_ids=cycle_ids,
            cycle_dirs=cycle_dirs,
            fc_current=fc_current,
            beta=beta,
            iteration=iteration,
            use_pressure_mode=use_pressure_mode,
            save_npy=getattr(args, "save_npy", False),
        )
        ess, entropy = _weight_stats(w_fit)

        ess_ratio = getattr(args, "ess_collapse_ratio", None)
        if ess_ratio is not None:
            ess_threshold = float(ess_ratio) * len(w_fit)
            if ess < ess_threshold:
                max_regen = int(getattr(args, "max_regen", 200))
                print(
                    f"WARNING: ESS collapse detected at cycle {iteration:03d} "
                    f"(ESS={ess:.2f} < threshold={ess_threshold:.2f})."
                )
                if regen_count >= max_regen:
                    raise RuntimeError(
                        f"ESS collapse persists and max regeneration count reached "
                        f"({regen_count}/{max_regen})."
                    )
                regen_count += 1
                mode = str(getattr(args, "ess_collapse_mode", "soft")).lower()
                if mode == "hard":
                    keep_ids = [iteration]
                    keep_dirs = [cycle_dir]
                    print(
                        f"INFO: ESS regeneration #{regen_count} (hard): "
                        "dropping history and using latest cycle only."
                    )
                else:
                    keep_n = max(2, int(np.ceil(len(cycle_ids) * 0.5)))
                    keep_n = min(keep_n, len(cycle_ids))
                    keep_ids = cycle_ids[-keep_n:]
                    keep_dirs = cycle_dirs[-keep_n:]
                    print(
                        f"INFO: ESS regeneration #{regen_count} (soft): "
                        f"keeping recent history cycles={keep_ids}."
                    )

                U_fit, F_fit, E_fit, w_fit, Stress_all, energy_lines = _assemble_fit_data(
                    cycle_ids=keep_ids,
                    cycle_dirs=keep_dirs,
                    fc_current=fc_current,
                    beta=beta,
                    iteration=iteration,
                    use_pressure_mode=use_pressure_mode,
                    save_npy=getattr(args, "save_npy", False),
                )
                cycle_ids = keep_ids
                ess, entropy = _weight_stats(w_fit)
                energy_lines.append(
                    f"# ess_regeneration mode={mode} threshold={ess_threshold:.2f}\n"
                )

        save_every = int(getattr(args, "save_every", 5))
        should_save_intermediate = (iteration % save_every == 0)
        if should_save_intermediate:
            _save_weights_qscaild(output_dir / f"weights_iter_{iteration:04d}", w_fit, save_npy=getattr(args, "save_npy", False))

        if args.include_third_order:
            print(f"Cycle {iteration:03d}: FC fit stage started (FC2+FC3).")
        else:
            print(f"Cycle {iteration:03d}: FC fit stage started (FC2 only).")
        t_fit_start = time.perf_counter()
        fc_fit, fc3_fit = fit_weighted_fc(
            U_fit, F_fit, w_fit, ph, include_third_order=args.include_third_order
        )
        t_fit_elapsed = time.perf_counter() - t_fit_start
        if args.include_third_order:
            if fc3_fit is None:
                print(f"Cycle {iteration:03d}: FC fit stage finished in {t_fit_elapsed:.2f} s (FC3 unavailable/fallback).")
            else:
                print(f"Cycle {iteration:03d}: FC fit stage finished in {t_fit_elapsed:.2f} s (FC3 fitted).")
        else:
            print(f"Cycle {iteration:03d}: FC fit stage finished in {t_fit_elapsed:.2f} s.")
        fc_next = (1.0 - args.mixing) * fc_fit + args.mixing * fc_current
        if fc3_fit is not None:
            if fc3_current is None:
                fc3_next = fc3_fit.copy()
            else:
                fc3_next = (1.0 - args.mixing) * fc3_fit + args.mixing * fc3_current
        else:
            fc3_next = fc3_current
        fc_diff_max = float(np.max(np.abs(fc_next - fc_current)))
        try:
            free_energy = compute_weighted_free_energy(
                U_fit, E_fit, fc_next, fc3_next, w_fit, ph, args.temperature, args.mesh, return_components=False
            )
            n_atom_primitive = len(ph.primitive)
            f_mev_atom = (free_energy * 1000.0) / n_atom_primitive
            with open(fe_hist_path, "a") as f:
                f.write(f"{iteration:6d}  {free_energy:16.8f}  {f_mev_atom:16.8f}\n")
            last_free_energy = free_energy
        except Exception:
            free_energy = np.nan
            f_mev_atom = np.nan

        if use_pressure_mode != "False" and Stress_all:
            stress_fit = np.concatenate(Stress_all, axis=0)
            stress_weighted = np.sum(stress_fit * w_fit[:, None], axis=0)
            kinetic_pressure_kbar = _compute_kinetic_pressure_kbar(U_fit, F_fit, w_fit, unitcell)
            unitcell, pressure_kbar, pressure_err = _update_cell_from_pressure(
                unitcell,
                use_pressure_mode,
                args.pressure_diag,
                args.pdiff,
                stress_weighted,
                kinetic_pressure_kbar=kinetic_pressure_kbar,
            )
            unitcell = _apply_mass_override_to_phonopy_atoms(unitcell, mass_map)
            pressure_converged = pressure_err < args.pdiff

            # Rebuild phonopy object for updated cell state.
            ph = Phonopy(unitcell, supercell_matrix=supercell_matrix, primitive_matrix="auto", symprec=tolerance_phonopy)
            ph.force_constants = fc_next
            base_supercell = ph.supercell
        elif use_pressure_mode != "False":
            pressure_converged = False

        if should_save_intermediate:
            write_FORCE_CONSTANTS(fc_fit, filename=str(output_dir / f"FORCE_CONSTANTS_FIT_{iteration}"))
            if fc3_fit is not None:
                if _HAS_WRITE_FC3 and write_fc3_to_hdf5 is not None:
                    write_fc3_to_hdf5(fc3_fit, filename=str(output_dir / f"FC3_FIT_{iteration}.hdf5"))
                else:
                    print("WARNING: write_fc3_to_hdf5 not available in this phonopy build; skipping FC3 save.")
            try:
                poscar_avg_path, adp_path, adp_cif_path = _save_weighted_average_and_adp(
                    output_dir=output_dir,
                    iteration=iteration,
                    ph=ph,
                    U_fit=U_fit,
                    w_fit=w_fit,
                    projection_axis_frac=getattr(args, "distribution_axis", None),
                )
                print(f"  Saved weighted average structure: {poscar_avg_path.name}")
                print(f"  Saved weighted ADP table: {adp_path.name}")
                print(f"  Saved weighted ADP CIF: {adp_cif_path.name}")
                print(
                    "  Saved element-average distributions: "
                    f"distribution_proj_iter_{iteration:04d}_element_*_avg.pdf/.dat"
                )
            except Exception as e:
                print(f"WARNING: failed to save weighted average/ADP at cycle {iteration}: {e}")

            if band_conf_path is not None:
                try:
                    ph.force_constants = fc_next
                    old_cwd = Path.cwd()
                    os.chdir(output_dir)
                    try:
                        _plot_iterative_band_structure(ph, band_conf_path, iteration, input_poscar_path.stem)
                        _promote_latest_band(output_dir, input_poscar_path.stem, iteration)
                    finally:
                        os.chdir(old_cwd)
                except Exception as e:
                    print(f"WARNING: failed to export band outputs at cycle {iteration}: {e}")
        write_FORCE_CONSTANTS(fc_current, filename=str(output_dir / "FORCE_CONSTANTS_PREVIOUS"))
        write_FORCE_CONSTANTS(fc_next, filename=str(output_dir / "FORCE_CONSTANTS_CURRENT"))
        (output_dir / "iteration").write_text(f"{iteration}\n")

        fit_line = (
            f"{iteration:8d} {len(cycle_ids):8d} {len(U_fit):12d} "
            f"{fc_diff_max:16.8e} {ess:10.6f} {entropy:10.6f} "
            f"{free_energy:16.8f} {f_mev_atom:18.8f}\n"
        )
        fit_line += f"# fit_elapsed_s {t_fit_elapsed:.6f}\n"
        _append_text(out_fit_paths, fit_line)
        _append_text(out_energy_paths, "".join(energy_lines))
        conv_line = (
            f"{iteration:8d} {fc_diff_max:16.8e} {ess:10.6f} {entropy:10.6f}"
        )
        if use_pressure_mode != "False":
            conv_line += f" {pressure_err:16.6f}"
        _append_text(out_conv_paths, conv_line + "\n")
        if use_pressure_mode != "False":
            _append_text(
                out_volume_paths,
                f"{iteration:8d} {pressure_kbar[0]:12.6f} {pressure_kbar[1]:12.6f} "
                f"{pressure_kbar[2]:12.6f} {pressure_err:16.6f}\n",
            )

        print(
            f"Cycle {iteration:03d}: "
            f"history={cycle_ids}, snapshots={len(U_fit)}, "
            f"fc_diff_max={fc_diff_max:.6e}, ESS={ess:.2f}"
        )
        fc_ratio = fc_diff_max / float(args.tolerance) if args.tolerance > 0 else np.inf
        print(
            "  FC convergence: "
            f"current={fc_diff_max:.6e}, target<{args.tolerance:.6e}, "
            f"ratio={fc_ratio:.2f}x"
        )
        if np.isfinite(f_mev_atom):
            print(f"  Free energy: {f_mev_atom:.6f} meV/atom")
        if should_save_intermediate:
            print(
                f"  Saved intermediate outputs at cycle {iteration} "
                f"(save-every={save_every})."
            )
            print(f"  Saved ensemble snapshot bundle: cycle_{iteration:03d}/")
        if use_pressure_mode != "False":
            p_ratio = pressure_err / float(args.pdiff) if args.pdiff > 0 else np.inf
            print(
                "  Pressure[kbar]="
                f"{pressure_kbar[0]:+.3f},{pressure_kbar[1]:+.3f},{pressure_kbar[2]:+.3f} "
                f"(err={pressure_err:.3f})"
            )
            print(
                "  Pressure convergence: "
                f"current_err={pressure_err:.6f}, target<{args.pdiff:.6f}, "
                f"ratio={p_ratio:.2f}x"
            )

        fc_current = fc_next
        fc3_current = fc3_next
        ph.force_constants = fc_current

        fc_converged = fc_diff_max < args.tolerance
        if use_pressure_mode == "False":
            converged = fc_converged
        else:
            converged = fc_converged and pressure_converged
        if converged:
            if use_pressure_mode == "False":
                final_reason = "fc_tolerance"
            else:
                final_reason = "fc_and_pressure_tolerance"
            converged = True
            break

    _save_weights_qscaild(output_dir / "weights_final", w_fit, save_npy=getattr(args, "save_npy", False))
    write_FORCE_CONSTANTS(fc_current, filename=str(output_dir / "FORCE_CONSTANTS_QSCAILD_final"))
    if fc3_current is not None:
        if _HAS_WRITE_FC3 and write_fc3_to_hdf5 is not None:
            write_fc3_to_hdf5(fc3_current, filename=str(output_dir / "FC3_QSCAILD_final.hdf5"))
        else:
            print("WARNING: write_fc3_to_hdf5 not available in this phonopy build; skipping final FC3 save.")
    if band_conf_path is not None:
        try:
            ph.force_constants = fc_current
            old_cwd = Path.cwd()
            os.chdir(output_dir)
            try:
                _plot_iterative_band_structure(ph, band_conf_path, "final", input_poscar_path.stem)
            finally:
                os.chdir(old_cwd)
            for ext in ("pdf", "dat", "yaml"):
                src = output_dir / f"band-{input_poscar_path.stem}_final.{ext}"
                if src.exists():
                    shutil.copyfile(src, output_dir / f"band.{ext}")
        except Exception as e:
            print(f"WARNING: failed to export final band outputs: {e}")
    _plot_free_energy_convergence(fe_hist_path, output_dir)
    _plot_sscha_history(output_dir, input_poscar_path.stem)
    if np.isfinite(last_free_energy):
        print(f"Final free energy (eV): {last_free_energy:.8f}")
    if converged:
        print(
            "qSCAILD converged "
            f"(reason={final_reason}) at cycle {iteration} "
            f"with fc_tol={args.tolerance}."
        )
    else:
        print(
            f"qSCAILD reached max cycles ({args.nsteps}) without convergence "
            f"(last reason={final_reason})."
        )
    _write_output_guide(output_dir, args, input_poscar_path)
    print("--- qSCAILD workflow finished ---")


def run_qscaild_workflow(args):
    if args.seed is not None:
        random.seed(args.seed)
        np.random.seed(args.seed)

    cmd = str(
        getattr(args, "phonopy_command", None)
        or getattr(args, "command", "")
        or ""
    ).lower()
    is_sscha_cmd = cmd in ("sscha", "run-sscha")
    banner = "SSCHA" if is_sscha_cmd else "qSCAILD"
    log_name = "macer_sscha.log" if is_sscha_cmd else "macer_qscaild.log"

    input_path = None
    output_dir = None
    try:
        # Resolve input first to derive deterministic default output prefix.
        tmp_output = Path.cwd()
        input_path = _prepare_input_poscar(args, tmp_output, convert_cif=False)
        output_dir = _build_output_dir(args, input_path, create=True)

        if getattr(args, "dry_run", False):
            # For dry-run, resolve CIF conversion if needed and skip expensive force evaluations.
            input_path = _prepare_input_poscar(args, Path.cwd(), convert_cif=True)
            args.output_dir = str(output_dir)
            print(f"--- Macer {banner} Workflow (direct/symfc) ---")
            print(f"Command: {' '.join(sys.argv)}")
            _run_qscaild_impl(args, input_path, output_dir)
            return

        # Ensure converted inputs are inside final output dir for real runs.
        input_path = _prepare_input_poscar(args, output_dir, convert_cif=True)
        args.output_dir = str(output_dir)
        log_file = output_dir / log_name
        orig_stdout = sys.stdout
        with Logger(str(log_file)) as lg:
            try:
                sys.stdout = lg
                print(f"--- Macer {banner} Workflow (direct/symfc) ---")
                print(f"Command: {' '.join(sys.argv)}")
                _run_qscaild_impl(args, input_path, output_dir)
            finally:
                try:
                    if output_dir is not None and input_path is not None:
                        _write_output_guide(output_dir, args, input_path)
                except Exception:
                    pass
                sys.stdout = orig_stdout
    except Exception as exc:
        print(f"qSCAILD run failed: {exc}")
        traceback.print_exc()
        raise




def run_sscha_workflow(args):
    """SSCHA entrypoint using direct symfc + qSCAILD-style cycle logic."""
    return run_qscaild_workflow(args)


def add_sscha_parser(subparsers):
    from macer import __version__
    from macer.cli.phonopy_main import MACER_LOGO

    parser = subparsers.add_parser(
        "sscha",
        aliases=["run-sscha"],
        help="SSCHA with qSCAILD-style cycle updates (direct/symfc)",
        description=MACER_LOGO + f"\nmacer_phonopy sscha (v{__version__}): "
                               "Direct SSCHA workflow using symfc fitting with qSCAILD-style "
                               "ensemble update and convergence scheme (no legacy backend).",
        epilog="""
  # 1) Validate setup only (no expensive force evaluation)
  macer phonopy sscha -p POSCAR -T 300 --dry-run

  # 2) Minimal finite-T run
  macer phonopy sscha -p POSCAR -T 300 --dim 3 3 3 --ff chgnet

  # 3) Explicit run controls + snapshots every 5 cycles
  macer phonopy sscha -p POSCAR -T 500 --dim 3 3 3 --nconf 10 --nsteps 30 --save-every 5 --ff chgnet

  # 4) Store outputs in a custom directory
  macer phonopy sscha -p POSCAR -T 500 --output-dir ./sscha-run-500K --ff chgnet

  # 5) Pressure-coupled volume optimization loop
  macer phonopy sscha -p POSCAR -T 500 --optimize-volume --pressure-diag 0 0 0 --pdiff 2.0 --nconf 10 --nsteps 30

  # 6) FC3 fitting + custom mesh
  macer phonopy sscha -p POSCAR -T 300 --include-third-order --mesh 9 9 9 --nconf 20 --nsteps 50

  # 7) Stabilize convergence with stronger IFC mixing + longer memory
  macer phonopy sscha -p POSCAR -T 300 --dim 3 3 3 --nconf 20 --nsteps 40 --mixing 0.3 --memory 0.6 --tolerance 0.01

  # 8) Tight final convergence pass (after a stable pre-run)
  macer phonopy sscha -p POSCAR -T 300 --dim 3 3 3 --nconf 30 --nsteps 60 --mixing 0.2 --memory 0.5 --tolerance 0.005
""",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    general_group = parser.add_argument_group("General & Input")
    general_group.add_argument("-p", "--poscar", required=False, default=None, help="Input crystal structure file (POSCAR).")
    general_group.add_argument("-c", "--cif", 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.")
    general_group.add_argument("--seed", type=int, help="Random seed for reproducibility.")

    mlff_group = 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", help="Path to the force field model file.")
    mlff_group.add_argument("--device", default=DEFAULT_DEVICE, choices=["cpu", "mps", "cuda"], help=f"Compute device (default: {DEFAULT_DEVICE}).")
    mlff_group.add_argument("--modal", type=str, default=None, help="Modal for SevenNet model, if required.")

    initial_group = parser.add_argument_group("Initial Harmonic FC Settings")
    initial_group.add_argument("--initial-fmax", type=float, default=5e-3, help="Force convergence for initial relaxation (eV/Å).")
    initial_group.add_argument("--initial-isif", type=int, default=3, help="VASP ISIF mode for initial relaxation (default: 3, use 0 to skip).")
    initial_group.add_argument("--dim", type=int, nargs="+", help="Supercell dimension (e.g., '2 2 2').")
    initial_group.add_argument("-l", "--min-length", type=float, default=15.0, help="Minimum supercell length if --dim is not set (Å).")
    initial_group.add_argument("--amplitude", type=float, default=0.01, help="Displacement amplitude for harmonic FC calculation (Å).")
    initial_group.add_argument("--pm", action="store_true", help="Use plus/minus displacements (always enabled).")
    initial_group.add_argument("--nodiag", action="store_true", help="Do not use diagonal displacements.")
    initial_group.add_argument("--symprec", type=float, default=1e-5, help="Symmetry tolerance (Å).")
    initial_group.add_argument("--tolerance-phonopy", type=float, default=5e-3, help="Symmetry tolerance for phonopy displacement/FC steps (default: 5e-3).")
    initial_group.add_argument("--read-initial-fc", type=str, help="Path to existing FORCE_CONSTANTS to skip initial calculation.")
    initial_group.add_argument("--initial-symmetry-off", dest="initial_use_symmetry", action="store_false", help="Disable FixSymmetry in initial relaxation.")

    ensemble_group = parser.add_argument_group("Reference Ensemble Settings")
    ensemble_group.add_argument("--reference-method", choices=["random", "md"], default="random", help="Method to generate reference ensemble (default: random).")
    ensemble_group.add_argument("--reference-n-samples", type=int, default=200, help=argparse.SUPPRESS)
    ensemble_group.add_argument("--reference-md-nsteps", type=int, default=200, help=argparse.SUPPRESS)
    ensemble_group.add_argument("--reference-md-nequil", type=int, default=100, help=argparse.SUPPRESS)
    ensemble_group.add_argument("--reference-md-tstep", type=float, default=1.0, help=argparse.SUPPRESS)
    ensemble_group.add_argument("--n-jobs", type=int, default=1, help=argparse.SUPPRESS)
    ensemble_group.add_argument("--md-thermostat", choices=["nve", "langevin"], default="langevin", help=argparse.SUPPRESS)
    ensemble_group.add_argument("--md-friction", type=float, default=0.01, help=argparse.SUPPRESS)
    ensemble_group.add_argument("--reference-ensemble", help=argparse.SUPPRESS)
    ensemble_group.add_argument("--no-save-reference-ensemble", action="store_true", help=argparse.SUPPRESS)
    ensemble_group.add_argument("--write-xdatcar", action="store_true", help=argparse.SUPPRESS)
    ensemble_group.add_argument("--xdatcar-step", type=int, default=50, help=argparse.SUPPRESS)

    compat_group = parser.add_argument_group("Reweighting / Compatibility Settings")
    compat_group.add_argument("-T", "--temperature", type=float, help="Target temperature in Kelvin.")
    compat_group.add_argument("--max-iter", type=int, default=200, help=argparse.SUPPRESS)
    compat_group.add_argument("--max-regen", type=int, default=200, help="Maximum ESS-regeneration attempts when collapse handling is enabled (default: 200).")
    compat_group.add_argument("--ess-collapse-ratio", type=float, default=None, help="Enable ESS collapse handling if ESS < ratio * snapshots (disabled by default).")
    compat_group.add_argument("--ess-collapse-mode", choices=["soft", "hard"], default="soft", help="ESS collapse handling mode: soft keeps recent history, hard keeps latest cycle only (default: soft).")
    compat_group.add_argument("--free-energy-conv", type=float, default=0.1, help=argparse.SUPPRESS)
    compat_group.add_argument("--mesh", type=int, nargs=3, default=[7, 7, 7], help="Q-point mesh for free-energy calculation (default: 7 7 7).")
    compat_group.add_argument("--fc-mixing-alpha", type=float, default=0.5, help=argparse.SUPPRESS)
    compat_group.add_argument("--include-third-order", action="store_true", help="Enable simultaneous fitting of 3rd-order force constants.")
    compat_group.add_argument("--mass", nargs="+", help="Specify atomic masses. Format: Symbol Mass Symbol Mass ... (e.g. --mass H 2.014 D 2.014)")
    compat_group.add_argument("--optimize-volume", action="store_true", help="Enable pressure-coupled volume optimization loop (default mode: cubic).")
    compat_group.add_argument("--max-volume-iter", type=int, default=10, help=argparse.SUPPRESS)
    compat_group.add_argument("--gamma-label", type=str, default="GM", help="Label for Gamma point in band plots.")
    compat_group.add_argument("--save-every", type=int, default=5, help="Save intermediate FORCE_CONSTANTS every N cycles (default: 5).")
    compat_group.add_argument("--no-plot-bands", dest="plot_bands", action="store_false", help="Do not plot band structures.")

    scheme_group = parser.add_argument_group("qSCAILD-Style Cycle Settings")
    scheme_group.add_argument("--nconf", type=int, default=30, help="Number of displaced configurations per cycle (default: 30).")
    scheme_group.add_argument("--nsteps", type=int, default=30, help="Maximum number of cycles (default: 30).")
    scheme_group.add_argument("--nfits", dest="nsteps", type=int, help=argparse.SUPPRESS)
    scheme_group.add_argument("--memory", type=float, default=0.5, help="History fraction used for multi-cycle reweighting fit (default: 0.5).")
    scheme_group.add_argument("--mixing", type=float, default=0.5, help="Linear mixing factor for IFC update (default: 0.5).")
    scheme_group.add_argument("--tolerance", type=float, default=0.02, help="IFC convergence tolerance (default: 0.02).")
    scheme_group.add_argument("--pdiff", type=float, default=2.0, help="Pressure convergence tolerance in kbar (default: 2.0).")
    scheme_group.add_argument("--grid", type=int, default=9, help="Q-grid size for covariance construction (default: 9).")
    scheme_group.add_argument("--use-smalldisp", action="store_true", help="Use small-displacement mode.")
    scheme_group.add_argument("--imaginary-freq", type=float, default=1.0, help="Imaginary-frequency handling parameter in THz.")
    scheme_group.add_argument("--use-pressure", choices=["False", "cubic", "tetragonal", "orthorhombic"], default="orthorhombic", help="Pressure-driven cell update mode (advanced override for --optimize-volume).")
    scheme_group.add_argument("--pressure-diag", type=float, nargs=3, metavar=("PX", "PY", "PZ"), default=[0.0, 0.0, 0.0], help="Target diagonal pressure components in kbar.")
    scheme_group.add_argument("--distribution-axis", type=float, nargs=3, metavar=("A", "B", "C"), default=None, help="Fractional direction vector for projected displacement distributions on primitive atoms (e.g., 1 2 1).")
    scheme_group.add_argument("--save-npy", action="store_true", help="Also save binary .npy arrays (default is dat-only outputs).")
    scheme_group.add_argument("--dry-run", action="store_true", help="Validate parser/settings and exit before expensive evaluations.")

    parser.set_defaults(plot_bands=True)
    parser.set_defaults(func=run_sscha_workflow)
    return parser
