#!/usr/bin/env python3
"""Granovetter (1978) threshold-model experiments for sim-lab ep02."""
from __future__ import annotations

import argparse
import json
import math
import platform
import statistics
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

import numpy as np


N = 100
MU = 25
SEEDS = list(range(1001, 1101))
CROWD_SIGMAS = list(range(8, 21))


def cascade(thresholds: Iterable[float]) -> dict[str, Any]:
    """Run simultaneous rounds until a round adds nobody."""
    values = [float(t) for t in thresholds]
    standing = [False] * len(values)
    stand_round = [0] * len(values)
    series: list[int] = []
    round_no = 0
    while True:
        count = sum(standing)
        newcomers = [i for i, threshold in enumerate(values)
                     if not standing[i] and count >= threshold]
        if not newcomers:
            break
        round_no += 1
        for i in newcomers:
            standing[i] = True
            stand_round[i] = round_no
        series.append(sum(standing))
    return {"final": sum(standing), "rounds": round_no, "series": series,
            "stand_round": stand_round}


def crowd_record(thresholds: list[float], changed_index: int | None,
                 from_value: int | None = None, to_value: int | None = None) -> dict[str, Any]:
    result = cascade(thresholds)
    return {
        "thresholds": thresholds,
        **result,
        "instigators": sum(t <= 0 for t in thresholds),
        "min_t": min(thresholds),
        "max_t": max(thresholds),
        "changed_index": changed_index,
        "from": from_value,
        "to": to_value,
        "changed_people": 0 if changed_index is None else 1,
    }


def normal_cdf(x: float) -> float:
    return 0.5 * (1.0 + math.erf(x / math.sqrt(2.0)))


def expected_equilibrium(sigma: float) -> float:
    if sigma == 0.0:
        return 0.0
    value = 0.0
    for _ in range(100_000):
        nxt = N * normal_cdf((value - MU) / sigma)
        if abs(nxt - value) < 1e-9:
            return nxt
        value = nxt
    raise RuntimeError(f"expected recurrence did not converge for sigma={sigma}")


def changed_final(thresholds: list[float], index: int) -> int:
    changed = thresholds.copy()
    changed[index] += 1
    return int(cascade(changed)["final"])


def run_experiments() -> dict[str, Any]:
    started = time.perf_counter()
    ladder_thresholds = list(range(N))
    ladder = crowd_record(ladder_thresholds, None)
    swap_thresholds = ladder_thresholds.copy()
    swap_thresholds[1] = 2
    swap = crowd_record(swap_thresholds, 1, 1, 2)

    swap_k_finals = []
    for k in range(N):
        values = ladder_thresholds.copy()
        values[k] = k + 1
        swap_k_finals.append(int(cascade(values)["final"]))
    swap_k = {"k": list(range(N)), "final": swap_k_finals,
              "all_equal_k": swap_k_finals == list(range(N)), "k_min": 0, "k_max": 99}

    expected_sigmas = [round(i * 0.1, 1) for i in range(201)]
    expected_re = [expected_equilibrium(sigma) for sigma in expected_sigmas]
    critical_index = next(i for i, value in enumerate(expected_re) if value >= 50)
    expected = {
        "sigma": expected_sigmas,
        "re": expected_re,
        "sigma_c": expected_sigmas[critical_index],
        "sigma_below": expected_sigmas[critical_index - 1],
        "re_below": expected_re[critical_index - 1],
        "re_at_c": expected_re[critical_index],
        "paper_sigma_c": 12.2,
    }

    by_sigma: dict[str, Any] = {}
    for sigma in CROWD_SIGMAS:
        finals: list[int] = []
        linchpin: list[bool] = []
        linchpin_orders: list[int] = []
        random_flip: list[bool] = []
        for seed in SEEDS:
            thresholds = np.random.default_rng(seed).normal(MU, sigma, N).tolist()
            baseline = int(cascade(thresholds)["final"])
            finals.append(baseline)
            has_linchpin = False
            if baseline >= 90:
                order = {index: rank for rank, index in enumerate(
                    sorted(range(N), key=lambda i: (thresholds[i], i)), start=1)}
                for index in range(N):
                    if changed_final(thresholds, index) <= 10:
                        has_linchpin = True
                        linchpin_orders.append(order[index])
            linchpin.append(has_linchpin)
            chosen = int(np.random.default_rng(seed + 500000).integers(0, N))
            random_flip.append(abs(changed_final(thresholds, chosen) - baseline) >= 50)
        full = sum(value >= 90 for value in finals)
        tiny_values = [value for value in finals if value <= 10]
        tiny = len(tiny_values)
        mid = N - full - tiny
        linchpin_crowds = sum(linchpin)
        by_sigma[str(sigma)] = {
            "n": N,
            "finals": finals,
            "mean_final": statistics.fmean(finals),
            "full": full,
            "tiny": tiny,
            "mid": mid,
            "full_share": full / N,
            "mid_share": mid / N,
            "tiny_mean_final": statistics.fmean(tiny_values) if tiny_values else None,
            "linchpin": linchpin,
            "linchpin_crowds": linchpin_crowds,
            "linchpin_share_of_full": linchpin_crowds / full if full else None,
            "linchpin_orders": linchpin_orders,
            "linchpin_median_order": statistics.median(linchpin_orders) if linchpin_orders else None,
            "random_flip": random_flip,
            "random_flip_share": sum(random_flip) / N,
        }
    split_sigma = min(CROWD_SIGMAS, key=lambda sigma: (abs(by_sigma[str(sigma)]["full_share"] - 0.5), sigma))
    split = dict(by_sigma[str(split_sigma)])
    split.update({
        "sigma": split_sigma,
        "rule": "full/n が 0.5 に最も近い sigma（同点は小さいほう）",
        "sigma_minus_expected_sigma_c": split_sigma - expected["sigma_c"],
    })
    return {
        "config": {
            "n_people": N, "mu": MU, "seeds": SEEDS, "n_per_condition": N,
            "full_min": 90, "tiny_max": 10, "sigma_grid_crowds": CROWD_SIGMAS,
            "sigma_step_expected": 0.1, "paper_sigma_c": 12.2, "paper_re_below": 6,
            "python": platform.python_version(), "numpy": np.__version__,
        },
        "groups": {"crowds": {f"s{sigma:02d}": {"n": N} for sigma in CROWD_SIGMAS}},
        "ladder": ladder, "swap": swap, "swap_k": swap_k, "expected": expected,
        "crowds": {"sigma": CROWD_SIGMAS, "by_sigma": by_sigma,
                    "max_mid_share": max(row["mid_share"] for row in by_sigma.values())},
        "split": split,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "runtime_s": time.perf_counter() - started,
    }


def metric_value(summary: dict[str, Any], metric: str) -> Any:
    parts = metric.split(".")
    if parts[:2] == ["crowds", "sigma8"]:
        value: Any = summary["crowds"]["by_sigma"]["8"]
        parts = parts[2:]
    elif parts[:2] == ["crowds", "sigma20"]:
        value = summary["crowds"]["by_sigma"]["20"]
        parts = parts[2:]
    else:
        value = summary
    for part in parts:
        value = value[part]
    return value


def prediction_hit(prediction: dict[str, Any], observed: Any) -> bool:
    if observed is None:
        return False
    kind = prediction["kind"]
    if kind == "eq":
        return observed == prediction["value"]
    if kind == "bool":
        return isinstance(observed, bool) and observed is prediction["value"]
    if kind == "range":
        return prediction["lo"] <= observed <= prediction["hi"]
    if kind == "max":
        return observed <= prediction["hi"]
    if kind == "min":
        return observed >= prediction["lo"]
    if kind == "in":
        return observed in prediction["values"]
    raise ValueError(f"unknown prediction kind: {kind}")


def observed_text(metric: str, observed: Any) -> str:
    if metric in {"ladder.final", "swap.final", "crowds.sigma8.full"}:
        return f"{observed}人" if not metric.endswith(".full") else f"{observed}組"
    if metric == "swap_k.all_equal_k":
        return "100通り一致" if observed else "不一致あり"
    if metric.endswith("share") or metric.endswith("max_mid_share"):
        return f"{observed * 100:.1f}%"
    if metric == "split.linchpin_median_order":
        return "該当なし" if observed is None else f"中央値 {observed:g}番目"
    if metric == "crowds.sigma20.mean_final" or metric == "split.tiny_mean_final":
        return "該当なし" if observed is None else f"{observed:.1f}人"
    if isinstance(observed, float):
        return f"{observed:.10g}"
    return str(observed)


def check_predictions(summary: dict[str, Any], prediction_path: Path) -> dict[str, Any]:
    source = json.loads(prediction_path.read_text(encoding="utf-8"))
    checks = []
    for prediction in source["predictions"]:
        observed = metric_value(summary, prediction["metric"])
        checks.append({
            "id": prediction["id"], "text": prediction["text"], "metric": prediction["metric"],
            "observed": observed, "observed_text": observed_text(prediction["metric"], observed),
            "hit": prediction_hit(prediction, observed),
        })
    return {"n": len(checks), "hits": sum(row["hit"] for row in checks), "checks": checks}


def dump_json(path: Path, value: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, sort_keys=True, indent=1) + "\n", encoding="utf-8")


def render_report(summary: dict[str, Any], checks: dict[str, Any]) -> str:
    lines = [
        "# No.02 しきい値モデルの結果", "",
        "## E1〜E4", "",
        "| 実験 | 結果 |", "|---|---:|",
        f"| E1 ladder final | {summary['ladder']['final']} |",
        f"| E2 swap final | {summary['swap']['final']} |",
        f"| E3 final == k | {sum(k == v for k, v in zip(summary['swap_k']['k'], summary['swap_k']['final']))}/100 |",
        f"| E4 σc | {summary['expected']['sigma_c']} |",
        f"| E4 直前（σ={summary['expected']['sigma_below']}） | {summary['expected']['re_below']:.10g}人 |",
        "", "## E5（各 n=100）", "",
        "| σ | full | tiny | mid | 平均 |", "|---:|---:|---:|---:|---:|",
    ]
    for sigma in summary["crowds"]["sigma"]:
        row = summary["crowds"]["by_sigma"][str(sigma)]
        lines.append(f"| {sigma} | {row['full']} | {row['tiny']} | {row['mid']} | {row['mean_final']:.2f} |")
    split = summary["split"]
    lines += [
        "", "## E6 と split", "",
        f"- split σ: {split['sigma']}（full {split['full']} / tiny {split['tiny']} / mid {split['mid']}、n={split['n']}）",
        f"- 要の人がいる組: {split['linchpin_crowds']}組（full {split['full']}組中）",
        f"- でたらめな1人の反転: {sum(split['random_flip'])}組（n={split['n']}）",
        "", "## 予想の照合", "",
        f"当たり {checks['hits']} / {checks['n']}。", "",
        "| ID | 予想 | 観測 | 判定 |", "|---|---|---:|---|",
    ]
    for row in checks["checks"]:
        lines.append(f"| {row['id']} | {row['text']} | {row['observed_text']} | {'当たり' if row['hit'] else '外れ'} |")
    return "\n".join(lines) + "\n"


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", type=Path, default=Path("sim/results/ep02"))
    parser.add_argument("--report", type=Path)
    args = parser.parse_args()
    summary = run_experiments()
    predictions = check_predictions(summary, Path(__file__).with_name("ep02_predictions.json"))
    dump_json(args.out / "summary.json", summary)
    dump_json(args.out / "predictions_check.json", predictions)
    if args.report:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(render_report(summary, predictions), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())