# pdb2reaction/tsopt.py

"""
Transition-state optimization using Hessian Guided Dimer (grad/dimer) or RS-I-RFO (hess/rsirfo).

Example:
    pdb2reaction tsopt -i ts_cand.pdb -q 0 -m 1 --opt-mode hess --out-dir ./result_tsopt/

For detailed documentation, see: docs/tsopt.md
"""

from __future__ import annotations

import gc
import logging
import sys
import textwrap
from pathlib import Path
from typing import Dict, Any, Optional, Tuple, List, Sequence

import click
import numpy as np
import torch
from ase import Atoms
from ase.io import write
import ase.units as units
import time

# ---------------- pysisyphus / pdb2reaction imports ----------------
from pysisyphus.helpers import geom_loader
from pysisyphus.optimizers.LBFGS import LBFGS
from pysisyphus.optimizers.exceptions import OptimizationError, ZeroStepLength
from pysisyphus.constants import BOHR2ANG, AMU2AU, AU2EV
from pysisyphus.calculators.Dimer import Dimer  # Dimer calculator (orientation-projected forces)

# RS-I-RFO optimizer (transition state)
from pysisyphus.tsoptimizers.RSIRFOptimizer import RSIRFOptimizer  # type: ignore

# local helpers from pdb2reaction
from .backends import create_calculator
from .defaults import (
    CALC_KW_DEFAULT,
    GEOM_KW_DEFAULT,
    UMA_CALC_KW,
    OPT_BASE_KW,
    TSOPT_MODE_ALIASES,
    OUT_DIR_TSOPT,
    HESSIAN_DIMER_CLI_KW,
    RSIRFO_KW,
)
from .utils import (
    resolve_freeze_atoms,
    apply_yaml_overrides,
    pretty_block,
    format_geom_for_echo,
    format_elapsed,
    normalize_choice,
    prepared_cli_input,
    set_convert_file_enabled,
    convert_xyz_like_outputs,
    strip_inherited_keys,
    cli_param_overridden,
)
from .cli_utils import resolve_yaml_sources, load_merged_yaml_cfg
from .freq import (
    _torch_device,
    _safe_masses_amu,
    _tr_orthonormal_basis,
    _mass_weighted_hessian,
    _calc_full_hessian_torch,
    _calc_energy,
    _write_mode_trj_and_pdb,
    _frequencies_cm_and_modes,
)

logger = logging.getLogger(__name__)


# ===================================================================
#               Mass-weighted projection & vib analysis
# ===================================================================
def _mw_projected_hessian_inplace(H: torch.Tensor,
                                  coords_bohr_t: torch.Tensor,
                                  masses_au_t: torch.Tensor,
                                  freeze_idx: Optional[List[int]] = None) -> torch.Tensor:
    """
    Mass-weight H in-place, optionally restrict to active DOF subspace (PHVA) and
    project out TR motions (in that subspace), also in-place. Symmetrizes after TR projection.
    Returns the (possibly reduced) Hessian to be diagonalized.
    """
    if H.dtype != torch.float64:
        H = H.to(dtype=torch.float64)
    dtype, device = H.dtype, H.device
    with torch.no_grad():
        N = coords_bohr_t.shape[0]
        if freeze_idx is not None and len(freeze_idx) > 0:
            frozen = set(int(i) for i in freeze_idx if 0 <= int(i) < N)
            active_idx = [i for i in range(N) if i not in frozen]
            if len(active_idx) == 0:
                raise RuntimeError("All atoms are frozen; no active DOF left for TR projection.")
            # mass-weight first
            H = _mass_weighted_hessian(H, masses_au_t.to(dtype=dtype, device=device))
            # take active DOF submatrix
            mask_dof = torch.ones(3 * N, dtype=torch.bool, device=device)
            for i in frozen:
                mask_dof[3 * i:3 * i + 3] = False
            H = H[mask_dof][:, mask_dof]
            # TR basis and projection in active subspace (in-place)
            coords_act = coords_bohr_t[active_idx, :].to(dtype=dtype, device=device)
            masses_act = masses_au_t[active_idx].to(dtype=dtype, device=device)
            Q, _ = _tr_orthonormal_basis(coords_act, masses_act)  # (3N_act, r)
            Qt = Q.T
            QtH = Qt @ H
            H.addmm_(Q, QtH, beta=1.0, alpha=-1.0)
            H.addmm_((QtH.T), Qt, beta=1.0, alpha=-1.0)
            QtHQ = QtH @ Q
            H.addmm_(Q @ QtHQ, Qt, beta=1.0, alpha=1.0)
            _t = H.T.clone()
            H.add_(_t).mul_(0.5)
            del _t
            del Q, Qt, QtH, QtHQ, mask_dof, coords_act, masses_act, active_idx, frozen
        else:
            # Full DOF: mass-weight + TR projection in-place
            H = _mass_weighted_hessian(H, masses_au_t.to(dtype=dtype, device=device))
            coords_bohr_t = coords_bohr_t.to(dtype=dtype, device=device)
            masses_au_t = masses_au_t.to(dtype=dtype, device=device)
            Q, _ = _tr_orthonormal_basis(coords_bohr_t, masses_au_t)  # (3N, r)
            Qt = Q.T
            QtH = Qt @ H
            H.addmm_(Q, QtH, beta=1.0, alpha=-1.0)
            H.addmm_(QtH.T, Qt, beta=1.0, alpha=-1.0)
            QtHQ = QtH @ Q
            H.addmm_(Q @ QtHQ, Qt, beta=1.0, alpha=1.0)
            _t = H.T.clone()
            H.add_(_t).mul_(0.5)
            del _t
            del Q, Qt, QtH, QtHQ
        if torch.cuda.is_available() and device.type == "cuda":
            torch.cuda.empty_cache()
        return H


def _omega2_to_freq_cm(omega2: torch.Tensor) -> float:
    """
    Convert a (possibly signed) omega^2 eigenvalue to cm^-1, matching freq.py.
    """
    s_new = (units._hbar * 1e10 / np.sqrt(units._e * units._amu) * np.sqrt(AU2EV) / BOHR2ANG)
    hnu = s_new * torch.sqrt(torch.abs(omega2))
    hnu = torch.where(omega2 < 0, -hnu, hnu)
    return float((hnu / units.invcm).item())

def _mode_direction_by_root(H: torch.Tensor,
                            coords_bohr_t: torch.Tensor,
                            masses_au_t: torch.Tensor,
                            root: int = 0,
                            freeze_idx: Optional[List[int]] = None) -> Tuple[np.ndarray, float]:
    """
    Get the eigenvector (Cartesian space) corresponding to the `root`-th most negative
    eigenvalue (root=0: most negative) of the mass-weighted, TR-projected Hessian.
    PHVA (active-subspace) is applied if freeze_idx is provided: frozen DOFs are zero.
    For root==0 the routine prefers LOBPCG and falls back to eigh as needed.
    """
    with torch.no_grad():
        # In-place: mass weight + (active-subspace) TR projection
        Hmw_proj = _mw_projected_hessian_inplace(H, coords_bohr_t, masses_au_t, freeze_idx=freeze_idx)

        # Solve eigenproblem in the (possibly reduced) space
        if int(root) == 0:
            try:
                w, v_mw_sub = torch.lobpcg(Hmw_proj, k=1, largest=False)
                u_mw_sub = v_mw_sub[:, 0]
                omega2 = w[0]
            except Exception:
                evals_f, evecs_f = torch.linalg.eigh(Hmw_proj)
                pick = int(torch.argmin(evals_f).item())
                u_mw_sub = evecs_f[:, pick]
                omega2 = evals_f[pick]
                del evals_f, evecs_f
        else:
            evals, evecs_mw = torch.linalg.eigh(Hmw_proj)  # ascending
            neg = (evals < 0)
            neg_inds = torch.nonzero(neg, as_tuple=False).view(-1)
            if neg_inds.numel() == 0:
                pick = int(torch.argmin(evals).item())
            else:
                k = max(0, min(int(root), neg_inds.numel() - 1))
                pick = int(neg_inds[k].item())
            u_mw_sub = evecs_mw[:, pick]
            omega2 = evals[pick]
            del evals, evecs_mw

        # Embed back to full 3N (frozen DOF as zeros) if we solved in subspace
        N = coords_bohr_t.shape[0]
        if freeze_idx and len(freeze_idx) > 0:
            frozen = set(int(i) for i in freeze_idx if 0 <= int(i) < N)
            mask_dof = torch.ones(3 * N, dtype=torch.bool, device=Hmw_proj.device)
            for i in frozen:
                mask_dof[3 * i:3 * i + 3] = False
            u_mw_full = torch.zeros(3 * N, dtype=Hmw_proj.dtype, device=Hmw_proj.device)
            u_mw_full[mask_dof] = u_mw_sub
            u_mw = u_mw_full
            del mask_dof, frozen
        else:
            u_mw = u_mw_sub

        # Convert mass-weighted → Cartesian & normalize
        masses_amu_t = (masses_au_t / AMU2AU).to(dtype=Hmw_proj.dtype, device=Hmw_proj.device)
        m3 = torch.repeat_interleave(masses_amu_t, 3)
        inv_sqrt_m = torch.sqrt(1.0 / m3)
        v = inv_sqrt_m * u_mw
        v = v / torch.linalg.norm(v)
        mode = v.reshape(-1, 3).detach().cpu().numpy()
        freq_cm = _omega2_to_freq_cm(omega2)

        del masses_amu_t, m3, inv_sqrt_m, v, u_mw, u_mw_sub, omega2
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return mode, freq_cm


def _calc_gradient(geom, uma_kwargs: dict) -> np.ndarray:
    """
    Return true Cartesian gradient (shape 3N,) in Hartree/Bohr.
    """
    calc = create_calculator(**uma_kwargs)
    geom.set_calculator(calc)
    g = np.array(geom.gradient, dtype=float).reshape(-1)
    geom.set_calculator(None)
    return g


# ===================================================================
#            Active-subspace helpers & Bofill update
# ===================================================================

def _active_indices(N: int, freeze_idx: Optional[List[int]]) -> List[int]:
    if not freeze_idx:
        return list(range(N))
    fz = set(int(i) for i in freeze_idx if 0 <= int(i) < N)
    return [i for i in range(N) if i not in fz]


def _active_mask_dof(N: int, freeze_idx: Optional[List[int]]) -> np.ndarray:
    mask = np.ones(3 * N, dtype=bool)
    if freeze_idx:
        for i in freeze_idx:
            if 0 <= int(i) < N:
                mask[3 * int(i):3 * int(i) + 3] = False
    return mask


def _extract_active_block(H_full: torch.Tensor, mask_dof: np.ndarray) -> torch.Tensor:
    """
    Return the active-DOF block as a torch.Tensor sharing device/dtype.
    """
    device = H_full.device
    m = torch.as_tensor(mask_dof, device=device, dtype=torch.bool)
    return H_full[m][:, m].clone()


def _representative_atoms_for_mode(mode_vec: torch.Tensor, flatten_k: int) -> np.ndarray:
    """
    Return indices of atoms with the largest displacements for a given mode.
    """
    vec = mode_vec.view(-1, 3)
    norms = torch.linalg.norm(vec, dim=1)
    k = min(int(flatten_k), vec.shape[0])
    if k <= 0:
        return np.zeros(0, dtype=int)
    topk = torch.topk(norms, k=k, largest=True)
    return topk.indices.detach().cpu().numpy()


def _select_flatten_targets_for_geom(
    freqs_cm: np.ndarray,
    modes: torch.Tensor,
    coords_bohr: np.ndarray,
    neg_freq_thresh_cm: float,
    root: int,
    flatten_sep_cutoff: float,
    flatten_k: int,
) -> List[int]:
    """
    Select a subset of imaginary modes to flatten for a geometry.
    """
    neg_idx_all = np.where(freqs_cm < -abs(neg_freq_thresh_cm))[0]
    if len(neg_idx_all) <= 1:
        return []

    order = np.argsort(freqs_cm[neg_idx_all])
    sorted_neg = neg_idx_all[order]
    root_clamped = max(0, min(int(root), len(order) - 1))
    primary_idx = sorted_neg[root_clamped]
    candidates = [int(i) for i in sorted_neg if int(i) != int(primary_idx)]
    if not candidates:
        return []

    coords_ang = torch.as_tensor(
        coords_bohr.reshape(-1, 3) * BOHR2ANG,
        dtype=modes.dtype,
        device=modes.device,
    )

    targets: List[int] = []
    reps_list: List[np.ndarray] = []

    for idx in candidates:
        rep = _representative_atoms_for_mode(modes[idx], flatten_k)
        if rep.size == 0:
            continue
        rep_coords = coords_ang[rep]
        if not reps_list:
            targets.append(idx)
            reps_list.append(rep)
            continue

        accept = True
        for prev_rep in reps_list:
            prev_coords = coords_ang[prev_rep]
            dmat = torch.cdist(rep_coords, prev_coords)
            min_dist = float(torch.min(dmat).item())
            if min_dist < float(flatten_sep_cutoff):
                accept = False
                break
        if accept:
            targets.append(idx)
            reps_list.append(rep)

    return targets


def _flatten_once_with_modes_for_geom(
    geom,
    masses_amu: np.ndarray,
    uma_kwargs: dict,
    freqs_cm: np.ndarray,
    modes: torch.Tensor,
    neg_freq_thresh_cm: float,
    flatten_amp_ang: float,
    flatten_sep_cutoff: float,
    flatten_k: int,
    root: int,
) -> bool:
    """
    Flatten extra imaginary modes for a geometry (single pass).
    """
    neg_idx_all = np.where(freqs_cm < -abs(neg_freq_thresh_cm))[0]
    if len(neg_idx_all) <= 1:
        return False

    targets = _select_flatten_targets_for_geom(
        freqs_cm,
        modes,
        geom.cart_coords,
        neg_freq_thresh_cm,
        root,
        flatten_sep_cutoff,
        flatten_k,
    )
    if not targets:
        return False

    mass_scale = np.sqrt(12.011 / masses_amu)[:, None]
    amp_bohr = float(flatten_amp_ang) / BOHR2ANG
    E_ref = _calc_energy(geom, uma_kwargs)

    m3 = np.repeat(masses_amu, 3).reshape(-1, 3)
    for idx in targets:
        v_mw = modes[idx].detach().cpu().numpy().reshape(-1, 3)
        v_cart = v_mw / np.sqrt(m3)
        v_cart /= np.linalg.norm(v_cart)

        disp = amp_bohr * mass_scale * v_cart
        ref = geom.cart_coords.reshape(-1, 3)

        plus = ref + disp
        minus = ref - disp

        geom.coords = plus.reshape(-1)
        E_plus = _calc_energy(geom, uma_kwargs)

        geom.coords = minus.reshape(-1)
        E_minus = _calc_energy(geom, uma_kwargs)

        use_plus = E_plus <= E_minus
        geom.coords = (plus if use_plus else minus).reshape(-1)
        E_keep = E_plus if use_plus else E_minus
        delta_e = E_keep - E_ref
        click.echo(
            f"[Flatten] mode={idx} freq={freqs_cm[idx]:+.2f} cm^-1 "
            f"E_disp={E_keep:.8f} Ha \u0394E={delta_e:+.8f} Ha"
        )

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return True


def _mw_tr_project_active_inplace(H: torch.Tensor,
                                  coords_act_t: torch.Tensor,
                                  masses_act_au_t: torch.Tensor) -> torch.Tensor:
    """
    Mass-weight & project TR in the *active* subspace (in-place; no explicit symmetrization).
    """
    with torch.no_grad():
        # mass-weight
        masses_amu_t = (masses_act_au_t / AMU2AU).to(dtype=H.dtype, device=H.device)
        m3 = torch.repeat_interleave(masses_amu_t, 3)
        inv_sqrt_m_col = torch.sqrt(1.0 / m3).view(1, -1)
        inv_sqrt_m_row = inv_sqrt_m_col.view(-1, 1)
        H.mul_(inv_sqrt_m_row)
        H.mul_(inv_sqrt_m_col)
        # TR basis & projection
        Q, _ = _tr_orthonormal_basis(coords_act_t, masses_act_au_t)  # (3N_act, r)
        Qt = Q.T
        QtH = Qt @ H
        H.addmm_(Q, QtH, beta=1.0, alpha=-1.0)
        H.addmm_(QtH.T, Qt, beta=1.0, alpha=-1.0)
        QtHQ = QtH @ Q
        H.addmm_(Q @ QtHQ, Qt, beta=1.0, alpha=1.0)
        del masses_amu_t, m3, inv_sqrt_m_col, inv_sqrt_m_row, Q, Qt, QtH, QtHQ
        return H




def _mode_direction_by_root_from_Hact(H: torch.Tensor,
                                      coords_bohr: np.ndarray,
                                      atomic_numbers: List[int],
                                      masses_au_t: torch.Tensor,
                                      active_idx: List[int],
                                      device: torch.device,
                                      root: int = 0) -> Tuple[np.ndarray, float]:
    """
    TS direction from the *active* Hessian block. Mass-weighting/TR are done in the
    active space. Result is embedded back to full 3N in Cartesian space.
    """
    with torch.no_grad():
        N = len(atomic_numbers)
        coords_act = torch.as_tensor(coords_bohr.reshape(-1, 3)[active_idx, :], dtype=H.dtype, device=device)
        masses_act_au = masses_au_t[active_idx].to(device=device, dtype=H.dtype)
        # mass-weight + TR in active space
        Hmw = H.clone()
        _mw_tr_project_active_inplace(Hmw, coords_act, masses_act_au)

        # eigenvector for requested root
        if int(root) == 0:
            try:
                w, V = torch.lobpcg(Hmw, k=1, largest=False)
                u_mw = V[:, 0]
                omega2 = w[0]
            except Exception:
                vals, vecs = torch.linalg.eigh(Hmw)
                pick = int(torch.argmin(vals).item())
                u_mw = vecs[:, pick]
                omega2 = vals[pick]
                del vals, vecs
        else:
            vals, vecs = torch.linalg.eigh(Hmw)
            neg = (vals < 0)
            neg_inds = torch.nonzero(neg, as_tuple=False).view(-1)
            if neg_inds.numel() == 0:
                pick = int(torch.argmin(vals).item())
            else:
                k = max(0, min(int(root), neg_inds.numel() - 1))
                pick = int(neg_inds[k].item())
            u_mw = vecs[:, pick]
            omega2 = vals[pick]
            del vals, vecs

        # Mass un-weight to Cartesian in the active space, then embed to full 3N
        masses_act_amu = (masses_act_au / AMU2AU).to(dtype=H.dtype, device=device)
        m3 = torch.repeat_interleave(masses_act_amu, 3)
        v_cart_act = u_mw / torch.sqrt(m3)
        v_cart_act = v_cart_act / torch.linalg.norm(v_cart_act)

        full = torch.zeros(3 * N, dtype=H.dtype, device=device)
        mask_dof = _active_mask_dof(N, list(set(range(N)) - set(active_idx)))
        mask_t = torch.as_tensor(mask_dof, dtype=torch.bool, device=device)
        full[mask_t] = v_cart_act
        mode = full.reshape(-1, 3).detach().cpu().numpy()
        freq_cm = _omega2_to_freq_cm(omega2)

        del coords_act, masses_act_au, masses_act_amu, m3, v_cart_act, full, mask_t, Hmw, u_mw, omega2
        if torch.cuda.is_available() and H.is_cuda:
            torch.cuda.empty_cache()
        return mode, freq_cm


def _bofill_update_active(H: torch.Tensor,
                          delta_act: np.ndarray,
                          g_new_act: np.ndarray,
                          g_old_act: np.ndarray,
                          eps: float = 1e-12) -> torch.Tensor:
    """
    Memory-efficient Bofill update on the *active* Cartesian Hessian block.
    Apply symmetric rank-1/2 updates directly in place using only the upper triangle
    index set (and mirror to the lower) to avoid allocating large temporaries.
    No explicit (H+H^T)/2 symmetrization step is performed.
    """
    device = H.device
    dtype = H.dtype

    # as torch vectors
    d = torch.as_tensor(delta_act, dtype=dtype, device=device).reshape(-1)
    g0 = torch.as_tensor(g_old_act, dtype=dtype, device=device).reshape(-1)
    g1 = torch.as_tensor(g_new_act, dtype=dtype, device=device).reshape(-1)
    y = g1 - g0

    # Use current symmetric H for matvec
    Hd = H @ d
    xi = y - Hd

    d_dot_xi = torch.dot(d, xi)
    d_norm2 = torch.dot(d, d)
    xi_norm2 = torch.dot(xi, xi)

    # guards
    if torch.abs(d_dot_xi) > eps:
        denom_ms = d_dot_xi
    else:
        sign = torch.sign(d_dot_xi)
        denom_ms = (sign if sign != 0 else torch.tensor(1.0, device=device)) * eps
    denom_psb_d4 = d_norm2 * d_norm2 if d_norm2 > eps else eps
    denom_psb_d2 = d_norm2 if d_norm2 > eps else eps
    denom_phi = d_norm2 * xi_norm2 if (d_norm2 > eps and xi_norm2 > eps) else (1.0)

    phi = 1.0 - (d_dot_xi * d_dot_xi) / denom_phi
    phi = torch.clamp(phi, 0.0, 1.0)

    # coefficients for rank updates
    alpha = (1.0 - phi) / denom_ms                      # for xi xi^T
    beta = -phi * (d_dot_xi / denom_psb_d4)             # for d d^T
    gamma = phi / denom_psb_d2                          # for d xi^T + xi d^T

    n = H.shape[0]
    iu0, iu1 = torch.triu_indices(n, n, device=device)
    is_diag = (iu0 == iu1)
    off = ~is_diag

    # Diagonal contributions (i == j): alpha*xi_i^2 + beta*d_i^2 + 2*gamma*d_i*xi_i
    if is_diag.any():
        idx = iu0[is_diag]
        diag_inc = (alpha * xi[idx] * xi[idx]
                    + beta * d[idx] * d[idx]
                    + 2.0 * gamma * d[idx] * xi[idx])
        # NOTE: use assignment so that advanced indexing actually updates H
        H[idx, idx] = H[idx, idx] + diag_inc

    # Off-diagonal (i < j): symmetric update
    if off.any():
        i = iu0[off]
        j = iu1[off]
        inc = (alpha * xi[i] * xi[j]
               + beta * d[i] * d[j]
               + gamma * (d[i] * xi[j] + xi[i] * d[j]))
        H[i, j] = H[i, j] + inc
        H[j, i] = H[j, i] + inc

    return H


def _imaginary_mode_indices_and_values(
    freqs_cm: np.ndarray,
    neg_freq_thresh_cm: float,
) -> Tuple[np.ndarray, List[float]]:
    """
    Return indices and values of frequencies regarded as imaginary under threshold.
    """
    thresh = abs(float(neg_freq_thresh_cm))
    neg_idx = np.where(freqs_cm < -thresh)[0]
    ims = [float(freqs_cm[i]) for i in neg_idx]
    return neg_idx, ims


def _echo_imaginary_modes(
    freqs_cm: np.ndarray,
    neg_freq_thresh_cm: float,
) -> np.ndarray:
    """
    Print and return imaginary-mode indices.
    """
    neg_idx, ims = _imaginary_mode_indices_and_values(freqs_cm, neg_freq_thresh_cm)
    click.echo(f"[Imaginary modes] n={len(neg_idx)}  ({ims})")
    return neg_idx


def _write_all_imaginary_modes(
    geom,
    freqs_cm: np.ndarray,
    modes: torch.Tensor,
    neg_freq_thresh_cm: float,
    masses_amu: Sequence[float],
    vib_dir: Path,
    prepared_input: Optional["PreparedInputStructure"] = None,
    ref_pdb: Optional[Path] = None,
    amplitude_ang: float = 0.8,
    n_frames: int = 20,
) -> int:
    """
    Write all imaginary modes as trajectory (and optional PDB) files.
    """
    neg_idx, _ = _imaginary_mode_indices_and_values(freqs_cm, neg_freq_thresh_cm)
    if len(neg_idx) == 0:
        return 0

    vib_dir.mkdir(parents=True, exist_ok=True)
    order = neg_idx[np.argsort(freqs_cm[neg_idx])]

    masses_amu_t = torch.as_tensor(masses_amu, dtype=modes.dtype, device=modes.device)
    m3 = torch.repeat_interleave(masses_amu_t, 3)
    freq_name_counts: Dict[str, int] = {}
    n_written = 0

    for mode_idx in order:
        i = int(mode_idx)
        freq = float(freqs_cm[i])
        mode_mw = modes[i]
        v_cart = (mode_mw / torch.sqrt(m3)).detach().cpu().numpy()
        norm = np.linalg.norm(v_cart)
        if norm <= 0.0:
            click.echo(
                f"[Mode write] WARNING: Skip mode {i} ({freq:+.2f} cm^-1) due to near-zero norm.",
                err=True,
            )
            continue
        v_cart = v_cart / norm

        freq_key = f"{freq:+.2f}"
        freq_name_counts[freq_key] = freq_name_counts.get(freq_key, 0) + 1
        serial = freq_name_counts[freq_key]
        suffix = "" if serial == 1 else f"_{serial:02d}"
        stem = f"imag_{freq:+.2f}cm-1{suffix}"

        out_trj = vib_dir / f"{stem}_trj.xyz"
        out_pdb = vib_dir / f"{stem}.pdb"
        _write_mode_trj_and_pdb(
            geom,
            v_cart,
            out_trj,
            amplitude_ang=amplitude_ang,
            n_frames=n_frames,
            comment=f"imag {freq:+.2f} cm-1",
            ref_pdb=ref_pdb,
            write_pdb=ref_pdb is not None,
            prepared_input=prepared_input,
            out_pdb=out_pdb,
        )
        n_written += 1

    del masses_amu_t, m3
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return n_written


# ===================================================================
#                   HessianDimer (Hessian Guided Dimer)
# ===================================================================

class HessianDimer:
    """
    Hessian Guided Dimer TS search with periodic Hessian updates.

    Extensions in this implementation:
      - `root` parameter: choose which imaginary mode to follow (0 = most negative).
      - Pass-through kwargs: `dimer_kwargs` and `lbfgs_kwargs` to tune internals.
      - Hard cap on total LBFGS steps across segments: `max_total_cycles`.
      - PHVA (active DOF subspace) + TR projection for mode picking,
        respecting `freeze_atoms`. For `root == 0` the implementation prefers LOBPCG.
      - The flatten loop can optionally use a Bofill-updated active Hessian block in the
        active DOF subspace (controlled by `flatten_loop_bofill`), but Bofill is applied
        *only* for the flatten displacements; after each dimer segment in the flatten
        loop, a fresh exact Hessian is recomputed.
      - Only **spatially separated** extra imaginary modes (based on representative
        atoms and a distance cutoff) are used for flattening to avoid overly large
        displacements when clustered imaginary modes are present.
      - UMA calculator kwargs accept `freeze_atoms` and `hessian_calc_mode` and
        default to returning a partial (active) Hessian when applicable.

    """

    def __init__(self,
                 fn: str,
                 out_dir: str = "./result_dimer",
                 thresh_loose: str = "gau_loose",
                 thresh: str = "baker",
                 update_interval_hessian: int = 500,
                 neg_freq_thresh_cm: float = 5.0,
                 flatten_amp_ang: float = 0.10,
                 flatten_max_iter: int = 0,
                 mem: int = 100000,
                 uma_kwargs: Optional[dict] = None,
                 device: str = "auto",
                 dump: bool = False,
                 #
                 root: int = 0,
                 dimer_kwargs: Optional[Dict[str, Any]] = None,
                 lbfgs_kwargs: Optional[Dict[str, Any]] = None,
                 max_total_cycles: int = 10000,
                 #
                 # Multi-mode flatten control
                 flatten_sep_cutoff: float = 0.0,
                 flatten_k: int = 10,
                 flatten_loop_bofill: bool = False,
                 #
                 # Propagate geometry kwargs so freeze-links and YAML geometry overrides
                 # also apply in grad mode.
                 geom_kwargs: Optional[Dict[str, Any]] = None,
                 prepared_input: Optional["PreparedInputStructure"] = None,
                 ) -> None:

        self.fn = fn
        self.out_dir = Path(out_dir); self.out_dir.mkdir(parents=True, exist_ok=True)
        self.vib_dir = self.out_dir / "vib"; self.vib_dir.mkdir(parents=True, exist_ok=True)
        # Use prepared_input.source_path if available (may be overridden by --ref-pdb),
        # otherwise fall back to fn directly.
        _src = (prepared_input.source_path if prepared_input is not None else Path(fn))
        self.ref_pdb: Optional[Path] = _src if _src.suffix.lower() == ".pdb" else None
        self.prepared_input = prepared_input

        self.thresh_loose = thresh_loose
        self.thresh = thresh
        self.update_interval_hessian = int(update_interval_hessian)
        self.neg_freq_thresh_cm = float(neg_freq_thresh_cm)
        self.flatten_amp_ang = float(flatten_amp_ang)
        self.flatten_max_iter = int(flatten_max_iter)
        self.mem = int(mem)
        self.root = int(root)
        self.dimer_kwargs = dict(dimer_kwargs or {})
        self.lbfgs_kwargs = dict(lbfgs_kwargs or {})
        self.max_total_cycles = int(max_total_cycles)

        # multi-mode flatten controls
        self.flatten_sep_cutoff = float(flatten_sep_cutoff)
        self.flatten_k = int(flatten_k)
        self.flatten_loop_bofill = bool(flatten_loop_bofill)

        # Track total cycles globally across all loops/segments
        self._cycles_spent = 0

        # UMA settings
        self.uma_kwargs = dict(charge=0, spin=1, model="uma-s-1p1",
                               task_name="omol", device="auto") if uma_kwargs is None else dict(uma_kwargs)

        # Geometry & masses (use provided geom kwargs so freeze_atoms etc. apply)
        gkw = dict(geom_kwargs or {})
        coord_type = gkw.pop("coord_type", GEOM_KW_DEFAULT["coord_type"])
        self.geom = geom_loader(fn, coord_type=coord_type, **gkw)
        self.masses_amu = _safe_masses_amu(self.geom.atomic_numbers)
        self.masses_au_t = torch.as_tensor(self.masses_amu * AMU2AU, dtype=torch.float32)

        # Preserve freeze list (for PHVA)
        self.freeze_atoms: List[int] = list(gkw.get("freeze_atoms", [])) if "freeze_atoms" in gkw else []

        # Device
        self.device = _torch_device(device)
        self.masses_au_t = self.masses_au_t.to(self.device)

        # temp file for Dimer orientation (N_raw)
        self.mode_path = self.out_dir / ".dimer_mode.dat"

        self.dump = bool(dump)
        self.optim_all_path = self.out_dir / "optimization_all_trj.xyz"

        self._raw_hessian_cache_cpu: Optional[torch.Tensor] = None
        self._raw_hessian_coords_cpu: Optional[np.ndarray] = None

    def _cache_raw_hessian_cpu(self, H: torch.Tensor) -> None:
        """
        Cache the *raw* Hessian on CPU for the current geometry.
        IMPORTANT:
          - Must be called before any in-place mass-weight / TR projection happens.
          - This cache is CPU-only by design (no persistent GPU cache).
        """
        self._raw_hessian_cache_cpu = H.detach().cpu().clone()
        self._raw_hessian_coords_cpu = self.geom.cart_coords.copy()

    def _sync_geom_hessian_cache(self, H: torch.Tensor) -> None:
        """
        Share a full-space raw Hessian with Geometry-level cache so RS-I-RFO/RFO
        can reuse the same Hessian on unchanged coordinates.

        Uses the CPU cache if available to avoid an extra GPU clone.
        """
        try:
            n_dof = int(self.geom.cart_coords.size)
            if tuple(H.shape) != (n_dof, n_dof):
                return
            # Prefer the CPU cache to avoid duplicating GPU memory
            if self._raw_hessian_cache_cpu is not None:
                H_cache = self._raw_hessian_cache_cpu
            else:
                H_cache = H.detach().cpu().clone()
            if hasattr(self.geom, "within_partial_hessian"):
                self.geom.within_partial_hessian = None
            self.geom.cart_hessian = H_cache
        except Exception:
            # Cache hand-off is best-effort; optimization must continue even if it fails.
            pass

    def _reuse_cached_hessian(self) -> Optional[torch.Tensor]:
        """
        If the cached geometry matches the current geometry, return the cached *raw*
        Hessian transferred to self.device. No persistent GPU caching is performed.

        Note:
          - On CPU, `.to('cpu')` can share storage; we clone to protect cache against
            later in-place ops (e.g., mass-weight/TR projection).
        """
        if self._raw_hessian_cache_cpu is None or self._raw_hessian_coords_cpu is None:
            return None
        if not np.array_equal(self.geom.cart_coords, self._raw_hessian_coords_cpu):
            return None
        H_dev = self._raw_hessian_cache_cpu.to(self.device)
        if self.device.type == "cpu":
            H_dev = H_dev.clone()
        return H_dev

    def _calc_full_hessian_cached(self, allow_reuse: bool) -> torch.Tensor:
        """
        Compute an exact Hessian, caching a CPU copy of the *raw* Hessian.
        If allow_reuse is True and the current geometry matches the cached one,
        reuse the cached Hessian (CPU→device transfer) instead of recalculating.
        """
        if allow_reuse:
            cached = self._reuse_cached_hessian()
            if cached is not None:
                click.echo("[Hessian] Reusing cached raw Hessian (0-step convergence).")
                self._sync_geom_hessian_cache(cached)
                return cached
        H = _calc_full_hessian_torch(self.geom, self.uma_kwargs, self.device)
        self._cache_raw_hessian_cpu(H)
        self._sync_geom_hessian_cache(H)
        return H

    # ----- One dimer segment for up to n_steps; returns (steps_done, converged) -----
    def _dimer_segment(self, threshold: str, n_steps: int) -> Tuple[int, bool]:
        # Dimer calculator using current mode as initial N
        calc_sp = create_calculator(**self.uma_kwargs)

        # Merge user dimer kwargs (but enforce N_raw & write_orientations)
        dimer_kwargs = dict(self.dimer_kwargs)
        dimer_kwargs.update({
            "calculator": calc_sp,
            "N_raw": str(self.mode_path),
            "write_orientations": False,
            "seed": 0,
            "mem": self.mem,
        })
        dimer = Dimer(**dimer_kwargs)

        self.geom.set_calculator(dimer)

        # LBFGS kwargs: enforce thresh/max_cycles/out_dir/dump; allow others
        lbfgs_kwargs = dict(self.lbfgs_kwargs)
        lbfgs_kwargs.update({
            "max_cycles": n_steps,
            "thresh": threshold,
            "out_dir": str(self.out_dir),
            "dump": self.dump,
        })
        opt = LBFGS(self.geom, **lbfgs_kwargs)
        opt.run()
        steps = opt.cur_cycle + 1
        converged = opt.is_converged
        self.geom.set_calculator(None)

        # Append to concatenated trajectory if dump enabled
        if self.dump:
            part_path = self.out_dir / "optimization_trj.xyz"
            if part_path.exists():
                with self.optim_all_path.open("a", encoding="utf-8") as f_all, \
                     part_path.open("r", encoding="utf-8") as f_part:
                    f_all.write(f_part.read())
        return steps, converged

    # ----- Loop dimer segments, updating mode from Hessian every interval -----
    def _dimer_loop(self, threshold: str) -> Tuple[int, bool]:
        """
        Run multiple LBFGS segments separated by periodic Hessian-based mode updates.
        Consumes from the global cycle budget self.max_total_cycles.

        Returns:
          (steps_in_this_call, zero_step_converged)
        where `zero_step_converged` is True iff the loop terminated by convergence
        without changing the geometry (i.e., 0-step convergence).
        """
        steps_in_this_call = 0
        zero_step_converged = False
        while True:
            remaining_global = max(0, self.max_total_cycles - self._cycles_spent)
            if remaining_global == 0:
                break
            steps_this = min(self.update_interval_hessian, remaining_global)
            coords_before = self.geom.cart_coords.copy()
            steps, ok = self._dimer_segment(threshold, steps_this)
            self._cycles_spent += steps
            steps_in_this_call += steps
            if ok:
                if np.array_equal(self.geom.cart_coords, coords_before):
                    zero_step_converged = True
                break
            # If budget exhausted after this segment, stop before doing a Hessian update
            if (self.max_total_cycles - self._cycles_spent) <= 0:
                break
            # Update mode from Hessian (respect freeze atoms via PHVA)
            H = self._calc_full_hessian_cached(allow_reuse=False)
            N = len(self.geom.atomic_numbers)
            coords_bohr_t = torch.as_tensor(self.geom.cart_coords.reshape(-1, 3),
                                            dtype=H.dtype, device=H.device)
            # full vs active-block Hessian
            if H.size(0) == 3 * N:
                mode_xyz, mode_freq_cm = _mode_direction_by_root(
                    H, coords_bohr_t, self.masses_au_t,
                    root=self.root, freeze_idx=self.freeze_atoms if len(self.freeze_atoms) > 0 else None
                )
            else:
                # partial (active) Hessian returned by UMA
                active_idx = _active_indices(N, self.freeze_atoms if len(self.freeze_atoms) > 0 else [])
                mode_xyz, mode_freq_cm = _mode_direction_by_root_from_Hact(
                    H, self.geom.cart_coords.reshape(-1, 3), self.geom.atomic_numbers,
                    self.masses_au_t, active_idx, self.device, root=self.root
                )
            click.echo(f"[Dimer mode] root={self.root} freq={mode_freq_cm:+.2f} cm^-1")
            np.savetxt(self.mode_path, mode_xyz, fmt="%.12f")
            del H, coords_bohr_t, mode_xyz, mode_freq_cm
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        return steps_in_this_call, zero_step_converged

    # ----- helpers for flatten selection -----
    def _representative_atoms_for_mode(self, mode_vec: torch.Tensor) -> np.ndarray:
        """
        Return indices of atoms with the largest displacements for a given mode.
        """
        vec = mode_vec.view(-1, 3)
        norms = torch.linalg.norm(vec, dim=1)
        k = min(self.flatten_k, vec.shape[0])
        if k <= 0:
            return np.zeros(0, dtype=int)
        topk = torch.topk(norms, k=k, largest=True)
        return topk.indices.detach().cpu().numpy()

    def _select_flatten_targets(self,
                                freqs_cm: np.ndarray,
                                modes: torch.Tensor) -> List[int]:
        """
        Select a subset of imaginary modes to flatten:
          * exclude the primary (TS) mode selected by `root`,
          * from the remaining imaginary modes, greedily pick modes whose
            representative atoms are separated by at least `flatten_sep_cutoff`
            (Å) from previously selected modes.
        """
        neg_idx_all = np.where(freqs_cm < -abs(self.neg_freq_thresh_cm))[0]
        if len(neg_idx_all) <= 1:
            return []

        # sort imaginary modes (most negative first)
        order = np.argsort(freqs_cm[neg_idx_all])
        sorted_neg = neg_idx_all[order]

        # primary mode index in freqs_cm (TS mode to keep)
        root_clamped = max(0, min(self.root, len(order) - 1))
        primary_idx = sorted_neg[root_clamped]

        # candidates = all other imaginary modes
        candidates = [int(i) for i in sorted_neg if int(i) != int(primary_idx)]
        if not candidates:
            return []

        # atomic coordinates in Å
        coords_ang = torch.as_tensor(
            self.geom.cart_coords.reshape(-1, 3) * BOHR2ANG,
            dtype=modes.dtype,
            device=modes.device,
        )

        targets: List[int] = []
        reps_list: List[np.ndarray] = []

        for idx in candidates:
            rep = self._representative_atoms_for_mode(modes[idx])
            if rep.size == 0:
                continue
            rep_coords = coords_ang[rep]  # (k, 3)
            if not reps_list:
                targets.append(idx)
                reps_list.append(rep)
                continue

            # check distance to all previously selected representative sets
            accept = True
            for prev_rep in reps_list:
                prev_coords = coords_ang[prev_rep]  # (k, 3)
                dmat = torch.cdist(rep_coords, prev_coords)
                min_dist = float(torch.min(dmat).item())
                if min_dist < self.flatten_sep_cutoff:
                    accept = False
                    break
            if accept:
                targets.append(idx)
                reps_list.append(rep)

        return targets

    def _flatten_once_with_modes(self, freqs_cm: np.ndarray, modes: torch.Tensor) -> bool:
        """
        Flatten using precomputed modes (mass-weighted, embedded).

        Only spatially separated extra imaginary modes are used, following a
        greedy selection similar to the PartialHessianDimer implementation.
        Modes are applied sequentially from the current geometry, so that each
        mode's displacement is decided from a 1D energy scan along that mode.
        """
        neg_idx_all = np.where(freqs_cm < -abs(self.neg_freq_thresh_cm))[0]
        if len(neg_idx_all) <= 1:
            return False

        # choose targets based on spatial separation
        targets = self._select_flatten_targets(freqs_cm, modes)
        if not targets:
            return False

        # mass scaling (C moves exactly flatten_amp_ang Å)
        mass_scale = np.sqrt(12.011 / self.masses_amu)[:, None]
        amp_bohr = self.flatten_amp_ang / BOHR2ANG

        # ensure energy reference is set up
        E_ref = _calc_energy(self.geom, self.uma_kwargs)

        # work in Bohr coordinates
        for idx in targets:
            # mode is currently mass-weighted & embedded to 3N
            v_mw = modes[idx].detach().cpu().numpy().reshape(-1, 3)
            # convert to Cartesian (Å-scale amplitude, but coords are Bohr)
            m3 = np.repeat(self.masses_amu, 3).reshape(-1, 3)
            v_cart = v_mw / np.sqrt(m3)
            v_cart /= np.linalg.norm(v_cart)

            disp = amp_bohr * mass_scale * v_cart  # Bohr
            ref = self.geom.cart_coords.reshape(-1, 3)

            plus = ref + disp
            minus = ref - disp

            self.geom.coords = plus.reshape(-1)
            E_plus = _calc_energy(self.geom, self.uma_kwargs)

            self.geom.coords = minus.reshape(-1)
            E_minus = _calc_energy(self.geom, self.uma_kwargs)

            # keep lower-energy side and continue from there
            use_plus = E_plus <= E_minus
            self.geom.coords = (plus if use_plus else minus).reshape(-1)
            E_keep = E_plus if use_plus else E_minus
            delta_e = E_keep - E_ref
            click.echo(
                f"[Flatten] mode={idx} freq={freqs_cm[idx]:+.2f} cm^-1 "
                f"E_disp={E_keep:.8f} Ha \u0394E={delta_e:+.8f} Ha"
            )

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return True

    # ----- Run full procedure -----
    def run(self) -> None:
        if self.dump and self.optim_all_path.exists():
            self.optim_all_path.unlink()

        N = len(self.geom.atomic_numbers)
        active_idx = _active_indices(N, self.freeze_atoms if len(self.freeze_atoms) > 0 else [])
        mask_dof = _active_mask_dof(N, self.freeze_atoms if len(self.freeze_atoms) > 0 else [])

        # (1) Initial Hessian → pick direction by `root`
        H = self._calc_full_hessian_cached(allow_reuse=False)
        coords_bohr_t = torch.as_tensor(self.geom.cart_coords.reshape(-1, 3),
                                        dtype=H.dtype, device=H.device)

        if H.size(0) == 3 * N:
            mode_xyz, mode_freq_cm = _mode_direction_by_root(
                H, coords_bohr_t, self.masses_au_t,
                root=self.root, freeze_idx=self.freeze_atoms if len(self.freeze_atoms) > 0 else None
            )
        else:
            click.echo("[CHECK] Using active-block Hessian from UMA (partial Hessian). Skip full-space TR check.")
            mode_xyz, mode_freq_cm = _mode_direction_by_root_from_Hact(
                H, self.geom.cart_coords.reshape(-1, 3), self.geom.atomic_numbers,
                self.masses_au_t, active_idx, self.device, root=self.root
            )
        click.echo(f"[Dimer mode] root={self.root} freq={mode_freq_cm:+.2f} cm^-1")
        np.savetxt(self.mode_path, mode_xyz, fmt="%.12f")
        del mode_xyz, coords_bohr_t, H, mode_freq_cm
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        # (2) Loose loop (or initial pass)
        if self.root != 0:
            click.echo("[WARNING] root != 0. Using this root only for the initial dimer loop.", err=True)
            click.echo(f"Dimer loop with initial direction from mode {self.root}...")
            self.root = 0
            self.thresh_loose = self.thresh
        else:
            click.echo("Loose dimer loop...")

        _steps_loose, zero_step_loose = self._dimer_loop(self.thresh_loose)

        # (3) Update mode & normal loop
        H = self._calc_full_hessian_cached(allow_reuse=zero_step_loose)
        coords_bohr_t = torch.as_tensor(self.geom.cart_coords.reshape(-1, 3),
                                        dtype=H.dtype, device=H.device)
        if H.size(0) == 3 * N:
            mode_xyz, mode_freq_cm = _mode_direction_by_root(
                H, coords_bohr_t, self.masses_au_t,
                root=self.root, freeze_idx=self.freeze_atoms if len(self.freeze_atoms) > 0 else None
            )
        else:
            click.echo("[CHECK] Using active-block Hessian from UMA (partial Hessian). Skip full-space TR check.")
            mode_xyz, mode_freq_cm = _mode_direction_by_root_from_Hact(
                H, self.geom.cart_coords.reshape(-1, 3), self.geom.atomic_numbers,
                self.masses_au_t, active_idx, self.device, root=self.root
            )
        click.echo(f"[Dimer mode] root={self.root} freq={mode_freq_cm:+.2f} cm^-1")
        np.savetxt(self.mode_path, mode_xyz, fmt="%.12f")
        del mode_xyz, coords_bohr_t, H, mode_freq_cm
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        click.echo("Normal dimer loop...")
        _steps_normal, zero_step_normal = self._dimer_loop(self.thresh)

        # (4) Flatten loop — exact Hessian each iteration & optional Bofill update
        if self.flatten_max_iter > 0:
            if self.flatten_loop_bofill:
                click.echo("Flatten loop with Bofill-updated active Hessian (flatten displacements only)...")
            else:
                click.echo("Flatten loop without Bofill updates (flatten displacements only)...")

            # (4.1) Evaluate one exact Hessian at the loop start and prepare the active block
            H = self._calc_full_hessian_cached(allow_reuse=zero_step_normal)
            if H.size(0) == 3 * N:
                # full → extract active
                H = _extract_active_block(H, mask_dof)  # torch (3N_act,3N_act)
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            # Flatten iterations with Hessian updates
            for it in range(self.flatten_max_iter):
                if (self.max_total_cycles - self._cycles_spent) <= 0:
                    break

                # (a) Frequencies & modes from the active Hessian
                freqs_cm, modes_embedded = _frequencies_cm_and_modes(
                    H.clone(),
                    self.geom.atomic_numbers,
                    self.geom.cart_coords.reshape(-1, 3),
                    self.device,
                    freeze_idx=self.freeze_atoms if len(self.freeze_atoms) > 0 else None,
                )
                n_imag = int(np.sum(freqs_cm < -abs(self.neg_freq_thresh_cm)))
                ims = [float(x) for x in freqs_cm if x < -abs(self.neg_freq_thresh_cm)]
                click.echo(f"[Imaginary modes] n={n_imag}  ({ims})")
                if n_imag <= 1:
                    break

                # (b) Flatten other imaginary modes
                if self.flatten_loop_bofill:
                    x_before_flat = self.geom.cart_coords.copy().reshape(-1)
                    g_before_flat = _calc_gradient(self.geom, self.uma_kwargs).reshape(-1)

                did_flatten = self._flatten_once_with_modes(freqs_cm, modes_embedded)
                if not did_flatten:
                    break

                if self.flatten_loop_bofill:
                    x_after_flat = self.geom.cart_coords.copy().reshape(-1)
                    g_after_flat = _calc_gradient(self.geom, self.uma_kwargs).reshape(-1)

                # (c) Bofill update using UMA gradients across the flatten displacement
                if self.flatten_loop_bofill:
                    delta_flat_full = x_after_flat - x_before_flat
                    delta_flat_act = delta_flat_full[mask_dof]
                    g_old_act = g_before_flat[mask_dof]
                    g_new_act = g_after_flat[mask_dof]
                    H = _bofill_update_active(H, delta_flat_act, g_new_act, g_old_act)

                # (d) Refresh dimer direction
                mode_xyz, mode_freq_cm = _mode_direction_by_root_from_Hact(
                    H, self.geom.cart_coords.reshape(-1, 3), self.geom.atomic_numbers,
                    self.masses_au_t, active_idx, self.device, root=self.root
                )
                click.echo(f"[Dimer mode] root={self.root} freq={mode_freq_cm:+.2f} cm^-1")
                np.savetxt(self.mode_path, mode_xyz, fmt="%.12f")

                # (e) Re-optimize with Dimer (consumes global cycle budget)
                _steps_flat, zero_step_flat = self._dimer_loop(self.thresh)

                # (f) After dimer optimization, recompute an exact Hessian for the next iteration
                H = self._calc_full_hessian_cached(allow_reuse=zero_step_flat)
                if H.size(0) == 3 * N:
                    H = _extract_active_block(H, mask_dof)
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

        # (5) Final outputs
        final_xyz = self.out_dir / "final_geometry.xyz"
        atoms_final = Atoms(self.geom.atoms, positions=(self.geom.cart_coords.reshape(-1, 3) * BOHR2ANG), pbc=False)
        write(final_xyz, atoms_final)

        # Final Hessian → imaginary mode animation
        H = self._calc_full_hessian_cached(allow_reuse=True)
        freqs_cm, modes = _frequencies_cm_and_modes(
            H,
            self.geom.atomic_numbers,
            self.geom.cart_coords.reshape(-1, 3),
            self.device,
            freeze_idx=self.freeze_atoms if len(self.freeze_atoms) > 0 else None,
        )

        del H
        neg_idx = _echo_imaginary_modes(freqs_cm, self.neg_freq_thresh_cm)
        if len(neg_idx) == 0:
            click.echo("[INFO] No imaginary mode found at the end (ν_min = %.2f cm^-1)." % (freqs_cm.min(),))
            del modes
        else:
            _write_all_imaginary_modes(
                self.geom,
                freqs_cm,
                modes,
                self.neg_freq_thresh_cm,
                self.masses_amu,
                self.vib_dir,
                prepared_input=self.prepared_input,
                ref_pdb=self.ref_pdb,
                amplitude_ang=0.8,
                n_frames=20,
            )
            del modes

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        click.echo(f"[DONE] Saved final geometry → {final_xyz}")
        click.echo(f"[DONE] Mode files → {self.vib_dir}")


# ===================================================================
#                         Defaults for CLI
# ===================================================================

# Geometry defaults
GEOM_KW = dict(GEOM_KW_DEFAULT)

CALC_KW = dict(UMA_CALC_KW)

# Optimizer base (common) — used by both RSIRFO and the inner LBFGS of Hessian Guided Dimer
OPT_BASE_KW_LOCAL = {**OPT_BASE_KW, "out_dir": OUT_DIR_TSOPT, "thresh": "baker"}

# hessian_dimer_KW kept as alias for backward compatibility with internal references
hessian_dimer_KW = HESSIAN_DIMER_CLI_KW

def _build_rsirfo_kwargs(
    opt_cfg: Dict[str, Any],
    rsirfo_cfg: Dict[str, Any],
    out_dir_path: Path,
) -> Dict[str, Any]:
    """Return the exact kwargs used for RSIRFO after normalization."""
    rs_args = dict(rsirfo_cfg)
    opt_base = dict(opt_cfg)
    opt_base["out_dir"] = str(out_dir_path)
    rs_args["out_dir"] = str(out_dir_path)

    roots = rs_args.get("roots", None)
    root_single = rs_args.get("root", None)
    if roots is None and root_single is not None:
        roots = [int(root_single)]
    if roots is None:
        roots = [0]
    rs_args["roots"] = [int(x) for x in roots]
    rs_args.pop("root", None)

    # Keep top-level opt knobs (max_cycles/dump/thresh/...) authoritative unless
    # rsirfo.* explicitly overrides them to a non-default value.
    for k in list(rs_args.keys()):
        if k in opt_base and k in RSIRFO_KW and rs_args[k] == RSIRFO_KW[k]:
            rs_args.pop(k, None)

    rsirfo_kwargs = {**opt_base, **rs_args}

    # RSIRFO ignores these DIIS-related knobs; drop them for clarity.
    for diis_kw in ("gediis", "gdiis", "gdiis_thresh", "gediis_thresh", "gdiis_test_direction", "adapt_step_func"):
        rsirfo_kwargs.pop(diis_kw, None)

    return rsirfo_kwargs


# ===================================================================
#                            CLI
# ===================================================================

@click.command(
    help="Transition state optimization (Dimer or RS-I-RFO).",
    context_settings={"help_option_names": ["-h", "--help"]},
)
@click.option(
    "-i", "--input",
    "input_path",
    type=click.Path(path_type=Path, exists=True, dir_okay=False),
    required=True,
    help="Input structure file (.pdb, .xyz, _trj.xyz, ...).",
)
@click.option(
    "-q",
    "--charge",
    type=int,
    required=False,
    help=(
        "Total charge. Required for non-.gjf inputs unless --ligand-charge is provided "
        "(PDB inputs or XYZ/GJF with --ref-pdb)."
    ),
)
@click.option(
    "--workers",
    type=int,
    default=CALC_KW["workers"],
    show_default=True,
    help="MLIP predictor workers; >1 spawns a parallel predictor (Hessian computation not supported with workers>1).",
)
@click.option(
    "--workers-per-node",
    "workers_per_node",
    type=int,
    default=CALC_KW["workers_per_node"],
    show_default=True,
    help="Workers per node when using a parallel MLIP predictor (workers>1).",
)
@click.option(
    "-l",
    "--ligand-charge",
    type=str,
    default=None,
    show_default=False,
    help=(
        "Total charge or per-resname mapping (e.g., GPP:-3,SAM:1) used to derive charge "
        "when -q is omitted (requires PDB input or --ref-pdb)."
    ),
)
@click.option("-m", "--multiplicity", "spin", type=int, default=None, show_default=False, help="Spin multiplicity (2S+1).")
@click.option(
    "--freeze-links/--no-freeze-links",
    "freeze_links",
    default=True,
    show_default=True,
    help="Freeze parent atoms of link hydrogens (PDB input or XYZ/GJF with --ref-pdb).",
)
@click.option(
    "--convert-files/--no-convert-files",
    "convert_files",
    default=True,
    show_default=True,
    help="Convert XYZ/TRJ outputs into PDB/GJF companions based on the input format.",
)
@click.option(
    "--ref-pdb",
    type=click.Path(path_type=Path, exists=True, dir_okay=False),
    default=None,
    help="Reference PDB topology to use when the input is XYZ/GJF (keeps XYZ coordinates).",
)
@click.option("--max-cycles", type=int, default=10000, show_default=True, help="Maximum number of optimization cycles.")
@click.option(
    "--flatten/--no-flatten",
    "flatten",
    default=False,
    show_default=True,
    help="Enable the extra-imaginary-mode flattening loop (grad: dimer loop, hess: post-RSIRFO).",
)
@click.option(
    "--opt-mode",
    type=click.Choice(["grad", "hess", "dimer", "rsirfo"], case_sensitive=False),
    default="hess",
    show_default=True,
    help="TS optimizer: 'grad'/'dimer' → Hessian Guided Dimer; 'hess'/'rsirfo' → RS-I-RFO.",
)
@click.option(
    "--dump/--no-dump",
    default=False,
    show_default=True,
    help="Write optimization trajectory to 'tsopt_trj.xyz' in the output directory.",
)
@click.option("-o", "--out-dir", type=str, default=OUT_DIR_TSOPT, show_default=True, help="Output directory.")
@click.option(
    "--thresh",
    type=str,
    default=None,
    show_default=False,
    help=(
        "Convergence preset for the active optimizer "
        "(gau_loose|gau|gau_tight|gau_vtight|baker|never). "
        "Defaults to 'baker' when not provided."
    ),
)
@click.option(
    "--config",
    "config_yaml",
    type=click.Path(path_type=Path, exists=True, dir_okay=False),
    default=None,
    help="Base YAML configuration file applied before explicit CLI options.",
)
@click.option(
    "--show-config/--no-show-config",
    "show_config",
    default=False,
    show_default=True,
    help="Print resolved configuration and continue execution.",
)
@click.option(
    "--dry-run/--no-dry-run",
    "dry_run",
    default=False,
    show_default=True,
    help="Validate options and print the execution plan without running TS optimization.",
)
@click.option(
    "--hessian-calc-mode",
    type=click.Choice(["FiniteDifference", "Analytical"], case_sensitive=False),
    default=None,
    help="Choose MLIP Hessian evaluation mode (used unless YAML sets calc.hessian_calc_mode). Defaults to 'FiniteDifference'.",
)
@click.option("-b", "--backend", type=click.Choice(["uma", "orb", "mace", "aimnet2"]), default="uma",
              show_default=True, help="MLIP backend.")
@click.option("--solvent", default="none", show_default=True,
              help="Implicit solvent name for xTB correction (e.g. 'water'). 'none' to disable.")
@click.option("--solvent-model", "solvent_model", default="alpb", type=click.Choice(["alpb", "cpcmx"]),
              show_default=True, help="xTB solvent model.")
@click.pass_context
def cli(
    ctx: click.Context,
    input_path: Path,
    charge: Optional[int],
    ligand_charge: Optional[str],
    workers: int,
    workers_per_node: int,
    spin: Optional[int],
    freeze_links: bool,
    convert_files: bool,
    ref_pdb: Optional[Path],
    max_cycles: int,
    flatten: bool,
    opt_mode: str,
    dump: bool,
    out_dir: str,
    thresh: Optional[str],
    config_yaml: Optional[Path],
    show_config: bool,
    dry_run: bool,
    hessian_calc_mode: Optional[str],
    backend: str,
    solvent: str,
    solvent_model: str,
) -> None:
    config_yaml, override_yaml, used_legacy_yaml = resolve_yaml_sources(
        config_yaml=config_yaml,
        override_yaml=None,
        args_yaml_legacy=None,
    )
    merged_yaml_cfg, config_layer_cfg, override_layer_cfg = load_merged_yaml_cfg(
        config_yaml=config_yaml,
        override_yaml=None,
    )

    set_convert_file_enabled(convert_files)
    with prepared_cli_input(
        input_path,
        ref_pdb=ref_pdb,
        charge=charge,
        spin=spin,
        ligand_charge=ligand_charge,
        prefix="[tsopt]",
    ) as (prepared_input, resolved_charge, resolved_spin):
        geom_input_path = prepared_input.geom_path
        source_path = prepared_input.source_path
        time_start = time.perf_counter()

        # --------------------------
        # 1) Assemble configuration (defaults < config < CLI(explicit) < override)
        # --------------------------
        geom_cfg = dict(GEOM_KW)
        calc_cfg = dict(CALC_KW)
        opt_cfg = dict(OPT_BASE_KW_LOCAL)
        import copy
        simple_cfg = copy.deepcopy(hessian_dimer_KW)
        # tsopt default: keep flatten loop off unless enabled via config or explicit --flatten.
        simple_cfg["flatten_max_iter"] = 0
        rsirfo_cfg = dict(RSIRFO_KW)

        apply_yaml_overrides(
            config_layer_cfg,
            [
                (geom_cfg, (("geom",),)),
                (calc_cfg, (("calc",),)),
                (opt_cfg, (("opt",),)),
                (simple_cfg, (("hessian_dimer",),)),
                (rsirfo_cfg, (("rsirfo",),)),
            ],
        )

        charge_value = calc_cfg.get("charge", resolved_charge)
        if charge_value is None:
            charge_value = resolved_charge
        calc_cfg["charge"] = int(charge_value)
        if cli_param_overridden(ctx, "charge"):
            calc_cfg["charge"] = int(resolved_charge)

        spin_value = calc_cfg.get("spin", resolved_spin)
        if spin_value is None:
            spin_value = resolved_spin
        calc_cfg["spin"] = int(spin_value)
        if cli_param_overridden(ctx, "spin"):
            calc_cfg["spin"] = int(resolved_spin)

        if cli_param_overridden(ctx, "workers"):
            calc_cfg["workers"] = int(workers)
        if cli_param_overridden(ctx, "workers_per_node"):
            calc_cfg["workers_per_node"] = int(workers_per_node)
        if cli_param_overridden(ctx, "backend"):
            calc_cfg["backend"] = backend
        if cli_param_overridden(ctx, "solvent"):
            calc_cfg["solvent"] = solvent
        if cli_param_overridden(ctx, "solvent_model"):
            calc_cfg["solvent_model"] = solvent_model
        if cli_param_overridden(ctx, "max_cycles"):
            opt_cfg["max_cycles"] = int(max_cycles)
        if cli_param_overridden(ctx, "dump"):
            opt_cfg["dump"] = bool(dump)
        if cli_param_overridden(ctx, "out_dir"):
            opt_cfg["out_dir"] = out_dir
        if cli_param_overridden(ctx, "thresh") and thresh is not None:
            opt_cfg["thresh"] = str(thresh)
            simple_cfg["thresh"] = str(thresh)
            rsirfo_cfg["thresh"] = str(thresh)
        if cli_param_overridden(ctx, "hessian_calc_mode") and hessian_calc_mode is not None:
            calc_cfg["hessian_calc_mode"] = str(hessian_calc_mode)
        if cli_param_overridden(ctx, "flatten"):
            if flatten:
                # --flatten explicitly enables flattening even when defaults/config disable it.
                default_flatten_iter = int(hessian_dimer_KW.get("flatten_max_iter", 0))
                if int(simple_cfg.get("flatten_max_iter", 0)) <= 0 and default_flatten_iter > 0:
                    simple_cfg["flatten_max_iter"] = default_flatten_iter
            else:
                simple_cfg["flatten_max_iter"] = 0

        apply_yaml_overrides(
            override_layer_cfg,
            [
                (geom_cfg, (("geom",),)),
                (calc_cfg, (("calc",),)),
                (opt_cfg, (("opt",),)),
                (simple_cfg, (("hessian_dimer",),)),
                (rsirfo_cfg, (("rsirfo",),)),
            ],
        )

        if "print_every" in opt_cfg:
            try:
                opt_print_every = int(opt_cfg["print_every"])
                if opt_print_every >= 1:
                    rsirfo_cfg.setdefault("print_every", opt_print_every)
                    simple_cfg.setdefault("lbfgs", {})
                    if isinstance(simple_cfg["lbfgs"], dict):
                        simple_cfg["lbfgs"].setdefault("print_every", opt_print_every)
            except (TypeError, ValueError):
                pass

        # Normalize freeze_atoms and optionally add link-parent indices for PDB inputs
        resolve_freeze_atoms(geom_cfg, source_path, freeze_links)

        # Pass freeze_atoms from geom → calc (so UMA knows active DOF for FD Hessian)
        if "freeze_atoms" in geom_cfg:
            calc_cfg["freeze_atoms"] = list(geom_cfg.get("freeze_atoms", []))
        calc_cfg["return_partial_hessian"] = True

        kind = normalize_choice(
            opt_mode,
            param="--opt-mode",
            alias_groups=TSOPT_MODE_ALIASES,
            allowed_hint="grad|hess|dimer|rsirfo",
        )
        out_dir_path = Path(opt_cfg["out_dir"]).resolve()

        # Pretty-print config summary
        click.echo(pretty_block("geom", format_geom_for_echo(geom_cfg)))
        click.echo(pretty_block("calc", format_geom_for_echo(calc_cfg)))
        echo_opt = {**opt_cfg, "out_dir": str(out_dir_path)}
        click.echo(pretty_block("opt", echo_opt))
        if kind == "dimer":
            sd_cfg_for_echo = dict(simple_cfg)
            sd_cfg_for_echo["dimer"] = dict(simple_cfg.get("dimer", {}))
            sd_cfg_for_echo["lbfgs"] = strip_inherited_keys(
                dict(simple_cfg.get("lbfgs", {})),
                echo_opt,
                mode="same",
            )
            click.echo(pretty_block("hessian_dimer", sd_cfg_for_echo))
        else:
            rsirfo_kwargs_for_echo = _build_rsirfo_kwargs(opt_cfg, rsirfo_cfg, out_dir_path)
            rsirfo_kwargs_for_echo = strip_inherited_keys(
                rsirfo_kwargs_for_echo,
                echo_opt,
                mode="same",
            )
            click.echo(pretty_block("rsirfo", rsirfo_kwargs_for_echo))

        if show_config:
            click.echo(
                pretty_block(
                    "yaml_layers",
                    {
                        "config": None if config_yaml is None else str(config_yaml),
                        "override_yaml": None if override_yaml is None else str(override_yaml),
                        "merged_keys": sorted(merged_yaml_cfg.keys()),
                    },
                )
            )
        if dry_run:
            click.echo(
                pretty_block(
                    "dry_run_plan",
                    {
                        "input_geometry": str(geom_input_path),
                        "output_dir": str(out_dir_path),
                        "optimizer_mode": str(kind),
                        "convert_files": bool(convert_files),
                        "freeze_links": bool(freeze_links),
                        "flatten": int(simple_cfg.get("flatten_max_iter", 0)) > 0,
                        "will_run_tsopt": True,
                        "will_write_summary": True,
                    },
                )
            )
            click.echo("[dry-run] Validation complete. TS optimization execution was skipped.")
            return

        # --------------------------
        # 2) Prepare geometry dir
        # --------------------------
        out_dir_path.mkdir(parents=True, exist_ok=True)

        # --------------------------
        # 3) Run
        # --------------------------
        try:
            if kind == "dimer":
                # Hessian Guided Dimer runner construction
                uma_kwargs_for_sd = dict(calc_cfg)
                runner = HessianDimer(
                    fn=str(geom_input_path),
                    out_dir=str(out_dir_path),
                    thresh_loose=simple_cfg.get("thresh_loose", "gau_loose"),
                    thresh=simple_cfg.get("thresh", "baker"),
                    update_interval_hessian=int(simple_cfg.get("update_interval_hessian", 500)),
                    neg_freq_thresh_cm=float(simple_cfg.get("neg_freq_thresh_cm", 5.0)),
                    flatten_amp_ang=float(simple_cfg.get("flatten_amp_ang", 0.10)),
                    flatten_max_iter=int(simple_cfg.get("flatten_max_iter", 0)),
                    mem=int(simple_cfg.get("mem", 100000)),
                    uma_kwargs=uma_kwargs_for_sd,
                    device=str(simple_cfg.get("device", calc_cfg.get("device", "auto"))),
                    dump=bool(opt_cfg["dump"]),
                    root=int(simple_cfg.get("root", 0)),
                    dimer_kwargs=dict(simple_cfg.get("dimer", {})),
                    lbfgs_kwargs=dict(simple_cfg.get("lbfgs", {})),
                    max_total_cycles=int(opt_cfg["max_cycles"]),
                    flatten_sep_cutoff=float(simple_cfg.get("flatten_sep_cutoff", 0.0)),
                    flatten_k=int(simple_cfg.get("flatten_k", 10)),
                    flatten_loop_bofill=bool(simple_cfg.get("flatten_loop_bofill", False)),
                    # Propagate geometry settings (freeze_atoms, coord_type, ...) to the HessianDimer runner
                    geom_kwargs=dict(geom_cfg),
                    prepared_input=prepared_input,
                )

                click.echo("====== TS optimization (Hessian Guided Dimer) started ======\n")
                runner.run()
                click.echo("====== TS optimization (Hessian Guided Dimer) finished ======\n")

                needs_pdb = source_path.suffix.lower() == ".pdb"
                needs_gjf = prepared_input.is_gjf
                ref_pdb = source_path.resolve() if needs_pdb else None

                final_xyz = out_dir_path / "final_geometry.xyz"
                if convert_xyz_like_outputs(
                    final_xyz,
                    prepared_input,
                    ref_pdb_path=ref_pdb,
                    out_pdb_path=out_dir_path / "final_geometry.pdb" if needs_pdb else None,
                    out_gjf_path=out_dir_path / "final_geometry.gjf" if needs_gjf else None,
                    context="final geometry",
                ):
                    click.echo("[convert] Wrote 'final_geometry' outputs.")

                if bool(opt_cfg.get("dump", False)) and needs_pdb:
                    all_trj = out_dir_path / "optimization_all_trj.xyz"
                    if all_trj.exists():
                        if convert_xyz_like_outputs(
                            all_trj,
                            prepared_input,
                            ref_pdb_path=ref_pdb,
                            out_pdb_path=out_dir_path / "optimization_all.pdb" if needs_pdb else None,
                            context="optimization trajectory",
                        ):
                            click.echo("[convert] Wrote 'optimization_all' outputs.")
                    else:
                        click.echo("[convert] WARNING: 'optimization_all_trj.xyz' not found; skipping conversion.", err=True)

            else:
                # RS-I-RFO (hess)
                #  - Build geometry now and attach UMA calculator
                coord_type = geom_cfg.get("coord_type", GEOM_KW_DEFAULT["coord_type"])
                coord_kwargs = dict(geom_cfg)
                coord_kwargs.pop("coord_type", None)
                geometry = geom_loader(geom_input_path, coord_type=coord_type, **coord_kwargs)

                calc = create_calculator(**calc_cfg)  # includes freeze_atoms / hessian_calc_mode / partial Hessian
                geometry.set_calculator(calc)

                # Construct RSIRFO optimizer
                rsirfo_kwargs = _build_rsirfo_kwargs(opt_cfg, rsirfo_cfg, out_dir_path)

                optimizer = RSIRFOptimizer(geometry, **rsirfo_kwargs)

                click.echo("====== TS optimization (RS-I-RFO) started ======\n")
                optimizer.run()
                click.echo("====== TS optimization (RS-I-RFO) finished ======\n")
                last_optimizer = optimizer

                # --- RSIRFO: count imaginary modes and optional flatten loop ---
                geometry.set_calculator(None)
                uma_kwargs_for_rsirfo = dict(calc_cfg)
                uma_kwargs_for_rsirfo["out_hess_torch"] = True
                device = _torch_device(calc_cfg.get("device", "auto"))

                def _attach_rsirfo_calc() -> None:
                    geometry.set_calculator(calc)

                def _calc_freqs_and_modes() -> Tuple[np.ndarray, torch.Tensor]:
                    H = _calc_full_hessian_torch(geometry, uma_kwargs_for_rsirfo, device)
                    from .hessian_cache import store as _hess_store
                    freeze_atoms = list(geom_cfg.get("freeze_atoms", []))
                    if freeze_atoms and H.shape[0] < 3 * len(geometry.atomic_numbers):
                        all_dofs = set(range(3 * len(geometry.atomic_numbers)))
                        frozen_dofs = set()
                        for _fi in freeze_atoms:
                            frozen_dofs.update([3 * _fi, 3 * _fi + 1, 3 * _fi + 2])
                        _active_dofs = sorted(all_dofs - frozen_dofs)
                    else:
                        _active_dofs = None
                    _hess_store("ts", H, active_dofs=_active_dofs)
                    freqs_local, modes_local = _frequencies_cm_and_modes(
                        H,
                        geometry.atomic_numbers,
                        geometry.cart_coords.reshape(-1, 3),
                        device,
                        freeze_idx=list(geom_cfg.get("freeze_atoms", [])) if len(geom_cfg.get("freeze_atoms", [])) > 0 else None,
                    )
                    del H
                    return freqs_local, modes_local

                freqs_cm, modes = _calc_freqs_and_modes()

                neg_freq_thresh_cm = float(simple_cfg.get("neg_freq_thresh_cm", 5.0))
                neg_mask = freqs_cm < -abs(neg_freq_thresh_cm)
                n_imag = int(np.sum(neg_mask))
                ims = [float(x) for x in freqs_cm if x < -abs(neg_freq_thresh_cm)]
                click.echo(f"[Imaginary modes] n={n_imag}  ({ims})")

                flatten_max_iter = int(simple_cfg.get("flatten_max_iter", 0))
                if flatten_max_iter > 0 and n_imag > 1:
                    click.echo("[flatten] Extra imaginary modes detected; starting RSIRFO flatten loop.")
                    masses_amu = _safe_masses_amu(geometry.atomic_numbers)
                    roots = rsirfo_kwargs.get("roots", [0])
                    main_root = int(roots[0]) if roots else 0
                    for it in range(flatten_max_iter):
                        click.echo(f"[flatten] RSIRFO iteration {it + 1}/{flatten_max_iter}")
                        did_flatten = _flatten_once_with_modes_for_geom(
                            geometry,
                            masses_amu,
                            uma_kwargs_for_rsirfo,
                            freqs_cm,
                            modes,
                            neg_freq_thresh_cm,
                            float(simple_cfg.get("flatten_amp_ang", 0.10)),
                            float(simple_cfg.get("flatten_sep_cutoff", 0.0)),
                            int(simple_cfg.get("flatten_k", 10)),
                            main_root,
                        )
                        if not did_flatten:
                            click.echo("[flatten] No eligible modes to flatten; stopping.")
                            break

                        _attach_rsirfo_calc()
                        optimizer = RSIRFOptimizer(geometry, **rsirfo_kwargs)
                        click.echo("====== TS optimization (RS-I-RFO) restarted ======\n")
                        optimizer.run()
                        click.echo("====== TS optimization (RS-I-RFO) finished ======\n")
                        last_optimizer = optimizer
                        geometry.set_calculator(None)

                        freqs_cm, modes = _calc_freqs_and_modes()
                        neg_mask = freqs_cm < -abs(neg_freq_thresh_cm)
                        n_imag = int(np.sum(neg_mask))
                        ims = [float(x) for x in freqs_cm if x < -abs(neg_freq_thresh_cm)]
                        click.echo(f"[Imaginary modes] n={n_imag}  ({ims})")
                        if n_imag <= 1:
                            break

                needs_pdb = source_path.suffix.lower() == ".pdb"
                needs_gjf = prepared_input.is_gjf
                ref_pdb = source_path.resolve() if needs_pdb else None
                final_xyz_path = last_optimizer.final_fn if isinstance(last_optimizer.final_fn, Path) else Path(last_optimizer.final_fn)
                if convert_xyz_like_outputs(
                    final_xyz_path,
                    prepared_input,
                    ref_pdb_path=ref_pdb,
                    out_pdb_path=out_dir_path / "final_geometry.pdb" if needs_pdb else None,
                    out_gjf_path=out_dir_path / "final_geometry.gjf" if needs_gjf else None,
                    context="final geometry",
                ):
                    click.echo("[convert] Wrote 'final_geometry' outputs.")

                if bool(opt_cfg.get("dump", False)) and needs_pdb:
                    trj_path = out_dir_path / "optimization_trj.xyz"
                    if trj_path.exists():
                        if convert_xyz_like_outputs(
                            trj_path,
                            prepared_input,
                            ref_pdb_path=ref_pdb,
                            out_pdb_path=out_dir_path / "optimization.pdb" if needs_pdb else None,
                            context="optimization trajectory",
                        ):
                            click.echo("[convert] Wrote 'optimization' outputs.")
                    else:
                        click.echo("[convert] WARNING: 'optimization_trj.xyz' not found; skipping conversion.", err=True)

                # --- RSIRFO: write all final imaginary modes ---
                neg_idx = _echo_imaginary_modes(freqs_cm, neg_freq_thresh_cm)
                if len(neg_idx) == 0:
                    click.echo("[INFO] No imaginary mode found at the end for RSIRFO.")
                else:
                    vib_dir = out_dir_path / "vib"
                    ref_pdb_mode = source_path if source_path.suffix.lower() == ".pdb" else None
                    _write_all_imaginary_modes(
                        geometry,
                        freqs_cm,
                        modes,
                        neg_freq_thresh_cm,
                        _safe_masses_amu(geometry.atomic_numbers).tolist(),
                        vib_dir,
                        prepared_input=prepared_input,
                        ref_pdb=ref_pdb_mode,
                        amplitude_ang=0.8,
                        n_frames=20,
                    )

                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

            # summary.md and key_* outputs are disabled.
            click.echo(format_elapsed("[time] Elapsed Time for TS Opt", time_start))

        except ZeroStepLength:
            click.echo("ERROR: Proposed step length dropped below the minimum allowed (ZeroStepLength).", err=True)
            sys.exit(2)
        except OptimizationError as e:
            click.echo(f"ERROR: Optimization failed — {e}", err=True)
            sys.exit(3)
        except KeyboardInterrupt:
            click.echo("Interrupted by user.", err=True)
            sys.exit(130)
        except Exception as e:
            import traceback
            tb = "".join(traceback.format_exception(type(e), e, e.__traceback__))
            click.echo("Unhandled error during optimization:\n" + textwrap.indent(tb, "  "), err=True)
            sys.exit(1)
        finally:
            # Release GPU memory (model + Hessian) so subsequent stages don't OOM
            calc = geometry = optimizer = last_optimizer = None
            freqs_cm = modes = None
            gc.collect()  # break cyclic refs inside torch.nn.Module
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
