from math import sqrt
from pathlib import Path
from typing import Literal, Optional

import numpy as np
from scipy.optimize import root_scalar

from pysisyphus.cos.ChainOfStates import ChainOfStates
from pysisyphus.Geometry import Geometry
from pysisyphus.helpers_pure import rms
# from pysisyphus.io.hessian import save_hessian
from pysisyphus.optimizers.guess_hessians import (
    get_guess_hessian,
    xtb_hessian,
    HessInit,
)
from pysisyphus.optimizers.hessian_updates import (
    bfgs_update,
    flowchart_update,
    damped_bfgs_update,
    bofill_update,
    ts_bfgs_update,
    ts_bfgs_update_org,
    ts_bfgs_update_revised,
)
from pysisyphus.optimizers.Optimizer import Optimizer
from pysisyphus.optimizers.exceptions import OptimizationError

from pysisyphus.helpers import array2string
import torch

def dummy_hessian_update(H, dx, dg):
    return np.zeros_like(H), "no"


HESS_UPDATE_FUNCS = {
    "none": dummy_hessian_update,
    None: dummy_hessian_update,
    False: dummy_hessian_update,
    "bfgs": bfgs_update,
    "damped_bfgs": damped_bfgs_update,
    "flowchart": flowchart_update,
    "bofill": bofill_update,
    "ts_bfgs": ts_bfgs_update,
    "ts_bfgs_org": ts_bfgs_update_org,
    "ts_bfgs_rev": ts_bfgs_update_revised,
}
HessUpdate = Literal[
    "none",
    None,
    False,
    "bfgs",
    "damped_bfgs",
    "flowchart",
    "bofill",
    "ts_bfgs",
    "ts_bfgs_org",
    "ts_bfgs_rev",
]


class HessianOptimizer(Optimizer):
    rfo_dict = {
        "min": (0, "min"),
        "max": (-1, "max"),
    }

    def __init__(
        self,
        geometry: Geometry,
        trust_radius: float = 0.5,
        trust_update: bool = True,
        trust_min: float = 0.1,
        trust_max: float = 1,
        max_energy_incr: Optional[float] = None,
        hessian_update: HessUpdate = "bfgs",
        hessian_init: HessInit = "fischer",
        hessian_recalc: Optional[int] = None,
        hessian_recalc_adapt: Optional[float] = None,
        hessian_xtb: bool = False,
        hessian_recalc_reset: bool = False,
        small_eigval_thresh: float = 1e-8,
        line_search: bool = False,
        alpha0: float = 1.0,
        max_micro_cycles: int = 25,
        rfo_overlaps: bool = False,
        **kwargs,
    ) -> None:
        """Baseclass for optimizers utilizing Hessian information.

        Parameters
        ----------
        geometry
            Geometry to be optimized.
        trust_radius
            Initial trust radius in whatever unit the optimization is carried out.
        trust_update
            Whether to update the trust radius throughout the optimization.
        trust_min
            Minimum trust radius.
        trust_max
            Maximum trust radius.
        max_energy_incr
            Maximum allowed energy increased after a faulty step. Optimization is
            aborted when the threshold is exceeded.
        hessian_update
            Type of Hessian update. Defaults to BFGS for minimizations and Bofill
            for saddle point searches.
        hessian_init
            Type of initial model Hessian.
        hessian_recalc
            Recalculate exact Hessian every n-th cycle instead of updating it.
        hessian_recalc_adapt
            Use a more flexible scheme to determine Hessian recalculation. Undocumented.
        hessian_xtb
            Recalculate the Hessian at the GFN2-XTB level of theory.
        hessian_recalc_reset
            Whether to skip Hessian recalculation after reset. Undocumented.
        small_eigval_thresh
            Threshold for small eigenvalues. Eigenvectors belonging to eigenvalues
            below this threshold are discardewd.
        line_search
            Whether to carry out a line search. Not implemented by a subclassing
            optimizers.
        alpha0
            Initial alpha for restricted-step (RS) procedure.
        max_micro_cycles
            Maximum number of RS iterations.
        rfo_overlaps
            Enable mode-following in RS procedure.

        Other Parameters
        ----------------
        **kwargs
            Keyword arguments passed to the Optimizer baseclass.
        """
        super().__init__(geometry, **kwargs)

        assert not issubclass(
            type(geometry), ChainOfStates
        ), "HessianOptimizer can't be used for and ChainOfStates objects!"

        self.trust_update = bool(trust_update)
        assert trust_min <= trust_max, "trust_min must be <= trust_max!"
        self.trust_min = float(trust_min)
        self.trust_max = float(trust_max)
        self.max_energy_incr = max_energy_incr
        # Constrain initial trust radius if trust_max > trust_radius
        self.trust_radius = min(trust_radius, trust_max)
        self.log(f"Initial trust radius: {self.trust_radius:.6f}")
        self.hessian_update = hessian_update
        self.hessian_update_func = HESS_UPDATE_FUNCS[hessian_update]
        self.hessian_init = hessian_init
        self.hessian_recalc = hessian_recalc
        self.hessian_recalc_adapt = hessian_recalc_adapt
        self.hessian_xtb = hessian_xtb
        self.hessian_recalc_reset = hessian_recalc_reset
        self.small_eigval_thresh = float(small_eigval_thresh)
        self.line_search = bool(line_search)
        # Restricted-step related
        self.alpha0 = alpha0
        self.max_micro_cycles = int(max_micro_cycles)
        assert max_micro_cycles >= 0
        self.rfo_overlaps = rfo_overlaps

        assert self.small_eigval_thresh > 0.0, "small_eigval_thresh must be > 0.!"
        if not self.restarted:
            self.hessian_recalc_in = None
            self.adapt_norm = None
            self.predicted_energy_changes = list()
        hessian_init_exists = Path(self.hessian_init).exists()
        if (
            # Allow actually calculated Hessians for all coordinate systems
            not hessian_init_exists
            and self.hessian_init not in ("calc", "xtb", "xtb1", "xtbff")
            # But disable model Hessian for Cartesian optimizations
            and self.geometry.coord_type in ("cart", "cartesian", "mwcartesian")
        ):
            self.hessian_init = "unit"
            self.log(
                f"Chosen initial (model) Hessian is incompatible with current "
                f"coord_type: {self.geometry.coord_type}!"
            )

        self._prev_eigvec_min = None
        self._prev_eigvec_max = None
        self._using_active_dofs = False
        self._active_dof_indices = None
        self.cur_H = None

    @property
    def using_active_dofs(self):
        return self._using_active_dofs

    @property
    def active_dof_indices(self):
        return self._active_dof_indices

    def _set_active_dofs(self, use_active):
        self._using_active_dofs = use_active
        if not use_active:
            self._active_dof_indices = None
            return
        if getattr(self.geometry, "within_partial_hessian", None) is not None:
            self._active_dof_indices = self.geometry.hess_active_dof_indices
            return
        # Fallback: infer active DOFs from calculator's Hessian-active atoms
        calc = getattr(self.geometry, "calculator", None)
        core = getattr(calc, "core", calc)
        hess_atoms = getattr(core, "hess_active_atoms", None)
        if hess_atoms is not None and len(hess_atoms) > 0:
            act = []
            for a in hess_atoms:
                base = 3 * int(a)
                act.extend([base, base + 1, base + 2])
            self._active_dof_indices = np.asarray(act, dtype=int)
            return
        self._active_dof_indices = self.geometry.active_dof_indices

    def active_from_full(self, vec):
        if not self.using_active_dofs:
            return vec
        inds = self._active_dof_indices
        if inds is None:
            return vec
        # Sanitize indices (drop negatives / out-of-bounds)
        try:
            inds = np.asarray(inds, dtype=int)
            if len(inds) > 0:
                if np.any(inds < 0):
                    inds = inds[inds >= 0]
                if len(inds) > 0:
                    max_valid = vec.shape[0] - 1
                    if np.any(inds > max_valid):
                        inds = inds[inds <= max_valid]
                if len(inds) > 0:
                    if np.min(inds) < 0 or np.max(inds) >= vec.shape[0]:
                        return vec
        except (ValueError, IndexError, TypeError):
            pass
        # Avoid double-slicing if vector is already in active space
        try:
            if vec.shape[0] == len(inds):
                return vec
        except (ValueError, IndexError, TypeError):
            pass
        try:
            if len(inds) > 0 and vec.shape[0] <= int(np.max(inds)):
                # Indices exceed vector length → already active (compact) vector.
                return vec
        except (ValueError, IndexError, TypeError):
            pass
        if isinstance(vec, torch.Tensor):
            if vec.device.type == "cuda":
                try:
                    vec_cpu = vec.detach().cpu().numpy()
                    return torch.as_tensor(vec_cpu[inds], dtype=vec.dtype, device=vec.device)
                except (ValueError, IndexError, TypeError):
                    return vec
            idx = torch.as_tensor(inds, dtype=torch.long, device=vec.device)
            return vec.index_select(0, idx)
        return vec[inds]

    def full_from_active(self, vec):
        if not self.using_active_dofs:
            return vec
        inds = self._active_dof_indices
        if inds is None:
            return vec
        if isinstance(vec, torch.Tensor):
            idx = torch.as_tensor(inds, dtype=torch.long, device=vec.device)
            full = torch.zeros(self.geometry.cart_coords.size, dtype=vec.dtype, device=vec.device)
            full.index_copy_(0, idx, vec)
            return full
        full = np.zeros(self.geometry.cart_coords.size, dtype=vec.dtype if hasattr(vec, "dtype") else float)
        full[inds] = vec
        return full

    def active_hessian(self, hessian):
        if not self.using_active_dofs:
            return hessian

        if getattr(self.geometry, "within_partial_hessian", None) is not None:
            act_n_dof = int(self.geometry.within_partial_hessian.get("active_n_dof", 0))
            if hessian.shape == (act_n_dof, act_n_dof):
                return hessian

        inds = self.active_dof_indices
        try:
            inds_arr = np.asarray(inds, dtype=int)
            if hessian.shape[0] == len(inds_arr) and len(inds_arr) > 0:
                if np.max(inds_arr) >= hessian.shape[0]:
                    # Likely already in active order; avoid double-slicing.
                    return hessian
        except (ValueError, IndexError, TypeError):
            pass
        try:
            inds = np.asarray(inds, dtype=int)
            if len(inds) > 0:
                max_valid = hessian.shape[0] - 1
                inds = inds[(inds >= 0) & (inds <= max_valid)]
        except (ValueError, IndexError, TypeError):
            pass
        if isinstance(hessian, torch.Tensor):
            if hessian.device.type == "cuda":
                try:
                    hess_cpu = hessian.detach().cpu().numpy()
                    return torch.as_tensor(hess_cpu[np.ix_(inds, inds)], dtype=hessian.dtype, device=hessian.device)
                except (ValueError, IndexError, TypeError):
                    return hessian
            idx = torch.as_tensor(inds, device=hessian.device, dtype=torch.int64)
            return hessian.index_select(0, idx).index_select(1, idx)
        return hessian[np.ix_(inds, inds)]

    def active_list(self, seq):
        if not self.using_active_dofs:
            return seq
        return [self.active_from_full(item) for item in seq]

    @property
    def prev_eigvec_min(self):
        return self._prev_eigvec_min

    @prev_eigvec_min.setter
    def prev_eigvec_min(self, prev_eigvec_min):
        if self.rfo_overlaps:
            self._prev_eigvec_min = prev_eigvec_min

    @property
    def prev_eigvec_max(self):
        return self._prev_eigvec_max

    @prev_eigvec_max.setter
    def prev_eigvec_max(self, prev_eigvec_max):
        if self.rfo_overlaps:
            self._prev_eigvec_max = prev_eigvec_max

    def reset(self):
        # Don't recalculate the hessian if we have to reset the optimizer
        hessian_init = self.hessian_init
        if (
            (not self.hessian_recalc_reset)
            and hessian_init == "calc"
            and self.geometry.coord_type != "cart"
        ):
            hessian_init = "fischer"
        self.prepare_opt(hessian_init)

    # def save_hessian(self):
    #     # Don't try to save Hessians of analytical potentials
    #     if self.geometry.is_analytical_2d:
    #         return

    #     h5_fn = self.get_path_for_fn(f"hess_calc_cyc_{self.cur_cycle}.h5")
    #     # Save the cartesian hessian, as it is independent of the
    #     # actual coordinate system that is used.
    #     save_hessian(
    #         h5_fn,
    #         self.geometry,
    #         self.geometry.cart_hessian,
    #         self.geometry.energy,
    #         self.geometry.calculator.mult,
    #     )
    #     self.log(f"Wrote calculated cartesian Hessian to '{h5_fn}'")

    def prepare_opt(self, hessian_init=None):
        if hessian_init is None:
            hessian_init = self.hessian_init

        self.H, hess_str = get_guess_hessian(self.geometry, hessian_init)
        if self.hessian_init != "calc" and self.geometry.is_analytical_2d:
            assert self.H.shape == (3, 3)
            self.H[2, 2] = 0.0

        msg = f"Using {hess_str} Hessian"
        if hess_str == "saved":
            msg += f" from '{hessian_init}'"
        self.log(msg)

        # # Dump to disk if hessian was calculated
        # if self.hessian_init == "calc":
        #     self.save_hessian()

        if (
            hasattr(self.geometry, "coord_type")
            and self.geometry.coord_type == "dlc"
            # Calculated Hessian is already in DLC
            and hessian_init != "calc"
        ):
            U = self.geometry.internal.U
            self.H = U.T.dot(self.H).dot(U)

        if self.hessian_recalc_adapt:
            self.adapt_norm = np.linalg.norm(self.geometry.forces)

        if self.hessian_recalc:
            # Already substract one, as we don't do a hessian update in
            # the first cycle.
            self.hessian_recalc_in = self.hessian_recalc - 1

    def _get_opt_restart_info(self):
        opt_restart_info = {
            "adapt_norm": self.adapt_norm,
            "H": self.H.tolist(),
            "hessian_recalc_in": self.hessian_recalc_in,
            "predicted_energy_changes": self.predicted_energy_changes,
        }
        return opt_restart_info

    def _set_opt_restart_info(self, opt_restart_info):
        self.adapt_norm = opt_restart_info["adapt_norm"]
        self.H = np.array(opt_restart_info["H"])
        self.hessian_recalc_in = opt_restart_info["hessian_recalc_in"]
        self.predicted_energy_changes = opt_restart_info["predicted_energy_changes"]

    def update_trust_radius(self):
        # The predicted change should be calculated at the end of optimize
        # of the previous cycle.
        assert (
            len(self.predicted_energy_changes) == len(self.forces) - 1
        ), "Did you forget to append to self.predicted_energy_changes?"
        self.log("Trust radius update")
        self.log(f"\tCurrent trust radius: {self.trust_radius:.6f}")
        predicted_change = self.predicted_energy_changes[-1]
        actual_change = self.energies[-1] - self.energies[-2]
        # Only report an unexpected increase if we actually predicted a
        # decrease.
        unexpected_increase = (actual_change > 0) and (predicted_change < 0)
        old_trust = self.trust_radius
        if unexpected_increase:
            self.log(f"Energy increased by {actual_change:.6f} au!")
            if self.max_energy_incr and (actual_change > self.max_energy_incr):
                raise OptimizationError("Actual energy change too high!")
        coeff = actual_change / predicted_change
        self.log(f"\tPredicted change: {predicted_change:.4e} au")
        self.log(f"\tActual change: {actual_change:.4e} au")
        self.log(f"\tCoefficient: {coeff:.2%}")
        step = self.steps[-1]
        last_step_norm = np.linalg.norm(step)
        self.set_new_trust_radius(coeff, last_step_norm)
        if unexpected_increase:
            self.table.print(
                f"Unexpected energy increase ({actual_change:.6f} au)! "
                f"Trust radius: old={old_trust:.4}, new={self.trust_radius:.4}"
            )

    def set_new_trust_radius(self, coeff, last_step_norm):
        # Nocedal, Numerical optimization Chapter 4, Algorithm 4.1

        # If actual and predicted energy change have different signs
        # coeff will be negative and lead to a decreased trust radius,
        # which is fine.
        if coeff < 0.25:
            self.trust_radius = max(self.trust_radius / 4, self.trust_min)
            self.log("\tDecreasing trust radius.")
        # Only increase trust radius if last step norm was at least 80% of it
        # See [5], Appendix, step size and direction control
        # elif coeff > 0.75 and (last_step_norm >= .8*self.trust_radius):
        #
        # Only increase trust radius if last step norm corresponded approximately
        # to the trust radius.
        elif coeff > 0.75 and abs(self.trust_radius - last_step_norm) <= 1e-3:
            self.trust_radius = min(self.trust_radius * 2, self.trust_max)
            self.log("\tIncreasing trust radius.")
        else:
            self.log(f"\tKeeping current trust radius at {self.trust_radius:.6f}")
            return
        self.log(f"\tUpdated trust radius: {self.trust_radius:.6f}")

    def update_hessian(self):
        # Compare current forces to reference forces to see if we shall recalc the
        # hessian.
        try:
            cur_norm = np.linalg.norm(self.forces[-1])
            ref_norm = self.adapt_norm / self.hessian_recalc_adapt
            recalc_adapt = cur_norm <= ref_norm
            self.log(
                "Check for adaptive Hessian recalculation: "
                f"{cur_norm:.6f} <= {ref_norm:.6f}, {recalc_adapt}"
            )
        except TypeError:
            recalc_adapt = False

        try:
            self.hessian_recalc_in = max(self.hessian_recalc_in - 1, 0)
            self.log(f"Recalculation of Hessian in {self.hessian_recalc_in} cycle(s).")
        except TypeError:
            self.hessian_recalc_in = None

        # Update reference norm if needed
        # TODO: Decide on whether to update the norm when the recalculation is
        # initiated by 'recalc'.
        if recalc_adapt:
            self.adapt_norm = cur_norm

        recalc = self.hessian_recalc_in == 0

        if recalc or recalc_adapt:
            # Free old Hessian from GPU before recalculating
            H_old = self.H
            self.H = None
            del H_old
            try:
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except ImportError:
                pass
            # Use xtb hessian
            self.log("Requested Hessian recalculation.")
            if self.hessian_xtb:
                self.H = xtb_hessian(self.geometry)
                key = "xtb"
            # Calculated hessian at actual level of theory
            else:
                self.H = self.geometry.hessian
                key = "exact"
                # self.save_hessian()
            if self.using_active_dofs:
                # Keep the optimizer Hessian in active-DOF space to avoid
                # shape mismatches during quasi-Newton updates.
                self.H = self.active_hessian(self.H)
            if not (self.cur_cycle == 0):
                self.log(f"Recalculated {key} Hessian in cycle {self.cur_cycle}.")
            # Reset counter. It is also reset when the recalculation was initiated
            # by the adaptive formulation.
            self.hessian_recalc_in = self.hessian_recalc
        # Simple hessian update
        else:
            dx = self.steps[-1]
            dg = -(self.forces[-1] - self.forces[-2])
            H_work = self.H
            if self.using_active_dofs:
                H_work = self.active_hessian(self.H)
                dx = self.active_from_full(dx)
                dg = self.active_from_full(dg)
            curv_cond = dx.dot(dg)
            if curv_cond < 0.0:
                self.log(
                    f"Curvature condition (s·y = {curv_cond:.4f} < 0) not satisfied!"
                )
            dH, key = self.hessian_update_func(H_work, dx, dg)
            self.H = H_work + dH
            self.log(f"Did {key} Hessian update.")

    def solve_rfo(self, rfo_mat, kind="min", prev_eigvec=None):
        # When using the restricted step variant of RFO the RFO matrix
        # may not be symmetric. Thats why we can't use eigh here.
        is_torch = isinstance(rfo_mat, torch.Tensor)

        if is_torch:
            if not torch.isfinite(rfo_mat).all():
                self.log("RFO matrix contains NaN/inf; sanitizing entries.")
                rfo_mat = torch.nan_to_num(
                    rfo_mat, nan=0.0, posinf=1e8, neginf=-1e8
                )
        else:
            if not np.isfinite(rfo_mat).all():
                self.log("RFO matrix contains NaN/inf; sanitizing entries.")
                rfo_mat = np.nan_to_num(rfo_mat, nan=0.0, posinf=1e8, neginf=-1e8)

        if is_torch:
            is_sym = torch.allclose(rfo_mat, rfo_mat.T)
        else:
            is_sym = np.allclose(rfo_mat, rfo_mat.T)

        if is_sym:
            try:
                eigenvalues, eigenvectors = (torch.linalg.eigh(rfo_mat) if is_torch else np.linalg.eigh(rfo_mat))
            except (torch._C._LinAlgError, np.linalg.LinAlgError):
                self.log("eigh failed; falling back to eig.")
                eigenvalues, eigenvectors = (torch.linalg.eig(rfo_mat) if is_torch else np.linalg.eig(rfo_mat))
                eigenvalues = eigenvalues.real
                eigenvectors = eigenvectors.real
        else:
            eigenvalues, eigenvectors = (torch.linalg.eig(rfo_mat) if is_torch else np.linalg.eig(rfo_mat))
            eigenvalues = eigenvalues.real
            eigenvectors = eigenvectors.real

        self.log("\tdiagonalized augmented Hessian")

        if isinstance(eigenvectors, torch.Tensor):
            sorted_inds = torch.argsort(eigenvalues)
        else:
            sorted_inds = np.argsort(eigenvalues)

        # Depending on wether we want to minimize (maximize) along
        # the mode(s) in the rfo mat we have to select the smallest
        # (biggest) eigenvalue and corresponding eigenvector.
        first_or_last, verbose = self.rfo_dict[kind]
        # Given sorted eigenvalue-indices (sorted_inds) use the first
        # (smallest eigenvalue) or the last (largest eigenvalue) index.
        if prev_eigvec is None:
            ind = sorted_inds[first_or_last]
        else:
            if isinstance(prev_eigvec, torch.Tensor):
                ovlps = prev_eigvec.matmul(eigenvectors)
            else:
                ovlps = np.array([prev_eigvec.dot(ev) for ev in eigenvectors.T])
            naive_ind = sorted_inds[first_or_last]
            ind = np.abs(ovlps).argmax() if isinstance(ovlps, np.ndarray) else torch.argmax(torch.abs(ovlps)).item()
            self.log(
                f"Overlap: {ind} ({eigenvalues[ind]:.6f}), "
                f"Naive: {naive_ind} ({eigenvalues[naive_ind]:.6f})"
            )
        follow_eigvec = eigenvectors.T[ind]
        if isinstance(follow_eigvec, torch.Tensor):
            step_nu = follow_eigvec.clone()
        else:
            step_nu = follow_eigvec.copy()
        nu = step_nu[-1]
        self.log(f"\tnu_{verbose}={nu:.8e}")
        # Scale eigenvector so that its last element equals 1. The
        # final is step is the scaled eigenvector without the last element.
        step = step_nu[:-1] / nu
        eigval = eigenvalues[ind]
        self.log(f"\teigenvalue_{verbose}={eigval:.8e}")
        return step, eigval, nu, follow_eigvec

    def solve_rfo_secular(self, eigvals, gradient, alpha=1.0, kind="min",
                          prev_eigvec=None, max_iter=50, tol=1e-12):
        """Solve the RFO eigenvalue problem via the secular equation.

        The augmented Hessian has arrowhead structure, so its eigenvalue
        problem reduces to: f(mu) = sum g_i^2/(alpha*mu - lam_i) - mu = 0,
        solvable in O(N) instead of O(N^3).

        Returns (step, eigval, nu, eigvec) on success, None on failure.
        """
        is_torch = isinstance(eigvals, torch.Tensor)

        # Convert all inputs to plain numpy/float for root-finding
        _np = lambda x: x.detach().cpu().numpy().copy() if isinstance(x, torch.Tensor) else np.asarray(x, dtype=np.float64)
        g = _np(gradient)
        lam = _np(eigvals)
        alpha = alpha.detach().cpu().item() if isinstance(alpha, torch.Tensor) else float(alpha)

        n = len(lam)
        g2 = g ** 2
        nz = g2 > 1e-30
        if not nz.any():
            step = np.zeros(n)
            eigvec = np.zeros(n + 1); eigvec[-1] = 1.0
            if is_torch:
                step = torch.zeros(n, device=eigvals.device, dtype=eigvals.dtype)
                eigvec = torch.zeros(n + 1, device=eigvals.device, dtype=eigvals.dtype)
                eigvec[-1] = 1.0
            return step, 0.0, 1.0, eigvec

        g2_nz, lam_nz = g2[nz], lam[nz]

        # Guard against alpha ≈ 0 (can arise from trust-radius adaptation)
        if abs(alpha) < 1e-14:
            return None

        def f_df(mu):
            d = alpha * mu - lam_nz
            return float(np.sum(g2_nz / d) - mu), float(-alpha * np.sum(g2_nz / d**2) - 1.0)

        _, verbose = self.rfo_dict[kind]

        # Bracket the root
        if kind == "min":
            pole = float(lam.min() / alpha)
            mu = pole - max(float(np.sqrt(g2_nz.sum())) / alpha, 1.0)
            for _ in range(20):
                if f_df(mu)[0] > 0: break
                mu = pole - 2.0 * (pole - mu)
            else:
                return None
            lo, hi = mu, pole - 1e-15 * max(abs(pole), 1.0)
        elif kind == "max":
            pole = float(lam.max() / alpha)
            mu = pole + max(float(np.sqrt(g2_nz.sum())) / alpha, 1.0)
            for _ in range(20):
                if f_df(mu)[0] < 0: break
                mu = pole + 2.0 * (mu - pole)
            else:
                return None
            lo, hi = pole + 1e-15 * max(abs(pole), 1.0), mu
        else:
            return None

        # Newton-Raphson with bisection safeguard
        mu_cur = (lo + hi) / 2.0
        for _ in range(max_iter):
            fval, dfval = f_df(mu_cur)
            if abs(fval) < tol:
                break
            mu_new = mu_cur - fval / dfval if abs(dfval) > 1e-30 else (lo + hi) / 2.0
            if mu_new <= lo or mu_new >= hi:
                mu_new = (lo + hi) / 2.0
            f_new = f_df(mu_new)[0]
            if f_new > 0: lo = mu_new
            else: hi = mu_new
            mu_cur = mu_new
        else:
            self.log(f"Secular equation did not converge in {max_iter} iters.")
            return None

        self.log(f"\teigenvalue_{verbose}={mu_cur:.8e} (secular)")
        self.log(f"\tnu_{verbose}={1.0:.8e}")

        # Compute step: s_i = g_i / (alpha * mu - lam_i)
        denom = alpha * mu_cur - lam
        step_np = np.where(nz, g / denom, 0.0)

        # Eigenvector for mode tracking
        eigvec_np = np.append(step_np, 1.0)
        eigvec_np /= np.linalg.norm(eigvec_np)

        # Mode tracking check
        if prev_eigvec is not None:
            prev_np = _np(prev_eigvec)
            if abs(float(np.dot(prev_np, eigvec_np))) < 0.3:
                self.log("Secular eigvec overlap too low; falling back.")
                return None

        if is_torch:
            step_np = torch.tensor(step_np, device=eigvals.device, dtype=eigvals.dtype)
            eigvec_np = torch.tensor(eigvec_np, device=eigvals.device, dtype=eigvals.dtype)

        return step_np, mu_cur, 1.0, eigvec_np

    def filter_small_eigvals(self, eigvals, eigvecs, mask=False):
        if isinstance(eigvals, torch.Tensor):
            small_inds = torch.abs(eigvals) < self.small_eigval_thresh
        else:
            small_inds = np.abs(eigvals) < self.small_eigval_thresh
        eigvals = eigvals[~small_inds]
        eigvecs = eigvecs[:, ~small_inds]
        small_num = sum(small_inds)
        self.log(
            f"Found {small_num} small eigenvalues in Hessian. Removed "
            "corresponding eigenvalues and eigenvectors."
        )
        # assert small_num <= 6, (
        #     "Expected at most 6 small eigenvalues in cartesian hessian "
        #     f"but found {small_num}!"
        # )
        if mask:
            return eigvals, eigvecs, small_inds
        else:
            return eigvals, eigvecs

    def log_negative_eigenvalues(self, eigvals, pre_str=""):
        neg_inds = eigvals < -self.small_eigval_thresh
        neg_eigval_str = array2string(eigvals[neg_inds], precision=6)
        self.log(f"{pre_str}Hessian has {neg_inds.sum()} negative eigenvalue(s).")
        self.log(f"\t{neg_eigval_str}")

    def housekeeping(self):
        """Calculate gradient and energy. Update trust radius and hessian
        if needed. Return energy, gradient and hessian for the current cycle."""
        gradient_full = self.geometry.gradient
        energy = self.geometry.energy
        self.energies.append(energy)
        self.log(f"    Energy: {energy: >12.6f} au")
        self.log(
            f"norm(grad): {np.linalg.norm(gradient_full): >12.6f} au / bohr (rad)"
        )
        self.log(
            f" rms(grad): {np.sqrt(np.mean(gradient_full**2)): >12.6f} au / bohr (rad)"
        )
        self.forces.append(-gradient_full)

        can_update = (
            # Allows gradient differences
            len(self.forces) > 1
            and (self.forces[-2].shape == gradient_full.shape)
            and len(self.coords) > 1
            # Coordinates may have been rebuilt. Take care of that.
            and (self.coords[-2].shape == self.coords[1].shape)
            and len(self.energies) > 1
        )
        if can_update:
            if self.trust_update:
                self.update_trust_radius()
            self.update_hessian()

        # Convert gradient to match H device/dtype AFTER update_hessian(),
        # so that hessian_recalc (which may replace self.H with a new tensor
        # on a different device) is accounted for.
        if isinstance(self.H, torch.Tensor):
            gradient_full = torch.from_numpy(gradient_full).to(
                self.H.device, self.H.dtype
            )

        H = self.H
        if self.geometry.internal:
            # Shift eigenvalues of orthogonal part to high values, so they
            # don't contribute to the actual step.
            H_proj = self.geometry.internal.project_hessian(self.H)
            # Symmetrize hessian, as the projection may break it?!
            H = (H_proj + H_proj.T) / 2

        if getattr(self.geometry, "within_partial_hessian", None) is not None:
            use_active = True
        elif (
            H.shape[0] != self.geometry.cart_coords.size
            and self.geometry.coord_type in ("cart", "cartesian", "mwcartesian")
        ):
            # Partial Hessian without explicit metadata: still use active slicing.
            # Only applies to Cartesian coordinate types; for internal coordinates
            # (e.g. DLC), the Hessian is naturally smaller than cart_coords.size.
            use_active = True
        else:
            use_active = (
                len(self.geometry.freeze_atoms) > 0
                and self.geometry.coord_type in ("cart", "cartesian", "mwcartesian")
                and H.shape[0] == self.geometry.cart_coords.size
            )
        self._set_active_dofs(use_active)

        H = self.active_hessian(H)
        if gradient_full.shape[0] == H.shape[0]:
            gradient = gradient_full
        else:
            gradient = self.active_from_full(gradient_full)

        if isinstance(H, torch.Tensor):
            eigvals, eigvecs = torch.linalg.eigh(H)
        else:
            eigvals, eigvecs = np.linalg.eigh(H)
        # Neglect small eigenvalues
        eigvals, eigvecs = self.filter_small_eigvals(eigvals, eigvecs)

        resetted = not can_update
        self.cur_H = H
        return energy, gradient, H, eigvals, eigvecs, resetted

    def get_augmented_hessian(self, eigvals, gradient, alpha=1.0):
        if isinstance(gradient, torch.Tensor):
            dim_ = eigvals.size(0) + 1
            H_aug = torch.zeros((dim_, dim_), device=gradient.device, dtype=gradient.dtype)
            H_aug[: dim_ - 1, : dim_ - 1] = torch.diag(eigvals / alpha)
        else:
            dim_ = eigvals.size + 1
            H_aug = np.zeros((dim_, dim_))
            H_aug[: dim_ - 1, : dim_ - 1] = np.diag(eigvals / alpha)
        H_aug[-1, :-1] = gradient
        H_aug[:-1, -1] = gradient

        H_aug[:-1, -1] /= alpha

        return H_aug

    def get_alpha_step(self, cur_alpha, rfo_eigval, step_norm, eigvals, gradient):
        # Derivative of the squared step w.r.t. alpha
        numer = gradient**2
        denom = (eigvals - rfo_eigval * cur_alpha) ** 3
        if isinstance(gradient, torch.Tensor):
            quot = torch.sum(numer / denom)
        else:
            quot = np.sum(numer / denom)
        self.log(f"quot={quot:.6f}")
        dstep2_dalpha = 2 * rfo_eigval / (1 + step_norm**2 * cur_alpha) * quot
        if isinstance(gradient, torch.Tensor):
            dstep2_valid = bool(
                torch.isfinite(dstep2_dalpha)
                & (torch.abs(dstep2_dalpha) > 1e-12)
            )
        else:
            dstep2_valid = np.isfinite(dstep2_dalpha) and abs(dstep2_dalpha) > 1e-12
        if not dstep2_valid:
            self.log(
                "alpha update skipped due to invalid derivative "
                f"(dstep2_dalpha={dstep2_dalpha})"
            )
            return 0.0
        self.log(f"analytic deriv.={dstep2_dalpha:.6f}")
        # Update alpha
        alpha_step = (
            2 * (self.trust_radius * step_norm - step_norm**2) / dstep2_dalpha
        )
        self.log(f"alpha_step={alpha_step:.4f}")
        min_alpha = 1e-8
        if (cur_alpha + alpha_step) <= min_alpha:
            self.log(
                "alpha update would make alpha non-positive; "
                f"clamping to min_alpha={min_alpha:.1e}"
            )
            alpha_step = min_alpha - cur_alpha
        return alpha_step

    def get_rs_step(self, eigvals, eigvecs, gradient, name="RS"):
        # Transform gradient to basis of eigenvectors
        if isinstance(eigvecs, torch.Tensor):
            if not isinstance(gradient, torch.Tensor):
                gradient = torch.as_tensor(
                    gradient, device=eigvecs.device, dtype=eigvecs.dtype
                )
            elif gradient.device != eigvecs.device:
                gradient = gradient.to(device=eigvecs.device)
            gradient_ = eigvecs.T @ gradient
        else:
            gradient_ = eigvecs.T.dot(gradient)

        alpha = self.alpha0
        for mu in range(self.max_micro_cycles):
            self.log(f"{name} micro cycle {mu:02d}, alpha={alpha:.6f}")
            # Try secular equation solver first (O(N) vs O(N^3))
            secular_result = self.solve_rfo_secular(
                eigvals, gradient_, alpha, kind="min",
                prev_eigvec=self.prev_eigvec_min,
            )
            if secular_result is not None:
                rfo_step_, eigval_min, nu, self.prev_eigvec_min = secular_result
            else:
                # Fallback to full eigendecomposition
                self.log("Secular solver failed; using full eigendecomposition.")
                H_aug = self.get_augmented_hessian(eigvals, gradient_, alpha)
                rfo_step_, eigval_min, nu, self.prev_eigvec_min = self.solve_rfo(
                    H_aug, "min", prev_eigvec=self.prev_eigvec_min
                )
            if isinstance(rfo_step_, torch.Tensor):
                rfo_norm_ = torch.linalg.norm(rfo_step_)
            else:
                rfo_norm_ = np.linalg.norm(rfo_step_)
            self.log(f"norm(rfo step)={rfo_norm_:.6f}")

            if rfo_norm_ <= 0:
                self.log(
                    "RFO step length is zero; falling back to trust-region Newton step."
                )
                step_ = self.get_newton_step_on_trust(
                    eigvals, eigvecs, gradient, transform=False
                )
                break

            if (rfo_norm_ < self.trust_radius) or abs(
                rfo_norm_ - self.trust_radius
            ) <= 1e-3:
                step_ = rfo_step_
                break

            alpha_step = self.get_alpha_step(
                alpha, eigval_min, rfo_norm_, eigvals, gradient_
            )
            alpha += alpha_step
            self.log("")
        # Otherwise, use trust region newton step
        else:
            self.log(
                "RS algorithm did not produce a desired step length in "
                f"{self.max_micro_cycles} micro cycles. Trying RFO with α=1.0."
            )
            secular_result = self.solve_rfo_secular(
                eigvals, gradient_, alpha=1.0, kind="min"
            )
            if secular_result is not None:
                rfo_step_, eigval_min, nu, _ = secular_result
            else:
                H_aug = self.get_augmented_hessian(eigvals, gradient_, alpha=1.0)
                rfo_step_, eigval_min, nu, _ = self.solve_rfo(H_aug, "min")
            if isinstance(rfo_step_, torch.Tensor):
                rfo_norm_ = torch.linalg.norm(rfo_step_)
            else:
                rfo_norm_ = np.linalg.norm(rfo_step_)

            # This should always be True if the above algorithm failed but we
            # keep this line nonetheless,  to make it more obvious.
            if rfo_norm_ > self.trust_radius:
                self.log(
                    f"Proposed RFO step with norm {rfo_norm_:.4f} is outside trust "
                    f"radius Δ={self.trust_radius:.4f}. "
                )
                step_ = self.get_newton_step_on_trust(
                    eigvals, eigvecs, gradient, transform=False
                )
                # Simple, downscaled RFO step
                # step_ = rfo_step_ / rfo_norm_ * self.trust_radius
            else:
                step_ = rfo_step_

        # Transform step back to original basis
        if isinstance(eigvecs, torch.Tensor):
            if isinstance(step_, torch.Tensor):
                pass
            else:
                step_ = torch.tensor(step_, device=eigvecs.device, dtype=eigvecs.dtype)
            step = eigvecs @ step_
            step = step.cpu().numpy()
        else:
            step = eigvecs.dot(step_)
        min_step_norm = getattr(self, "min_step_norm", 0.0)
        step_norm = np.linalg.norm(step)
        if step_norm <= min_step_norm:
            self.log(
                "RFO step length below minimum threshold; "
                "falling back to trust-region Newton step."
            )
            step = self.get_newton_step_on_trust(eigvals, eigvecs, gradient)
        if not np.isfinite(step).all():
            self.log(
                "RFO step contains NaN/inf; falling back to trust-region Newton step."
            )
            step = self.get_newton_step_on_trust(eigvals, eigvecs, gradient)
            if not np.isfinite(step).all():
                raise ValueError(
                    "Fallback Newton step still contains NaN/inf; "
                    "aborting to avoid corrupting coordinates."
                )
        return step

    @staticmethod
    def get_shifted_step_trans(eigvals, gradient_trans, shift):
        return -gradient_trans / (eigvals + shift)

    @staticmethod
    def get_newton_step(eigvals, eigvecs, gradient):
        if isinstance(eigvecs, torch.Tensor):
            eigvals = eigvals.to(eigvecs.device, dtype=eigvecs.dtype)
            gradient = gradient.to(eigvecs.device, dtype=eigvecs.dtype)
            return (eigvecs @ (eigvecs.T @ gradient / eigvals)).cpu().numpy()
        else:
            return eigvecs.dot(eigvecs.T.dot(gradient) / eigvals)

    def get_newton_step_on_trust(self, eigvals, eigvecs, gradient, transform=True):
        """Step on trust-radius.

        See Nocedal 4.3 Iterative solutions of the subproblem
        """
        if isinstance(eigvals, torch.Tensor):
            eigvals = eigvals.cpu().numpy()

        min_ind = eigvals.argmin()
        min_eigval = eigvals[min_ind]
        pos_definite = bool((eigvals > 0.0).all())
        if isinstance(eigvecs, torch.Tensor):
            if not isinstance(gradient, torch.Tensor):
                gradient = torch.tensor(
                    gradient, device=eigvecs.device, dtype=eigvecs.dtype
                )
            else:
                gradient = gradient.to(device=eigvecs.device, dtype=eigvecs.dtype)
            gradient_trans = eigvecs.T @ gradient
            gradient_trans = gradient_trans.cpu().numpy()
        else:
            gradient_trans = eigvecs.T.dot(gradient)

        # This will be also be True when we come close to a minimizer,
        # but then the Hessian will also be positive definite and a
        # simple Newton step will be used.
        hard_case = abs(gradient_trans[min_ind]) <= 1e-6
        self.log(f"Smallest eigenvalue: {min_eigval:.6f}")
        self.log(f"Positive definite Hessian: {pos_definite}")
        self.log(f"Hard case: {hard_case}")

        def get_step(shift):
            return -gradient_trans / (eigvals + shift)

        # Unshifted Newton step
        newton_step_trans = get_step(0.0)
        newton_norm = np.linalg.norm(newton_step_trans)

        def on_trust_radius_lin(step):
            return 1 / self.trust_radius - 1 / np.linalg.norm(step)

        def finalize_step(shift):
            step = get_step(shift)
            if transform:
                if isinstance(eigvecs, torch.Tensor):
                    step = torch.tensor(step, device=eigvecs.device, dtype=eigvecs.dtype)
                    return (eigvecs @ step).cpu().numpy()
                else:
                    return eigvecs.dot(step)
            return step

        # Simplest case. Positive definite Hessian and predicted step is
        # already in trust radius.
        if pos_definite and newton_norm <= self.trust_radius:
            self.log("Using unshifted Newton step.")
            if isinstance(eigvecs, torch.Tensor):
                newton_step_trans = torch.tensor(
                    newton_step_trans, device=eigvecs.device, dtype=eigvecs.dtype
                )
                return (eigvecs @ newton_step_trans).cpu().numpy()
            else:
                return eigvecs.dot(newton_step_trans)

        # If the Hessian is not positive definite or if the step is too
        # long we have to determine the shift parameter lambda.
        rs_kwargs = {
            "f": lambda shift: on_trust_radius_lin(get_step(shift)),
            "xtol": 1e-3,
            # Would otherwise be chosen automatically, but we set it
            # here explicitly for verbosity.
            "method": "brentq",
        }

        def root_search(bracket):
            rs_kwargs.update(
                {
                    "bracket": bracket,
                    "x0": bracket[0] + 1e-3,
                }
            )
            res = root_scalar(**rs_kwargs)
            return res

        BRACKET_END = 1e10
        if not hard_case:
            bracket_start = 0.0 if pos_definite else -min_eigval + 1e-2
            bracket = (bracket_start, BRACKET_END)
            try:
                res = root_search(bracket)
                assert res.converged
                return finalize_step(res.root)
            # ValueError may be raised when the function values for the
            # initial bracket have the same sign. If so, we continue with
            # treating it as a hard case.
            except ValueError:
                pass

        # Now we would try the bracket (-b2, -b1). The resulting step should have
        # a suitable length, but the (shifted) Hessian would have an incorrect
        # eigenvalue spectrum (not positive definite). To solve this we use a
        # different formula to calculate the step.
        mask = np.ones_like(gradient_trans)
        mask[min_ind] = 0
        mask = mask.astype(bool)
        without_min = gradient_trans[mask] / (eigvals[mask] - min_eigval)
        tau_sq = self.trust_radius**2 - (without_min**2).sum()
        if tau_sq >= 0.0:
            tau = sqrt(tau_sq)
            step_trans = [tau] + (-without_min).tolist()
        else:
            # Hard case. Search in open interval (endpoints not included)
            # (-min_eigval, inf).
            bracket = (-min_eigval + 1e-6, BRACKET_END)
            try:
                res = root_search(bracket)
                if res.converged:
                    return finalize_step(res.root)
            except ValueError:
                pass
            # Fallback: clamp tau to 0 so the step excludes the
            # minimum-eigenvalue component but remains valid.
            self.log("Hard case fallback: tau clamped to 0.")
            tau = 0.0
            step_trans = [tau] + (-without_min).tolist()

        if not transform:
            return step_trans

        if isinstance(eigvecs, torch.Tensor):
            step_trans = torch.tensor(step_trans, device=eigvecs.device, dtype=eigvecs.dtype)
            return (eigvecs @ step_trans).cpu().numpy()
        else:
            return eigvecs.dot(step_trans)

    @staticmethod
    def quadratic_model(gradient, hessian, step):
        if isinstance(gradient, torch.Tensor):
            step = torch.tensor(step, device=gradient.device, dtype=gradient.dtype)
            return (step @ gradient + 0.5 * step @ hessian @ step).cpu().numpy()
        else:
            step = np.asarray(step).ravel()
            return step.dot(gradient) + 0.5 * step.dot(hessian).dot(step)

    @staticmethod
    def rfo_model(gradient, hessian, step):
        return HessianOptimizer.quadratic_model(gradient, hessian, step) / (
            1 + step.dot(step)
        )

    def get_step_func(self, eigvals, gradient, grad_rms_thresh=1e-2):
        positive_definite = (eigvals < 0).sum() == 0
        gradient_small = rms(gradient) < grad_rms_thresh

        if self.adapt_step_func and gradient_small and positive_definite:
            return self.get_newton_step_on_trust, self.quadratic_model
        # RFO fallback
        else:
            return self.get_rs_step, self.rfo_model
