# -*- coding: utf-8 -*-
"""稳定性(AMMI/Finlay-Wilkinson)正式 tc 套件：TestClient 走真实 API + 落库 detail_json 断言。

fixture（镜像 e2e_advanced_stability）：3 组合 × 3 站点 cell means
  C0=[0,6,12]（高响应 → FW b≈2.0、flag_stable=False）
  C1=[2,5,8] （稳定   → FW b≈1.0、flag_stable=True，|b-1|≤2·se(b)）
  C2=[4,4,4] （低响应 → FW b≈0.0、flag_stable=False）
每组合 2 株、每株在 3 站点各有观察值（evaluate_year=2023/2024/2025）。

经真实 HTTP 端点 POST /api/v1/bre/statistics/stability/run 断言：
  [1] site 维度：env_dim=site、engine_version 落库、methods=["ammi","finlay"]、
      FW b 斜坡（C0>1.3 / C1≈1±0.05 / C2<0.5）+ flag_stable 语义（C1=True 其余 False）、
      r²≈1（常量响应基因型为 None 合法）；AMMI n=9、rank 按 ASV 升序、IPC1 占比≥IPC2、互作 SS>0
  [2] year 维度：env_dim=year 且 finlay 落库
  [3] 门禁：非法 methods / 非法 gxe_env → 409
  [4] list / detail 端点命中本批次
依赖: Redis + PG 正常（TestClient 走真实 lifespan）。运行后自动清理。
"""
import os
os.environ["ENVIRONMENT"] = "dev"
os.environ["PYTHONUTF8"] = "1"

import sys, asyncio  # noqa: E402
sys.path.insert(0, r"d:\dpb\dpb\backend")

import main  # noqa: E402
from fastapi.testclient import TestClient  # noqa: E402
from sqlalchemy import delete, select  # noqa: E402

from app.core.database import create_async_engine_and_session  # noqa: E402
from app.api.v1.module_system.user.model import UserModel  # noqa: E402 (注册 mapper)
from app.api.v1.module_bre.trait.model import TraitModel  # noqa: E402
from app.api.v1.module_bre.target.model import TargetModel  # noqa: E402
from app.api.v1.module_bre.trial.model import TrialModel  # noqa: E402
from app.api.v1.module_bre.trial_study.model import TrialStudyModel  # noqa: E402
from app.api.v1.module_bre.germplasm.model import BreedingGermplasmModel  # noqa: E402
from app.api.v1.module_bre.cross_combination.model import CrossCombinationModel  # noqa: E402
from app.api.v1.module_bre.tree.model import TreeModel  # noqa: E402
from app.api.v1.module_bre.trait_observation.model import TraitObservationModel  # noqa: E402
from app.api.v1.module_bre.statistics.model import StatisticsJobModel, StabilityResultModel  # noqa: E402
from scripts.breeding_stats import stability  # noqa: E402

create_app = main.create_app
TOKEN = None
ok, fail = 0, 0

PREFIX = "TCSTAB"
trait_code = f"tR_{PREFIX}"
tokens: dict[str, list[int]] = {
    "trait": [], "target": [], "germ": [], "combo": [], "tree": [], "obs": [], "study": [],
}
result_ids: list[int] = []
FIX: dict = {}


def check(name, cond, detail=""):
    global ok, fail
    if cond:
        ok += 1
        print(f"  [ok] {name} {detail}")
    else:
        fail += 1
        print(f"  [FAIL] {name} {detail}")


def login(client):
    global TOKEN
    d = {"username": "super", "password": "123456", "grant_type": "password", "login_type": "PC端"}
    r = client.post("/api/v1/system/auth/login", data=d)
    b = r.json()
    if r.status_code == 200 and b.get("code") == 0:
        TOKEN = b["data"]["access_token"]
        return
    key = client.get("/api/v1/system/auth/captcha/get").json()["data"]["key"]
    client.post("/api/v1/system/auth/captcha/slider/complete", json={"captcha_key": key})
    d["captcha_key"] = key
    r = client.post("/api/v1/system/auth/login", data=d)
    b = r.json()
    assert r.status_code == 200 and b.get("code") == 0, f"LOGIN FAIL {r.status_code} {b}"
    TOKEN = b["data"]["access_token"]


def auth():
    return {"Authorization": f"Bearer {TOKEN}"}


async def _wipe() -> None:
    """启动前清理本前缀残留（防上次进程被杀/建 fixture 中途失败留下的脏数据）。"""
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            result_rows = list((await db.execute(select(StabilityResultModel.id).where(
                StabilityResultModel.trait_code == trait_code))).scalars())
            if result_rows:
                await db.execute(delete(StabilityResultModel).where(
                    StabilityResultModel.id.in_(result_rows)))
            jobs = list((await db.execute(select(StatisticsJobModel.id).where(
                StatisticsJobModel.job_type == "STABILITY",
                StatisticsJobModel.params_json.op("->>")("trait_code") == trait_code))).scalars())
            if jobs:
                await db.execute(delete(StatisticsJobModel).where(StatisticsJobModel.id.in_(jobs)))
            study_ids = list((await db.execute(select(TrialStudyModel.id).where(
                TrialStudyModel.study_name.like(f"站%_{PREFIX}")))).scalars())
            if study_ids:
                await db.execute(delete(TrialStudyModel).where(TrialStudyModel.id.in_(study_ids)))
            await db.execute(delete(TrialModel).where(TrialModel.trial_name == f"区域试验{PREFIX}"))
            combo_ids = list((await db.execute(select(CrossCombinationModel.id).where(
                CrossCombinationModel.combination_code.like(f"C%_{PREFIX}")))).scalars())
            if combo_ids:
                await db.execute(delete(CrossCombinationModel).where(
                    CrossCombinationModel.id.in_(combo_ids)))
            await db.execute(delete(BreedingGermplasmModel).where(
                BreedingGermplasmModel.cultivar_name.in_([f"FM_{PREFIX}", f"MM_{PREFIX}"])))
            await db.execute(delete(TraitModel).where(TraitModel.trait_code == trait_code))
            await db.execute(delete(TargetModel).where(TargetModel.target_name == f"目标{PREFIX}"))
            await db.commit()
            print(f"[preclean] stability tc 前缀残留已清（result={len(result_rows)} study={len(study_ids)}）")
    finally:
        await engine.dispose()


async def _build_fixture() -> None:
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            trait = TraitModel(trait_code=trait_code, trait_name=f"坐果率{PREFIX}", data_type="numeric",
                               unit="%", is_core="1", direction="desc", into_ebv="1",
                               default_h2=0.4, created_id=1)
            db.add(trait)
            await db.flush()
            tokens["trait"].append(trait.id)

            target = TargetModel(target_name=f"目标{PREFIX}", created_id=1)
            db.add(target)
            await db.flush()
            tokens["target"].append(target.id)

            trial = TrialModel(trial_name=f"区域试验{PREFIX}", start_year=2023, end_year=2025, created_id=1)
            db.add(trial)
            await db.flush()

            studies = []
            for i in range(3):
                s = TrialStudyModel(trial_id=trial.id, study_name=f"站{i + 1}_{PREFIX}",
                                    year=2023 + i, created_id=1)
                db.add(s)
                studies.append(s)
            await db.flush()
            tokens["study"] = [s.id for s in studies]

            fm = BreedingGermplasmModel(cultivar_name=f"FM_{PREFIX}", can_be_female=True, created_id=1)
            mm = BreedingGermplasmModel(cultivar_name=f"MM_{PREFIX}", can_be_male=True, created_id=1)
            db.add_all([fm, mm])
            await db.flush()
            tokens["germ"] += [fm.id, mm.id]

            # 3 组合，每组合 2 株；cell 均值目标：
            # C0=[0,6,12]  C1=[2,5,8]  C2=[4,4,4]（对应 3 站点）
            targets = {0: [0.0, 6.0, 12.0], 1: [2.0, 5.0, 8.0], 2: [4.0, 4.0, 4.0]}
            combos: list[CrossCombinationModel] = []
            trees: list[TreeModel] = []
            for k in range(3):
                combo = CrossCombinationModel(
                    combination_code=f"C{k}_{PREFIX}", bre_target_id=target.id,
                    female_parent_id=fm.id, male_parent_id=mm.id,
                    design_type="full_diallel", created_id=1)
                db.add(combo)
                await db.flush()
                combos.append(combo)
                tokens["combo"].append(combo.id)
                for j in range(2):
                    t = TreeModel(combination_id=combo.id, tree_no=f"{PREFIX}-{k}{j}", status="alive",
                                  stage="seedling", generation="F1", planted_date="2022-03-10",
                                  created_id=1)
                    db.add(t)
                    trees.append(t)
            await db.flush()
            tokens["tree"] = [t.id for t in trees]

            obs_rows = []
            for k, combo in enumerate(combos):
                pair = [trees[2 * k], trees[2 * k + 1]]
                for si, (study, yv) in enumerate(zip(studies, (2023, 2024, 2025))):
                    v = targets[k][si]
                    for t in pair:
                        o = TraitObservationModel(
                            tree_id=t.id, combination_id=combo.id, trait_id=trait.id,
                            value_numeric=v, evaluate_year=yv, trial_study_id=study.id, created_id=1)
                        db.add(o)
                        obs_rows.append(o)
            await db.flush()
            tokens["obs"] = [o.id for o in obs_rows]
            FIX["combo_ids"] = [c.id for c in combos]
            await db.commit()
    finally:
        await engine.dispose()


async def _verify_and_cleanup() -> None:
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            if result_ids:
                await db.execute(delete(StabilityResultModel).where(
                    StabilityResultModel.id.in_(result_ids)))
            jobs = list((await db.execute(select(StatisticsJobModel.id).where(
                StatisticsJobModel.job_type == "STABILITY",
                StatisticsJobModel.params_json.op("->>")("trait_code") == trait_code))).scalars())
            if jobs:
                await db.execute(delete(StatisticsJobModel).where(StatisticsJobModel.id.in_(jobs)))
            if tokens["obs"]:
                await db.execute(delete(TraitObservationModel).where(
                    TraitObservationModel.id.in_(tokens["obs"])))
            if tokens["tree"]:
                await db.execute(delete(TreeModel).where(TreeModel.id.in_(tokens["tree"])))
            if tokens["combo"]:
                await db.execute(delete(CrossCombinationModel).where(
                    CrossCombinationModel.id.in_(tokens["combo"])))
            if tokens["germ"]:
                await db.execute(delete(BreedingGermplasmModel).where(
                    BreedingGermplasmModel.id.in_(tokens["germ"])))
            if tokens["study"]:
                await db.execute(delete(TrialStudyModel).where(TrialStudyModel.id.in_(tokens["study"])))
            await db.execute(delete(TrialModel).where(TrialModel.trial_name == f"区域试验{PREFIX}"))
            if tokens["trait"]:
                await db.execute(delete(TraitModel).where(TraitModel.id.in_(tokens["trait"])))
            if tokens["target"]:
                await db.execute(delete(TargetModel).where(TargetModel.id.in_(tokens["target"])))
            await db.commit()
            print(f"[cleanup] stability tc 数据已清（result={len(result_ids)} obs={len(tokens['obs'])}）")
    finally:
        await engine.dispose()


def main_() -> None:
    asyncio.run(_wipe())
    asyncio.run(_build_fixture())
    trait_id = tokens["trait"][0]
    try:
        with TestClient(create_app()) as client:
            login(client)
            H = auth()

            def run(payload):
                r = client.post("/api/v1/bre/statistics/stability/run", json=payload, headers=H)
                b = r.json()
                check(f"[HTTP] stability/run 200", r.status_code == 200,
                      f"{r.status_code} {str(r.text)[:150]}")
                return b.get("data") if b.get("code") == 0 else None

            # ================= [1] site 维度 =================
            sid = run({"trait_id": trait_id, "trait_code": trait_code,
                       "gxe_env": "site", "methods": ["ammi", "finlay"]})
            result_ids.append(sid)
            r = client.get(f"/api/v1/bre/statistics/stability/{sid}", headers=H)
            b = r.json()
            check("[1] detail 200 + id 命中", r.status_code == 200 and b.get("data", {}).get("id") == sid,
                  f"{r.status_code}")
            d0 = b["data"]
            check("[1] env_dim=site + engine_version 落库",
                  d0["env_dim"] == "site" and d0["engine_version"] == stability.ENGINE_VERSION,
                  f"env={d0['env_dim']} ver={d0['engine_version']}")
            d = d0["detail_json"]
            check("[1] methods=['ammi','finlay']", d["methods"] == ["ammi", "finlay"], f"{d['methods']}")
            fw = d["finlay"]
            rows = {r_["genotype"]: r_ for r_ in fw["rows"]}
            k = [f"f{cid}" for cid in FIX["combo_ids"]]
            check("[1] FW 3 行（3 组合）", len(rows) == 3 and all(g in rows for g in k), f"{list(rows)}")
            bmap = {g: rows[g]["b"] for g in k}
            check("[1] FW b 斜坡：C0>1.3 / C1≈1.0±0.05 / C2<0.5",
                  bmap[k[0]] > 1.3 and abs(bmap[k[1]] - 1.0) < 0.05 and bmap[k[2]] < 0.5,
                  f"{ {g: round(bmap[g], 3) for g in k} }")
            check("[1] flag_stable 语义：C1=True、C0/C2=False",
                  rows[k[0]]["flag_stable"] is False and rows[k[1]]["flag_stable"] is True
                  and rows[k[2]]["flag_stable"] is False,
                  f"{ {g: rows[g]['flag_stable'] for g in k} }")
            check("[1] r²≈1（常量响应基因型为 None 合法）",
                  all(r_["r2"] is None or r_["r2"] >= 0.99 for r_ in fw["rows"]),
                  f"{ {g: rows[g]['r2'] for g in k} }")
            am = d["ammi"]
            check("[1] AMMI n=9 格子", am["n"] == 9, f"n={am['n']}")
            check("[1] AMMI rank 按 ASV 升序 + 互作 SS>0",
                  len(am["rows"]) == 3 and all(r_["rank"] == i + 1 for i, r_ in enumerate(am["rows"]))
                  and am["rows"][0]["asv"] <= am["rows"][-1]["asv"] and am["ss_interaction"] > 0,
                  f"asv={[round(r_['asv'], 3) for r_ in am['rows']]} ss={round(am['ss_interaction'], 3)}")
            check("[1] IPC1 解释比例 ≥ IPC2",
                  am["ipc_variance"][0]["proportion"] >= am["ipc_variance"][1]["proportion"],
                  f"[{round(am['ipc_variance'][0]['proportion'], 3)}, {round(am['ipc_variance'][1]['proportion'], 3)}]")

            # ================= [2] year 维度 =================
            sid2 = run({"trait_id": trait_id, "trait_code": trait_code,
                        "gxe_env": "year", "methods": ["ammi", "finlay"]})
            result_ids.append(sid2)
            r2 = client.get(f"/api/v1/bre/statistics/stability/{sid2}", headers=H).json()["data"]
            check("[2] year 维度：env_dim=year + finlay 落库",
                  r2["env_dim"] == "year" and r2["detail_json"]["finlay"]
                  and r2["detail_json"]["ammi"]["n"] == 9,
                  f"env={r2['env_dim']} ammi_n={r2['detail_json']['ammi']['n']}")

            # ================= [3] 门禁 =================
            r = client.post("/api/v1/bre/statistics/stability/run",
                            json={"trait_id": trait_id, "trait_code": trait_code,
                                  "gxe_env": "site", "methods": ["bad"]}, headers=H)
            check("[3] 非法 methods → 409", r.status_code == 409, f"{r.status_code}")
            r = client.post("/api/v1/bre/statistics/stability/run",
                            json={"trait_id": trait_id, "trait_code": trait_code,
                                  "gxe_env": "bad", "methods": ["ammi", "finlay"]}, headers=H)
            check("[3] 非法 gxe_env → 409", r.status_code == 409, f"{r.status_code}")

            # ================= [4] list / detail =================
            r = client.get("/api/v1/bre/statistics/stability", headers=H)
            lst = r.json()["data"]
            check("[4] list 含本批次 + detail 命中",
                  any(x["id"] == sid for x in lst) and
                  client.get(f"/api/v1/bre/statistics/stability/{sid}", headers=H).json()["data"]["id"] == sid,
                  f"list 共 {len(lst)} 条")
    finally:
        asyncio.run(_verify_and_cleanup())

    print(f"\n===== stability tc 套件：ok={ok} fail={fail} =====")


if __name__ == "__main__":
    main_()
    sys.exit(1 if fail else 0)