"""Grounded extraction on multimodal contexts: an extract field whose value is in a media
item says where, as text extraction says where with ``meta.span``. No schema annotation:
a field's mode follows from its type (SPEC.md §3); whether an extract field points into
text or into media follows from where its pointer lands in the context.

The span head (:class:`deg.model.SpanHead`) already scores every context token, the media
**soft tokens** included:

* image: the ~260 pooled soft tokens are a ``gw x gh`` grid in raster order (Gemma 4 pools
  3x3 patches of 16 px of the aspect-preserving resize, :func:`image_grid`). ``start`` is
  the box's **top-left** cell and ``end`` its **bottom-right** cell, so ``start <= end``
  in raster order (the span head's own order) and two pointers give a rectangle;
  :func:`cells_to_box` maps them back to pixels of the source image.
* audio: one soft token per :data:`AUDIO_MS` ms; ``start``/``end`` are the first and last
  token of the window, :func:`tokens_to_time` gives seconds.

:func:`landing` decides text vs media vs null from the start pointer's probability mass.
A media landing has no text to copy, so the value is **read from the region**: the
region (padded, :func:`crop_box`) is cut from the source and encoded again by the frozen
towers as the only media item of a one-field x-generate fill (:func:`read_schema`), and
the text is parsed by the field's mode like a copied span. At the image budget of a whole
receipt one soft token covers ~48x48 resized px (a printed word is often smaller); a crop
at the same budget is read at several times the resolution, which a region-masked
attention bias cannot add, and it needs no change to the fork kernels or decode graphs.
Crop-read records come from gold boxes (:mod:`deg.datagen.mm_ground`).
"""

from __future__ import annotations

import hashlib
import math
from dataclasses import dataclass
from typing import Any

import torch
from torch import Tensor

PATCH = 16
POOL = 3
SOFT_BUDGETS = (70, 140, 280, 560, 1120)  # Gemma 4 image processor's supported budgets
AUDIO_MS = 40  # processor_config audio_ms_per_token
CROP_PAD = 0.15  # crop margin: share of the box side, each side
CROP_MIN_PX = 48  # crops are at least this wide/high in source pixels
AUDIO_PAD_S = 0.5


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


def resized_size(width: int, height: int, max_soft_tokens: int, patch: int = PATCH,
                 pool: int = POOL) -> tuple[int, int]:
    """(target width, target height) of Gemma 4's aspect-preserving resize (HF
    ``get_aspect_ratio_preserving_size``): the largest multiple of ``pool * patch`` per side
    with at most ``max_soft_tokens * pool**2`` patches."""
    max_patches = max_soft_tokens * pool * pool
    factor = math.sqrt(max_patches * patch * patch / (height * width))
    side = pool * patch
    th = int(math.floor(factor * height / side)) * side
    tw = int(math.floor(factor * width / side)) * side
    if th == 0 and tw == 0:
        raise ValueError(f"image {width}x{height} resizes to 0x0")
    longest = (max_patches // (pool * pool)) * side
    if th == 0:
        th, tw = side, min(int(math.floor(width / height)) * side, longest)
    elif tw == 0:
        tw, th = side, min(int(math.floor(height / width)) * side, longest)
    return tw, th


@dataclass(frozen=True)
class Grid:
    """Pooled soft-token grid of one image: ``gw x gh`` cells in raster order, each
    ``cell_w x cell_h`` source pixels."""

    gw: int
    gh: int
    width: int  # source image
    height: int

    @property
    def n(self) -> int:
        return self.gw * self.gh

    @property
    def cell_w(self) -> float:
        return self.width / self.gw

    @property
    def cell_h(self) -> float:
        return self.height / self.gh


def image_grid(width: int, height: int, max_soft_tokens: int) -> Grid:
    tw, th = resized_size(width, height, max_soft_tokens)
    return Grid(tw // (PATCH * POOL), th // (PATCH * POOL), int(width), int(height))


def grid_for(width: int, height: int, n: int, prefer: int | None = None) -> Grid:
    """The grid of an image whose embedding has ``n`` soft tokens (the budget that produced
    it: ``prefer`` first, then every supported one)."""
    for budget in ((prefer,) if prefer else ()) + SOFT_BUDGETS:
        g = image_grid(width, height, budget)
        if g.n == n:
            return g
    raise ValueError(f"no soft-token budget gives {n} tokens for a {width}x{height} image")


def cells_to_box(i: int, j: int, g: Grid) -> list[float]:
    """Top-left cell ``i`` and bottom-right cell ``j`` (raster indices) -> pixel box
    ``[x0, y0, x1, y1]`` in the source image."""
    (y0, x0), (y1, x1) = divmod(i, g.gw), divmod(j, g.gw)
    x0, x1 = min(x0, x1), max(x0, x1)
    y0, y1 = min(y0, y1), max(y0, y1)
    return [round(x0 * g.cell_w, 1), round(y0 * g.cell_h, 1),
            round(min(g.width, (x1 + 1) * g.cell_w), 1),
            round(min(g.height, (y1 + 1) * g.cell_h), 1)]


def box_to_cells(box: Any, g: Grid) -> tuple[int, int]:
    """Pixel box -> (top-left cell, bottom-right cell): the cells its corners fall in."""
    x0, y0, x1, y1 = (float(v) for v in box)
    eps = 1e-6

    def col(x: float) -> int:
        return min(g.gw - 1, max(0, int(x / g.cell_w)))

    def row(y: float) -> int:
        return min(g.gh - 1, max(0, int(y / g.cell_h)))

    c0, r0 = col(x0), row(y0)
    c1, r1 = col(max(x0, x1 - eps)), row(max(y0, y1 - eps))
    return r0 * g.gw + c0, r1 * g.gw + c1


def tokens_to_time(i: int, j: int, seconds: float | None = None, ms: int = AUDIO_MS
                   ) -> list[float]:
    s, e = i * ms / 1000.0, (j + 1) * ms / 1000.0
    if seconds is not None:
        s, e = min(s, seconds), min(e, seconds)
    return [round(s, 3), round(e, 3)]


def time_to_tokens(window: Any, n: int, ms: int = AUDIO_MS) -> tuple[int, int]:
    s, e = float(window[0]), float(window[1])
    i = min(n - 1, max(0, int(s * 1000 // ms)))
    j = min(n - 1, max(i, int(math.ceil(e * 1000 / ms)) - 1))
    return i, j


# --------------------------------------------------------------------------- media in context


@dataclass(frozen=True)
class MediaSpan:
    """One media item of an encoded context: its soft tokens are context tokens
    ``[lo, lo + n)``; ``grid`` (image) or ``seconds`` (audio) map pointers to the source."""

    id: str
    kind: str
    lo: int
    n: int
    grid: Grid | None = None
    seconds: float | None = None


def media_spans(ectx: Any, parts: list[Any], max_soft_tokens: int | None = None
                ) -> list[MediaSpan]:
    """The media of ``ectx`` (rendered from ``parts``) with their geometry. An image part
    without a ``size`` cannot be mapped to pixels and is left out."""
    from deg.model import Media

    media = [p for p in parts if isinstance(p, Media)]
    out = []
    for k, ((start, emb), part) in enumerate(zip(ectx.media, media)):
        n, lo = int(emb.shape[0]), start - ectx.ctx_start
        mid = part.id or f"m{k}"
        if part.kind == "image":
            if part.size is None:
                continue
            out.append(MediaSpan(mid, "image", lo, n, grid_for(*part.size, n, max_soft_tokens)))
        else:
            out.append(MediaSpan(mid, "audio", lo, n, seconds=part.seconds))
    return out


def located(span: MediaSpan, i: int, j: int) -> dict[str, Any]:
    """Pointer cells (item-relative) -> the response's ground entry."""
    out: dict[str, Any] = {"media": span.id, "kind": span.kind, "tokens": [i, j]}
    if span.kind == "image":
        assert span.grid is not None
        out["box"] = cells_to_box(i, j, span.grid)
    else:
        out["time"] = tokens_to_time(i, j, span.seconds)
    return out


def ground_target(span: MediaSpan, gold: dict[str, Any]) -> tuple[int, int] | None:
    """A gold ``{box}`` / ``{time}`` -> context-token (start, end) of ``span``."""
    if span.kind == "image" and gold.get("box") is not None and span.grid is not None:
        i, j = box_to_cells(gold["box"], span.grid)
    elif span.kind == "audio" and gold.get("time") is not None:
        i, j = time_to_tokens(gold["time"], span.n)
    else:
        return None
    return span.lo + i, span.lo + j


def soft_mask(spans: list[MediaSpan], T: int, device: Any = None) -> Tensor:
    """[T + 1] bool: the soft tokens of ``spans`` and the null slot."""
    m = torch.zeros(T + 1, dtype=torch.bool, device=device)
    for s in spans:
        m[s.lo:s.lo + s.n] = True
    m[T] = True
    return m


def landing(start: Tensor, spans: list[MediaSpan], temperature: Any = 1.0) -> str:
    """Where an extract pointer lands: ``"media"`` when the start pointer's probability
    mass on the media soft tokens beats both the text tokens' and null's, else
    ``"text"`` (the text path, null included, decides as without media)."""
    T = start.shape[0] - 1
    ls = torch.log_softmax(start.float() / temperature, -1)
    soft = soft_mask(spans, T, start.device)
    soft[T] = False
    text = ~soft
    text[T] = False
    m = torch.logsumexp(ls[soft], 0)
    t = torch.logsumexp(ls[text], 0) if bool(text.any()) else torch.tensor(float("-inf"))
    return "media" if float(m) > max(float(t), float(ls[T])) else "text"


def text_only(logits: Tensor, spans: list[MediaSpan]) -> Tensor:
    """Span logits with the media soft tokens masked (the text path of a media context)."""
    soft = soft_mask(spans, logits.shape[0] - 1, logits.device)
    soft[-1] = False
    return logits.masked_fill(soft, float("-inf"))


def decode_ground(start: Tensor, end: Tensor, spans: list[MediaSpan], temperature: Any = 1.0
                  ) -> tuple[dict[str, Any] | None, float]:
    """The best (top-left, bottom-right) pair inside one media item, or None when null
    wins; returns (entry, its probability). Logits are renormalized over the soft tokens
    of ``spans`` plus null. Image pairs need ``col(end) >= col(start)`` and ``row(end) >=
    row(start)``; audio pairs ``end >= start``."""
    T = start.shape[0] - 1
    keep = soft_mask(spans, T, start.device)
    ls = torch.log_softmax((start.float() / temperature).masked_fill(~keep, float("-inf")), -1)
    le = torch.log_softmax((end.float() / temperature).masked_fill(~keep, float("-inf")), -1)
    best, best_score = None, float(ls[T] + le[T])
    for s in spans:
        a, b = ls[s.lo:s.lo + s.n], le[s.lo:s.lo + s.n]
        score = a[:, None] + b[None, :]
        idx = torch.arange(s.n, device=score.device)
        if s.kind == "image":
            assert s.grid is not None
            col, row = idx % s.grid.gw, idx // s.grid.gw
            ok = (col[None, :] >= col[:, None]) & (row[None, :] >= row[:, None])
        else:
            ok = idx[None, :] >= idx[:, None]
        score = score.masked_fill(~ok, float("-inf"))
        v, f = score.flatten().max(0)
        if float(v) > best_score:
            best_score = float(v)
            best = located(s, int(f) // s.n, int(f) % s.n)
    return best, min(1.0, math.exp(best_score))


# --------------------------------------------------------------------------- crop re-read


def crop_box(box: Any, width: int, height: int, pad: float = CROP_PAD,
             min_px: int = CROP_MIN_PX) -> tuple[int, int, int, int]:
    """Integer crop ``(x0, y0, x1, y1)`` around ``box``: ``pad`` of each side added, at
    least ``min_px`` per side, clipped to the image."""
    x0, y0, x1, y1 = (float(v) for v in box)
    w, h = max(x1 - x0, 1.0), max(y1 - y0, 1.0)
    px, py = max(pad * w, (min_px - w) / 2), max(pad * h, (min_px - h) / 2)
    return (max(0, int(math.floor(x0 - px))), max(0, int(math.floor(y0 - py))),
            min(int(width), int(math.ceil(x1 + px))), min(int(height), int(math.ceil(y1 + py))))


def crop_window(window: Any, seconds: float | None, pad: float = AUDIO_PAD_S
                ) -> tuple[float, float]:
    s, e = float(window[0]) - pad, float(window[1]) + pad
    return max(0.0, s), (e if seconds is None else min(seconds, e))


def crop_sha(sha: str, region: Any) -> str:
    """Content key of a crop: its source's sha256 and the integer/ms region (train crops
    from gold boxes and inference crops from predicted boxes share the media cache)."""
    return hashlib.sha256(f"{sha}:{list(region)}".encode()).hexdigest()


def crop_item(item: dict[str, Any], entry: dict[str, Any]) -> dict[str, Any]:
    """An in-memory media item (``data``: PIL image / float32 waveform) of ``entry``'s
    region of ``item`` (a record media item with ``path``)."""
    from deg.media import SAMPLE_RATE, item_sha, load_audio, load_image

    sha = item_sha(item)
    if item["kind"] == "image":
        im = load_image(item["path"])
        region = crop_box(entry["box"], im.width, im.height)
        return {"id": "crop", "kind": "image", "data": im.crop(region),
                "sha256": crop_sha(sha, region), "width": region[2] - region[0],
                "height": region[3] - region[1]}
    x = load_audio(item["path"])
    s, e = crop_window(entry["time"], len(x) / SAMPLE_RATE)
    region = (int(round(s * 1000)), int(round(e * 1000)))
    clip = x[int(s * SAMPLE_RATE):int(e * SAMPLE_RATE)]
    return {"id": "crop", "kind": "audio", "data": clip, "sha256": crop_sha(sha, region),
            "duration_s": len(clip) / SAMPLE_RATE}


READ_MAX_LENGTH = 200


def read_schema(fld: Any) -> tuple[dict[str, Any], str]:
    """The one-field schema a region is read with (inference and crop-read training
    records alike): the extract leaf's key as an x-generate string with the leaf's own
    title/description; the read text is then parsed by the leaf's mode."""
    key = next((p for p in reversed(fld.path) if isinstance(p, str)), "value")
    leaf: dict[str, Any] = {"type": "string", "x-generate": True,
                            "maxLength": int(fld.constraints.get("maxLength", READ_MAX_LENGTH))}
    for k in ("title", "description"):
        if getattr(fld, k):
            leaf[k] = getattr(fld, k)
    return {"type": "object", "properties": {key: leaf}, "required": [key],
            "additionalProperties": False}, key


# --------------------------------------------------------------------------- metrics


def box_iou(a: Any, b: Any) -> float:
    ix = max(0.0, min(a[2], b[2]) - max(a[0], b[0]))
    iy = max(0.0, min(a[3], b[3]) - max(a[1], b[1]))
    inter = ix * iy
    area = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter / area if area > 0 else 0.0


def box_cover(pred: Any, gold: Any) -> float:
    """Share of the gold box inside the predicted one (containment)."""
    ix = max(0.0, min(pred[2], gold[2]) - max(pred[0], gold[0]))
    iy = max(0.0, min(pred[3], gold[3]) - max(pred[1], gold[1]))
    g = (gold[2] - gold[0]) * (gold[3] - gold[1])
    return ix * iy / g if g > 0 else 0.0


def center_in(pred: Any, gold: Any) -> bool:
    cx, cy = (gold[0] + gold[2]) / 2, (gold[1] + gold[3]) / 2
    return pred[0] <= cx <= pred[2] and pred[1] <= cy <= pred[3]


def time_iou(a: Any, b: Any) -> float:
    inter = max(0.0, min(a[1], b[1]) - max(a[0], b[0]))
    union = max(a[1], b[1]) - min(a[0], b[0])
    return inter / union if union > 0 else 0.0


def time_cover(pred: Any, gold: Any) -> float:
    inter = max(0.0, min(pred[1], gold[1]) - max(pred[0], gold[0]))
    g = gold[1] - gold[0]
    return inter / g if g > 0 else 0.0