"""tkminfit.py

scipy.optimize.minimize ベースの最小二乗共通ランナー。

既存の tknlsq.nonlinear_lsq は scipy.optimize.least_squares 用。
このモジュールは、Nelder-Mead など minimize 系メソッドを使いたい
アプリ用の薄い共通層。

主な用途:
- optid=1 の非線形パラメータ最適化
- optid_lin=1 の線形ブロック最適化
- variable projection: 非線形パラメータごとに線形パラメータを内部で解く
- 数値Jacobianから共分散と標準誤差を推定
"""

from __future__ import annotations

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

import tklsq.tklsq as tkcore
import tklsq.tkparamio as tkparamio
from tklsq.tknlsq import numerical_jacobian, covariance_from_jacobian


ArrayLike = Sequence[float] | np.ndarray
ParamValues = Dict[str, float]


@dataclass
class MinimizeLSQResult:
    """minimize ベース最小二乗の結果。"""

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

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


def pack_values(values: Mapping[str, float], names: Sequence[str]) -> np.ndarray:
    """名前リストに従って値をベクトル化する。"""
    return np.array([float(values[n]) for n in names], dtype=float)


def unpack_values(
    p_free: ArrayLike,
    base_values: Mapping[str, float],
    free_names: Sequence[str],
) -> Dict[str, float]:
    """free vector を base_values に戻す。"""
    out = {k: float(v) for k, v in base_values.items()}
    for name, value in zip(free_names, np.asarray(p_free, dtype=float).reshape(-1)):
        out[name] = float(value)
    return out


def residuals_to_objective(
    residuals: ArrayLike,
    *,
    penalty: float = 0.0,
) -> float:
    """RSS + penalty を返す。"""
    r = np.asarray(residuals, dtype=float).reshape(-1)
    r = np.where(np.isfinite(r), r, 1.0e100)
    return float(r @ r + penalty)


def solve_linear_block(
    y: ArrayLike,
    params: Mapping[str, Mapping[str, object]],
    p_current: Mapping[str, float],
    design_matrix_func: Callable[[Mapping[str, float], Sequence[str]], np.ndarray],
    *,
    lin_names: Optional[Sequence[str]] = None,
    weights: Optional[ArrayLike] = None,
) -> Tuple[Dict[str, float], object]:
    """現在の非線形パラメータのもとで、線形パラメータを線形最小二乗で解く。

    design_matrix_func は以下の形を想定する。

        X = design_matrix_func(p_current, lin_names)

    たとえばアプリ側で x をクロージャとして持たせる。

        def design_matrix_func(p, lin_names):
            return make_X(x, p, lin_names)
    """
    if lin_names is None:
        lin_names = tkparamio.linear_param_names(params)
    lin_names = list(lin_names)

    if not lin_names:
        return {k: float(v) for k, v in p_current.items()}, None

    y_arr = np.asarray(y, dtype=float).reshape(-1)
    X = np.asarray(design_matrix_func(p_current, lin_names), dtype=float)

    if X.ndim != 2:
        raise ValueError("design_matrix_func must return a 2D array")
    if X.shape[0] != y_arr.size:
        raise ValueError("design matrix row count and y size mismatch")
    if X.shape[1] != len(lin_names):
        raise ValueError("design matrix column count must match lin_names")

    lin_res = tkcore.linear_lsq(X, y_arr, weights=weights)

    p_new = {k: float(v) for k, v in p_current.items()}
    for name, beta in zip(lin_names, lin_res.beta):
        p_new[name] = float(beta)

    return p_new, lin_res


def estimate_covariance_for_params(
    residual_func: Callable[[Mapping[str, float]], ArrayLike],
    params_fit: Mapping[str, float],
    free_names: Sequence[str],
    *,
    rel_step: float = 1.0e-6,
    abs_step: float = 1.0e-12,
) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray], Dict[str, Optional[float]], Optional[float], int, str]:
    """任意の residual_func(params) について、free_names の共分散を推定する。"""
    free_names = list(free_names)
    p_free = pack_values(params_fit, free_names)

    def r_of_free(v):
        p = unpack_values(v, params_fit, free_names)
        r = np.asarray(residual_func(p), dtype=float).reshape(-1)
        return np.where(np.isfinite(r), r, 1.0e100)

    residuals = r_of_free(p_free)
    J = numerical_jacobian(r_of_free, p_free, rel_step=rel_step, abs_step=abs_step)

    dof = residuals.size - len(free_names)
    cov, stderr_free, sigma2, warning = covariance_from_jacobian(residuals, J, dof=dof)

    stderr: Dict[str, Optional[float]] = {name: None for name in params_fit.keys()}
    if stderr_free is not None:
        for name, se in zip(free_names, stderr_free):
            stderr[name] = float(se)

    return J, cov, stderr_free, stderr, sigma2, dof, warning


def minimize_lsq(
    residual_func: Callable[[Mapping[str, float]], ArrayLike],
    params: Mapping[str, Mapping[str, object]],
    *,
    free_names: Optional[Sequence[str]] = None,
    method: str = "Nelder-Mead",
    use_penalty: bool = True,
    maxiter: Optional[int] = None,
    options: Optional[dict] = None,
    callback: Optional[Callable[[int, Dict[str, float], float], None]] = None,
    print_interval: int = 0,
    rel_step: float = 1.0e-6,
    abs_step: float = 1.0e-12,
) -> MinimizeLSQResult:
    """scipy.optimize.minimize で RSS + penalty を最小化する。

    residual_func:
        residual_func(params_dict) -> residual vector

    params:
        tkparamio.read_param_csv() の戻り値。
    """
    try:
        from scipy.optimize import minimize
    except Exception as exc:  # pragma: no cover
        raise ImportError("minimize_lsq requires scipy. Install with: pip install scipy") from exc

    base_values = tkparamio.values_from_params(params)

    if free_names is None:
        free_names = tkparamio.nonlinear_param_names(params)
    free_names = list(free_names)

    x0 = pack_values(base_values, free_names)
    counter = {"n": 0}

    def residual_from_vector(v):
        p = unpack_values(v, base_values, free_names)
        r = np.asarray(residual_func(p), dtype=float).reshape(-1)
        return np.where(np.isfinite(r), r, 1.0e100)

    def objective(v):
        p = unpack_values(v, base_values, free_names)
        r = residual_from_vector(v)
        penalty = tkparamio.bounds_penalty(params, p, names=free_names) if use_penalty else 0.0
        return residuals_to_objective(r, penalty=penalty)

    def scipy_callback(v):
        counter["n"] += 1
        obj = objective(v)
        p = unpack_values(v, base_values, free_names)

        if print_interval > 0 and counter["n"] % print_interval == 0:
            print(f"[iter {counter['n']:5d}] objective={obj:.10g}, params={p}")

        if callback is not None:
            callback(counter["n"], p, obj)

    if options is None:
        options = {}
    if maxiter is not None and "maxiter" not in options:
        options = dict(options)
        options["maxiter"] = maxiter

    raw = minimize(
        objective,
        x0,
        method=method,
        callback=scipy_callback,
        options=options,
    )

    p_fit = unpack_values(raw.x, base_values, free_names)
    residuals = residual_from_vector(raw.x)
    RSS = float(residuals @ residuals)

    J, cov, stderr_free, stderr, sigma2, dof, warning = estimate_covariance_for_params(
        residual_func,
        p_fit,
        free_names,
        rel_step=rel_step,
        abs_step=abs_step,
    )

    return MinimizeLSQResult(
        params=p_fit,
        params_free=pack_values(p_fit, free_names),
        names=list(params.keys()),
        free_names=free_names,
        optimized_names=free_names,
        linear_names=[],
        residuals=residuals,
        jacobian=J,
        cov_free=cov,
        stderr_free=stderr_free,
        stderr=stderr,
        sigma2_resid=sigma2,
        dof=dof,
        N=residuals.size,
        p_free=len(free_names),
        RSS=RSS,
        objective=float(raw.fun),
        success=bool(raw.success),
        message=str(raw.message),
        nfev=getattr(raw, "nfev", None),
        nit=getattr(raw, "nit", None),
        warning=warning,
        raw_result=raw,
    )


def variable_projection_lsq(
    y: ArrayLike,
    params: Mapping[str, Mapping[str, object]],
    model_func: Callable[[Mapping[str, float]], ArrayLike],
    design_matrix_func: Callable[[Mapping[str, float], Sequence[str]], np.ndarray],
    *,
    nonlin_names: Optional[Sequence[str]] = None,
    lin_names: Optional[Sequence[str]] = None,
    method: str = "Nelder-Mead",
    use_penalty: bool = True,
    weights: Optional[ArrayLike] = None,
    maxiter: Optional[int] = None,
    options: Optional[dict] = None,
    callback: Optional[Callable[[int, Dict[str, float], float], None]] = None,
    print_interval: int = 0,
    rel_step: float = 1.0e-6,
    abs_step: float = 1.0e-12,
) -> MinimizeLSQResult:
    """非線形パラメータ探索中に、線形パラメータを毎回線形最小二乗で解く。

    model_func:
        y_fit = model_func(params_dict)

    design_matrix_func:
        X = design_matrix_func(params_dict, lin_names)

    どちらも x はアプリ側のクロージャで保持する想定。
    """
    try:
        from scipy.optimize import minimize
    except Exception as exc:  # pragma: no cover
        raise ImportError("variable_projection_lsq requires scipy. Install with: pip install scipy") from exc

    y_arr = np.asarray(y, dtype=float).reshape(-1)
    base_values = tkparamio.values_from_params(params)

    if nonlin_names is None:
        nonlin_names = tkparamio.nonlinear_param_names(params)
    if lin_names is None:
        lin_names = tkparamio.linear_param_names(params)

    nonlin_names = list(nonlin_names)
    lin_names = list(lin_names)

    x0 = pack_values(base_values, nonlin_names)
    counter = {"n": 0}
    last_linear_result = {"res": None}

    def complete_params_from_nonlin(v):
        p = unpack_values(v, base_values, nonlin_names)
        p2, lin_res = solve_linear_block(
            y_arr,
            params,
            p,
            design_matrix_func,
            lin_names=lin_names,
            weights=weights,
        )
        last_linear_result["res"] = lin_res
        return p2

    def residual_from_params(p):
        y_fit = np.asarray(model_func(p), dtype=float).reshape(-1)
        if y_fit.size != y_arr.size:
            raise ValueError("model_func output size and y size mismatch")
        r = y_arr - y_fit
        return np.where(np.isfinite(r), r, 1.0e100)

    def objective(v):
        p = complete_params_from_nonlin(v)
        r = residual_from_params(p)
        penalty_names = list(dict.fromkeys(nonlin_names + lin_names))
        penalty = tkparamio.bounds_penalty(params, p, names=penalty_names) if use_penalty else 0.0
        return residuals_to_objective(r, penalty=penalty)

    def scipy_callback(v):
        counter["n"] += 1
        p = complete_params_from_nonlin(v)
        obj = objective(v)

        if print_interval > 0 and counter["n"] % print_interval == 0:
            shown = {name: p[name] for name in list(dict.fromkeys(nonlin_names + lin_names))}
            print(f"[iter {counter['n']:5d}] objective={obj:.10g}, params={shown}")

        if callback is not None:
            callback(counter["n"], p, obj)

    if options is None:
        options = {}
    if maxiter is not None and "maxiter" not in options:
        options = dict(options)
        options["maxiter"] = maxiter

    raw = minimize(
        objective,
        x0,
        method=method,
        callback=scipy_callback,
        options=options,
    )

    p_fit = complete_params_from_nonlin(raw.x)
    residuals = residual_from_params(p_fit)
    RSS = float(residuals @ residuals)

    free_names = list(dict.fromkeys(nonlin_names + lin_names))

    def residual_for_cov(p):
        return residual_from_params(p)

    J, cov, stderr_free, stderr, sigma2, dof, warning = estimate_covariance_for_params(
        residual_for_cov,
        p_fit,
        free_names,
        rel_step=rel_step,
        abs_step=abs_step,
    )

    return MinimizeLSQResult(
        params=p_fit,
        params_free=pack_values(p_fit, free_names),
        names=list(params.keys()),
        free_names=free_names,
        optimized_names=nonlin_names,
        linear_names=lin_names,
        residuals=residuals,
        jacobian=J,
        cov_free=cov,
        stderr_free=stderr_free,
        stderr=stderr,
        sigma2_resid=sigma2,
        dof=dof,
        N=residuals.size,
        p_free=len(free_names),
        RSS=RSS,
        objective=float(raw.fun),
        success=bool(raw.success),
        message=str(raw.message),
        nfev=getattr(raw, "nfev", None),
        nit=getattr(raw, "nit", None),
        warning=warning,
        raw_result=raw,
        linear_result=last_linear_result["res"],
    )