"""
IF&PA benchmark: Monte Carlo comparison of the preferential median with common
alternatives under chosen conditions. Reproduces the site's applicability
configurator (benchmark.js): same model, same seeded generator.

Usage:
    python benchmark.py                       # default conditions
    python benchmark.py --preset overconfident
    python benchmark.py --m 12 --contamination 0.2 --shift 8 --outlier-u 0.3
    python benchmark.py --list                # show the presets

Model (units of sigma; the true value is 0). For each of m results:
    u_k  stated uncertainty: 1, or log-uniform on [0.5, 2] if heterogeneous;
    x_k = understatement * u_k * noise + dark * N(0, 1),
          noise ~ normal / uniform / Student t3 scaled to unit variance;
    with probability `contamination` the result is an outlier:
          x_k += sign * shift, u_k *= outlier_u.

Needs ifpa.py in the same folder. Pure Python.
"""
import argparse
import math

from ifpa import ESTIMATORS, compare_estimators, preferential_median

M32 = 0xFFFFFFFF

PRESETS = {
    "clean":         dict(m=10, kernel="normal",  heterogeneous=True,  understatement=1, dark=0,   contamination=0,    shift=0,  one_sided=True,  outlier_u=1),
    "overconfident": dict(m=10, kernel="normal",  heterogeneous=True,  understatement=1, dark=0,   contamination=0.15, shift=6,  one_sided=True,  outlier_u=0.3),
    "faults":        dict(m=8,  kernel="normal",  heterogeneous=False, understatement=1, dark=0,   contamination=0.2,  shift=10, one_sided=False, outlier_u=1),
    "dark":          dict(m=10, kernel="normal",  heterogeneous=True,  understatement=1, dark=1.5, contamination=0,    shift=0,  one_sided=True,  outlier_u=1),
    "bounded":       dict(m=15, kernel="uniform", heterogeneous=False, understatement=1, dark=0,   contamination=0.1,  shift=9,  one_sided=True,  outlier_u=1),
    "heavy":         dict(m=12, kernel="t3",      heterogeneous=True,  understatement=1, dark=0,   contamination=0,    shift=0,  one_sided=True,  outlier_u=1),
    "small":         dict(m=4,  kernel="normal",  heterogeneous=True,  understatement=1, dark=0,   contamination=0.25, shift=6,  one_sided=True,  outlier_u=1),
    "precise":       dict(m=10, kernel="uniform", heterogeneous=True,  understatement=1, dark=0,   contamination=0,    shift=0,  one_sided=True,  outlier_u=1),
    "skewed":        dict(m=10, kernel="skewed",  heterogeneous=True,  understatement=1, dark=0,   contamination=0.2,  shift=6,  one_sided=True,  outlier_u=0.3),
    "fewlabs":       dict(m=5, kernel="normal",  heterogeneous=True,   understatement=1, dark=0,   contamination=0.2,  shift=6,  one_sided=True,   outlier_u=0.1),
    "twosided":      dict(m=10, kernel="normal",  heterogeneous=True,   understatement=1, dark=0,   contamination=0.2,  shift=6,  one_sided=False,  outlier_u=0.3),
    "heavyconf":     dict(m=10, kernel="t3",      heterogeneous=True,   understatement=1, dark=0,   contamination=0.2,  shift=6,  one_sided=True,   outlier_u=0.1),
    "allunder":      dict(m=10, kernel="normal",  heterogeneous=True,   understatement=2, dark=0,   contamination=0.3,  shift=6,  one_sided=True,   outlier_u=0.1),
    "digital":       dict(m=20, kernel="uniform", heterogeneous=False,  understatement=1, dark=0,   contamination=0.2,  shift=6,  one_sided=False,  outlier_u=1),
}
DEFAULTS = dict(m=10, kernel="normal", heterogeneous=True, understatement=1, dark=0,
                contamination=0.1, shift=6, one_sided=True, outlier_u=1, replicates=500, seed=1,
                pm_k=2.0, pm_grid="fine")
KERNELS = ["normal", "uniform", "triangular", "laplace", "t3", "skewed"]
GRID_STEP = {"fine": 0.05, "medium": 0.2}   # шаг сетки ДАЗ в долях медианы u; "auto" — сетка IF&PA-L


def _imul(a, b):
    return (a * b) & M32


def rng(seed):
    """mulberry32 — the same generator as in benchmark.js."""
    state = [seed & M32]

    def draw():
        state[0] = (state[0] + 0x6D2B79F5) & M32
        t = state[0]
        t = _imul(t ^ (t >> 15), t | 1)
        t = t ^ ((t + _imul(t ^ (t >> 7), t | 61)) & M32)
        return ((t ^ (t >> 14)) & M32) / 4294967296.0
    return draw


def normal(r):
    u1, u2 = r(), r()
    return math.sqrt(-2.0 * math.log(1.0 - u1)) * math.cos(2.0 * math.pi * u2)


SKEW_S = 0.6   # логнормальное: параметр формы; ошибка приведена к нулевому среднему и единичной дисперсии


def noise(r, kernel):
    """Погрешность с нулевым средним и единичной дисперсией; число случайных чисел фиксировано для ядра."""
    if kernel == "uniform":
        return (2.0 * r() - 1.0) * math.sqrt(3.0)
    if kernel == "triangular":                      # сумма двух равномерных, [-sqrt6; sqrt6]
        return (r() + r() - 1.0) * math.sqrt(6.0)
    if kernel == "laplace":                         # острый пик и экспоненциальные хвосты
        a = r() - 0.5
        return -(1.0 / math.sqrt(2.0)) * math.copysign(1.0, a) * math.log(1.0 - 2.0 * abs(a))
    if kernel == "t3":
        z = normal(r)
        chi = normal(r) ** 2 + normal(r) ** 2 + normal(r) ** 2
        return z / math.sqrt(chi / 3.0) / math.sqrt(3.0)
    if kernel == "skewed":                          # логнормальная, скошенная вправо
        s2 = SKEW_S * SKEW_S
        return (math.exp(SKEW_S * normal(r)) - math.exp(s2 / 2.0)) / math.sqrt((math.exp(s2) - 1.0) * math.exp(s2))
    return normal(r)


def pm_estimate(x, u, p):
    """ПМ с настройками конфигуратора: интервалы x ± k·u;
    сетка — как в IF&PA-L («auto»), средняя или мелкая (шаг 0,2 или 0,05 медианы u)."""
    uu = [p["pm_k"] * v for v in u]
    n = None
    if p["pm_grid"] in GRID_STEP:
        lo = min(a - b for a, b in zip(x, uu))
        hi = max(a + b for a, b in zip(x, uu))
        su = sorted(u)
        med = su[len(u) // 2] if len(u) % 2 else (su[len(u) // 2 - 1] + su[len(u) // 2]) / 2
        h = GRID_STEP[p["pm_grid"]] * med
        if h > 0 and hi > lo:
            n = min(4001, max(11, int(math.ceil((hi - lo) / h)) + 1))
    # мелкая сетка: несмежные победители — реальные разные области, дробить сетку дальше бессмысленно
    r = preferential_median(x, uu, n=n, n_max=n if n else 201)
    return r["pm"], r["u_pm"]


def dataset(p, r):
    """One simulated data set; the order of random draws matches benchmark.js."""
    x, u = [], []
    for _ in range(p["m"]):
        a = r()
        uk = math.exp(math.log(0.5) + a * math.log(4.0)) if p["heterogeneous"] else 1.0
        xk = p["understatement"] * uk * noise(r, p["kernel"])
        xk += p["dark"] * normal(r)
        is_out = r() < p["contamination"]
        sign = -1.0 if r() < 0.5 else 1.0
        if is_out:
            xk += (1.0 if p["one_sided"] else sign) * p["shift"]
            uk *= p["outlier_u"]
        x.append(xk)
        u.append(uk)
    return x, u


def run(**params):
    p = dict(DEFAULTS)
    p.update(params)
    r = rng(p["seed"])
    errors = {n: [] for n in ESTIMATORS}
    covered = {n: 0 for n in ESTIMATORS}
    counted = {n: 0 for n in ESTIMATORS}
    for _ in range(p["replicates"]):
        x, u = dataset(p, r)
        est = compare_estimators(x, u)
        est[0] = dict(name=ESTIMATORS[0], value=0, u=0)
        est[0]["value"], est[0]["u"] = pm_estimate(x, u, p)
        for e in est:
            errors[e["name"]].append(e["value"])          # the true value is 0
            if e["u"] is not None and math.isfinite(e["u"]):
                counted[e["name"]] += 1
                covered[e["name"]] += abs(e["value"]) <= 2 * e["u"]
    rows = []
    for name in ESTIMATORS:
        e = errors[name]
        a = sorted(abs(v) for v in e)
        pct = lambda q: a[max(0, math.ceil(q * len(a)) - 1)]  # noqa: E731
        rows.append(dict(name=name, rmse=math.sqrt(sum(v * v for v in e) / len(e)),
                         mae=pct(0.5), p90=pct(0.9), q25=pct(0.25), q75=pct(0.75),
                         coverage=covered[name] / counted[name] if counted[name] else None))
    rows.sort(key=lambda row: row["rmse"])
    return p, rows


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--preset", choices=sorted(PRESETS))
    ap.add_argument("--list", action="store_true", help="show presets and exit")
    ap.add_argument("--m", type=int)
    ap.add_argument("--kernel", choices=KERNELS)
    ap.add_argument("--pm-k", type=float, help="interval for the PM: x ± k·u (default 2)")
    ap.add_argument("--pm-grid", choices=["fine", "medium", "auto"], help="RAV grid for the PM: step 0.05 or 0.2 of median u, or IF&PA-L (default fine)")
    ap.add_argument("--equal-u", action="store_true", help="all stated uncertainties equal")
    ap.add_argument("--understatement", type=float)
    ap.add_argument("--dark", type=float)
    ap.add_argument("--contamination", type=float)
    ap.add_argument("--shift", type=float)
    ap.add_argument("--two-sided", action="store_true")
    ap.add_argument("--outlier-u", type=float)
    ap.add_argument("--replicates", type=int)
    ap.add_argument("--seed", type=int)
    a = ap.parse_args()
    if a.list:
        for k, v in PRESETS.items():
            print(f"{k:14s} {v}")
        return
    params = dict(PRESETS[a.preset]) if a.preset else {}
    for key, val in [("m", a.m), ("kernel", a.kernel), ("understatement", a.understatement),
                     ("dark", a.dark), ("contamination", a.contamination), ("shift", a.shift),
                     ("outlier_u", a.outlier_u), ("replicates", a.replicates), ("seed", a.seed),
                     ("pm_k", a.pm_k), ("pm_grid", a.pm_grid)]:
        if val is not None:
            params[key] = val
    if a.equal_u:
        params["heterogeneous"] = False
    if a.two_sided:
        params["one_sided"] = False
    p, rows = run(**params)
    print("Conditions:", {k: p[k] for k in DEFAULTS})
    print(f"\n{'Estimator':34s} {'RMSE':>7s} {'median|e|':>10s} {'P90|e|':>8s} {'cover ±2u':>10s}")
    for row in rows:
        cov = "" if row["coverage"] is None else f"{row['coverage']:.0%}"
        mark = "  <-- IF&PA" if row["name"].startswith("Preferential") else ""
        print(f"{row['name']:34s} {row['rmse']:7.3f} {row['mae']:10.3f} {row['p90']:8.3f} {cov:>10s}{mark}")
    print("\nErrors in units of sigma (typical measurement standard deviation); true value 0.")


if __name__ == "__main__":
    main()