"""tknlsq.py

非線形最小二乗の小さな汎用ラッパー。

設計方針
--------
- アプリ固有の物理モデルは外に置く
- ここでは residual_func(params) -> residual vector を最小化する
- 数値 Jacobian と cov ≈ s^2 (J^T J)^(-1) を共通化する
- 固定パラメータを扱いやすくする

SciPy がある場合は scipy.optimize.least_squares を使う。
SciPy がない環境では ImportError を出す。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
import math
import numpy as np

ArrayLike = Union[Sequence[float], np.ndarray]
ParamsLike = Union[Sequence[float], Mapping[str, float]]


@dataclass
class NonlinearLSQResult:
    """非線形最小二乗の結果。"""

    params: Union[np.ndarray, Dict[str, float]]
    params_free: np.ndarray
    names: List[str]
    free_names: List[str]
    fixed_names: List[str]
    residuals: np.ndarray
    jacobian: Optional[np.ndarray]
    cov_free: Optional[np.ndarray]
    stderr_free: Optional[np.ndarray]
    stderr: Union[Optional[np.ndarray], Dict[str, Optional[float]]]
    sigma2_resid: Optional[float]
    dof: int
    N: int
    p_free: int
    RSS: float
    cost: float
    success: bool
    message: str
    nfev: Optional[int] = None
    njev: Optional[int] = None
    status: Optional[int] = None
    warning: str = ""
    raw_result: object = field(default=None, repr=False)

    @property
    def sigma_resid(self) -> Optional[float]:
        if self.sigma2_resid is None:
            return None
        return math.sqrt(max(float(self.sigma2_resid), 0.0))


def numerical_jacobian(
    fun_vec: Callable[[np.ndarray], ArrayLike],
    p: ArrayLike,
    *,
    rel_step: float = 1e-6,
    abs_step: float = 1e-12,
) -> np.ndarray:
    """中心差分でベクトル関数の Jacobian を計算する。

    fun_vec(p) -> residual vector, J[i, j] = d residual_i / d p_j
    """
    p = np.asarray(p, dtype=float).reshape(-1)
    f0 = np.asarray(fun_vec(p), dtype=float).reshape(-1)
    J = np.zeros((f0.size, p.size), dtype=float)

    for j in range(p.size):
        dp = rel_step * (abs(p[j]) + 1.0) + abs_step
        p_plus = p.copy(); p_plus[j] += dp
        p_minus = p.copy(); p_minus[j] -= dp
        f_plus = np.asarray(fun_vec(p_plus), dtype=float).reshape(-1)
        f_minus = np.asarray(fun_vec(p_minus), dtype=float).reshape(-1)
        if f_plus.size != f0.size or f_minus.size != f0.size:
            raise ValueError("fun_vec output size changed during numerical differentiation")
        J[:, j] = (f_plus - f_minus) / (2.0 * dp)
    return J


def covariance_from_jacobian(
    residuals: ArrayLike,
    J: np.ndarray,
    *,
    dof: Optional[int] = None,
    rcond: float = 1e-12,
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[float], str]:
    """残差 Jacobian からパラメータ共分散を近似する。

    cov ≈ s^2 (J^T J)^(-1), s^2 = RSS / dof

    dof <= 0 の場合、cov と stderr は None を返す。
    """
    r = np.asarray(residuals, dtype=float).reshape(-1)
    J = np.asarray(J, dtype=float)
    if J.ndim != 2:
        raise ValueError("J must be 2D")
    if J.shape[0] != r.size:
        raise ValueError("J rows and residual size mismatch")

    if dof is None:
        dof = r.size - J.shape[1]
    RSS = float(r.T @ r)
    if dof <= 0:
        return None, None, None, (
            f"WARNING: residual degrees of freedom N-p = {r.size}-{J.shape[1]} = {dof}; "
            "covariance and parameter errors are disabled."
        )

    sigma2 = RSS / dof
    JTJ = J.T @ J
    warning = ""
    try:
        cov0 = np.linalg.inv(JTJ)
    except np.linalg.LinAlgError:
        cov0 = np.linalg.pinv(JTJ, rcond=rcond)
        warning = "WARNING: J^T J is singular; pseudo-inverse was used for covariance."

    cov = sigma2 * cov0
    stderr = np.sqrt(np.maximum(np.diag(cov), 0.0))
    return cov, stderr, float(sigma2), warning


def _normalize_p0(p0: ParamsLike, names: Optional[Sequence[str]]) -> Tuple[np.ndarray, List[str], bool]:
    """p0 を配列化する。戻り値の bool は mapping 入力だったかどうか。"""
    if isinstance(p0, Mapping):
        if names is None:
            names = list(p0.keys())
        names = list(names)
        vals = np.array([float(p0[n]) for n in names], dtype=float)
        return vals, names, True

    vals = np.asarray(p0, dtype=float).reshape(-1)
    if names is None:
        names = [f"p{i}" for i in range(vals.size)]
    names = list(names)
    if len(names) != vals.size:
        raise ValueError("len(names) must match len(p0)")
    return vals, names, False


def _normalize_fixed(
    names: Sequence[str],
    p0_full: np.ndarray,
    fixed: Optional[Union[Sequence[str], Mapping[str, float]]],
) -> Tuple[np.ndarray, List[str], List[str], Dict[str, float]]:
    names = list(names)
    fixed_values: Dict[str, float] = {}
    if fixed is None:
        fixed_names: List[str] = []
    elif isinstance(fixed, Mapping):
        fixed_names = list(fixed.keys())
        fixed_values = {str(k): float(v) for k, v in fixed.items()}
    else:
        fixed_names = [str(x) for x in fixed]
        fixed_values = {n: float(p0_full[names.index(n)]) for n in fixed_names}

    for n in fixed_names:
        if n not in names:
            raise ValueError(f"fixed parameter {n!r} is not in names")

    p0_mod = p0_full.copy()
    for n, v in fixed_values.items():
        p0_mod[names.index(n)] = v

    free_names = [n for n in names if n not in fixed_names]
    return p0_mod, free_names, fixed_names, fixed_values


def _build_full_vector(
    p_free: np.ndarray,
    p_base: np.ndarray,
    names: Sequence[str],
    free_names: Sequence[str],
) -> np.ndarray:
    p_full = np.asarray(p_base, dtype=float).copy()
    for n, v in zip(free_names, p_free):
        p_full[list(names).index(n)] = float(v)
    return p_full


def _vector_to_output(p_full: np.ndarray, names: Sequence[str], as_mapping: bool):
    if as_mapping:
        return {n: float(v) for n, v in zip(names, p_full)}
    return p_full.copy()


def _stderr_to_output(
    stderr_free: Optional[np.ndarray],
    names: Sequence[str],
    free_names: Sequence[str],
    fixed_names: Sequence[str],
    as_mapping: bool,
):
    if as_mapping:
        out: Dict[str, Optional[float]] = {n: None for n in names}
        if stderr_free is not None:
            for n, se in zip(free_names, stderr_free):
                out[n] = float(se)
        for n in fixed_names:
            out[n] = 0.0
        return out

    if stderr_free is None:
        return None
    arr = np.full(len(names), np.nan, dtype=float)
    for n, se in zip(free_names, stderr_free):
        arr[list(names).index(n)] = float(se)
    for n in fixed_names:
        arr[list(names).index(n)] = 0.0
    return arr


def _normalize_bounds(
    bounds,
    names: Sequence[str],
    free_names: Sequence[str],
) -> Tuple[np.ndarray, np.ndarray]:
    if bounds is None:
        return np.full(len(free_names), -np.inf), np.full(len(free_names), np.inf)

    if isinstance(bounds, Mapping):
        lo = []
        hi = []
        for n in free_names:
            b = bounds.get(n, (-np.inf, np.inf))
            lo.append(float(b[0]))
            hi.append(float(b[1]))
        return np.asarray(lo, dtype=float), np.asarray(hi, dtype=float)

    lo, hi = bounds
    lo = np.asarray(lo, dtype=float).reshape(-1)
    hi = np.asarray(hi, dtype=float).reshape(-1)
    if lo.size == len(names) and hi.size == len(names):
        idx = [list(names).index(n) for n in free_names]
        return lo[idx], hi[idx]
    if lo.size == len(free_names) and hi.size == len(free_names):
        return lo, hi
    raise ValueError("bounds must match either all parameters or free parameters")


def nonlinear_lsq(
    residual_func: Callable[[Union[np.ndarray, Dict[str, float]]], ArrayLike],
    p0: ParamsLike,
    *,
    names: Optional[Sequence[str]] = None,
    fixed: Optional[Union[Sequence[str], Mapping[str, float]]] = None,
    bounds=None,
    loss: str = "linear",
    max_nfev: Optional[int] = None,
    rel_step: float = 1e-6,
    abs_step: float = 1e-12,
    use_result_jac: bool = True,
    scipy_kwargs: Optional[dict] = None,
) -> NonlinearLSQResult:
    """非線形最小二乗を実行する。

    residual_func は residual_func(params) -> 1D residual vector を返す関数。

    p0 が dict の場合、residual_func には dict が渡される。
    p0 が list/array の場合、residual_func には np.ndarray が渡される。

    fixed は固定パラメータ名のリスト、または {name: value} の辞書。
    """
    try:
        from scipy.optimize import least_squares
    except Exception as exc:  # pragma: no cover
        raise ImportError("nonlinear_lsq requires scipy. Install with: pip install scipy") from exc

    p0_full, names, as_mapping = _normalize_p0(p0, names)
    p_base, free_names, fixed_names, _fixed_values = _normalize_fixed(names, p0_full, fixed)
    p0_free = np.array([p_base[names.index(n)] for n in free_names], dtype=float)
    lb, ub = _normalize_bounds(bounds, names, free_names)

    def call_residual_full(p_full: np.ndarray) -> np.ndarray:
        params = _vector_to_output(p_full, names, as_mapping)
        r = np.asarray(residual_func(params), dtype=float).reshape(-1)
        if not np.all(np.isfinite(r)):
            # least_squares dislikes non-finite residuals. Replace by large residuals.
            r = np.where(np.isfinite(r), r, 1e100)
        return r

    def call_residual_free(p_free: np.ndarray) -> np.ndarray:
        p_full = _build_full_vector(p_free, p_base, names, free_names)
        return call_residual_full(p_full)

    if scipy_kwargs is None:
        scipy_kwargs = {}

    if len(free_names) == 0:
        r = call_residual_full(p_base)
        RSS = float(r.T @ r)
        return NonlinearLSQResult(
            params=_vector_to_output(p_base, names, as_mapping),
            params_free=np.array([], dtype=float),
            names=list(names),
            free_names=[],
            fixed_names=list(fixed_names),
            residuals=r,
            jacobian=None,
            cov_free=None,
            stderr_free=None,
            stderr=_stderr_to_output(None, names, free_names, fixed_names, as_mapping),
            sigma2_resid=None,
            dof=r.size,
            N=r.size,
            p_free=0,
            RSS=RSS,
            cost=0.5 * RSS,
            success=True,
            message="all parameters are fixed",
            warning="All parameters are fixed; optimization and covariance are skipped.",
            raw_result=None,
        )

    res = least_squares(
        call_residual_free,
        p0_free,
        bounds=(lb, ub),
        loss=loss,
        max_nfev=max_nfev,
        **scipy_kwargs,
    )

    p_full = _build_full_vector(res.x, p_base, names, free_names)
    residuals = call_residual_free(res.x)
    N = residuals.size
    p = len(free_names)
    dof = N - p
    RSS = float(residuals.T @ residuals)

    if use_result_jac and getattr(res, "jac", None) is not None:
        J = np.asarray(res.jac, dtype=float)
    else:
        J = numerical_jacobian(call_residual_free, res.x, rel_step=rel_step, abs_step=abs_step)

    cov, stderr_free, sigma2, cov_warning = covariance_from_jacobian(residuals, J, dof=dof)
    stderr = _stderr_to_output(stderr_free, names, free_names, fixed_names, as_mapping)

    return NonlinearLSQResult(
        params=_vector_to_output(p_full, names, as_mapping),
        params_free=np.asarray(res.x, dtype=float),
        names=list(names),
        free_names=list(free_names),
        fixed_names=list(fixed_names),
        residuals=residuals,
        jacobian=J,
        cov_free=cov,
        stderr_free=stderr_free,
        stderr=stderr,
        sigma2_resid=sigma2,
        dof=dof,
        N=N,
        p_free=p,
        RSS=RSS,
        cost=float(res.cost),
        success=bool(res.success),
        message=str(res.message),
        nfev=getattr(res, "nfev", None),
        njev=getattr(res, "njev", None),
        status=getattr(res, "status", None),
        warning=cov_warning,
        raw_result=res,
    )


def delta_method_variance(
    output_func: Callable[[np.ndarray], ArrayLike],
    p: ArrayLike,
    cov: np.ndarray,
    *,
    rel_step: float = 1e-6,
    abs_step: float = 1e-12,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """出力量 y=f(p) の分散をデルタ法で計算する。

    Returns
    -------
    y0 : ndarray, shape (N,)
    var_y : ndarray, shape (N,)
    grads : ndarray, shape (N, M)
    """
    p = np.asarray(p, dtype=float).reshape(-1)
    cov = np.asarray(cov, dtype=float)
    if cov.shape != (p.size, p.size):
        raise ValueError("cov shape must match p")

    def f(v):
        return np.asarray(output_func(v), dtype=float).reshape(-1)

    y0 = f(p)
    grads = numerical_jacobian(f, p, rel_step=rel_step, abs_step=abs_step)
    var_y = np.einsum("ni,ij,nj->n", grads, cov, grads)
    var_y = np.maximum(var_y, 0.0)
    return y0, var_y, grads


def delta_method_band(
    output_func: Callable[[np.ndarray], ArrayLike],
    p: ArrayLike,
    cov: np.ndarray,
    *,
    nsigma: float = 1.0,
    rel_step: float = 1e-6,
    abs_step: float = 1e-12,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """デルタ法で y, y_low, y_high, sigma_y を返す。"""
    y0, var_y, _grads = delta_method_variance(
        output_func, p, cov, rel_step=rel_step, abs_step=abs_step
    )
    sigma = np.sqrt(var_y)
    return y0, y0 - nsigma * sigma, y0 + nsigma * sigma, sigma


def result_free_values(result: NonlinearLSQResult) -> np.ndarray:
    """NonlinearLSQResult から free parameter values を返す。"""
    return np.asarray(result.params_free, dtype=float)


def result_free_cov(result: NonlinearLSQResult) -> Optional[np.ndarray]:
    """NonlinearLSQResult から free parameter covariance を返す。"""
    return result.cov_free