import json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F

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


def softmax_extreme(x, tau):
    return tau * torch.logsumexp(x / tau, dim=-1)


def softmin_extreme(x, tau):
    return -tau * torch.logsumexp(-x / tau, dim=-1)


def rncoa_loss(j, obstacle_min, obstacle_max, gamma, M=2.0, tau=0.05,
               lambda_gamma=0.01, lambda_c=10.0):
    """The authoritative displayed L_RNCOA formula, including its signs."""
    smax = softmax_extreme(j, tau)
    smin = softmin_extreme(j, tau)
    g1, g2 = gamma[..., 0], gamma[..., 1]
    v1 = F.relu(obstacle_min - M*g1 - smax)
    v2 = F.relu(smin - obstacle_max - M*g2)
    coupling = F.relu(g1 + g2 - 1.0)
    return (v1 + v2 + lambda_gamma*(g1+g2) + lambda_c*coupling,
            {"smax": smax, "smin": smin, "v1": v1, "v2": v2,
             "coupling": coupling})


def independent_overlap_loss(j, obstacle_min, obstacle_max):
    """Per-vertex overlap hinge; it does not encode whole-body separation."""
    return torch.minimum(j - obstacle_min, obstacle_max - j).clamp_min(0).sum(-1)


def exact_body_collision(j, lo, hi):
    return bool((j.max().item() >= lo) and (j.min().item() <= hi))


def analytic_check():
    cases = {
        "left_safe": torch.tensor([[-1.4, -1.1, -1.3, -1.2]]),
        "collision": torch.tensor([[-0.3, 0.1, 0.3, -0.1]]),
        "right_safe": torch.tensor([[1.1, 1.3, 1.2, 1.4]]),
        "straddling": torch.tensor([[-1.2, -0.2, 0.2, 1.2]]),
    }
    lo, hi = 0.0, 1.0
    rows = {}
    for name, j in cases.items():
        gamma = torch.zeros((1, 2))
        total, parts = rncoa_loss(j, lo, hi, gamma)
        indep = independent_overlap_loss(j, lo, hi)
        # Residuals of the two displayed inequalities at gamma=0.
        residuals = [float((lo - j.max()).item()), float((j.min() - hi).item())]
        rows[name] = {"rncoa_loss": float(total.item()),
                      "independent_loss": float(indep.item()),
                      "collision": exact_body_collision(j[0], lo, hi),
                      "constraint_residuals": residuals,
                      "rncoa_v1": float(parts["v1"].item()),
                      "rncoa_v2": float(parts["v2"].item())}
    x = torch.tensor([[[-1.0, 0.0, 2.0, 3.0]]])
    errors = []
    for tau in [0.2, 0.1, 0.05, 0.01]:
        errors.append({"tau": tau,
                       "smax_error": abs(float(softmax_extreme(x, tau))-3.0),
                       "smin_error": abs(float(softmin_extreme(x, tau))-(-1.0))})
    return rows, errors


def sampled_confusion(n=10000):
    """Random intervals test whether zero stated loss identifies collision-free bodies."""
    rng = np.random.default_rng(SEED)
    # Four vertices are a translated rectangular body in one obstacle coordinate.
    centers = rng.uniform(-1.5, 2.5, n)
    halfwidth = 0.55
    j = torch.tensor(np.stack([centers-halfwidth, centers-halfwidth,
                               centers+halfwidth, centers+halfwidth], axis=1), dtype=torch.float32)
    with torch.no_grad():
        loss, _ = rncoa_loss(j, 0., 1., torch.zeros(n, 2), M=2., tau=.01,
                             lambda_gamma=0., lambda_c=10.)
    collision = ((j.min(1).values <= 1.) & (j.max(1).values >= 0.)).numpy()
    zero = (loss.numpy() < 1e-6)
    return {"n": n, "collision_rate": float(collision.mean()),
            "zero_loss_rate_on_collisions": float(zero[collision].mean()),
            "zero_loss_rate_on_safe": float(zero[~collision].mean()),
            "false_negative_count": int((zero & collision).sum()),
            "false_positive_count": int((~zero & ~collision).sum())}


def tiny_optimization(steps=180):
    results = {}
    for method in ["independent", "rncoa"]:
        torch.manual_seed(SEED)
        p = torch.nn.Parameter(torch.tensor([[0.45, 0.55, 0.65, 0.35]]))
        opt = torch.optim.Adam([p], lr=0.025)
        initial = p.detach().clone()
        for _ in range(steps):
            opt.zero_grad()
            if method == "independent":
                loss = independent_overlap_loss(p, 0., 1.).mean()
            else:
                loss, _ = rncoa_loss(p, 0., 1., torch.zeros((1,2)), M=2., tau=.05,
                                     lambda_gamma=0., lambda_c=10.)
                loss = loss.mean()
            loss.backward(); opt.step()
        final = p.detach()[0]
        eval_loss = (independent_overlap_loss(final[None],0.,1.) if method == "independent"
                     else rncoa_loss(final[None],0.,1.,torch.zeros(1,2),lambda_gamma=0.)[0])
        results[method] = {"initial_training_loss": float(loss.detach() * 0 +
                         (independent_overlap_loss(initial,0.,1.) if method=="independent" else
                          rncoa_loss(initial,0.,1.,torch.zeros(1,2),lambda_gamma=0.)[0]).item()),
            "final_eval_loss": float(eval_loss.item()),
            "final_vertices": [round(float(x), 5) for x in final],
            "exact_collision": exact_body_collision(final, 0., 1.)}
    return results


def main():
    analytic, smooth = analytic_check()
    out = {"seed": SEED, "analytic_cases": analytic,
           "smooth_extreme_errors": smooth,
           "sampled_confusion": sampled_confusion(),
           "optimization": tiny_optimization(),
           "interpretation": "With obstacle interval [0,1], the displayed inequalities encode the wrong disjunctive direction for collision avoidance: gamma=0 is feasible for bodies inside/straddling the obstacle and infeasible for bodies wholly outside."
           }
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()