"""单因素随机效应 ANOVA + 广义遗传力 H²（纯 Python，numpy + fdist）。

模型：y_ij = μ + a_i + e_ij，a_i ~ N(0, σ²g)，e_ij ~ N(0, σ²e)。
分组（group）= 家系/组合，用于估计家系方差组分：
    n₀ = (N - Σ n_i²/N)/(k-1)        加权平均每组样本量
    σ²e = MSW                         组内均方 = 环境/误差方差
    σ²g = (MSB - MSW)/n₀              组间方差组分（家系遗传方差）
    H²  = σ²g/(σ²g + σ²e)             广义遗传力（单株 basis）
    F   = MSB/MSW，p = 1-P(F≤f)

鲁棒性：组数/样本不足、组内方差为 0（全组同值）、n₀≤0 时均不崩溃，
按合理默认（H²=0、F=1、p=0.5）返回并给 warning。
"""
from __future__ import annotations

import numpy as np

from . import fdist


def solve(rows: list[dict]) -> dict:
    """rows: [{"group": str|int, "value": float}]。

    返回：方差分析表字段 + σ²g/σ²e + H² + F/p + 各组均值 + warning。
    """
    if not rows:
        raise ValueError("无分组表型数据")
    y = np.array([float(r["value"]) for r in rows], dtype=float)
    g = np.array([r["group"] for r in rows])
    if len(y) < 2 or len(np.unique(g)) < 2:
        raise ValueError("至少需要 2 个组且总样本 ≥ 2 才能做方差分析")
    n = len(y)
    overall = float(y.mean())

    group_keys = sorted(set(g))
    k = len(group_keys)
    sizes: list[int] = []
    means: list[float] = []
    for gv in group_keys:
        mask = g == gv
        sizes.append(int(mask.sum()))
        means.append(float(y[mask].mean()))

    ss_total = float(((y - overall) ** 2).sum())
    ss_between = float(sum(sz * (m - overall) ** 2 for sz, m in zip(sizes, means)))
    ss_within = ss_total - ss_between
    ms_between = ss_between / (k - 1) if k > 1 else 0.0
    ms_within = ss_within / (n - k) if n > k else 0.0
    n0 = (n - sum(sz * sz for sz in sizes) / n) / (k - 1) if k > 1 else 0.0
    sigma_e = ms_within
    sigma_g = (ms_between - ms_within) / n0 if n0 > 0 else 0.0
    h2 = sigma_g / (sigma_g + sigma_e) if (sigma_g + sigma_e) > 1e-12 else 0.0
    h2 = float(np.clip(h2, 0.0, 1.0))
    f = ms_between / ms_within if ms_within > 0 else 1.0
    p = fdist.f_pvalue(f, k - 1, n - k) if (k > 1 and n > k and ms_within > 0) else 0.5

    warning = None
    if ms_within == 0:
        warning = "组内方差为 0（各组观测完全一致），方差组分不可分，H² 无信息。"
    elif sigma_g <= 1e-12:
        warning = "家系方差组分 ≈ 0，组间差异不显著，H² 接近 0。"
    return {
        "n_groups": k,
        "n": n,
        "overall_mean": round(overall, 4),
        "ss_total": round(ss_total, 4),
        "ss_between": round(ss_between, 4),
        "ss_within": round(ss_within, 4),
        "ms_between": round(ms_between, 4),
        "ms_within": round(ms_within, 4),
        "sigma_g": round(sigma_g, 4),
        "sigma_e": round(sigma_e, 4),
        "h2_broad": h2,
        "f": round(float(f), 4),
        "p_value": round(p, 6),
        "groups": [
            {"group": str(gv), "n": sz, "mean": round(m, 4)}
            for gv, sz, m in zip(group_keys, sizes, means)
        ],
        "warning": warning,
    }


def solve_rcbd(rows: list[dict]) -> dict:
    """RCBD 双因素 ANOVA：y_ij = μ + a_i(组合/家系) + b_j(区组) + e_ij。

    区组效应序贯析出（Type I：grand mean → group → block），遗传差异为区组校正后
    （SS_group 先于区组入模型、SS_block 为其后增量、残差为两者剔除后）。
    输出区组行 SS_block/MS_block/F_block/p_block 与区组均值。区组数不足或无残余
    自由度时按数据现状返回 warning，不崩溃。

    rows: [{"group": str|int, "block": str|int, "value": float}]
    """
    if not rows:
        raise ValueError("无 RCBD 表型数据")
    y = np.array([float(r["value"]) for r in rows], dtype=float)
    g = np.array([str(r["group"]) for r in rows])
    b = np.array([str(r["block"]) for r in rows])
    n = len(y)
    if n < 2:
        raise ValueError("至少需要 2 个观测")
    g_keys = sorted(set(g))
    b_keys = sorted(set(b))
    kg, kb = len(g_keys), len(b_keys)
    if kg < 2:
        raise ValueError("至少需要 2 个组合")
    if kb < 2:
        raise ValueError("至少需要 2 个区组才能做 RCBD 分析")

    overall = float(y.mean())

    def design(*factors: tuple[list, np.ndarray]) -> np.ndarray:
        # 截距 + 每个因子 drop 第一个水平的哑变量 → 满秩序贯设计阵
        cols = [np.ones(n)]
        for keys, arr in factors:
            for kv in keys[1:]:
                cols.append((arr == kv).astype(float))
        return np.column_stack(cols)

    X1 = np.ones((n, 1))
    Xg = design((g_keys, g))
    Xgb = design((g_keys, g), (b_keys, b))

    def rss(X: np.ndarray) -> float:
        beta, *_ = np.linalg.lstsq(X, y, rcond=None)
        resid = y - X @ beta
        return float(resid @ resid)

    ss_total = float(((y - overall) ** 2).sum())
    ss_group = max(rss(X1) - rss(Xg), 0.0)   # 组合效应（区组校正前）
    ss_block = max(rss(Xg) - rss(Xgb), 0.0)  # 区组效应（组合剔除后析出）
    ss_error = max(rss(Xgb), 0.0)            # 残余误差
    df_group = kg - 1
    df_block = kb - 1
    df_error = n - kg - kb + 1
    ms_group = ss_group / df_group if df_group > 0 else 0.0
    ms_block = ss_block / df_block if df_block > 0 else 0.0
    ms_error = ss_error / df_error if df_error > 0 else 0.0
    sizes = [int((g == gv).sum()) for gv in g_keys]
    means = [float(y[g == gv].mean()) for gv in g_keys]
    n0 = (n - sum(sz * sz for sz in sizes) / n) / (kg - 1) if kg > 1 else 0.0
    sigma_e = ms_error
    sigma_g = (ms_group - ms_error) / n0 if n0 > 0 else 0.0
    h2 = sigma_g / (sigma_g + sigma_e) if (sigma_g + sigma_e) > 1e-12 else 0.0
    h2 = float(np.clip(h2, 0.0, 1.0))
    f_group = ms_group / ms_error if ms_error > 0 else 1.0
    f_block = ms_block / ms_error if ms_error > 0 else 1.0
    p_group = fdist.f_pvalue(f_group, df_group, df_error) if (df_error > 0 and ms_error > 0) else 0.5
    p_block = fdist.f_pvalue(f_block, df_block, df_error) if (df_error > 0 and ms_error > 0) else 0.5

    warning = None
    if df_error <= 0:
        warning = "区组设计过饱和：无残余误差自由度（每组合每区组仅 1 株），无法估计误差项。"
    elif ms_error == 0:
        warning = "残余误差为 0（区组+组合完全解释变异），F 值无信息。"
    elif sigma_g <= 1e-12:
        warning = "家系方差组分 ≈ 0，组间差异不显著，H² 接近 0。"
    return {
        "n_groups": kg,
        "n_blocks": kb,
        "n": n,
        "overall_mean": round(overall, 4),
        "ss_total": round(ss_total, 4),
        "ss_between": round(ss_group, 4),
        "ss_block": round(ss_block, 4),
        "ss_within": round(ss_error, 4),
        "ms_between": round(ms_group, 4),
        "ms_block": round(ms_block, 4),
        "ms_within": round(ms_error, 4),
        "sigma_g": round(sigma_g, 4),
        "sigma_e": round(sigma_e, 4),
        "h2_broad": h2,
        "f": round(float(f_group), 4),
        "p_value": round(p_group, 6),
        "f_block": round(float(f_block), 4),
        "p_block": round(p_block, 6),
        "groups": [
            {"group": str(gv), "n": sz, "mean": round(m, 4)}
            for gv, sz, m in zip(g_keys, sizes, means)
        ],
        "blocks": [
            {"block": str(bv), "n": int((b == bv).sum()), "mean": round(float(y[b == bv].mean()), 4)}
            for bv in b_keys
        ],
        "warning": warning,
    }