"""tkfitdiag.py

フィット結果の診断ユーティリティ。

線形/非線形を問わず、パラメータ共分散行列が得られた後に使う。
主な用途:
- 共分散行列 -> 標準偏差・相関係数行列
- 共分散行列の固有値/固有ベクトルから、不確かな方向を調べる
- J^T J の固有値/条件数から、同定しにくいパラメータ結合を調べる
- 相対誤差・強相関・不確かな固有方向への寄与から固定候補を提案する
"""

from __future__ import annotations

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


@dataclass
class EigenSummary:
    """固有ベクトルの主要成分サマリ。"""

    rank: int
    eigenvalue: float
    components: List[Tuple[str, float]]


@dataclass
class FixCandidate:
    """固定または外部拘束の候補。"""

    param: str
    score: float
    value: float
    stderr: float
    relerr: Optional[float]
    reasons: List[str] = field(default_factory=list)


@dataclass
class FitDiagnostics:
    """パラメータ共分散に基づく診断結果。"""

    names: List[str]
    values: np.ndarray
    stderr: np.ndarray
    relerr: np.ndarray
    cov: np.ndarray
    corr: np.ndarray
    eig_cov_values_desc: np.ndarray
    eig_cov_vectors_desc: np.ndarray
    eig_cov_summary: List[EigenSummary]
    jtj: Optional[np.ndarray] = None
    eig_jtj_values_asc: Optional[np.ndarray] = None
    eig_jtj_vectors_asc: Optional[np.ndarray] = None
    cond_jtj: Optional[float] = None
    warning: str = ""


def covariance_to_correlation(cov: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """共分散行列から相関係数行列と標準偏差を計算する。"""
    cov = np.asarray(cov, dtype=float)
    if cov.ndim != 2 or cov.shape[0] != cov.shape[1]:
        raise ValueError("cov must be a square matrix")

    stderr = np.sqrt(np.maximum(np.diag(cov), 0.0))
    denom = np.outer(stderr, stderr)
    with np.errstate(invalid="ignore", divide="ignore"):
        corr = cov / denom
    corr[denom == 0] = np.nan
    np.fill_diagonal(corr, 1.0)
    return corr, stderr


def eigen_sorted_symmetric(A: np.ndarray, descending: bool = True) -> Tuple[np.ndarray, np.ndarray]:
    """対称行列の固有値・固有ベクトルをソートして返す。"""
    A = np.asarray(A, dtype=float)
    vals, vecs = np.linalg.eigh(A)
    idx = np.argsort(vals)
    if descending:
        idx = idx[::-1]
    return vals[idx], vecs[:, idx]


def summarize_eigenvectors(
    eigenvalues: Sequence[float],
    eigenvectors: np.ndarray,
    names: Sequence[str],
    *,
    topk: int = 3,
    compk: int = 3,
) -> List[EigenSummary]:
    """固有ベクトルの主要成分を読みやすくまとめる。"""
    names = list(names)
    eigenvalues = np.asarray(eigenvalues, dtype=float)
    eigenvectors = np.asarray(eigenvectors, dtype=float)
    out: List[EigenSummary] = []

    n_show = min(topk, len(eigenvalues))
    for i in range(n_show):
        v = eigenvectors[:, i]
        order = np.argsort(np.abs(v))[::-1]
        comps = [(names[j], float(v[j])) for j in order[: min(compk, len(order))]]
        out.append(EigenSummary(rank=i + 1, eigenvalue=float(eigenvalues[i]), components=comps))
    return out


def diagnose_covariance(
    names: Sequence[str],
    values: Sequence[float],
    cov: np.ndarray,
    *,
    jacobian: Optional[np.ndarray] = None,
    topk_eigen: int = 3,
    compk_eigen: int = 4,
) -> FitDiagnostics:
    """共分散行列からフィット診断を作る。

    Parameters
    ----------
    names, values
        パラメータ名と値。
    cov
        パラメータ共分散行列。
    jacobian
        残差の Jacobian。渡すと J^T J の条件数と最小固有方向も評価する。
    """
    names = list(names)
    values = np.asarray(values, dtype=float)
    cov = np.asarray(cov, dtype=float)

    if cov.shape != (len(names), len(names)):
        raise ValueError("cov shape does not match names")
    if values.size != len(names):
        raise ValueError("values size does not match names")

    corr, stderr = covariance_to_correlation(cov)
    tiny = 1e-300
    relerr = stderr / (np.abs(values) + tiny)

    eig_vals, eig_vecs = eigen_sorted_symmetric(cov, descending=True)
    eig_summary = summarize_eigenvectors(
        eig_vals, eig_vecs, names, topk=topk_eigen, compk=compk_eigen
    )

    jtj = None
    eig_jtj_vals = None
    eig_jtj_vecs = None
    cond_jtj = None
    warning = ""

    if jacobian is not None:
        J = np.asarray(jacobian, dtype=float)
        if J.ndim != 2 or J.shape[1] != len(names):
            warning = "WARNING: jacobian shape does not match parameter count; J^T J skipped."
        else:
            jtj = J.T @ J
            eig_jtj_vals, eig_jtj_vecs = eigen_sorted_symmetric(jtj, descending=False)
            pos = eig_jtj_vals[eig_jtj_vals > 0]
            if len(pos) == len(eig_jtj_vals):
                cond_jtj = float(eig_jtj_vals[-1] / eig_jtj_vals[0])
            else:
                cond_jtj = math.inf
                warning = "WARNING: J^T J has zero or negative eigenvalues; problem may be rank deficient."

    return FitDiagnostics(
        names=names,
        values=values,
        stderr=stderr,
        relerr=relerr,
        cov=cov,
        corr=corr,
        eig_cov_values_desc=eig_vals,
        eig_cov_vectors_desc=eig_vecs,
        eig_cov_summary=eig_summary,
        jtj=jtj,
        eig_jtj_values_asc=eig_jtj_vals,
        eig_jtj_vectors_asc=eig_jtj_vecs,
        cond_jtj=cond_jtj,
        warning=warning,
    )


def propose_fix_candidates(
    names: Sequence[str],
    values: Sequence[float],
    stderr: Sequence[float],
    corr: np.ndarray,
    *,
    eig_cov_values: Optional[Sequence[float]] = None,
    eig_cov_vectors: Optional[np.ndarray] = None,
    corr_thr: float = 0.95,
    relerr_thr: float = 0.5,
    topn: int = 3,
) -> List[FixCandidate]:
    """固定/外部拘束候補をヒューリスティックに提案する。

    判断材料:
    - 相対標準誤差 stderr / |value|
    - 他パラメータとの強相関
    - 共分散最大固有値方向への寄与

    注意: これは数値的な「決まりにくさ」の診断であり、最終判断は物理的妥当性、
    独立測定、文献値の有無と合わせて行う。
    """
    names = list(names)
    values = np.asarray(values, dtype=float)
    stderr = np.asarray(stderr, dtype=float)
    corr = np.asarray(corr, dtype=float)

    if len(names) != values.size or values.size != stderr.size:
        raise ValueError("names, values, stderr sizes must match")
    if corr.shape != (len(names), len(names)):
        raise ValueError("corr shape does not match names")

    tiny = 1e-300
    relerr = stderr / (np.abs(values) + tiny)
    score = np.nan_to_num(relerr, nan=0.0, posinf=1e300, neginf=0.0).astype(float)
    reasons: Dict[str, List[str]] = {n: [] for n in names}

    for i, n in enumerate(names):
        if np.isfinite(relerr[i]):
            if relerr[i] >= relerr_thr:
                reasons[n].append(f"relative stderr = {relerr[i]:.3g} (>= {relerr_thr})")
            else:
                reasons[n].append(f"relative stderr = {relerr[i]:.3g}")
        else:
            reasons[n].append("relative stderr is not finite")

    for i in range(len(names)):
        for j in range(i + 1, len(names)):
            cij = corr[i, j]
            if not np.isfinite(cij):
                continue
            if abs(cij) >= corr_thr:
                bonus = (abs(cij) - corr_thr) * 5.0 + 0.5
                score[i] += bonus
                score[j] += bonus
                reasons[names[i]].append(f"strong corr with {names[j]}: {cij:+.3f} (>= {corr_thr})")
                reasons[names[j]].append(f"strong corr with {names[i]}: {cij:+.3f} (>= {corr_thr})")

    if eig_cov_values is not None and eig_cov_vectors is not None:
        vecs = np.asarray(eig_cov_vectors, dtype=float)
        if vecs.ndim == 2 and vecs.shape[0] == len(names) and vecs.shape[1] > 0:
            v = vecs[:, 0]
            w = np.abs(v) / (np.max(np.abs(v)) + tiny)
            for i, n in enumerate(names):
                if w[i] >= 0.5:
                    score[i] += 0.7 * w[i]
                    reasons[n].append(
                        f"dominant in most-uncertain eigen-direction: |v|/max={w[i]:.2f}"
                    )

    items = [
        FixCandidate(
            param=names[i],
            score=float(score[i]),
            value=float(values[i]),
            stderr=float(stderr[i]),
            relerr=float(relerr[i]) if np.isfinite(relerr[i]) else None,
            reasons=reasons[names[i]],
        )
        for i in range(len(names))
    ]
    items.sort(key=lambda x: x.score, reverse=True)

    filtered = [
        item
        for item in items
        if (item.relerr is not None and item.relerr >= relerr_thr)
        or any("strong corr" in r for r in item.reasons)
    ]
    if not filtered:
        filtered = items
    return filtered[: min(topn, len(filtered))]


def propose_fix_candidates_from_diagnostics(
    diag: FitDiagnostics,
    *,
    corr_thr: float = 0.95,
    relerr_thr: float = 0.5,
    topn: int = 3,
) -> List[FixCandidate]:
    """FitDiagnostics から固定候補を提案するショートカット。"""
    return propose_fix_candidates(
        diag.names,
        diag.values,
        diag.stderr,
        diag.corr,
        eig_cov_values=diag.eig_cov_values_desc,
        eig_cov_vectors=diag.eig_cov_vectors_desc,
        corr_thr=corr_thr,
        relerr_thr=relerr_thr,
        topn=topn,
    )


def format_fix_candidates(candidates: Sequence[FixCandidate]) -> str:
    """固定候補をコンソール表示しやすい文字列にする。"""
    lines: List[str] = []
    for k, c in enumerate(candidates, 1):
        rel = "NA" if c.relerr is None else f"{c.relerr:.3g}"
        lines.append(
            f"[{k}] {c.param}: value={c.value:.6g}, stderr={c.stderr:.3g}, "
            f"relerr={rel}, score={c.score:.3g}"
        )
        for r in c.reasons:
            lines.append(f"    - {r}")
    return "\n".join(lines)


def diagnostics_to_dict(diag: FitDiagnostics) -> dict:
    """JSON保存しやすい辞書に変換する。"""
    return {
        "names": diag.names,
        "values": diag.values.tolist(),
        "stderr": diag.stderr.tolist(),
        "relerr": diag.relerr.tolist(),
        "cov": diag.cov.tolist(),
        "corr": diag.corr.tolist(),
        "eig_cov": {
            "eigenvalues_desc": diag.eig_cov_values_desc.tolist(),
            "eigenvectors_desc": diag.eig_cov_vectors_desc.tolist(),
            "summary": [
                {"rank": s.rank, "eigenvalue": s.eigenvalue, "components": s.components}
                for s in diag.eig_cov_summary
            ],
        },
        "eig_jtj": None
        if diag.eig_jtj_values_asc is None
        else {
            "eigenvalues_asc": diag.eig_jtj_values_asc.tolist(),
            "eigenvectors_asc": diag.eig_jtj_vectors_asc.tolist(),
            "cond": diag.cond_jtj,
        },
        "warning": diag.warning,
    }