# tkminfit.py callback-compatible edition
#
# この版では callback インターフェースを scipy.optimize.minimize
# と同じ callback(q) に変更しています。
#
# NOTE:
#   callback に渡される q は scipy.optimize.minimize 内部の
#   最適化ベクトルです。
#
#   log parameter や normalized parameter を使う場合は、
#   アプリ側で unpack / inverse transform を行ってください。
#
# 変更点:
#   callback(counter, params, objective)
#       ↓
#   callback(q)
#
# backward compatibility:
#   callback=None は従来通り使用可能。

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]:
    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:
    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]:
    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,
):
    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[[np.ndarray], None]] = None,
    print_interval: int = 0,
    rel_step: float = 1.0e-6,
    abs_step: float = 1.0e-12,
) -> MinimizeLSQResult:

    try:
        from scipy.optimize import minimize
    except Exception as exc:
        raise ImportError(
            "minimize_lsq requires 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

        if (
            print_interval > 0
            and counter["n"] % print_interval == 0
        ):
            obj = objective(v)
            p = unpack_values(v, base_values, free_names)

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

        if callback is not None:
            # NOTE:
            # v is the internal optimization vector used by
            # scipy.optimize.minimize.
            #
            # Applications using transformed parameters
            # (e.g. log-scaled parameters) should unpack /
            # inverse-transform v on the application side.
            callback(v)

    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[[np.ndarray], None]] = None,
    print_interval: int = 0,
    rel_step: float = 1.0e-6,
    abs_step: float = 1.0e-12,
) -> MinimizeLSQResult:

    try:
        from scipy.optimize import minimize
    except Exception as exc:
        raise ImportError(
            "variable_projection_lsq requires 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

        if (
            print_interval > 0
            and counter["n"] % print_interval == 0
        ):
            p = complete_params_from_nonlin(v)
            obj = objective(v)

            shown = {
                name: p[name]
                for name in list(
                    dict.fromkeys(nonlin_names + lin_names)
                )
            }

            print(
                f"[iter {counter['n']:5d}] "
                f"objective={obj:.10g}, params={shown}"
            )

        if callback is not None:
            # NOTE:
            # v contains only nonlinear optimization variables.
            #
            # Linear parameters are internally solved by
            # variable projection and are NOT included in v.
            callback(v)

    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"],
    )