"""基因组选择：VanRaden G 矩阵 + GBLUP + ssGBLUP（仅 numpy，无 R）。

- build_g_matrix：dosage(0/1/2) 样本×标记矩阵 → VanRaden G（method1=ZZ'/2Σp(1-p)、
  method2=ZDZ'，D=1/[2p(1-p)]），缺格按列均值(2p)填充；MAF 过滤 + blend 岭回归保证可逆。
- solve_gblup：仅基因型个体入模型，MME 用 G⁻¹ 替代 A⁻¹；方差组分用剖面 REML
  （黄金分割最大化精确 REML 对数似然，与 blup._solve_gxe 同模板）。
- solve_ssgblup：单步法 H⁻¹ = A⁻¹ + [[0,0],[0,G⁻¹ − A22⁻¹]]（Aguilar et al. 2010；
  A22 = 基因型个体在 A 中的子阵），基因型+非基因型个体都出 EBV。

可靠性 = 1 − PEV/σ²a，PEV = C22[ii]·σ²e（C 为 MME 系数矩阵逆的加性块）。
"""
from __future__ import annotations

import math
import random

import numpy as np

from scripts.breeding_stats import blup

ENGINE_VERSION = "1.4.0"  # 1.4.0: 非线性基因组预测 rrBLUP + BayesB; 1.3.0: kfold split=family; 1.2.0: SSR 多等位 + EMMAX GRM 复用
_RIDGE = 1e-6  # A22⁻¹ 病态加小岭

_GBLUP_KEYS = ["ebv", "reliability", "h2", "sigma_a", "sigma_e",
               "n_obs", "n_individuals", "n_markers", "converged", "warning"]


def build_g_matrix(genotypes: dict[str, dict[str, float]],
                   markers: list[str], *, maf_min: float = 0.05,
                   method: int = 2, blend: float = 0.02):
    """VanRaden G 矩阵。

    genotypes: {sample: {marker: dosage(0/1/2)}}；markers: 有序标记列表。
    返回 (G 按 sorted(genotypes) 排序, meta{G 统计量})。
    """
    samples = sorted(genotypes)
    n = len(samples)
    if n < 2 or not markers:
        raise ValueError("基因型样本或标记数不足（需 n≥2、m≥1）")
    m = len(markers)
    M = np.full((n, m), np.nan)
    n_missing = 0
    for i, s in enumerate(samples):
        g = genotypes[s]
        for j, mk in enumerate(markers):
            v = g.get(mk)
            if v is None or (isinstance(v, float) and math.isnan(v)):
                n_missing += 1
            else:
                M[i, j] = float(v)
    p = np.zeros(m)
    for j in range(m):
        col = M[:, j]
        valid = col[~np.isnan(col)]
        if len(valid) == 0:
            p[j] = 0.5
        else:
            p[j] = float(np.mean(valid)) / 2.0
        M[:, j] = np.where(np.isnan(M[:, j]), 2.0 * p[j], M[:, j])
    maf = np.minimum(p, 1.0 - p)
    keep = maf >= maf_min
    if int(keep.sum()) == 0:
        raise ValueError(f"MAF≥{maf_min} 的标记为 0，请降低 maf_min 或检查标记等位频率")
    M = M[:, keep]
    p = p[keep]
    m_keep = int(keep.sum())
    Z = M - 2.0 * p[None, :]
    if method == 1:
        denom = 2.0 * float(np.sum(p * (1.0 - p)))
        G = (Z @ Z.T) / denom if denom > 0 else np.eye(n)
    else:
        D = 1.0 / (2.0 * p * (1.0 - p))
        G = Z @ np.diag(D) @ Z.T
    if blend and blend > 0:
        diag_mean = float(np.mean(np.diag(G)))
        G = (1.0 - blend) * G + blend * diag_mean * np.eye(n)
    off = G[np.triu_indices(n, 1)] if n > 1 else np.array([0.0])
    return G, {
        "n_individuals": n,
        "n_markers": m,
        "m_after_maf": m_keep,
        "n_missing": n_missing,
        "maf_min": maf_min,
        "method": method,
        "blend": blend,
        "diag_mean": round(float(np.mean(np.diag(G))), 6),
        "diag_range": [round(float(np.min(np.diag(G))), 6), round(float(np.max(np.diag(G))), 6)],
        "offdiag_range": [round(float(np.min(off)), 6), round(float(np.max(off)), 6)],
    }


def _scale_g_to_a22(G: np.ndarray, A22: np.ndarray) -> np.ndarray:
    """把基因关系阵 G 缩放到系谱关系阵 A22 的尺度（Christensen et al. 调整）。

    VanRaden method=2 的 G 对角线≈标记数 m（远大于 A22 的 ≈1），若直接组 H，
    基因型块会被 G 主导、h²/EBV 尺度偏置。用两参数对齐 G*=a·G+b 使
    对角均值/非对角均值与 A22 一致；G 退化（对角≈非对角）时跳过并原样返回。
    """
    n = len(A22)
    if n < 2:
        return G
    diag_g = np.mean(np.diag(G))
    diag_a = np.mean(np.diag(A22))
    off_g = float(np.mean(G[np.triu_indices(n, 1)]))
    off_a = float(np.mean(A22[np.triu_indices(n, 1)]))
    denom = diag_g - off_g
    if abs(denom) < 1e-12:
        return G
    a = (diag_a - off_a) / denom
    b = off_a - a * off_g
    return a * G + b


def _reml_ll(ratio: float, ZGZ: np.ndarray, X: np.ndarray, y: np.ndarray) -> float:
    """V = ZGZ·ratio + I（Ve=1 尺度），精确 REML 对数似然（slogdet + y'Py）。"""
    m = len(y)
    V = ZGZ * ratio + np.eye(m)
    try:
        VinvX = np.linalg.solve(V, X)
        Vinvy = np.linalg.solve(V, y)
        XtVinvX = X.T @ VinvX
        XtVinvY = X.T @ Vinvy
        yPy = float(y @ Vinvy) - float(XtVinvY @ np.linalg.solve(XtVinvX, XtVinvY))
        _, lv = np.linalg.slogdet(V)
        _, lx = np.linalg.slogdet(XtVinvX)
        return -0.5 * (float(lv) + float(lx) + yPy)
    except np.linalg.LinAlgError:
        return -np.inf


def _profile_ve(V: np.ndarray, X: np.ndarray, y: np.ndarray, df: int) -> float:
    """REML 剖面 σ²e = y'Py/df（V 固定尺度下）。"""
    VinvX = np.linalg.solve(V, X)
    Vinvy = np.linalg.solve(V, y)
    XtVinvX = X.T @ VinvX
    XtVinvY = X.T @ Vinvy
    yPy = float(y @ Vinvy) - float(XtVinvY @ np.linalg.solve(XtVinvX, XtVinvY))
    return yPy / df if df > 0 else float("nan")


def _profile_solve(ZGZ: np.ndarray, X: np.ndarray, y: np.ndarray,
                   Ginv: np.ndarray, n_all: int, z: np.ndarray):
    """共享剖面 REML + MME 求解。返回 (ebv, rel, h2, sigma_a, sigma_e, n_iter)。"""
    m = len(y)
    df = m - X.shape[1]
    # 剖面参数为 log10(σ²a/σ²e)。h²=ratio/(1+ratio)，故 h²∈[H2_MIN,H2_MAX]
    # 对应 ratio∈[H2_MIN/(1−H2_MIN), H2_MAX/(1−H2_MAX)]（直接用 log10(H2_MAX)≈−0.0004
    # 会把 σ²a/σ²e 封顶 0.999，令 h² 上限只有 0.5）。
    lo, hi = math.log10(blup.H2_MIN / (1.0 - blup.H2_MIN)), math.log10(blup.H2_MAX / (1.0 - blup.H2_MAX))

    def _ll(r: float) -> float:
        return _reml_ll(10.0 ** r, ZGZ, X, y)

    r_opt, n_iter = blup._golden_max(_ll, lo, hi, tol=1e-4, max_iter=blup.MAX_PROFILE_EVALS)
    ratio = 10.0 ** r_opt
    V = ZGZ * ratio + np.eye(m)
    ve = _profile_ve(V, X, y, df)
    va = ratio * ve if not math.isnan(ve) else 1.0
    lam = ve / va if va > 0 else 1.0

    Zmat = np.zeros((n_all, m))
    Zmat[z, np.arange(m)] = 1.0
    C = np.block([
        [X.T @ X, X.T @ Zmat.T],
        [Zmat @ X, Zmat @ Zmat.T + lam * Ginv],
    ])
    rhs = np.concatenate([X.T @ y, Zmat @ y])
    try:
        sol = np.linalg.solve(C, rhs)
    except np.linalg.LinAlgError:
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None,
                "sigma_e": None, "n_iter": n_iter, "converged": False,
                "warning": "MME 系数矩阵奇异（基因型共线/样本不足），求解失败"}
    C_inv = np.linalg.inv(C)
    C22inv = C_inv[1:, 1:]
    rel = np.clip(1.0 - np.diag(C22inv) * ve / va, 0.0, 1.0) if va > 0 else np.zeros(n_all)
    h2 = va / (va + ve) if (va + ve) > 0 else None
    return {"ebv": sol[1:], "reliability": rel, "h2": h2, "sigma_a": float(va),
            "sigma_e": float(ve) if not math.isnan(ve) else None,
            "n_iter": n_iter, "converged": True, "warning": None}


def solve_gblup(individuals: list[str], phenos: dict[str, float], G: np.ndarray,
                *, tol: float = blup.TOL):
    """GBLUP：仅基因型个体。individuals 为 G 行序（个体 id）；phenos 子集有观测。"""
    n = len(individuals)
    obs = [i for i in individuals if i in phenos]
    m = len(obs)
    if m < 3:
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
                "n_obs": m, "n_individuals": n, "n_markers": None,
                "converged": False, "warning": "表型观测不足 3，无法估计方差组分"}
    idx = {i: k for k, i in enumerate(individuals)}
    z = np.array([idx[i] for i in obs])
    y = np.array([float(phenos[i]) for i in obs])
    X = np.ones((m, 1))
    ZGZ = G[np.ix_(z, z)]
    try:
        Ginv = np.linalg.inv(G)
    except np.linalg.LinAlgError:
        Ginv = np.linalg.pinv(G)
    res = _profile_solve(ZGZ, X, y, Ginv, n, z)
    out = dict(res)
    ebv = {individuals[i]: float(out["ebv"][i]) for i in range(n)}
    rel = {individuals[i]: float(out["reliability"][i]) for i in range(n)}
    out.update(ebv=ebv, reliability=rel, n_obs=m, n_individuals=n,
               n_markers=None if not isinstance(G, np.ndarray) else G.shape[0])
    if out["warning"] is None and out.get("h2") is not None and (out["h2"] >= blup.H2_MAX or out["h2"] <= blup.H2_MIN):
        out["warning"] = "遗传力估计逼近边界，剖面 REML 未完全收敛（样本小/信噪比低）"
    return out


def _build_h(pedigree: list[dict], genotyped: list[str], G: np.ndarray,
             *, ridge: float = _RIDGE, extra_base: list[str] | None = None):
    """ssGBLUP 合并关系阵（H 协方差 + H⁻¹ 仅进 MME，Aguilar et al. 2010）。

    返回 (order, idx, H, Hinv)。系谱中缺失的基因型个体按 base 补入（与 phenos 兜底一致）。
    """
    base: list[int] = []
    non_base: list[tuple[int, int | None, int | None]] = []
    seen: set[int] = set()
    for rec in pedigree:
        i = rec["individual"]
        if i in seen:
            continue
        seen.add(i)
        d, s = rec.get("dam"), rec.get("sire")
        if d is None and s is None:
            base.append(i)
        else:
            non_base.append((i, d, s))
    if extra_base:
        for g in extra_base:
            if g not in seen:
                seen.add(g)
                base.append(g)
    for g in genotyped:  # 无表型/无系谱基因型树也按 base 补入 H（不再 raise）
        if g not in seen:
            seen.add(g)
            base.append(g)
    if not seen:
        raise ValueError("系谱与基因型均为空")
    order, idx = blup._order_pedigree(base, non_base)
    n = len(order)
    parent_of = {i: (d, s) for (i, d, s) in non_base}
    A, Ainv = blup._build_ainv(order, len(base), parent_of)

    gen_indices = [idx[g] for g in genotyped]
    if len(gen_indices) < 2:
        raise ValueError("基因型个体不足 2，无法 ssGBLUP")
    A22 = A[np.ix_(gen_indices, gen_indices)]
    G = _scale_g_to_a22(G, A22)
    try:
        A22inv = np.linalg.inv(A22 + ridge * np.eye(len(gen_indices)))
        Ginv = np.linalg.inv(G)
    except np.linalg.LinAlgError:
        A22inv = np.linalg.pinv(A22)
        Ginv = np.linalg.pinv(G)
    Hinv = Ainv.copy()
    Hinv[np.ix_(gen_indices, gen_indices)] += Ginv - A22inv

    # H 合并关系阵（协方差；H⁻¹=A⁻¹+[[0,0],[0,G⁻¹−A22⁻¹]]，Aguilar et al. 2010）
    # H = A + [[A12·A22⁻¹·D·A22⁻¹·A21, A12·A22⁻¹·D], [D·A22⁻¹·A21, D]]，D=G−A22
    gen_set = set(gen_indices)
    non_gen = [i for i in range(n) if i not in gen_set]
    if non_gen:
        D = G - A22
        L = A[np.ix_(non_gen, gen_indices)] @ A22inv
        LD = L @ D
        H = A.copy()
        H[np.ix_(non_gen, non_gen)] += LD @ L.T
        H[np.ix_(non_gen, gen_indices)] += LD
        H[np.ix_(gen_indices, non_gen)] += LD.T
        H[np.ix_(gen_indices, gen_indices)] = G
    else:
        H = G
    return order, idx, H, Hinv


def solve_ssgblup(pedigree: list[dict], phenos: dict[str, float],
                  G: np.ndarray, genotyped: list[str], *, tol: float = blup.TOL,
                  ridge: float = _RIDGE):
    """ssGBLUP：单步法。genotyped 为有序列表（与 G 行对应）。全系谱个体出 EBV。"""
    if not pedigree and not phenos:
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
                "n_obs": 0, "n_individuals": 0, "n_markers": None,
                "converged": False, "warning": "系谱与表型均为空"}
    try:
        order, idx, H, Hinv = _build_h(
            pedigree, genotyped, G, ridge=ridge, extra_base=list(phenos))
    except ValueError as e:
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
                "n_obs": 0, "n_individuals": len({r["individual"] for r in pedigree}),
                "n_markers": None, "converged": False, "warning": str(e)}
    n = len(order)

    obs = [i for i in order if i in phenos]
    m = len(obs)
    if m < 3:
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
                "n_obs": m, "n_individuals": n, "n_markers": None,
                "converged": False, "warning": "表型观测不足 3，无法估计方差组分"}
    z = np.array([idx[i] for i in obs])
    y = np.array([float(phenos[i]) for i in obs])
    X = np.ones((m, 1))
    ZGZ = H[np.ix_(z, z)]  # V 用 H 协方差块（H⁻¹ 仅进 MME 系数矩阵）
    res = _profile_solve(ZGZ, X, y, Hinv, n, z)
    out = dict(res)
    ebv = {order[i]: float(out["ebv"][i]) for i in range(n)}
    rel = {order[i]: float(out["reliability"][i]) for i in range(n)}
    out.update(ebv=ebv, reliability=rel, n_obs=m, n_individuals=n, n_markers=len(genotyped))
    if out["warning"] is None and out.get("h2") is not None and (out["h2"] >= blup.H2_MAX or out["h2"] <= blup.H2_MIN):
        out["warning"] = "遗传力估计逼近边界，剖面 REML 未完全收敛"
    return out


def _build_dosage_matrix(genotypes: dict[str, dict[str, float]],
                         markers: list[str], *, maf_min: float = 0.05):
    """dosage dict → 有序样本×标记剂量矩阵（与 build_g_matrix 同 MAF 过滤口径）。

    返回 (X[n×q], 保留标记列表, 等位频率 p[q])。X 已按列均值(≈2p)填充缺失。
    rrBLUP/BayesB 直接消费 X（Xc = X − 2p 即 GBLUP 的 Z）。
    """
    samples = sorted(genotypes)
    n = len(samples)
    m = len(markers)
    M = np.full((n, m), np.nan)
    for i, s in enumerate(samples):
        g = genotypes[s]
        for j, mk in enumerate(markers):
            v = g.get(mk)
            if v is not None and not (isinstance(v, float) and math.isnan(v)):
                M[i, j] = float(v)
    p = np.zeros(m)
    for j in range(m):
        col = M[:, j]
        valid = col[~np.isnan(col)]
        p[j] = 0.5 if len(valid) == 0 else float(np.mean(valid)) / 2.0
        M[:, j] = np.where(np.isnan(M[:, j]), 2.0 * p[j], M[:, j])
    keep = np.minimum(p, 1.0 - p) >= maf_min
    if int(keep.sum()) == 0:
        raise ValueError(f"MAF≥{maf_min} 的标记为 0，请降低 maf_min 或检查标记等位频率")
    return M[:, keep], [mk for mk, k in zip(markers, keep) if k], p[keep]


def solve_rrblup(individuals: list[str], phenos: dict[str, float],
                 X: np.ndarray, marker_order: list[str], p: np.ndarray,
                 *, tol: float = blup.TOL):
    """RR-BLUP（岭回归逐标记）：y = μ + Xu + e，u=(Xc'Xc+λI)⁻¹Xc'yc。

    与 GBLUP 数学对偶（G = XcXc'/c 时 EBV 一致）：方差组分仍用剖面 REML
    （_profile_solve 同一求解器）估 σa²/σe² → σu²=σa²/c（c=2Σp(1-p)）→
    λ=σe²/σu²。输出 EBV（基因型个体，含表型截距）+ marker effects（GBLUP
    无的解释能力）+ 与 GBLUP 同构的可靠性（PEV 近似）。
    """
    n = len(individuals)
    obs = [i for i in individuals if i in phenos]
    m = len(obs)
    if m < 3:
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
                "n_obs": m, "n_individuals": n, "n_markers": None, "marker_effects": {},
                "converged": False, "warning": "表型观测不足 3，无法估计方差组分"}
    idx = {i: k for k, i in enumerate(individuals)}
    z = np.array([idx[i] for i in obs])
    y = np.array([float(phenos[i]) for i in obs])
    Xc = X - 2.0 * p[None, :]  # 中心化剂量 = GBLUP 的 Z
    q = Xc.shape[1]
    c = 2.0 * float(np.sum(p * (1.0 - p)))
    if c <= 0:
        c = 1.0
    G = Xc @ Xc.T / c
    try:
        Ginv = np.linalg.inv(G)
    except np.linalg.LinAlgError:
        Ginv = np.linalg.pinv(G)
    res = _profile_solve(G[np.ix_(z, z)], np.ones((m, 1)), y, Ginv, n, z)
    va, ve = res.get("sigma_a"), res.get("sigma_e")
    if not res.get("converged") or va is None or ve is None or va <= 0 or ve <= 0:
        return {"ebv": {}, "reliability": {}, "h2": res.get("h2"), "sigma_a": va,
                "sigma_e": ve, "n_obs": m, "n_individuals": n, "n_markers": q,
                "marker_effects": {},
                "converged": False, "warning": res.get("warning") or "方差组分估计失败，rrBLUP 未收敛"}
    lam = ve / (va / c) if va > 0 else 1.0  # λ = σe²/σu² = σe²·c/σa²
    yc = y - float(y.mean())
    Xo = Xc[z]
    try:
        u = np.linalg.solve(Xo.T @ Xo + lam * np.eye(q), Xo.T @ yc)
    except np.linalg.LinAlgError:
        u = np.linalg.pinv(Xo.T @ Xo + lam * np.eye(q)) @ (Xo.T @ yc)
    mu = float(y.mean())
    ebv = {individuals[i]: float(Xc[i] @ u + mu) for i in range(n)}
    rel = res["reliability"]
    out = {"ebv": ebv, "reliability": {individuals[i]: float(rel[i]) for i in range(n)},
           "h2": res["h2"], "sigma_a": float(va), "sigma_e": float(ve),
           "n_obs": m, "n_individuals": n, "n_markers": q,
           "marker_effects": {marker_order[j]: float(u[j]) for j in range(q)},
           "converged": True, "warning": res.get("warning")}
    return out


def _bayesb_gibbs(Xo: np.ndarray, yc: np.ndarray, *, n_iter: int, burn_in: int,
                  seed: int, pi: float, df_s: float, scale0: float,
                  df_e: float, scale_e0: float) -> tuple[np.ndarray, float, float, int]:
    """BayesB 逐标记 Gibbs 核心：给定中心化剂量观测子阵 Xo 与中心化表型 yc。

    返回 (u_hat 后验均值, σe² 末段, SSE 末段, 收敛样本数 n_save)。
    numpy 固定 seed（default_rng）保证同数据同参数逐位可复现。
    """
    q = Xo.shape[1]
    m = len(yc)
    rng = np.random.default_rng(seed)
    u = np.zeros(q)
    e = yc.astype(float).copy()
    sigma_e2 = float(np.var(yc))
    if sigma_e2 <= 0:
        sigma_e2 = 1.0
    u_acc = np.zeros(q)
    n_save = 0
    last_sse = float("inf")
    for it in range(n_iter):
        for j in range(q):
            e_star = e + Xo[:, j] * u[j]
            xjxj = float(Xo[:, j] @ Xo[:, j])
            xjy = float(Xo[:, j] @ e_star)
            if rng.uniform() > pi:  # 1−π 概率标记入模型
                s2 = (df_s * scale0 + u[j] * u[j]) / (rng.chisquare(df_s + 1))
                if s2 <= 0 or xjxj <= 0:
                    u[j] = 0.0
                else:
                    denom = xjxj + sigma_e2 / s2
                    u[j] = rng.normal(xjy / denom, math.sqrt(sigma_e2 / denom))
            else:
                u[j] = 0.0
            e = e_star - Xo[:, j] * u[j]
        sse = float(e @ e)
        sigma_e2 = (df_e * scale_e0 + sse) / rng.chisquare(df_e + m)
        if sigma_e2 <= 0:
            sigma_e2 = 1e-8
        if it >= burn_in:
            u_acc += u
            n_save += 1
        last_sse = sse
    if n_save == 0:
        raise ValueError("后验样本为空（burn_in 覆盖全部迭代）")
    return u_acc / n_save, sigma_e2, last_sse, n_save


def solve_bayesb(individuals: list[str], phenos: dict[str, float],
                 X: np.ndarray, marker_order: list[str], p: np.ndarray,
                 *, n_iter: int = 2000, burn_in: int = 500, seed: int = 20260805,
                 pi: float = 0.95, df_s: float = 5.0, scale0: float = 0.001,
                 df_e: float = 5.0, scale_e0: float = 1.0,
                 tol: float = blup.TOL):
    """BayesB（Meuwissen et al. 2001）：贝叶斯可变选择，逐标记 Gibbs 抽样。

    模型 y = μ + Σ Xc_j·u_j + e；每标记 δ_j ~ Bernoulli(1−π) 概率入模型，
    u_j | δ=1 ~ N(0, σ²j)、σ²j ~ inv-χ²(df_s, scale)；σe² ~ inv-χ²(df_e, scale_e0)。
    固定 seed（numpy default_rng）保证可复现（MLOps 铁律：同数据同参数重跑逐位一致）。
    输出 EBV（后验均值，基因型个体，含表型截距）+ marker effects + 收敛诊断。
    可靠性未单独估计（后验方差才是正确不确定性度量），按 0 保守处理（note 说明）。
    """
    n = len(individuals)
    obs = [i for i in individuals if i in phenos]
    m = len(obs)
    if m < 3:
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
                "n_obs": m, "n_individuals": n, "n_markers": None, "marker_effects": {},
                "converged": False, "warning": "表型观测不足 3，无法估计方差组分"}
    if n_iter <= burn_in:
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
                "n_obs": m, "n_individuals": n, "n_markers": None, "marker_effects": {},
                "converged": False, "warning": "n_iter 必须大于 burn_in"}
    idx = {i: k for k, i in enumerate(individuals)}
    z = np.array([idx[i] for i in obs])
    y = np.array([float(phenos[i]) for i in obs])
    Xc = X - 2.0 * p[None, :]
    q = Xc.shape[1]
    yc = y - float(y.mean())
    Xo = Xc[z]

    try:
        u_hat, sigma_e2, sse, n_save = _bayesb_gibbs(
            Xo, yc, n_iter=n_iter, burn_in=burn_in, seed=seed, pi=pi,
            df_s=df_s, scale0=scale0, df_e=df_e, scale_e0=scale_e0)
    except Exception as exc:  # noqa: BLE001
        return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
                "n_obs": m, "n_individuals": n, "n_markers": q, "marker_effects": {},
                "converged": False, "warning": f"BayesB 抽样异常: {exc!s}"}
    mu = float(yc.mean())
    ebv = {individuals[i]: float(Xc[i] @ u_hat + mu) for i in range(n)}
    note = (f"BayesB π={pi}·{n_iter}iter·{burn_in}burn-in·seed={seed}（后验均值，"
            f"收敛样本 {n_save}；σe² 末段 SSE={sse:.3f}）")
    return {"ebv": ebv,
            "reliability": {individuals[i]: 0.0 for i in range(n)},
            "h2": None, "sigma_a": None, "sigma_e": float(sigma_e2),
            "n_obs": m, "n_individuals": n, "n_markers": q,
            "marker_effects": {marker_order[j]: float(u_hat[j]) for j in range(q)},
            "converged": True, "warning": note}


def kfold_cv_genomic(phenos: dict[str, float], G: np.ndarray, genotyped: list[str],
                     *, k: int = 5, seed: int = 20260804, method: str = "gblup",
                     pedigree: list[dict] | None = None,
                     ridge: float = _RIDGE,
                     split: str = "random",
                     family_of: dict[str, str] | None = None,
                     genotypes: dict[str, dict[str, float]] | None = None,
                     markers: list[str] | None = None,
                     n_iter: int = 1000, burn_in: int = 300, pi: float = 0.95,
                     df_s: float = 5.0, scale0: float = 0.001,
                     df_e: float = 5.0, scale_e0: float = 1.0) -> dict:
    """GS k 折交叉验证（实证预测准确度：留出 GEBV vs 表型 pearson/RMSE）。

    与 blup.kfold_cv 区别：关系矩阵为 VanRaden G（GBLUP）或单步 H（ssGBLUP），
    折划分默认按**基因型个体随机分层**（固定种子；GS 同世代样本无系谱家系树，
    随机分层即标准做法，区别于 ABLUP 的 sire 家系留出）。若样本确有家系结构
    （同父半同胞），随机分层会把同家系亲缘个体摊入 train/test → 泄漏亲缘、
    高估准确度，此时应传 `split="family"` + `family_of` 走**家系阻塞折**
    （同家系个体同折，照 blup.kfold_cv 的 sire 家系留出模式）；无家系个体自成一族。
    家系数 < 折数 k 时折叠为 n_families 折并附 warning。

    method 分派：
    - `gblup`：VanRaden G（折内 _profile_solve，掩蔽留出经 G⁻¹ 交叉关系预测）。
    - `ssgblup`：单步 H（需 pedigree）。
    - `rrblup`：岭回归逐标记（需 genotypes+markers 剂量矩阵；折内 REML 定 λ，
      EBV=Xc·u+μ 直接由标记基因型预测，不经关系矩阵逆）。
    - `bayesb`：贝叶斯可变选择（需 genotypes+markers；折内 Gibbs，固定 seed 可复现）。
    rrblup/bayesb 缺 genotypes/markers 时返回含 warning 的空结果。

    掩蔽留出：个体保留在 G/H 中，仅掩蔽其表型（_profile_solve 的 MME 非观测个体
    经 G⁻¹/H⁻¹ 交叉关系预测）。G/H/Hinv 只建一次，逐折只换 z_train 子阵。
    剂量矩阵（Xc）同样只建一次，逐折只取观测行。

    输出 dict 与 blup.kfold_cv 逐键对齐 + `n_families`。cand=genotyped∩phenos（可验证个体）。
    """
    method = (method or "gblup").lower()
    if method == "ssgblup":
        order, idx, H, Hinv = _build_h(
            pedigree or [], genotyped, G, ridge=ridge, extra_base=list(phenos))
        individuals = order
        zmat_full = H
        relmat_inv = Hinv
    elif method in ("rrblup", "bayesb"):
        if genotypes is None or markers is None:
            return {
                "k": k, "n_total": 0, "n_individuals": 0,
                "folds": [], "mean_pearson": None, "mean_rmse": None,
                "pooled_pearson": None, "pooled_rmse": None, "cv_accuracy": None,
                "h2": None, "n_families": None,
                "warning": f"{method} 需要 genotypes+markers 剂量矩阵，无法交叉验证",
            }
        X_full, marker_order, p_full = _build_dosage_matrix(genotypes, markers)
        Xc_full = X_full - 2.0 * p_full[None, :]
        marker_c = 2.0 * float(np.sum(p_full * (1.0 - p_full)))
        if marker_c <= 0:
            marker_c = 1.0
        individuals = sorted(genotypes)
        idx = {i: k for k, i in enumerate(individuals)}
        # _profile_solve 需全 n×n Ginv（MME 随机效应块为 Zmat@Zmat.T + λ·Ginv_full），
        # 折内只取 zmat_full[np.ix_(z,z)] 观测子阵（= XoXo'/c），与 GBLUP 分支同构。
        zmat_full = Xc_full @ Xc_full.T / marker_c
        try:
            relmat_inv = np.linalg.inv(zmat_full)
        except np.linalg.LinAlgError:
            relmat_inv = np.linalg.pinv(zmat_full)
    else:
        individuals = list(genotyped)
        idx = {i: k for k, i in enumerate(individuals)}
        zmat_full = G
        try:
            relmat_inv = np.linalg.inv(G)
        except np.linalg.LinAlgError:
            relmat_inv = np.linalg.pinv(G)
    n = len(individuals)
    cand = sorted([i for i in genotyped if i in phenos])
    n_total = len(cand)
    if n_total < 2:
        return {
            "k": k, "n_total": n_total, "n_individuals": n,
            "folds": [], "mean_pearson": None, "mean_rmse": None,
            "pooled_pearson": None, "pooled_rmse": None, "cv_accuracy": None,
            "h2": None, "warning": "可验证个体（有基因型+表型）不足 2，无法交叉验证",
        }
    k = max(2, min(int(k or 5), 10))
    rng = random.Random(seed)
    n_families: int | None = None
    fam_warning: str | None = None
    split = (split or "random").lower()
    if split == "family":
        fam_of = {ind: (family_of or {}).get(ind, ("self", ind)) for ind in cand}
        fams = sorted({f for f in fam_of.values()})
        n_families = len(fams)
        if n_families < 2:
            return {
                "k": k, "n_total": n_total, "n_individuals": n,
                "folds": [], "mean_pearson": None, "mean_rmse": None,
                "pooled_pearson": None, "pooled_rmse": None, "cv_accuracy": None,
                "h2": None, "n_families": n_families,
                "warning": "可验证个体仅构成 1 个家系，无法家系阻塞交叉验证（请改用随机分层）",
            }
        if n_families < k:
            k = n_families
            fam_warning = f"家系数 {n_families} < 请求折数，折叠为 {k} 折（family-blocked）"
        rng.shuffle(fams)
        fam_bin = {f: i % k for i, f in enumerate(fams)}
        bin_of = {c: fam_bin[fam_of[c]] for c in cand}
    else:
        order_shuf = list(cand)
        rng.shuffle(order_shuf)
        bin_of = {c: i % k for i, c in enumerate(order_shuf)}

    folds: list[dict] = []
    all_x: list[float] = []
    all_y: list[float] = []
    for fold in range(k):
        test_set = {c for c, b in bin_of.items() if b == fold}
        train_phenos = {i: v for i, v in phenos.items() if i not in test_set}
        train_obs = [i for i in individuals if i in train_phenos]
        if not test_set or len(train_obs) < 2:
            folds.append({
                "fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
                "n_eval": 0, "pearson": None, "rmse": None,
                "error": "训练集个体不足 2 或测试集为空",
            })
            continue
        z = np.array([idx[i] for i in train_obs])
        y = np.array([float(train_phenos[i]) for i in train_obs])
        ebv: dict[str, float] | None = None
        fold_h2: float | None = None
        res: dict | None = None
        try:
            if method in ("rrblup", "bayesb"):
                Xo = Xc_full[z]
                yc = y - float(y.mean())
                if method == "rrblup":
                    res = _profile_solve(zmat_full[np.ix_(z, z)], np.ones((len(train_obs), 1)),
                                         y, relmat_inv, n, z)
                    va, ve = res.get("sigma_a"), res.get("sigma_e")
                    if (res.get("converged") and va is not None and ve is not None
                            and va > 0 and ve > 0):
                        lam = ve / (va / marker_c)
                        q = Xc_full.shape[1]
                        try:
                            u = np.linalg.solve(Xo.T @ Xo + lam * np.eye(q), Xo.T @ yc)
                        except np.linalg.LinAlgError:
                            u = np.linalg.pinv(Xo.T @ Xo + lam * np.eye(q)) @ (Xo.T @ yc)
                        mu = float(y.mean())
                        ebv = np.array([float(Xc_full[i] @ u + mu) for i in range(n)])
                        fold_h2 = res.get("h2")
                else:
                    u_hat, _sigma_e2, _sse, _n_save = _bayesb_gibbs(
                        Xo, yc, n_iter=n_iter, burn_in=burn_in, seed=seed, pi=pi,
                        df_s=df_s, scale0=scale0, df_e=df_e, scale_e0=scale_e0)
                    mu = float(y.mean())
                    ebv = np.array([float(Xc_full[i] @ u_hat + mu) for i in range(n)])
                    res = {"h2": None, "converged": True}
            else:
                X = np.ones((len(train_obs), 1))
                res = _profile_solve(zmat_full[np.ix_(z, z)], X, y, relmat_inv, n, z)
                if res.get("converged"):
                    ebv = res["ebv"]
                    fold_h2 = res.get("h2")
        except Exception as e:  # noqa: BLE001
            folds.append({
                "fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
                "n_eval": 0, "pearson": None, "rmse": None,
                "error": f"求解失败: {e!s}",
            })
            continue
        if not res or not res.get("converged") or ebv is None:
            folds.append({
                "fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
                "n_eval": 0, "pearson": None, "rmse": None,
                "error": (res or {}).get("warning") or "求解失败",
            })
            continue
        xs = [float(ebv[idx[c]]) for c in sorted(test_set)]
        ys = [float(phenos[c]) for c in sorted(test_set)]
        if len(xs) < 2:
            folds.append({
                "fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
                "n_eval": len(xs), "pearson": None, "rmse": None,
                "error": "测试集有效个体不足 2，无法计算 pearson",
            })
            continue
        pearson = blup._pearson(list(xs), list(ys))
        rmse = math.sqrt(sum((a - b) ** 2 for a, b in zip(xs, ys)) / len(xs))
        all_x.extend(xs)
        all_y.extend(ys)
        folds.append({
            "fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
            "n_eval": len(xs), "pearson": round(pearson, 4),
            "rmse": round(rmse, 6),
            "h2": float(fold_h2) if fold_h2 is not None else None,
        })

    valid = [f for f in folds if f.get("pearson") is not None]
    mean_pearson = round(sum(f["pearson"] for f in valid) / len(valid), 4) if valid else None
    rmse_vals = [f["rmse"] for f in valid if f.get("rmse") is not None]
    mean_rmse = round(sum(rmse_vals) / len(rmse_vals), 6) if rmse_vals else None
    pooled_pearson = round(blup._pearson(all_x, all_y), 4) if len(all_x) >= 2 else None
    pooled_rmse = (round(math.sqrt(sum((a - b) ** 2 for a, b in zip(all_x, all_y)) / len(all_x)), 6)
                   if len(all_x) else None)
    h2_vals = [f["h2"] for f in valid if f.get("h2") is not None]
    h2_mean = round(sum(h2_vals) / len(h2_vals), 4) if h2_vals else None
    errors = [f.get("error") for f in folds if f.get("error")]
    warning_parts = []
    if fam_warning:
        warning_parts.append(fam_warning)
    if errors:
        warning_parts.append(f"{len(errors)} 折失败: " + "；".join(errors))
    warning = "；".join(warning_parts) if warning_parts else None
    return {
        "k": k, "n_total": n_total, "n_individuals": n,
        "folds": folds,
        "mean_pearson": mean_pearson, "mean_rmse": mean_rmse,
        "pooled_pearson": pooled_pearson, "pooled_rmse": pooled_rmse,
        "cv_accuracy": mean_pearson, "h2": h2_mean,
        "n_families": n_families, "warning": warning,
    }