import math
import pytest
import torch
from glint.box import (BoxHead, BoxLossWeights, box_answer, box_iou, box_loss, box_metrics, decode_box, decode_edges,
                       dfl_loss, generalized_iou, max_prob_confidence, order_box, two_hot)

K = 64


def logits_from_target(target, bins=K):
    return two_hot(target, bins).clamp_min(1e-12).log()


# ---------------------------------------------------------------- bins encode / decode
@pytest.mark.parametrize("bins", [2, 16, 64])
def test_two_hot_round_trip(bins):
    torch.manual_seed(0)
    t = torch.rand(100, 4)
    t[0] = torch.tensor([0.0, 1.0, 0.0, 1.0])                   # exact endpoints
    t[1] = torch.tensor([0.5, 1 / (bins - 1), 1 - 1e-7, min(1.0, 2 / (bins - 1))])
    th = two_hot(t, bins)
    assert th.shape == (100, 4, bins)
    assert torch.allclose(th.sum(-1), torch.ones(100, 4))
    assert ((th > 0).sum(-1) <= 2).all()                         # mass only on two adjacent bins
    assert torch.allclose(decode_edges(logits_from_target(t, bins)), t, atol=1e-5)


def test_two_hot_clamps_out_of_range():
    t = torch.tensor([[-0.2, 1.3, 0.25, 0.75]])
    assert torch.allclose(decode_edges(logits_from_target(t)), t.clamp(0, 1), atol=1e-5)


# ---------------------------------------------------------------- DFL
def test_dfl_matches_manual_and_is_minimal_at_target():
    torch.manual_seed(0)
    t = torch.tensor([[0.1, 0.4, 0.33, 0.9]])
    z = torch.randn(1, 4, K, requires_grad=True)
    loss = dfl_loss(z, t)
    manual = 0.0
    for e in range(4):
        x = float(t[0, e]) * (K - 1); lo = int(math.floor(x)); w = x - lo
        lp = z[0, e].log_softmax(-1)
        manual += -((1 - w) * lp[lo] + w * lp[lo + 1])
    assert torch.allclose(loss, manual / 4, atol=1e-5)
    loss.backward()
    assert torch.isfinite(z.grad).all() and z.grad.abs().sum() > 0
    # loss at the two-hot target logits == entropy of the two-hot distribution (its minimum)
    th = two_hot(t, K)
    ent = -(th * th.clamp_min(1e-12).log()).sum(-1).mean()
    assert torch.allclose(dfl_loss(logits_from_target(t), t), ent, atol=1e-4)
    assert dfl_loss(logits_from_target(t), t) < dfl_loss(torch.zeros(1, 4, K), t)


def test_dfl_exact_bin_target_is_one_hot_ce():
    t = torch.full((1, 4), 10 / (K - 1))
    z = torch.randn(1, 4, K)
    assert torch.allclose(dfl_loss(z, t), -z.log_softmax(-1)[..., 10].mean(), atol=1e-5)


# ---------------------------------------------------------------- IoU / GIoU
def test_iou_giou_known_values():
    a = torch.tensor([0.0, 0.5, 0.0, 0.5])                       # top, bottom, left, right
    assert torch.allclose(box_iou(a, a), torch.tensor(1.0))
    assert torch.allclose(generalized_iou(a, a), torch.tensor(1.0))
    b = torch.tensor([0.25, 0.75, 0.25, 0.75])                   # overlap 0.25x0.25
    assert torch.allclose(box_iou(a, b), torch.tensor(0.0625 / (0.25 + 0.25 - 0.0625)), atol=1e-6)
    c = torch.tensor([0.5, 1.0, 0.5, 1.0])                       # touching corner, hull = 1
    assert torch.allclose(box_iou(a, c), torch.tensor(0.0))
    assert torch.allclose(generalized_iou(a, c), torch.tensor(-0.5), atol=1e-6)
    d = torch.tensor([0.0, 0.1, 0.9, 1.0])
    assert generalized_iou(a, d) < 0 and generalized_iou(a, d) > -1


def test_giou_gradient_for_disjoint_boxes():
    p = torch.tensor([0.0, 0.1, 0.0, 0.1], requires_grad=True)
    g = torch.tensor([0.8, 0.9, 0.8, 0.9])
    (1 - generalized_iou(p, g)).backward()
    assert p.grad.abs().sum() > 0                                # IoU alone would have zero gradient here


def test_iou_broadcast():
    a = torch.rand(5, 4).sort(-1).values[:, [0, 1, 2, 3]]
    a = order_box(a)
    assert box_iou(a, a).shape == (5,)
    assert torch.allclose(box_iou(a, a), torch.ones(5), atol=1e-5)


# ---------------------------------------------------------------- ordering
def test_order_box_swaps_inverted_edges():
    e = torch.tensor([[0.8, 0.2, 0.9, 0.1], [0.1, 0.2, 0.3, 0.4]])
    assert torch.equal(order_box(e), torch.tensor([[0.2, 0.8, 0.1, 0.9], [0.1, 0.2, 0.3, 0.4]]))
    z = logits_from_target(torch.tensor([[0.7, 0.3, 0.6, 0.2]]))
    t, b, l, r = decode_box(z)[0].tolist()
    assert t <= b and l <= r


def test_box_answer_jev_shape_and_strict_order():
    a = box_answer(torch.tensor([0.7, 0.3, 0.5, 0.5]), 1.3)
    assert a["type"] == "box" and set(a["box"]) == {"top", "bottom", "left", "right"}
    assert a["box"]["top"] < a["box"]["bottom"] and a["box"]["left"] < a["box"]["right"]
    assert a["box"]["top"] == 0.3 and a["box"]["bottom"] == 0.7
    assert a["confidence"] == 1.0
    z = box_answer([1.0, 1.0, 0.0, 0.0], 0.5)["box"]              # degenerate at the image border
    assert 0 <= z["top"] < z["bottom"] <= 1 and 0 <= z["left"] < z["right"] <= 1


# ---------------------------------------------------------------- loss / metrics
def test_box_loss_parts_and_backward():
    torch.manual_seed(0)
    z = torch.randn(3, 4, K, requires_grad=True); c = torch.zeros(3, requires_grad=True)
    tgt = torch.tensor([[0.1, 0.2, 0.3, 0.6], [0.5, 0.9, 0.1, 0.2], [0.0, 1.0, 0.0, 1.0]])
    loss, parts = box_loss(z, c, tgt, BoxLossWeights())
    assert set(parts) == {"dfl", "l1", "giou", "conf", "iou"}
    loss.backward()
    assert torch.isfinite(z.grad).all() and z.grad.abs().sum() > 0 and c.grad.abs().sum() > 0
    w0, _ = box_loss(z, c, tgt, BoxLossWeights(dfl=0, l1=0, giou=0, conf=0))
    assert float(w0.detach()) == 0.0
    perfect, pp = box_loss(logits_from_target(tgt), torch.full((3,), 10.0), tgt)
    assert pp["l1"] < 1e-4 and pp["giou"] < 1e-3 and pp["iou"] > 0.999


def test_max_prob_confidence_range():
    z = torch.randn(2, 4, K)
    c = max_prob_confidence(z)
    assert c.shape == (2,) and ((c > 0) & (c <= 1)).all()
    assert torch.allclose(max_prob_confidence(logits_from_target(torch.rand(2, 4))), torch.ones(2), atol=1e-5)


def test_box_metrics():
    m = box_metrics([0.9, 0.6, 0.3, 0.8], [0.9, 0.7, 0.2, 0.8], [0.5] * 4)
    assert m["n_box_questions"] == 4
    assert m["mean_iou"] == pytest.approx(0.65)
    assert m["iou_ge_0.5"] == 0.75 and m["iou_ge_0.75"] == 0.5
    assert 0 <= m["ece_15bin_conf_vs_iou50"] <= 1 and "ece_15bin_maxprob_vs_iou50" in m
    assert box_metrics([], []) == {"n_box_questions": 0}


# ---------------------------------------------------------------- head forward / backward
class _Rec:
    def __init__(self, spans, coords=None):
        from glint.schema import EncodedBoxQuestion
        self.box_questions = tuple(EncodedBoxQuestion(f"q{i}", s) for i, s in enumerate(spans))
        self.image_coords = coords


def test_box_head_forward_backward_random_hidden():
    torch.manual_seed(0)
    H, L = 48, 40
    head = BoxHead(hidden_size=H, width=32, layers=2, heads=4, feedforward=64, bins=K)
    hidden = torch.randn(2, L, H, requires_grad=True)
    mask = torch.ones(2, L, dtype=torch.long); mask[1, 30:] = 0
    coords = (tuple(range(2, 14)), tuple(((r + .5) / 3, (c + .5) / 4) for r in range(3) for c in range(4)))
    recs = [_Rec([(20, 25), (26, 30)], coords), _Rec([(5, 9)])]
    out = head(hidden, mask, recs)
    assert out[0].logits.shape == (2, 4, K) and out[0].conf_logit.shape == (2,)
    assert out[1].logits.shape == (1, 4, K) and out[0].question_ids == ("q0", "q1")
    tgt = torch.tensor([[0.1, 0.3, 0.2, 0.8], [0.6, 0.7, 0.1, 0.4]])
    loss, _ = box_loss(out[0].logits, out[0].conf_logit, tgt)
    loss2, _ = box_loss(out[1].logits, out[1].conf_logit, tgt[:1])
    (loss + loss2).backward()
    assert hidden.grad is not None and torch.isfinite(hidden.grad).all()
    assert hidden.grad[1, 30:].abs().sum() == 0                  # padding is ignored
    for name, p in head.named_parameters():
        assert p.grad is not None, name
    assert head(hidden, mask, [_Rec([])])[0] is None


def test_box_head_questions_do_not_interact():
    torch.manual_seed(0)
    head = BoxHead(hidden_size=16, width=16, layers=1, heads=2, feedforward=32, bins=8).eval()
    hidden = torch.randn(1, 20, 16); mask = torch.ones(1, 20, dtype=torch.long)
    both = head(hidden, mask, [_Rec([(2, 5), (8, 12)])])[0]
    one = head(hidden, mask, [_Rec([(2, 5)])])[0]
    assert torch.allclose(both.logits[0], one.logits[0], atol=1e-5)


def test_box_head_can_overfit_and_bf16():
    torch.manual_seed(0)
    head = BoxHead(hidden_size=16, width=32, layers=1, heads=4, feedforward=64, bins=K)
    hidden = torch.randn(1, 12, 16); mask = torch.ones(1, 12, dtype=torch.long)
    rec = [_Rec([(3, 6)])]; tgt = torch.tensor([[0.42, 0.47, 0.18, 0.83]])
    opt = torch.optim.AdamW(head.parameters(), lr=3e-3)
    for _ in range(150):
        o = head(hidden, mask, rec)[0]; loss, parts = box_loss(o.logits, o.conf_logit, tgt)
        opt.zero_grad(); loss.backward(); opt.step()
    assert parts["iou"] > 0.9
    with torch.no_grad():
        assert float(torch.sigmoid(head(hidden, mask, rec)[0].conf_logit)) > 0.7   # confidence tracks IoU
    hb = head.to(torch.bfloat16)
    ob = hb(hidden.bfloat16(), mask, rec)[0]
    assert ob.logits.dtype == torch.bfloat16 and box_iou(decode_box(ob.logits)[0], tgt[0]) > 0.8


def test_box_head_edges_do_not_collapse():
    # regression: with a tiny edge embedding all 4 edges decoded to the same value (the L1-optimal mean of the targets)
    torch.manual_seed(1)
    head = BoxHead(hidden_size=16, width=32, layers=1, heads=4, feedforward=64, bins=K)
    hidden = torch.randn(4, 12, 16); mask = torch.ones(4, 12, dtype=torch.long)
    tgt = torch.tensor([[0.48, 0.54, 0.50, 0.86], [0.53, 0.59, 0.08, 0.92], [0.20, 0.30, 0.10, 0.40], [0.70, 0.95, 0.60, 0.90]])
    opt = torch.optim.AdamW(head.parameters(), lr=1e-3)
    for step in range(400):
        i = step % 4
        o = head(hidden[i:i + 1], mask[i:i + 1], [_Rec([(3, 6)])])[0]
        loss, _ = box_loss(o.logits, o.conf_logit, tgt[i:i + 1])
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        out = torch.stack([decode_box(head(hidden[i:i + 1], mask[i:i + 1], [_Rec([(3, 6)])])[0].logits)[0] for i in range(4)])
    assert box_iou(out, tgt).mean() > 0.5