"""Generate reproducible charts for the public Convee A/B testing guide.

Run from ``Convee BE``:

    MPLCONFIGDIR=/tmp/convee-matplotlib .venv/bin/python \
        -m simulations.generate_help_ab_charts

The simulation compares three procedures under a true A/A null:

* one ordinary two-sided pooled z-test at a preselected final horizon;
* repeatedly applying the same p <= 0.05 rule and stopping on the first hit;
* Convee's pairwise Dirichlet-multinomial e-process.

It also demonstrates family-wise error in A/A/A/A and creates one treatment
uplift trajectory. The vectorized e-value formula is checked against the
production implementation before charts are written.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.special import gammaln, ndtr

from src.sequential_analysis import pair_log_e_value


CONVEE = "#6871E4"
CONVEE_DARK = "#2A33C0"
RED = "#EF4444"
AMBER = "#F59E0B"
GREEN = "#22C55E"
INK = "#172033"
MUTED = "#687386"
GRID = "#E6E9F0"
PAPER = "#FFFFFF"


def _log_beta(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return gammaln(a) + gammaln(b) - gammaln(a + b)


def pair_log_e_vectorized(
    control_conversions: np.ndarray,
    treatment_conversions: np.ndarray,
    prior_concentration: float = 100.0,
) -> np.ndarray:
    prior = prior_concentration / 2
    return (
        _log_beta(prior + control_conversions, prior + treatment_conversions)
        - _log_beta(np.asarray(prior), np.asarray(prior))
        - (control_conversions + treatment_conversions) * math.log(0.5)
    )


def _assert_matches_production() -> None:
    for control, treatment in ((0, 0), (10, 15), (50, 80), (50, 100)):
        vectorized = float(pair_log_e_vectorized(
            np.asarray(control), np.asarray(treatment)
        ))
        production = pair_log_e_value(
            [control, treatment], [0.5, 0.5], 100.0
        )
        if not math.isclose(vectorized, production, abs_tol=1e-10):
            raise AssertionError(
                f"e-value mismatch for {control}/{treatment}: "
                f"{vectorized} != {production}"
            )


def _ordinary_p_values(
    control: np.ndarray,
    treatment: np.ndarray,
    visitors_per_arm: np.ndarray,
) -> np.ndarray:
    n = visitors_per_arm[None, :]
    pooled = (control + treatment) / (2 * n)
    variance = pooled * (1 - pooled) * (2 / n)
    difference = treatment / n - control / n
    z = np.divide(
        difference,
        np.sqrt(variance),
        out=np.zeros_like(difference, dtype=float),
        where=variance > 0,
    )
    return 2 * ndtr(-np.abs(z))


def simulate_aa(
    *, runs: int, looks: int, visitors_step: int, baseline: float, seed: int
) -> dict:
    rng = np.random.default_rng(seed)
    increments = rng.binomial(
        visitors_step, baseline, size=(runs, looks, 2)
    )
    cumulative = np.cumsum(increments, axis=1)
    visitors = np.arange(1, looks + 1) * visitors_step
    p_values = _ordinary_p_values(
        cumulative[:, :, 0], cumulative[:, :, 1], visitors
    )
    ordinary_repeated = np.minimum.accumulate(p_values, axis=1) <= 0.05
    log_e = pair_log_e_vectorized(cumulative[:, :, 0], cumulative[:, :, 1])
    sequential = np.maximum.accumulate(log_e, axis=1) >= math.log(20)
    return {
        "visitors": visitors,
        "ordinary_repeated": ordinary_repeated.mean(axis=0),
        "sequential": sequential.mean(axis=0),
        "fixed_final": float((p_values[:, -1] <= 0.05).mean()),
        "ordinary_final": float(ordinary_repeated[:, -1].mean()),
        "sequential_final": float(sequential[:, -1].mean()),
    }


def simulate_aaaa(
    *, runs: int, looks: int, visitors_step: int, baseline: float, seed: int
) -> dict:
    rng = np.random.default_rng(seed)
    increments = rng.binomial(
        visitors_step, baseline, size=(runs, looks, 4)
    )
    cumulative = np.cumsum(increments, axis=1)
    pair_logs = np.stack([
        pair_log_e_vectorized(cumulative[:, :, 0], cumulative[:, :, arm])
        for arm in range(1, 4)
    ], axis=2)
    pair_max = np.maximum.accumulate(pair_logs, axis=1)
    any_uncorrected = np.any(pair_max >= math.log(20), axis=2)
    any_bonferroni = np.any(pair_max >= math.log(60), axis=2)
    visitors = np.arange(1, looks + 1) * visitors_step
    return {
        "visitors": visitors,
        "uncorrected": any_uncorrected.mean(axis=0),
        "bonferroni": any_bonferroni.mean(axis=0),
        "uncorrected_final": float(any_uncorrected[:, -1].mean()),
        "bonferroni_final": float(any_bonferroni[:, -1].mean()),
    }


def simulate_uplift_path(
    *, looks: int, visitors_step: int, seed: int
) -> dict:
    selected = None
    selected_seed = None
    for candidate in range(seed, seed + 10_000):
        rng = np.random.default_rng(candidate)
        increments_a = rng.binomial(visitors_step, 0.05, size=looks)
        increments_b = rng.binomial(visitors_step, 0.08, size=looks)
        control = np.cumsum(increments_a)
        treatment = np.cumsum(increments_b)
        log_e = pair_log_e_vectorized(control, treatment)
        crossing = np.flatnonzero(log_e >= math.log(20))
        if len(crossing) and 8 <= crossing[0] <= looks - 8:
            selected = (control, treatment, log_e)
            selected_seed = candidate
            break
    if selected is None:
        raise RuntimeError("could not find a readable uplift trajectory")
    control, treatment, log_e = selected
    return {
        "seed": selected_seed,
        "visitors": np.arange(1, looks + 1) * visitors_step * 2,
        "control": control,
        "treatment": treatment,
        "current_e": np.exp(log_e),
        "max_e": np.exp(np.maximum.accumulate(np.maximum(log_e, 0))),
        "crossing_look": int(np.flatnonzero(log_e >= math.log(20))[0]),
    }


def _style_axis(axis) -> None:
    axis.set_facecolor(PAPER)
    axis.grid(axis="y", color=GRID, linewidth=1)
    axis.spines[["top", "right", "left"]].set_visible(False)
    axis.spines["bottom"].set_color(GRID)
    axis.tick_params(colors=MUTED, labelsize=10, length=0)
    axis.set_axisbelow(True)


def _save(fig, path: Path) -> None:
    fig.savefig(path, dpi=180, bbox_inches="tight", facecolor=PAPER)
    plt.close(fig)


def chart_peeking(result: dict, output: Path) -> None:
    fig, axis = plt.subplots(figsize=(10.5, 6.2))
    axis.set_facecolor(PAPER)
    axis.set_xlim(-0.7, 10)
    axis.set_ylim(-0.8, 5.7)
    axis.axis("off")

    rows = [
        ("Один p-value в запланированном финале", result["fixed_final"], INK),
        ("Обычный p-value проверяли 30 раз", result["ordinary_final"], RED),
        ("Метод Convee проверяли постоянно", result["sequential_final"], CONVEE),
    ]
    for row_index, (label, rate, color) in enumerate(rows):
        y_base = 4.7 - row_index * 2.15
        false_worlds = round(rate * 100)
        axis.text(0, y_base + 0.65, label, color=INK, fontsize=12, fontweight="bold")
        axis.text(
            9.8, y_base + 0.65,
            f"{rate * 100:.2f}%  ≈  {false_worlds} из 100",
            color=color, fontsize=12, fontweight="bold", ha="right",
        )
        for index in range(100):
            column = index % 20
            line = index // 20
            axis.scatter(
                column * 0.5, y_base - line * 0.23,
                s=22, marker="o",
                color=color if index < false_worlds else "#E7EAF0",
                edgecolors="none",
            )

    _save(fig, output / "peeking-aa-false-positives.webp")


def chart_multiple_variants(result: dict, output: Path) -> None:
    fig, axis = plt.subplots(figsize=(10.5, 5.7))
    _style_axis(axis)
    x = result["visitors"] * 4
    axis.plot(
        x, result["uncorrected"] * 100,
        color=AMBER, linewidth=3, label="3 сравнения, каждому оставить 5%",
    )
    axis.plot(
        x, result["bonferroni"] * 100,
        color=CONVEE, linewidth=3, label="Convee: общий лимит 5% делится на 3",
    )
    axis.axhline(5, color=INK, linestyle="--", linewidth=1.5, label="Общий лимит 5%")
    axis.set_ylim(0, max(7, result["uncorrected"].max() * 125))
    axis.set_xlabel("Всего посетителей A + B + C + D", color=MUTED, labelpad=12)
    axis.set_ylabel("Хотя бы один ложный вывод, %", color=MUTED, labelpad=12)
    axis.legend(frameon=False, loc="upper left", fontsize=10)
    _save(fig, output / "abn-family-error.webp")


def chart_e_path(result: dict, output: Path) -> None:
    fig, axis = plt.subplots(figsize=(10.5, 5.7))
    _style_axis(axis)
    crossing = result["crossing_look"]
    x = result["visitors"][:crossing + 1]
    current = result["current_e"][:crossing + 1]

    axis.axhspan(0, 5, color="#F7F8FA", zorder=0)
    axis.axhspan(5, 20, color="#FFF8E8", alpha=0.85, zorder=0)
    axis.axhspan(20, 30, color="#ECFDF5", alpha=0.9, zorder=0)
    axis.plot(
        x, current, color=CONVEE, linewidth=3, marker="o", markersize=5,
        markerfacecolor=PAPER, markeredgewidth=2, label="Накопленные доказательства",
    )
    axis.axhline(20, color=GREEN, linestyle="--", linewidth=2)
    axis.scatter(
        [x[-1]], [current[-1]], s=95, color=GREEN,
        edgecolor=PAPER, linewidth=2, zorder=5,
    )
    axis.annotate(
        "Результат подтверждён\nA: 74 заказа · B: 121 заказ",
        xy=(x[-1], current[-1]),
        xytext=(-165, 28), textcoords="offset points",
        color=INK, fontsize=10, linespacing=1.5,
        arrowprops={"arrowstyle": "->", "color": MUTED},
    )
    axis.text(x[0], 2.5, "Данных мало", color=MUTED, fontsize=10, va="center")
    axis.text(x[0], 12.5, "Есть сигнал, но решать рано", color="#A16207", fontsize=10, va="center")
    axis.text(x[0], 24.8, "Зона подтверждённого результата", color="#047857", fontsize=10, va="center")
    axis.set_xlim(0, x[-1] * 1.04)
    axis.set_ylim(0, 30)
    axis.set_xlabel("Всего посетителей A + B", color=MUTED, labelpad=12)
    axis.set_ylabel("Сила доказательств", color=MUTED, labelpad=12)
    _save(fig, output / "e-value-path.webp")


def _serializable(result: dict) -> dict:
    serialized = {}
    for key, value in result.items():
        if isinstance(value, np.ndarray):
            serialized[key] = value.tolist()
        elif isinstance(value, np.generic):
            serialized[key] = value.item()
        else:
            serialized[key] = value
    return serialized


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    default_output = (
        Path(__file__).resolve().parents[2]
        / "Convee FE" / "public" / "help-assets" / "ab-testing"
    )
    parser.add_argument("--output", type=Path, default=default_output)
    parser.add_argument("--runs", type=int, default=20_000)
    args = parser.parse_args()
    if args.runs < 1_000:
        raise SystemExit("Use at least 1,000 runs for a public chart")

    _assert_matches_production()
    output = args.output.resolve()
    output.mkdir(parents=True, exist_ok=True)

    aa = simulate_aa(
        runs=args.runs, looks=30, visitors_step=250,
        baseline=0.05, seed=2026072201,
    )
    aaaa = simulate_aaaa(
        runs=args.runs, looks=200, visitors_step=1_000,
        baseline=0.05, seed=2026072202,
    )
    uplift = simulate_uplift_path(
        looks=80, visitors_step=100, seed=2026072203,
    )

    chart_peeking(aa, output)
    chart_multiple_variants(aaaa, output)
    chart_e_path(uplift, output)

    summary = {
        "simulation_version": "help_ab_testing_v1",
        "runs": args.runs,
        "baseline_conversion_rate": 0.05,
        "seed_aa": 2026072201,
        "seed_aaaa": 2026072202,
        "aa": {
            "looks": 30,
            "visitors_per_arm_per_look": 250,
            "fixed_horizon_false_positive_rate": aa["fixed_final"],
            "repeated_p_value_false_positive_rate": aa["ordinary_final"],
            "convee_false_positive_rate": aa["sequential_final"],
        },
        "aaaa": {
            "looks": 200,
            "visitors_per_arm_per_look": 1_000,
            "uncorrected_family_false_positive_rate": aaaa["uncorrected_final"],
            "convee_bonferroni_family_false_positive_rate": aaaa["bonferroni_final"],
        },
        "uplift_path": {
            "seed": uplift["seed"],
            "crossing_total_visitors": int(uplift["visitors"][uplift["crossing_look"]]),
            "control_conversions_at_crossing": int(uplift["control"][uplift["crossing_look"]]),
            "treatment_conversions_at_crossing": int(uplift["treatment"][uplift["crossing_look"]]),
            "e_value_at_crossing": float(uplift["current_e"][uplift["crossing_look"]]),
        },
    }
    (output / "simulation-results.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (output / "simulation.py").write_text(
        Path(__file__).read_text(encoding="utf-8"),
        encoding="utf-8",
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()