import json, random
import numpy as np
import torch
import torch.nn.functional as F

SEED = 1215
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)


def js(p, q):
    m = (p + q) / 2
    return 0.5 * ((p * (p.log() - m.log())).sum(-1) +
                  (q * (q.log() - m.log())).sum(-1))


def exact_math_check():
    cases = [({0, 1}, {0}, True), ({0}, {1}, False),
             ({0, 2}, {2}, True), ({1}, {2}, False)]
    intersection_ok = all(bool(a & b) == expected for a, b, expected in cases)
    p = torch.tensor([[.9, .1], [.9, .1], [.2, .8]], dtype=torch.float64)
    return {
        "intersection_equivalence": intersection_ok,
        "js_zero_equal": float(js(p[0], p[1])),
        "js_positive_unequal": float(js(p[0], p[2])),
        "js_symmetric_error": abs(float(js(p[0], p[2]) - js(p[2], p[0]))),
        "js_upper_bound_log2": float(np.log(2) - js(p[0], p[2])),
    }


def safety_structure_check():
    feasible = [{0}, {0, 1}]
    empty = [{0}, {1}]
    return {
        "feasible_intersection": sorted(set.intersection(*map(set, feasible))),
        "empty_intersection": sorted(set.intersection(*map(set, empty))),
        "empty_intersection_rate": float(len(set.intersection(*map(set, empty))) == 0),
    }


def train(consistency_weight, safety_weight, steps=1200):
    # Histories 0,1 share observation o_feasible and have compatible safety.
    # Histories 2,3 share o_empty and have contradictory safety requirements.
    # Imitation labels intentionally conflict within both observation classes.
    logits = torch.nn.Parameter(torch.zeros(4, 2))
    opt = torch.optim.Adam([logits], lr=0.035)
    labels = torch.tensor([0, 1, 0, 1])
    margin = 0.90
    for _ in range(steps):
        q = logits.softmax(-1)
        task = F.cross_entropy(logits, labels)
        consistency = (js(q[0], q[1]) + js(q[2], q[3])) / 2
        # C(x,U) is 1 iff U is safe for x. For the feasible group, action 0
        # is the sole common safe action, so its group-average margin is q_bar[0].
        q_bar_feasible = q[:2].mean(0)
        robust_hinge = F.relu(margin - q_bar_feasible[0])
        loss = task + consistency_weight * consistency + safety_weight * robust_hinge
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        q = logits.softmax(-1)
        kappa = float((js(q[0], q[1]) + js(q[2], q[3])) / 2)
        feasible_violation = float(1 - q[:2, 0].mean())
        # Empty intersection has max_U min_x C(x,U)=0 exactly, independent of q.
        empty_robust_margin = 0.0
        empty_certificate = True
        return {
            "kappa": kappa, "feasible_violation": feasible_violation,
            "empty_robust_margin": empty_robust_margin,
            "empty_intersection_certificate": empty_certificate,
            "task_ce": float(F.cross_entropy(logits, labels)),
            "q": q.numpy().round(4).tolist(),
        }


def main():
    math = exact_math_check()
    structure = safety_structure_check()
    # First sweep the JS coefficient at zero safety weight: predicted kappa decline.
    consistency_weights = [0.0, 0.02, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0]
    consistency_sweep = {
        str(a): train(a, 0.0) for a in consistency_weights
    }
    # Then sweep robust-safety weight at fixed consistency: predicted feasible
    # violation decline, while the empty intersection remains certified impossible.
    safety_weights = [0.0, 0.02, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0]
    safety_sweep = {
        str(b): train(1.0, b) for b in safety_weights
    }
    baseline = consistency_sweep["0.0"]
    kappas = [consistency_sweep[str(a)]["kappa"] for a in consistency_weights]
    violations = [safety_sweep[str(b)]["feasible_violation"] for b in safety_weights]
    checks = {
        "P1_js_zero_is_exact": math["js_zero_equal"] == 0.0,
        "P2_kappa_reduced_high_consistency": kappas[-1] < kappas[0],
        "P2_monotone_kappa_fraction": float(np.mean(np.diff(kappas) <= 1e-5)),
        "P3_feasible_violation_reduced_high_safety": violations[-1] < violations[0],
        "P3_empty_intersection_always_certificate": all(
            safety_sweep[str(b)]["empty_intersection_certificate"] and
            safety_sweep[str(b)]["empty_robust_margin"] == 0.0 for b in safety_weights),
    }
    result = {
        "seed": SEED, "math_check": math, "safety_structure": structure,
        "baseline_task_only": baseline,
        "consistency_sweep": consistency_sweep,
        "safety_sweep": safety_sweep,
        "prediction_checks": checks,
        "prediction_summary": {
            "P1": "JS is exactly zero for equal distributions and positive for unequal ones.",
            "P2": "Increasing the JS coefficient reduces observational disagreement kappa.",
            "P3": "Increasing robust-safety weight reduces feasible-group unsafe probability; empty intersections stay explicitly infeasible.",
        },
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()