"""Jev-1.14-style `box` questions: bounding box of a UI element as fractions of the image.

Answer format (same as Jev 1.14):
    {"type": "box", "box": {"top": .., "bottom": .., "left": .., "right": ..}, "confidence": 0..1}
top/bottom are fractions of the image height, left/right fractions of the width, origin top-left,
left < right and top < bottom. There is no "not found" answer: presence is a separate noul question
(convention: `<box qid>_visible`).

Edge layout of every tensor in this module: (..., 4) = (top, bottom, left, right).

Head: for each edge, logits over K bins (distribution focal loss style, GFL/YOLOv8).
Bin i has centre i / (K - 1), so 0 and 1 are exactly representable. Decoded edge = expectation
over bin centres; ordering is enforced at decode with min/max per axis.
Confidence: a learned sigmoid head regressing the IoU with gold (BCE against the detached IoU).
"""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Any
import torch
import torch.nn.functional as F

BOX_EDGES = ("top", "bottom", "left", "right")
DEFAULT_BINS = 64
FOURIER_BANDS = 8


def is_box_question(question: dict[str, Any]) -> bool:
    return str(question.get("type")) == "box"


# ---------------------------------------------------------------------------- bins

def bin_centers(bins: int, device=None, dtype=torch.float32) -> torch.Tensor:
    return torch.linspace(0.0, 1.0, bins, device=device, dtype=dtype)


def two_hot(target: torch.Tensor, bins: int) -> torch.Tensor:
    """Soft label on the two bins adjacent to each target value in [0, 1]. (...,) -> (..., bins)."""
    t = target.float().clamp(0.0, 1.0) * (bins - 1)
    lo = t.floor().long().clamp(max=bins - 2)
    w_hi = t - lo.float()
    out = torch.zeros(*target.shape, bins, device=target.device)
    out.scatter_(-1, lo.unsqueeze(-1), (1.0 - w_hi).unsqueeze(-1))
    out.scatter_add_(-1, (lo + 1).unsqueeze(-1), w_hi.unsqueeze(-1))
    return out


def decode_edges(logits: torch.Tensor) -> torch.Tensor:
    """Expectation over bin centres. (..., 4, K) logits -> (..., 4) edges (unordered)."""
    p = logits.float().softmax(-1)
    return (p * bin_centers(logits.shape[-1], logits.device)).sum(-1)


def order_box(edges: torch.Tensor) -> torch.Tensor:
    """Enforce top <= bottom and left <= right by swapping (min/max per axis; differentiable)."""
    t, b, l, r = edges.unbind(-1)
    return torch.stack([torch.minimum(t, b), torch.maximum(t, b), torch.minimum(l, r), torch.maximum(l, r)], -1)


def decode_box(logits: torch.Tensor) -> torch.Tensor:
    return order_box(decode_edges(logits))


def max_prob_confidence(logits: torch.Tensor) -> torch.Tensor:
    """Fallback confidence: mean over the 4 edges of the sum of the two most likely adjacent bins."""
    p = logits.float().softmax(-1)
    pair = p[..., :-1] + p[..., 1:]
    return pair.max(-1).values.mean(-1)


# ---------------------------------------------------------------------------- geometry

def box_iou(a: torch.Tensor, b: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
    """IoU of (top, bottom, left, right) boxes, broadcasting over leading dims."""
    ih = (torch.minimum(a[..., 1], b[..., 1]) - torch.maximum(a[..., 0], b[..., 0])).clamp(min=0)
    iw = (torch.minimum(a[..., 3], b[..., 3]) - torch.maximum(a[..., 2], b[..., 2])).clamp(min=0)
    inter = ih * iw
    area_a = (a[..., 1] - a[..., 0]).clamp(min=0) * (a[..., 3] - a[..., 2]).clamp(min=0)
    area_b = (b[..., 1] - b[..., 0]).clamp(min=0) * (b[..., 3] - b[..., 2]).clamp(min=0)
    return inter / (area_a + area_b - inter + eps)


def generalized_iou(a: torch.Tensor, b: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
    """GIoU in [-1, 1] of (top, bottom, left, right) boxes (both must be ordered)."""
    ih = (torch.minimum(a[..., 1], b[..., 1]) - torch.maximum(a[..., 0], b[..., 0])).clamp(min=0)
    iw = (torch.minimum(a[..., 3], b[..., 3]) - torch.maximum(a[..., 2], b[..., 2])).clamp(min=0)
    inter = ih * iw
    area_a = (a[..., 1] - a[..., 0]) * (a[..., 3] - a[..., 2])
    area_b = (b[..., 1] - b[..., 0]) * (b[..., 3] - b[..., 2])
    union = area_a + area_b - inter
    hull = ((torch.maximum(a[..., 1], b[..., 1]) - torch.minimum(a[..., 0], b[..., 0]))
            * (torch.maximum(a[..., 3], b[..., 3]) - torch.minimum(a[..., 2], b[..., 2])))
    return inter / (union + eps) - (hull - union) / (hull + eps)


# ---------------------------------------------------------------------------- losses

def dfl_loss(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Distribution focal loss: soft-label CE on the two bins adjacent to each edge, mean over edges.
    logits (..., 4, K), target (..., 4) in [0, 1]."""
    logp = logits.float().log_softmax(-1)
    return -(two_hot(target, logits.shape[-1]) * logp).sum(-1).mean()


@dataclass
class BoxLossWeights:
    dfl: float = 1.0
    l1: float = 5.0
    giou: float = 2.0
    conf: float = 1.0


def box_loss(logits: torch.Tensor, conf_logit: torch.Tensor, target: torch.Tensor,
             weights: BoxLossWeights | None = None) -> tuple[torch.Tensor, dict[str, float]]:
    """Loss for one or more box questions. logits (N, 4, K), conf_logit (N,), target (N, 4) ordered gold.
    Returns mean over N and the (float) parts."""
    w = weights or BoxLossWeights()
    target = target.float()
    pred = decode_box(logits)
    dfl = dfl_loss(logits, target)
    l1 = (pred - target).abs().mean()
    giou = (1.0 - generalized_iou(pred, target)).mean()
    iou = box_iou(pred, target).detach()
    conf = F.binary_cross_entropy_with_logits(conf_logit.float(), iou)
    loss = w.dfl * dfl + w.l1 * l1 + w.giou * giou + w.conf * conf
    return loss, {"dfl": float(dfl.detach()), "l1": float(l1.detach()), "giou": float(giou.detach()),
                  "conf": float(conf.detach()), "iou": float(iou.mean())}


# ---------------------------------------------------------------------------- answers / metrics

def box_answer(box: torch.Tensor | list[float], confidence: float, min_size: float = 1e-3, digits: int = 4) -> dict[str, Any]:
    """Jev-1.14 box answer. Guarantees 0 <= top < bottom <= 1 and 0 <= left < right <= 1."""
    t, b, l, r = [float(v) for v in (box.tolist() if torch.is_tensor(box) else box)]
    t, b = sorted((min(max(t, 0.0), 1.0), min(max(b, 0.0), 1.0)))
    l, r = sorted((min(max(l, 0.0), 1.0), min(max(r, 0.0), 1.0)))

    def widen(lo, hi):
        if hi - lo >= min_size:
            return lo, hi
        c = min(max((lo + hi) / 2, min_size / 2), 1 - min_size / 2)
        return c - min_size / 2, c + min_size / 2
    t, b = widen(t, b); l, r = widen(l, r)
    rd = lambda v: round(v, digits)
    return {"type": "box", "box": {"top": rd(t), "bottom": rd(b), "left": rd(l), "right": rd(r)},
            "confidence": rd(min(max(float(confidence), 0.0), 1.0))}


def ece(confs, hits, bins=15) -> float:
    tot, n = 0.0, len(confs)
    for k in range(bins):
        lo, hi = k / bins, (k + 1) / bins
        idx = [i for i, c in enumerate(confs) if lo < c <= hi or (k == 0 and c == 0)]
        if idx:
            tot += len(idx) / n * abs(sum(hits[i] for i in idx) / len(idx) - sum(confs[i] for i in idx) / len(idx))
    return tot


def box_metrics(ious: list[float], confs: list[float], maxprob_confs: list[float] | None = None, bins: int = 15) -> dict[str, Any]:
    n = len(ious)
    if not n:
        return {"n_box_questions": 0}
    hit50 = [int(i >= 0.5) for i in ious]
    out = {"n_box_questions": n, "mean_iou": sum(ious) / n, "iou_ge_0.5": sum(hit50) / n,
           "iou_ge_0.75": sum(int(i >= 0.75) for i in ious) / n, "mean_confidence": sum(confs) / n,
           f"ece_{bins}bin_conf_vs_iou50": ece(confs, hit50, bins)}
    if maxprob_confs:
        out[f"ece_{bins}bin_maxprob_vs_iou50"] = ece(maxprob_confs, hit50, bins)
    return out


# ---------------------------------------------------------------------------- head

def fourier_features(yx: torch.Tensor, bands: int = FOURIER_BANDS) -> torch.Tensor:
    """(N, 2) coords in [0, 1] -> (N, 4 * bands) sin/cos features."""
    freqs = (2.0 ** torch.arange(bands, device=yx.device, dtype=torch.float32)) * math.pi
    ang = yx.float().unsqueeze(-1) * freqs                                 # (N, 2, bands)
    return torch.cat([ang.sin(), ang.cos()], -1).flatten(1)


@dataclass
class BoxPrediction:
    logits: torch.Tensor        # (Q, 4, K) per box question of the record
    conf_logit: torch.Tensor    # (Q,)
    question_ids: tuple[str, ...]


class BoxHead(torch.nn.Module):
    """Box head on the backbone's last hidden states.

    Query per box question = projection of the same pooled question representation the joint schema
    head uses (mean of the normalised hidden states over the instruction span) + the global (last-token)
    vector, layer-normed and expanded to 4 edge queries (top/bottom/left/right) with a learned edge
    embedding of the same (unit) scale, so the edges are distinguishable from the first step.
    The edge queries run through pre-norm transformer decoder layers (self-attention between the 4 edges
    + cross-attention to the projected sequence). Image tokens of the memory get an explicit 2D Fourier
    position encoding (fractional row/col of the merged patch) so the head can read *where* evidence is.
    Each edge query -> K bin logits through its own (per-edge) scorer; the 4 edge queries + per-edge (max prob, entropy) -> IoU logit.
    """

    def __init__(self, hidden_size: int, width: int = 256, layers: int = 2, heads: int = 4,
                 feedforward: int = 1024, bins: int = DEFAULT_BINS, fourier_bands: int = FOURIER_BANDS,
                 dropout: float = 0.0) -> None:
        super().__init__()
        self.bins, self.fourier_bands = bins, fourier_bands
        self.hidden_norm = torch.nn.LayerNorm(hidden_size)
        self.memory_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.position_projection = torch.nn.Linear(4 * fourier_bands, width)
        self.question_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.global_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.query_norm = torch.nn.LayerNorm(width)
        self.edge_embedding = torch.nn.Parameter(torch.randn(4, width))
        self.layers = torch.nn.ModuleList([
            torch.nn.TransformerDecoderLayer(d_model=width, nhead=heads, dim_feedforward=feedforward, dropout=dropout,
                                             activation="gelu", batch_first=True, norm_first=True)
            for _ in range(layers)])
        self.output_norm = torch.nn.LayerNorm(width)
        bound = 1.0 / math.sqrt(width)                                          # nn.Linear default init, one scorer per edge
        self.bin_weight = torch.nn.Parameter(torch.empty(4, width, bins).uniform_(-bound, bound))
        self.bin_bias = torch.nn.Parameter(torch.empty(4, bins).uniform_(-bound, bound))
        self.confidence_scorer = torch.nn.Sequential(
            torch.nn.Linear(4 * width + 8, width), torch.nn.GELU(), torch.nn.Dropout(dropout), torch.nn.Linear(width, 1))

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, records: list[Any]) -> list[BoxPrediction | None]:
        results: list[BoxPrediction | None] = []
        normalized_hidden = self.hidden_norm(hidden_states)
        for batch_index, record in enumerate(records):
            box_questions = getattr(record, "box_questions", ()) or ()
            if not box_questions:
                results.append(None)
                continue
            sequence_length = int(attention_mask[batch_index].sum().item())
            sequence_hidden = normalized_hidden[batch_index, :sequence_length]
            memory = self.memory_projection(sequence_hidden)
            coords = getattr(record, "image_coords", None)
            if coords is not None and len(coords[0]):
                index = torch.tensor(coords[0], dtype=torch.long, device=memory.device)
                yx = torch.tensor(coords[1], dtype=torch.float32, device=memory.device)
                position = self.position_projection(fourier_features(yx, self.fourier_bands).to(memory.dtype))
                memory = memory.index_add(0, index, position)
            question_vectors = torch.stack([sequence_hidden[s:e].mean(0) for s, e in (q.question_span for q in box_questions)])
            base = self.query_norm(self.question_projection(question_vectors) + self.global_projection(sequence_hidden[-1]).unsqueeze(0))
            queries = (base.unsqueeze(1) + self.edge_embedding.to(base.dtype)).reshape(1, -1, base.shape[-1])
            # edge queries of different questions must not attend to each other
            q_index = torch.arange(len(box_questions), device=memory.device).repeat_interleave(4)
            self_mask = q_index.unsqueeze(0) != q_index.unsqueeze(1)
            for layer in self.layers:
                queries = layer(queries, memory.unsqueeze(0), tgt_mask=self_mask)
            queries = self.output_norm(queries[0]).view(len(box_questions), 4, -1)
            logits = torch.einsum("qew,ewk->qek", queries, self.bin_weight.to(queries.dtype)) + self.bin_bias.to(queries.dtype)  # (Q, 4, K)
            p = logits.detach().float().softmax(-1)
            stats = torch.cat([p.max(-1).values, -(p * p.clamp_min(1e-9).log()).sum(-1) / math.log(self.bins)], -1)
            conf_logit = self.confidence_scorer(torch.cat([queries.flatten(1), stats.to(queries.dtype)], -1)).squeeze(-1)
            results.append(BoxPrediction(logits, conf_logit, tuple(q.question_id for q in box_questions)))
        return results