"""纯 Python 配合力分析（GCA / SCA），按交配设计分支。

仅依赖 numpy，无 R / 外部运行时。

模型 A —— 双列（full_diallel / partial_diallel，Griffing 对称双亲，每组合一条均值观测）：
    y_ij = μ + g_i + g_j + s_ij
    - μ：总体均值；
    - g_i、g_j：亲本一般配合力（GCA，施加 Σg=0 约束）；
    - s_ij：组合特殊配合力（SCA）= 残差 y_ij - (μ + g_i + g_j)。

模型 B —— line×tester（line_tester / NCII）与 NCIII（nciii）两因素模型：
    y_ij = μ + l_i + t_j + (lt)_ij
    - l_i：line（母本）一般配合力，t_j：tester（父本）一般配合力（各施加 Σ=0 约束）；
    - (lt)_ij：line×tester 互作 = 特殊配合力 SCA（残差）。
    NCIII 即 tester 恰好为 2 个的同一模型（测交种 T1/T2）。

与旧 R(lme4) 随机效应版的语义一致（GCA 归并到亲本级、SCA 归并到组合级），
但用最小二乘 BLUE（经典 Griffing / 经典 NCII），确定、可复现、无 R 依赖。
"""
from __future__ import annotations

import numpy as np

from . import fdist

DIALLEL_DESIGNS = {"full_diallel", "partial_diallel"}
TWO_FACTOR_DESIGNS = {"line_tester", "ncii", "nciii"}
DESIGNS = DIALLEL_DESIGNS | TWO_FACTOR_DESIGNS


def solve(rows: list[dict], design_type: str = "full_diallel") -> dict:
    """运行配合力 GCA / SCA，按交配设计分支。

    参数:
        rows: [{"combo": 组合标识(str), "parent1": 母本种质id(int),
                "parent2": 父本种质id(int), "value": 组合表型均值(float)}]
        design_type: full_diallel(完全双列,默认) / partial_diallel(部分双列) /
                     line_tester(line×tester NCII 两因素) / nciii(NCIII 测交)。
                     NCII/NCIII 中 parent1=line(母本)、parent2=tester(父本)，两角色集合须不相交。

    返回:
        gca: {parent_id(str): float} 各亲本一般配合力（Σg=0）；
        sca: {combo: {"parent1": int, "parent2": int, "sca": float}} 组合特殊配合力；
        gca_se: {parent_id(str): float} GCA 估计标准误；
        anova: 配合力变异分解（GCA vs SCA / line vs tester vs 互作），
               F/p 给出「对应配合力是否显著」。NCII/NCIII 附 roles 标注 line/tester。
    """
    if not rows:
        raise ValueError("无组合表型数据")
    dt = (design_type or "full_diallel").lower()
    if dt not in DESIGNS:
        raise ValueError(f"未知交配设计: {design_type}（可用: {sorted(DESIGNS)}）")
    if dt in TWO_FACTOR_DESIGNS:
        return _solve_two_factor(rows, dt)
    return _solve_diallel(rows, dt)


def _solve_diallel(rows: list[dict], design_type: str) -> dict:
    """模型 A：Griffing 对称双亲。full_diallel 与 partial_diallel 同一最小二乘。"""
    parents = sorted({r["parent1"] for r in rows} | {r["parent2"] for r in rows})
    if len(parents) < 2:
        raise ValueError("至少需要 2 个不同亲本才能估计配合力")

    n_p = len(parents)
    p_idx = {p: i for i, p in enumerate(parents)}
    n = len(rows)
    last = parents[-1]

    # 设计阵：截距 + 每个亲本一列（自交时亲本同格计数 2）。X 保持原样用于预测。
    X = np.zeros((n, 1 + n_p))
    X[:, 0] = 1.0
    for k, r in enumerate(rows):
        for p in (r["parent1"], r["parent2"]):
            X[k, 1 + p_idx[p]] += 1.0
    # 施加 Σg=0：在副本上用最后一亲本列吸收其余亲本列，得到可识别的设计阵。
    Xd = X[:, :n_p].copy()
    for j in range(n_p - 1):
        Xd[:, 1 + j] -= X[:, 1 + (n_p - 1)]

    y = np.array([float(r["value"]) for r in rows])
    beta, *_ = np.linalg.lstsq(Xd, y, rcond=None)
    mu = float(beta[0])
    g = np.zeros(n_p)
    g[:-1] = beta[1:]
    g[-1] = -float(g[:-1].sum())

    pred = X @ np.concatenate([[mu], g])
    gca = {str(p): round(float(g[p_idx[p]]), 6) for p in parents}
    sca: dict[str, dict] = {}
    sca_vals = np.zeros(n)
    for k, r in enumerate(rows):
        s = float(y[k] - pred[k])
        sca_vals[k] = s
        sca[r["combo"]] = {
            "parent1": int(r["parent1"]),
            "parent2": int(r["parent2"]),
            "sca": round(s, 6),
        }

    # 配合力方差分解：SS_gca（亲本列贡献）、SS_sca（残差，无重复观测时作误差项）。
    df_gca = n_p - 1
    df_sca = n - 1 - df_gca
    ss_total = float((y - y.mean()) @ (y - y.mean()))
    ss_gca = float((pred - y.mean()) @ (pred - y.mean()))
    ss_sca = float(sca_vals @ sca_vals)
    ms_gca = ss_gca / df_gca if df_gca > 0 else 0.0
    ms_sca = ss_sca / df_sca if df_sca > 0 else 0.0
    f_gca = ms_gca / ms_sca if ms_sca > 0 else 1.0
    p_gca = fdist.f_pvalue(f_gca, df_gca, df_sca) if (df_sca > 0 and ms_sca > 0) else 0.5
    # GCA 标准误：SE = sqrt(MSe / n_i)，n_i = 亲本出现的观测次数；MSe 用 MS_sca 近似。
    gca_se: dict[str, float] = {}
    warnings: list[str] = []
    for p in parents:
        cnt = sum(1 for r in rows if r["parent1"] == p or r["parent2"] == p)
        gca_se[str(p)] = round(float(np.sqrt(ms_sca / cnt)) if ms_sca > 0 and cnt > 0 else 0.0, 6)
        if cnt == 1:
            warnings.append(f"亲本 {p} 仅出现在 1 个组合，GCA 与 SCA 部分混杂（设计不完备）")

    return {
        "gca": gca,
        "sca": sca,
        "gca_se": gca_se,
        "anova": {
            "design": design_type,
            "df_gca": df_gca,
            "df_sca": df_sca,
            "ss_gca": round(ss_gca, 4),
            "ss_sca": round(ss_sca, 4),
            "ms_gca": round(ms_gca, 4),
            "ms_sca": round(ms_sca, 4),
            "f_gca": round(float(f_gca), 4),
            "p_gca": round(p_gca, 6),
            "warnings": warnings,
        },
    }


def _solve_two_factor(rows: list[dict], design_type: str) -> dict:
    """模型 B：line×tester 两因素（NCII）/ NCIII。

    parent1=line(母本)、parent2=tester(父本)。line GCA 与 tester GCA 分别估计，
    SCA = line×tester 互作（残差）。NCIII 要求 tester 恰好 2 个。
    """
    lines = sorted({r["parent1"] for r in rows})
    testers = sorted({r["parent2"] for r in rows})
    if not lines or not testers:
        raise ValueError("line×tester 设计需母本(line)与父本(tester)亲本")
    overlap = set(lines) & set(testers)
    if overlap:
        raise ValueError(
            f"line×tester 设计父本角色重叠: {sorted(overlap)}——同一亲本不能既作 line 又作 tester"
        )
    if design_type == "nciii" and len(testers) != 2:
        raise ValueError(f"NCIII 设计需恰好 2 个测交种(tester)，当前 {len(testers)} 个")

    n_l, n_t = len(lines), len(testers)
    l_idx = {p: i for i, p in enumerate(lines)}
    t_idx = {p: i for i, p in enumerate(testers)}
    n = len(rows)

    # 设计阵：截距 + 每 line 一列 + 每 tester 一列（全哑变量）。X 保持原样用于预测。
    X = np.zeros((n, 1 + n_l + n_t))
    X[:, 0] = 1.0
    for k, r in enumerate(rows):
        X[k, 1 + l_idx[r["parent1"]]] = 1.0
        X[k, 1 + n_l + t_idx[r["parent2"]]] = 1.0
    # 施加 Σl=0、Σt=0：选取截距 + 前 n_l-1 个 line 列 + 前 n_t-1 个 tester 列（基线 line/tester 列
    # 被吸收进截距），再在副本上用基线列吸收其余列，得到可识别的 (1+(n_l-1)+(n_t-1)) 列设计阵。
    # 注意 tester 列位于 1+n_l 起，不能整段切片。
    sel = [0] + list(range(1, n_l)) + list(range(1 + n_l, 1 + n_l + n_t - 1))
    Xd = X[:, sel].copy()
    for j in range(n_l - 1):
        Xd[:, 1 + j] -= X[:, 1 + (n_l - 1)]
    for j in range(n_t - 1):
        Xd[:, 1 + (n_l - 1) + j] -= X[:, 1 + n_l + (n_t - 1)]

    y = np.array([float(r["value"]) for r in rows])
    beta, *_ = np.linalg.lstsq(Xd, y, rcond=None)
    mu = float(beta[0])
    g_l = np.zeros(n_l)
    g_l[:-1] = beta[1:n_l]
    g_l[-1] = -float(g_l[:-1].sum())
    g_t = np.zeros(n_t)
    g_t[:-1] = beta[n_l:n_l + n_t - 1]
    g_t[-1] = -float(g_t[:-1].sum())

    pred = X @ np.concatenate([[mu], g_l, g_t])
    gca: dict[str, float] = {}
    for i, p in enumerate(lines):
        gca[str(p)] = round(float(g_l[i]), 6)
    for j, p in enumerate(testers):
        gca[str(p)] = round(float(g_t[j]), 6)
    sca: dict[str, dict] = {}
    sca_vals = np.zeros(n)
    for k, r in enumerate(rows):
        s = float(y[k] - pred[k])
        sca_vals[k] = s
        sca[r["combo"]] = {
            "parent1": int(r["parent1"]),
            "parent2": int(r["parent2"]),
            "sca": round(s, 6),
        }

    # 顺序方差分解：SS_line（截距→line）、SS_tester（→tester，扣除 line）、SS_sca（残差=互作）。
    ss_total = float((y - y.mean()) @ (y - y.mean()))
    rss_line = _residual_ss(y, _design(X, n_l))
    rss_tester = _residual_ss(y, Xd)  # 全模型(含 line+tester)残差 = Σsca²
    ss_line = ss_total - rss_line
    ss_tester = rss_line - rss_tester
    ss_sca = float(sca_vals @ sca_vals)
    df_line = n_l - 1
    df_tester = n_t - 1
    df_sca = n - 1 - df_line - df_tester
    ms_line = ss_line / df_line if df_line > 0 else 0.0
    ms_tester = ss_tester / df_tester if df_tester > 0 else 0.0
    ms_sca = ss_sca / df_sca if df_sca > 0 else 0.0
    f_line = ms_line / ms_sca if ms_sca > 0 else 1.0
    f_tester = ms_tester / ms_sca if ms_sca > 0 else 1.0
    p_line = fdist.f_pvalue(f_line, df_line, df_sca) if (df_sca > 0 and ms_sca > 0) else 0.5
    p_tester = fdist.f_pvalue(f_tester, df_tester, df_sca) if (df_sca > 0 and ms_sca > 0) else 0.5
    # GCA 标准误：SE = sqrt(MSe / n_i)，n_i = 亲本出现次数（line 每 tester 组 n_t 次等）。
    gca_se: dict[str, float] = {}
    for p in lines:
        cnt = sum(1 for r in rows if r["parent1"] == p)
        gca_se[str(p)] = round(float(np.sqrt(ms_sca / cnt)) if ms_sca > 0 and cnt > 0 else 0.0, 6)
    for p in testers:
        cnt = sum(1 for r in rows if r["parent2"] == p)
        gca_se[str(p)] = round(float(np.sqrt(ms_sca / cnt)) if ms_sca > 0 and cnt > 0 else 0.0, 6)

    return {
        "gca": gca,
        "sca": sca,
        "gca_se": gca_se,
        "anova": {
            "design": design_type,
            "roles": {"lines": lines, "testers": testers},
            "df_line": df_line,
            "df_tester": df_tester,
            "df_sca": df_sca,
            "ss_line": round(ss_line, 4),
            "ss_tester": round(ss_tester, 4),
            "ss_sca": round(ss_sca, 4),
            "ms_line": round(ms_line, 4),
            "ms_tester": round(ms_tester, 4),
            "ms_sca": round(ms_sca, 4),
            "f_line": round(float(f_line), 4),
            "f_tester": round(float(f_tester), 4),
            "p_line": round(p_line, 6),
            "p_tester": round(p_tester, 6),
        },
    }


def _design(X: np.ndarray, n_l: int) -> np.ndarray:
    """截距 + 全 line 哑变量的子设计阵（SS_line 用）。"""
    return X[:, : 1 + n_l].copy()


def _residual_ss(y: np.ndarray, A: np.ndarray) -> float:
    """最小二乘残差平方和 RSS = y^T y - y^T A (A^T A)^-1 A^T y（lstsq 求拟合格）。"""
    beta, *_ = np.linalg.lstsq(A, y, rcond=None)
    pred = A @ beta
    return float((y - pred) @ (y - pred))