import json
from pathlib import Path
import numpy as np
from sklearn.linear_model import LogisticRegression

SEED = 3048


def primary_action(x):
    # Aggressive controller: fast near the goal, but saturation makes it brittle
    # from boundary states under disturbances.
    return (-0.85 * x if abs(x) < 0.85 else 0.75 * np.sign(x))


def fallback_action(x):
    # Conservative controller: slower, with smaller action excursions.
    return np.clip(-0.42 * x, -0.55, 0.55)


def rollout(x0, policy, noise_std, horizon=12, rng=None):
    rng = np.random.default_rng() if rng is None else rng
    x = float(x0)
    reached = False
    failed = False
    for _ in range(horizon):
        x = x + float(policy(x)) + rng.normal(0.0, noise_std)
        if abs(x) >= 2.0:
            failed = True
            break
        if abs(x) <= 0.16:
            reached = True
            break
    return reached and not failed, failed, x


def features(xs):
    xs = np.asarray(xs)
    return np.column_stack([np.abs(xs), xs * xs])


def make_critic_data(n, noise, policy, seed):
    rng = np.random.default_rng(seed)
    xs = rng.uniform(-1.95, 1.95, size=n)
    successes = np.array([rollout(x, policy, noise, rng=rng)[0] for x in xs], dtype=int)
    return features(xs), successes


def train_critics():
    # P estimates primary-policy success probability.
    xp, yp = make_critic_data(8000, 0.42, primary_action, SEED)
    # V is failure risk under the conservative fallback, so lower is safer.
    xv, yv_success = make_critic_data(8000, 0.62, fallback_action, SEED + 1)
    yv_failure = 1 - yv_success
    p = LogisticRegression(C=10.0, solver="lbfgs", random_state=SEED).fit(xp, yp)
    v = LogisticRegression(C=10.0, solver="lbfgs", random_state=SEED + 1).fit(xv, yv_failure)
    return p, v


def gate(p_hat, v_hat, eta_p=0.55, eta_r=0.20):
    # Execute primary iff both paper conditions pass.
    return (p_hat >= eta_p) & (v_hat <= eta_r)


def evaluate(mode, p, v, n=5000, noise=0.48, seed=SEED + 20):
    rng = np.random.default_rng(seed)
    xs = np.clip(rng.normal(0.0, 1.45, n), -1.95, 1.95)
    p_hat = p.predict_proba(features(xs))[:, 1]
    v_hat = v.predict_proba(features(xs))[:, 1]
    if mode == "primary":
        use_primary = np.ones(n, dtype=bool)
    elif mode == "p_gate":
        use_primary = p_hat >= 0.55
    elif mode == "dual_gate":
        use_primary = gate(p_hat, v_hat)
    else:
        raise ValueError(mode)
    success = np.zeros(n, dtype=bool)
    failed = np.zeros(n, dtype=bool)
    for i, x in enumerate(xs):
        pol = primary_action if use_primary[i] else fallback_action
        success[i], failed[i], _ = rollout(x, pol, noise, rng=rng)
    return {
        "catastrophic_failure_rate": float(failed.mean()),
        "goal_success_rate": float(success.mean()),
        "primary_coverage": float(use_primary.mean()),
        "mean_p_selected": float(p_hat[use_primary].mean()) if use_primary.any() else 0.0,
        "mean_v_selected": float(v_hat[use_primary].mean()) if use_primary.any() else 0.0,
    }


def verify_math():
    rng = np.random.default_rng(SEED)
    p, v = rng.random(10000), rng.random(10000)
    a = gate(p, v, 0.55, 0.20)
    direct = (p >= 0.55) & (v <= 0.20)
    tighter = gate(p, v, 0.70, 0.10)
    return {"formula_exact": bool(np.array_equal(a, direct)),
            "stricter_threshold_subset": bool(np.all(~tighter | a)),
            "selected_fraction": float(a.mean())}


def main():
    math = verify_math()
    p, v = train_critics()
    results = {m: evaluate(m, p, v) for m in ("primary", "p_gate", "dual_gate")}
    out = {"math_check": math, "results": results,
           "config": {"seed": SEED, "train_noise_p": 0.42, "train_noise_v": 0.62,
                      "eval_noise": 0.48, "eta_p": 0.55, "eta_r": 0.20}}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()