"""Golden-set subsampling simulation: gate decisions on R golden sets vs the population truth."""
import numpy as np

from . import gates, stats

EXACT = ["bm25", "bge", "minilm", "e5"]


def scenarios(cfg):
    """(family, A, B) system pairs; A is the baseline, B the candidate."""
    seeds = cfg["hnsw"]["seeds"]
    pairs = [("noop", f"hnsw_ef{ef}_s{seeds[0]}", f"hnsw_ef{ef}_s{s}") for ef in (16, 128) for s in seeds[1:]]
    pairs += [("natural", "bge", "minilm"), ("natural", "bge", "e5"), ("natural", "bge", "bm25"),
              ("natural", "hnsw_ef128_s0", "hnsw_ef16_s0")]
    pairs += [("natural", "bge", f"del{round(r * 100):02d}") for r in cfg["deletion_rates"]]
    pairs += [("controlled", "bge", f"ctl{round(d * 100):02d}") for d in cfg["controlled_deltas"]]
    return pairs


def controlled(base, deltas, rng):
    """Zero all metrics of randomly ordered queries until population recall@5 drops by each delta."""
    N = len(base["recall@5"])
    order = rng.permutation(N)
    order = order[base["recall@5"][order] > 0]
    cum = np.cumsum(base["recall@5"][order]) / N
    out = {}
    for delta in deltas:
        chosen = order[: np.searchsorted(cum, delta) + 1]
        out[f"ctl{round(delta * 100):02d}"] = {m: np.where(np.isin(np.arange(N), chosen), 0.0, v) for m, v in base.items()}
    return out


def _contrast(rows, key, x):
    mean, lo, hi = stats.mean_ci(x)
    rows.append({**key, "mean": mean, "lo": lo, "hi": hi})


def simulate(name, ds_i, systems, cfg, R=None):
    sim, seed = cfg["simulation"], cfg["seed"]
    R, B, alpha = R or sim["R"], sim["B"], sim["alpha"]
    N = len(systems["bge"]["recall@5"])
    split = np.random.default_rng([seed, ds_i, 99]).permutation(N)
    truths = {"population": (np.arange(N), np.arange(N)), "split_half": (split[: N // 2], split[N // 2:])}
    rows_t, rows_p, rows_c = [], [], []

    for ti, (truth, (pool, tidx)) in enumerate(truths.items()):
        for n in [n for n in sim["n"] if n <= N // 2]:
            rng = np.random.default_rng([seed, ds_i, ti, n])
            S = rng.permuted(np.tile(pool, (R, 1)), axis=1)[:, :n]
            Wa = stats.bootstrap_weights(n, B, rng).astype(np.float32)
            Wb = stats.bootstrap_weights(n, B, rng).astype(np.float32)
            flips = rng.choice(np.array([-1.0, 1.0], np.float32), size=(B, n))
            for metric in sim["metrics"]:
                for sysname in EXACT:
                    v = systems[sysname][metric]
                    true, est = v[tidx].mean(), v[S].mean(1)
                    for d in sim["d"]:
                        rows_t.append(dict(dataset=name, truth=truth, system=sysname, metric=metric, n=n, d=d,
                                           true_value=true, truth_pass=d >= 0,
                                           k_pass=int(gates.g1_threshold(v[S], true - d).sum()), R=R))
                family = {}
                for fam, A, Bn in scenarios(cfg):
                    a, b = systems[A][metric], systems[Bn][metric]
                    XA, XB = a[S].astype(np.float32), b[S].astype(np.float32)
                    delta, sd_d = (b - a)[tidx].mean(), (b - a)[pool].std(ddof=1)
                    lb_p = stats.lower_bound((XB - XA) @ Wa.T, alpha)
                    lb_u = stats.lower_bound(XB @ Wb.T - XA @ Wa.T, alpha)
                    key = dict(dataset=name, truth=truth, family=fam, A=A, B=Bn, metric=metric, n=n, R=R)
                    base = {**key, "delta_true": delta, "sd_d": sd_d}
                    g1 = {h: gates.g1_threshold(XB, a[tidx].mean() - h) for h in sim["headroom"]}
                    for h, x in g1.items():
                        rows_p.append({**base, "gate": "G1", "param": h, "k_pass": int(x.sum())})
                    for m in sim["margins"]:
                        dec = {"G2": gates.g2_naive_delta(XA, XB, m), "G3": lb_p >= -m, "G3u": lb_u >= -m,
                               "G3signflip": gates.g3_sign_flip(XA, XB, m, flips, alpha)}
                        for g, x in dec.items():
                            rows_p.append({**base, "gate": g, "param": m, "k_pass": int(x.sum()),
                                           "analytic_pass": stats.noninferiority_pass_prob(delta, m, sd_d, n, alpha)
                                           if g == "G3" else np.nan})
                        h2 = g1[max(sim["headroom"])].astype(float) - dec["G3"]
                        h3 = (~dec["G3u"]).astype(float) - (~dec["G3"])
                        _contrast(rows_c, {**key, "param": m, "contrast": "G1pass-G3pass"}, h2)
                        _contrast(rows_c, {**key, "param": m, "contrast": "G3ufail-G3fail"}, h3)
                        if fam == "noop":
                            family[m] = family.get(m, 0) + h3
                n_noop = sum(f == "noop" for f, _, _ in scenarios(cfg))
                for m, total in family.items():
                    _contrast(rows_c, dict(dataset=name, truth=truth, family="noop", A="*", B="*", metric=metric,
                                           n=n, R=R, param=m, contrast="G3ufail-G3fail"), total / n_noop)
    return rows_t, rows_p, rows_c