"""稳定性分析：AMMI（加性主效应乘性互作 SVD）与 Finlay-Wilkinson（环境指数回归）。

仅依赖 numpy。输入为基因型×环境两因素均值表 {genotype: {env: mean}}（缺格 NaN）。

- Finlay-Wilkinson：每基因型对该环境指数（该环境全部基因型均值）回归 y_gj = b·I_j + a。
  b≈1（|b-1| ≤ 2·se(b)）且残差小 → 稳定（与总体响应一致）；b>1 高响应（好环境下增益更大）、
  b<1 低响应。
- AMMI：两因素主效应分解（μ+g+e）后残差矩阵 SVD → IPC1/IPC2 得分、
  ASV（Purchase et al. 2000：两轴权重按各自 SS 归一后合成）、Wricke ecovalence
  （Σ_j (y_ij − y_i. − y_.j + y..)²）。

稳定排名：AMMI 按 ASV 升序（越小越稳定，rank=1 最稳）；FW 由 flag_stable 直接标注。
"""
from __future__ import annotations

import math

import numpy as np

ENGINE_VERSION = "1.0.0"


def _table_to_arrays(table: dict[str, dict[str, float]]):
    genotypes = sorted(table)
    envs: list[str] = []
    seen: set[str] = set()
    for g in genotypes:
        for e in table[g]:
            if e not in seen:
                seen.add(e)
                envs.append(e)
    envs.sort()
    Y = np.full((len(genotypes), len(envs)), np.nan)
    for gi, g in enumerate(genotypes):
        for e, v in table[g].items():
            Y[gi, envs.index(e)] = float(v)
    return Y, genotypes, envs


def finlay_wilkinson(table: dict) -> dict:
    """Finlay-Wilkinson：每基因型对环境指数回归，返回 b/截距/R²/残差MS/稳定标记。"""
    Y, genotypes, envs = _table_to_arrays(table)
    env_index = np.nanmean(Y, axis=0)  # 环境指数 = 该环境全基因型均值
    rows: list[dict] = []
    for gi, g in enumerate(genotypes):
        y = Y[gi]
        mask = (~np.isnan(y)) & (~np.isnan(env_index))
        n = int(mask.sum())
        if n < 2:
            rows.append({"genotype": g, "n_env": n, "b": None, "intercept": None,
                         "r2": None, "se_b": None, "dev_ms": None,
                         "flag_stable": None, "note": "环境数不足 2，无法回归"})
            continue
        xs = env_index[mask]
        ys = y[mask]
        A = np.column_stack([np.ones(n), xs])
        coef, *_ = np.linalg.lstsq(A, ys, rcond=None)
        intercept, b = float(coef[0]), float(coef[1])
        resid = ys - A @ coef
        dof = n - 2
        dev_ms = float(resid @ resid / dof) if dof > 0 else float("nan")
        ss_tot = float(((ys - ys.mean()) ** 2).sum())
        r2 = 1.0 - float(resid @ resid) / ss_tot if ss_tot > 0 else float("nan")
        var_x = float(((xs - xs.mean()) ** 2).sum())
        se_b = math.sqrt(dev_ms / var_x) if var_x > 0 and not math.isnan(dev_ms) else float("nan")
        y_scale = float(np.ptp(ys)) if len(ys) > 1 else 0.0
        if math.isnan(se_b):
            flag_stable = False
        elif dev_ms <= 1e-12 * max(1.0, y_scale ** 2):
            # 完美拟合（残差≈0）时 se_b 无信息量；b 精确已知，b≈1 才视为稳定
            flag_stable = abs(b - 1.0) <= 1e-6
        else:
            flag_stable = abs(b - 1.0) <= 2.0 * se_b
        rows.append({
            "genotype": g, "n_env": n, "b": round(b, 4), "intercept": round(intercept, 4),
            "r2": round(r2, 4) if not math.isnan(r2) else None,
            "se_b": round(se_b, 4) if not math.isnan(se_b) else None,
            "dev_ms": round(dev_ms, 6) if not math.isnan(dev_ms) else None,
            "flag_stable": bool(flag_stable),
        })
    return {
        "method": "finlay_wilkinson",
        "genotypes": genotypes, "environments": envs,
        "env_index": [None if math.isnan(v) else round(float(v), 4) for v in env_index],
        "rows": rows,
    }


def ammi(table: dict, n_components: int = 2) -> dict:
    """AMMI：两因素主效应分解 + 残差 SVD，输出 IPC 得分/ASV/ecovalence 与稳定排名。"""
    Y, genotypes, envs = _table_to_arrays(table)
    g, e = Y.shape
    if g < 2 or e < 2:
        return {"method": "ammi", "genotypes": genotypes, "environments": envs,
                "rows": [], "ipc_variance": [], "n": 0,
                "warning": "基因型或环境数不足 2，无法 AMMI"}
    has_missing = bool(np.isnan(Y).any())
    grand = float(np.nanmean(Y))
    row_mean = np.nanmean(Y, axis=1)
    col_mean = np.nanmean(Y, axis=0)
    resid = np.nan_to_num(Y - row_mean[:, None] - col_mean[None, :] + grand, nan=0.0)
    U, S, _Vt = np.linalg.svd(resid, full_matrices=False)
    k = min(g, e)
    ss_total = float(np.sum(resid ** 2))
    ipc_var = []
    for i in range(k):
        prop = (S[i] ** 2) / ss_total if ss_total > 0 else 0.0
        ipc_var.append({"ipc": i + 1, "singular_value": round(float(S[i]), 6),
                        "ss": round(float(S[i] ** 2), 6), "proportion": round(prop, 4)})
    n_keep = min(n_components, k)
    g_ipc = U[:, :n_keep] * S[:n_keep]  # 基因型 IPC 得分（sqrt(λ)·U）
    ss_ipc1 = float((g_ipc[:, 0] ** 2).sum()) if n_keep >= 1 else 0.0
    ss_ipc2 = float((g_ipc[:, 1] ** 2).sum()) if n_keep >= 2 else 0.0
    rows: list[dict] = []
    for gi in range(g):
        ipc1 = float(g_ipc[gi, 0]) if n_keep >= 1 else 0.0
        ipc2 = float(g_ipc[gi, 1]) if n_keep >= 2 else 0.0
        # ASV（Purchase 2000）：IPC1 按 SS 比例缩放到 IPC2 尺度，与 IPC2 正交合成
        if n_keep >= 2 and ss_ipc1 > 0:
            asv = math.sqrt((ipc1 * ss_ipc2 / ss_ipc1) ** 2 + ipc2 ** 2)
        else:
            asv = abs(ipc1)
        ecoval = float(np.nansum((Y[gi] - row_mean[gi] - col_mean + grand) ** 2))
        rows.append({
            "genotype": genotypes[gi],
            "ipc1": round(ipc1, 4),
            "ipc2": round(ipc2, 4) if n_keep >= 2 else None,
            "asv": round(asv, 4),
            "ecovalence": round(ecoval, 6),
        })
    rows.sort(key=lambda r: r["asv"])
    for rank, r in enumerate(rows, 1):
        r["rank"] = rank
    warning = None
    if has_missing:
        warning = "两因素表存在缺格，SVD 前缺格按 0 填充，IPC/ASV 估计仅供参考"
    return {
        "method": "ammi", "genotypes": genotypes, "environments": envs,
        "rows": rows, "ipc_variance": ipc_var, "n": g * e,
        "ss_interaction": round(ss_total, 6), "warning": warning,
    }


def stability_analysis(table: dict, *, methods: tuple[str, ...] = ("ammi", "finlay"),
                       n_components: int = 2) -> dict:
    """稳定性综合分析：AMMI + Finlay-Wilkinson 可选。"""
    out: dict = {"genotypes": sorted(table), "methods": list(methods)}
    if "finlay" in methods:
        out["finlay"] = finlay_wilkinson(table)
    if "ammi" in methods:
        out["ammi"] = ammi(table, n_components=n_components)
    return out