"""Design and reproducibility tools for SHWRP-2026-015.

This script creates a randomized field-trial allocation, blank data templates,
planning calculations, official 2024 cost benchmarks, and original figures.
It does not create or impute field outcomes.
"""

from __future__ import annotations

import json
from itertools import product
from pathlib import Path
from statistics import NormalDist

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


ROOT = Path(__file__).resolve().parent
OUT = ROOT / "outputs"
FIG = OUT / "figures"
OUT.mkdir(parents=True, exist_ok=True)
FIG.mkdir(parents=True, exist_ok=True)

SEED = 20260717
SITES = ["北部試驗點", "中部試驗點", "東部試驗點"]
SEASONS = ["一期作", "二期作"]
BLOCKS = [1, 2, 3, 4]

FACTORS = {
    "nitrogen": {
        "N100": "土壤檢測與地方推薦量的100%",
        "N80": "同一推薦量的80%",
    },
    "water": {
        "CSF": "連續淺水；依生育期與地方指引調整",
        "AWD": "安全間歇灌溉；田間水管約-15 cm再灌水，抽穗前後維持水層",
    },
    "density": {
        "D22": "30 x 15 cm，約22.2叢/m²",
        "D14": "30 x 24 cm，約13.9叢/m²",
    },
}

OUTCOME_COLUMNS = [
    "grain_yield_kg_ha_14pct",
    "irrigation_water_m3_ha",
    "water_productivity_kg_m3",
    "productive_panicles_m2",
    "filled_grains_panicle",
    "thousand_grain_weight_g",
    "lodging_pct",
    "disease_severity_pct",
    "brown_rice_pct",
    "head_rice_pct",
    "chalkiness_pct",
    "protein_pct_dry_basis",
    "amylose_pct_dry_basis",
    "taste_panel_score",
    "instrumental_taste_value",
    "direct_cost_ntd_ha",
    "full_cost_ntd_ha",
    "cost_ntd_kg",
]


def treatment_catalog() -> pd.DataFrame:
    rows = []
    for idx, (nitrogen, water, density) in enumerate(
        product(FACTORS["nitrogen"], FACTORS["water"], FACTORS["density"]), start=1
    ):
        rows.append(
            {
                "treatment_id": f"T{idx:02d}",
                "nitrogen": nitrogen,
                "nitrogen_definition": FACTORS["nitrogen"][nitrogen],
                "water": water,
                "water_definition": FACTORS["water"][water],
                "density": density,
                "density_definition": FACTORS["density"][density],
            }
        )
    return pd.DataFrame(rows)


def randomized_allocation(treatments: pd.DataFrame) -> pd.DataFrame:
    rng = np.random.default_rng(SEED)
    rows = []
    for site in SITES:
        for season in SEASONS:
            for block in BLOCKS:
                order = rng.permutation(treatments["treatment_id"].to_numpy())
                for plot_position, treatment_id in enumerate(order, start=1):
                    rows.append(
                        {
                            "site_code": site,
                            "season": season,
                            "block": block,
                            "plot_position": plot_position,
                            "plot_id": f"{site[:2]}-{season}-{block:02d}-{plot_position:02d}",
                            "treatment_id": treatment_id,
                        }
                    )
    allocation = pd.DataFrame(rows).merge(treatments, on="treatment_id", how="left")
    for column in OUTCOME_COLUMNS:
        allocation[column] = np.nan
    return allocation


def measurement_schedule() -> pd.DataFrame:
    rows = [
        ("整地前", "土壤基線", "pH、有機質、EC、有效性磷鉀、土壤質地", "每試驗點混合樣；區集差異大時按區集採樣"),
        ("插秧日", "處理確認", "品種、秧苗日齡、每叢苗數、行株距、基肥投入", "拍照並記錄批號"),
        ("插秧後14日", "成活與早期生育", "缺株率、株高、葉色/SPAD、雜草覆蓋", "固定樣方"),
        ("有效分蘗終期", "群體建立", "分蘗數、葉色/SPAD、灌溉量、病蟲害", "開始或維持水分處理"),
        ("幼穗形成期", "生殖轉換", "株高、葉色/SPAD、氮肥投入、灌溉量", "記錄追肥與穗肥"),
        ("抽穗前後各7日", "關鍵水分期", "抽穗日、病害、灌溉量、倒伏", "AWD組維持安全水層"),
        ("成熟期", "產量構成", "穗數、稔實率、每穗粒數、千粒重、倒伏", "排除邊行"),
        ("收穫與乾燥", "產量", "小區鮮重、水分、14%標準產量、乾燥能耗", "同批設備與程序"),
        ("收穫後", "米質與食味", "糙米率、完整米率、白堊質、蛋白質、直鏈澱粉、食味", "盲碼、同條件碾米與炊飯"),
        ("收穫後", "經濟結果", "種苗、肥料、農藥、人工、機械、能源、用水與地租", "分直接成本與完整成本"),
    ]
    return pd.DataFrame(rows, columns=["stage", "domain", "measurements", "quality_control"])


def evidence_matrix() -> pd.DataFrame:
    rows = [
        ("氮肥", "產量", "非線性；不足與過量皆可能不利", "臺灣結果受品種、期作與基礎肥力影響"),
        ("氮肥", "蛋白質", "較高氮量通常提高", "蛋白質不是食味的唯一決定因子"),
        ("氮肥", "食味", "過量穗肥可能不利", "須以官能與理化指標共同評估"),
        ("安全AWD", "灌溉用水", "預期降低", "必須量測實際灌溉量，不能只記灌水次數"),
        ("安全AWD", "產量", "管理得當時可維持", "土壤、供水、氣候與抽穗期管理可改變結果"),
        ("栽植密度", "產量構成", "改變穗數與每穗粒數的平衡", "不存在跨品種通用的單一最佳密度"),
        ("栽植密度", "倒伏與病害", "可能透過冠層結構改變", "需現地量測，不能由密度直接推定"),
        ("期作", "產量與米質", "具有顯著背景差異的可能", "應納入固定效果或分層估計"),
    ]
    return pd.DataFrame(rows, columns=["factor", "outcome", "evidence_direction", "boundary"])


def official_cost_benchmark() -> pd.DataFrame:
    # MOA Agricultural Statistics Yearbook 2024, pp. 374-375.
    rows = [
        ("一期粳稻", 7083, 107787, 132167, 1522, 1866, 12617, 2366),
        ("二期粳稻", 5409, 113399, 134542, 2096, 2487, 13727, 1729),
    ]
    return pd.DataFrame(
        rows,
        columns=[
            "crop_season",
            "sample_yield_kg_ha",
            "primary_cost_ntd_ha",
            "secondary_cost_ntd_ha",
            "primary_cost_ntd_100kg",
            "secondary_cost_ntd_100kg",
            "fertilizer_cost_ntd_ha",
            "irrigation_charge_ntd_ha",
        ],
    )


def planning_precision() -> pd.DataFrame:
    alpha = 0.05
    power = 0.80
    z_alpha = NormalDist().inv_cdf(1 - alpha / 2)
    z_power = NormalDist().inv_cdf(power)
    rows = []
    for sites, seasons in [(1, 1), (1, 2), (2, 2), (3, 2)]:
        total = sites * seasons * len(BLOCKS) * 8
        per_level = total // 2
        mde_d = (z_alpha + z_power) * np.sqrt(2 / per_level)
        rows.append(
            {
                "sites": sites,
                "seasons": seasons,
                "total_plot_seasons": total,
                "observations_per_main_effect_level": per_level,
                "approx_mde_standardized_difference": round(float(mde_d), 3),
                "assumption": "獨立、平衡、雙尾alpha=.05、power=.80；未計區集與重複量測相關",
            }
        )
    return pd.DataFrame(rows)


def analysis_plan() -> dict:
    return {
        "status": "prospective protocol; no field outcomes collected",
        "primary_outcomes": [
            "grain_yield_kg_ha_14pct",
            "instrumental_taste_value and blinded taste_panel_score",
            "irrigation_water_m3_ha",
            "full_cost_ntd_kg",
        ],
        "model": "mixed model with nitrogen*water*density, season, site, and interactions; block nested in site-season as random intercept",
        "estimands": "main effects and prespecified two-way interactions with 95% confidence intervals",
        "missing_data": "report reasons; no outcome imputation in primary analysis",
        "multiple_testing": "four co-primary domains reported separately; secondary outcomes labeled exploratory",
        "decision_rule": "no single treatment is declared best unless yield non-inferiority and prespecified quality, water, and cost thresholds are all met",
        "non_inferiority_margin": "must be fixed with agronomists before planting; not set by observed data",
        "random_seed": SEED,
    }


def make_charts(
    treatments: pd.DataFrame,
    allocation: pd.DataFrame,
    schedule: pd.DataFrame,
    evidence: pd.DataFrame,
    costs: pd.DataFrame,
    precision: pd.DataFrame,
) -> None:
    font_file = ROOT.parent.parent / "public" / "fonts" / "NotoSansTC-Regular.ttf"
    if font_file.exists():
        from matplotlib import font_manager

        font_manager.fontManager.addfont(str(font_file))
        plt.rcParams["font.family"] = font_manager.FontProperties(fname=str(font_file)).get_name()
    else:
        plt.rcParams["font.family"] = ["Microsoft JhengHei", "DejaVu Sans"]
    plt.rcParams["axes.unicode_minus"] = False
    plt.rcParams["figure.facecolor"] = "#fbf7f0"
    brown, green, gold, muted = "#6f3519", "#315c4c", "#c28b42", "#8c8177"

    fig, ax = plt.subplots(figsize=(12, 6.7))
    matrix = np.arange(8).reshape(2, 4)
    ax.imshow(matrix, cmap="YlGn", aspect="auto", alpha=.72)
    for i, row in treatments.iterrows():
        y, x = divmod(i, 4)
        ax.text(x, y, f"{row.treatment_id}\n{row.nitrogen} · {row.water} · {row.density}", ha="center", va="center", fontsize=12, weight="bold", color="#211a15")
    ax.set_xticks([]); ax.set_yticks([])
    ax.set_title("圖1　2 × 2 × 2 因子試驗的八種處理", fontsize=17, weight="bold", color=brown, pad=18)
    ax.text(1.5, 1.75, "每個試驗點、每一期作設4區集；每區集包含全部8處理", ha="center", fontsize=12, color=green)
    fig.tight_layout(); fig.savefig(FIG / "fig1_factorial_design.png", dpi=220, bbox_inches="tight"); plt.close(fig)

    counts = allocation.groupby(["site_code", "season"]).size().unstack()
    fig, ax = plt.subplots(figsize=(11, 6.6))
    counts.plot(kind="bar", stacked=True, color=[gold, green], ax=ax)
    ax.set_title("圖2　確認性方案共192個小區期作觀測", fontsize=17, weight="bold", color=brown)
    ax.set_ylabel("小區期作數"); ax.set_xlabel("")
    ax.tick_params(axis="x", rotation=0)
    ax.legend(title="期作", frameon=False)
    ax.grid(axis="y", alpha=.2)
    fig.tight_layout(); fig.savefig(FIG / "fig2_site_season_allocation.png", dpi=220); plt.close(fig)

    fig, ax = plt.subplots(figsize=(10.8, 6.5))
    ax.plot(precision["total_plot_seasons"], precision["approx_mde_standardized_difference"], marker="o", linewidth=2.5, color=green)
    for _, row in precision.iterrows():
        ax.annotate(f"d≈{row.approx_mde_standardized_difference:.2f}", (row.total_plot_seasons, row.approx_mde_standardized_difference), xytext=(0, 10), textcoords="offset points", ha="center")
    ax.set_title("圖3　規模越大，可辨識的主效應越小", fontsize=17, weight="bold", color=brown)
    ax.set_xlabel("總小區期作數"); ax.set_ylabel("近似最小可辨識標準化差異 d")
    ax.grid(alpha=.22)
    ax.text(.01, -.18, "規劃估計：雙尾α=.05、power=.80、平衡且獨立；正式分析須納入區集與跨期相關。", transform=ax.transAxes, fontsize=9.5, color=muted)
    fig.tight_layout(); fig.savefig(FIG / "fig3_power_planning.png", dpi=220, bbox_inches="tight"); plt.close(fig)

    direction_order = ["預期降低", "管理得當時可維持", "非線性；不足與過量皆可能不利", "較高氮量通常提高", "過量穗肥可能不利", "改變穗數與每穗粒數的平衡", "可能透過冠層結構改變", "具有顯著背景差異的可能"]
    encoded = evidence["evidence_direction"].map({value: i for i, value in enumerate(direction_order)}).to_numpy().reshape(-1, 1)
    fig, ax = plt.subplots(figsize=(12, 7.2))
    ax.imshow(encoded, cmap="YlGn", aspect="auto", alpha=.7)
    ax.set_xticks([0], ["證據方向（不是效果量）"])
    ax.set_yticks(range(len(evidence)), [f"{r.factor} → {r.outcome}" for r in evidence.itertuples()])
    for i, row in evidence.iterrows():
        ax.text(0, i, row.evidence_direction, ha="center", va="center", fontsize=10.5, weight="bold")
    ax.set_title("圖4　既有證據支持提出方向性假說，但不足以預填田間結果", fontsize=16, weight="bold", color=brown, pad=16)
    fig.tight_layout(); fig.savefig(FIG / "fig4_evidence_direction.png", dpi=220, bbox_inches="tight"); plt.close(fig)

    fig, ax = plt.subplots(figsize=(12, 7.4))
    y = np.arange(len(schedule))
    ax.scatter(np.zeros_like(y), y, s=120, color=green, zorder=3)
    ax.vlines(0, y.min(), y.max(), color=gold, linewidth=3)
    for idx, row in schedule.iterrows():
        ax.text(.08, idx, f"{row.stage}｜{row.domain}", va="center", fontsize=11, weight="bold", color="#211a15")
        ax.text(.62, idx, row.measurements, va="center", fontsize=9.3, color=muted)
    ax.set_xlim(-.1, 1.9); ax.set_ylim(len(schedule)-.4, -.6); ax.axis("off")
    ax.set_title("圖5　從整地前到收穫後的預定量測鏈", fontsize=17, weight="bold", color=brown, pad=18)
    fig.tight_layout(); fig.savefig(FIG / "fig5_measurement_timeline.png", dpi=220, bbox_inches="tight"); plt.close(fig)

    fig, axes = plt.subplots(1, 2, figsize=(12, 6.6))
    axes[0].bar(costs["crop_season"], costs["sample_yield_kg_ha"], color=[green, gold])
    axes[0].set_title("樣本平均產量（kg/ha）"); axes[0].grid(axis="y", alpha=.2)
    width = .35; x = np.arange(len(costs))
    axes[1].bar(x-width/2, costs["primary_cost_ntd_100kg"], width, label="第一種成本", color=gold)
    axes[1].bar(x+width/2, costs["secondary_cost_ntd_100kg"], width, label="第二種成本", color=green)
    axes[1].set_xticks(x, costs["crop_season"]); axes[1].set_title("每100公斤生產成本（NT$）")
    axes[1].legend(frameon=False); axes[1].grid(axis="y", alpha=.2)
    fig.suptitle("圖6　農業部2024年粳稻成本與產量基準", fontsize=17, weight="bold", color=brown)
    fig.text(.5, .01, "官方樣本統計只作外部基準；不得代替本試驗逐筆成本紀錄。", ha="center", color=muted, fontsize=9.5)
    fig.tight_layout(rect=[0, .04, 1, .94]); fig.savefig(FIG / "fig6_official_cost_benchmark.png", dpi=220); plt.close(fig)


def main() -> None:
    treatments = treatment_catalog()
    allocation = randomized_allocation(treatments)
    schedule = measurement_schedule()
    evidence = evidence_matrix()
    costs = official_cost_benchmark()
    precision = planning_precision()

    treatments.to_csv(OUT / "treatment_catalog.csv", index=False, encoding="utf-8-sig")
    allocation.drop(columns=OUTCOME_COLUMNS).to_csv(OUT / "randomization_schedule.csv", index=False, encoding="utf-8-sig")
    allocation.to_csv(OUT / "field_data_template.csv", index=False, encoding="utf-8-sig")
    schedule.to_csv(OUT / "measurement_schedule.csv", index=False, encoding="utf-8-sig")
    evidence.to_csv(OUT / "evidence_matrix.csv", index=False, encoding="utf-8-sig")
    costs.to_csv(OUT / "official_2024_cost_benchmark.csv", index=False, encoding="utf-8-sig")
    precision.to_csv(OUT / "planning_precision.csv", index=False, encoding="utf-8-sig")
    (OUT / "analysis_plan.json").write_text(json.dumps(analysis_plan(), ensure_ascii=False, indent=2), encoding="utf-8")
    make_charts(treatments, allocation, schedule, evidence, costs, precision)

    summary = {
        "report_number": "SHWRP-2026-015",
        "status": "prospective protocol",
        "treatments": len(treatments),
        "pilot_plot_seasons": 32,
        "confirmatory_plot_seasons": len(allocation),
        "sites": len(SITES),
        "seasons": len(SEASONS),
        "blocks_per_site_season": len(BLOCKS),
        "field_outcomes_present": False,
        "random_seed": SEED,
    }
    (OUT / "analysis_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
