import json
from pathlib import Path
import numpy as np


def survival_from_hazard(h):
    h = np.asarray(h, dtype=float)
    S = np.ones(len(h) + 1)
    for n in range(len(h)):
        S[n + 1] = S[n] * (1.0 - h[n])
    return S, h * S[:-1]


def make_sequence(rng, n, heavy_tail=True):
    vals, durations, regime = [], [], 0
    while len(vals) < n:
        if heavy_tail:
            d = int(np.clip(np.floor(2.0 * (1.0 + rng.pareto(1.7))), 2, 80))
        else:
            d = int(rng.geometric(1 / 12.0))
        durations.append(d)
        vals.extend([regime] * d)
        regime = 1 - regime
    return np.asarray(vals[:n], dtype=np.int64), durations


def duration_hazard(durations, max_age):
    d = np.asarray(durations)
    return np.array([
        np.mean(d == a + 1) / max(np.mean(d >= a + 1), 1e-12)
        for a in range(max_age)
    ]).clip(1e-5, 1 - 1e-5)


def tokenwise_router(x, noise, rng):
    obs = x + rng.normal(0, noise, len(x))
    return (obs > 0.5).astype(np.int64)


def age_router(x, noise, rng, hazard, max_age, evidence_gain=5.0):
    obs = x + rng.normal(0, noise, len(x))
    out = np.empty(len(x), dtype=np.int64)
    regime = int(obs[0] > 0.5)
    age = 0
    calls, switches = 1, 0
    out[0] = regime
    for t in range(1, len(x)):
        candidate = int(obs[t] > 0.5)
        base = hazard[min(age, max_age - 1)]
        disagreement = abs(obs[t] - 0.5) if candidate != regime else 0.0
        logit = np.log(base / (1.0 - base)) + evidence_gain * disagreement
        p_switch = 1.0 / (1.0 + np.exp(-np.clip(logit, -30, 30)))
        if rng.random() < p_switch:
            calls += 1
            if candidate != regime:
                regime = candidate
                switches += 1
                age = 0
            else:
                age += 1
        else:
            age += 1
        out[t] = regime
    return out, calls, switches


def score(pred, x, n):
    return {"accuracy": float(np.mean(pred == x)),
            "switches": int(np.sum(pred[1:] != pred[:-1])),
            "switch_rate": float(np.mean(pred[1:] != pred[:-1])),
            "router_calls": int(n)}


def evaluate(seed=7, n=30000, noise=0.65):
    rng = np.random.default_rng(seed)
    x, true_durations = make_sequence(rng, n, heavy_tail=True)
    train_rng = np.random.default_rng(seed + 100)
    _, train_durations = make_sequence(train_rng, 200000, heavy_tail=True)
    max_age = 80
    hazards = duration_hazard(train_durations, max_age)
    b = tokenwise_router(x, noise, np.random.default_rng(seed + 1))
    a, calls, _ = age_router(x, noise, np.random.default_rng(seed + 2), hazards, max_age)
    bm, am = score(b, x, n), score(a, x, calls)
    h = duration_hazard(true_durations, max_age)
    S, p = survival_from_hazard(h)
    identity_error = float(abs(np.sum(p) + S[-1] - 1.0))
    d = np.asarray(true_durations)
    implied_mean = float(np.sum((np.arange(max_age) + 1) * p) + (max_age + 1) * S[-1])
    result = {
        "settings": {"seed": seed, "n": n, "noise": noise},
        "math": {"max_survival_identity_error": identity_error,
                 "empirical_mean_duration": float(np.mean(d)),
                 "implied_mean_duration_truncated_tail": implied_mean,
                 "hazard_first_12": hazards[:12].tolist()},
        "baseline": bm, "age_conditioned": am,
        "relative_switch_reduction": float(1 - am["switches"] / max(bm["switches"], 1)),
        "relative_router_call_reduction": float(1 - calls / n),
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))
    return result


if __name__ == "__main__":
    evaluate()