"""Per-pixel language targets (label_text v1: SigLIP2-base text embeddings of 624 concepts) for splatted features.

Source: point4d-datasets tools/label_text/v1 (README there). Every semantic class / object category maps to one
concept row of `targets [624, 768]` (L2-normalised google/siglip2-base-patch16-224 text embeddings); row -1 means
IGNORE (unlabelled, void, background, ...), and ids or names missing from a table are IGNORE too.
  * table key "semantic id": rows = lut[semantic]
  * table key "category name": rows = classes[instances[str(instance)]["name"]]["row"]
  * no semantic channel but a semantic-id table (data-012 in the HF mirror): read instances[...]["semantic_id"]
Loss (as recommended there): contrastive over all concepts, logits = cos(pred, targets) / tau with a learnable tau
initialised at the fitted 0.009; plain cosine regression separates classes poorly.
"""

import json
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

DEFAULT_DIR = "/scratch/data/point4d-datasets/tools/label_text/v1"
IGNORE_ROW = -1


class LabelText:
    def __init__(self, root=DEFAULT_DIR):
        self.root = Path(root)
        z = np.load(self.root / "label_embeddings_v1.npz")
        self.targets = z["targets"].astype(np.float32)                    # (C, 768), L2-normalised
        self.concepts = [str(c) for c in z["concept_names"]]
        self.tables = {}
        for path in self.root.glob("*_label_text.json"):
            table = json.loads(path.read_text())
            classes = table.get("classes") or {}
            if not classes:
                continue                                                  # e.g. ADT: empty for now -> all IGNORE
            semantic = table.get("key", "").startswith("semantic")
            entry = dict(semantic=semantic, rows={k: int(v.get("row", IGNORE_ROW)) for k, v in classes.items()})
            if semantic:
                ids = [int(k) for k in classes]
                lut = np.full(max(ids) + 1, IGNORE_ROW, np.int32)
                for k, row in entry["rows"].items():
                    lut[int(k)] = row
                entry["lut"] = lut
            self.tables[table.get("dataset", path.name.removesuffix("_label_text.json"))] = entry

    def rows(self, dataset, semantic=None, instance=None, instances=None):
        """(H, W) int32 concept rows for one frame, or None when the dataset has no usable channel/table."""
        table = self.tables.get(dataset)
        if table is None:
            return None
        if table["semantic"]:
            ids = semantic
            if ids is None and instance is not None and instances:              # semantic id via instances.json
                lut = np.full(int(instance.max()) + 1, -1, np.int64)
                for k, v in instances.items():
                    if int(k) <= instance.max() and v.get("semantic_id") is not None:
                        lut[int(k)] = int(v["semantic_id"])
                ids = lut[instance]
            if ids is None:
                return None
            ids = ids.astype(np.int64)
            lut = table["lut"]
            ok = (ids >= 0) & (ids < len(lut))
            return np.where(ok, lut[np.clip(ids, 0, len(lut) - 1)], IGNORE_ROW).astype(np.int32)
        if instance is None or not instances:
            return None
        lut = np.full(int(instance.max()) + 1, IGNORE_ROW, np.int32)
        for k, v in instances.items():
            name = v.get("name")
            if name is not None and int(k) <= instance.max():
                lut[int(k)] = table["rows"].get(name, IGNORE_ROW)
        return lut[instance]


class LanguageDecoder(nn.Module):
    """Splatted per-pixel latent (lang_dim) -> 768-d SigLIP2 space; contrastive logits over the concept targets."""

    def __init__(self, lang_dim, targets, tau=0.009):
        super().__init__()
        self.proj = nn.Sequential(nn.Linear(lang_dim, 256), nn.GELU(), nn.Linear(256, targets.shape[1]))
        self.register_buffer("targets", F.normalize(torch.as_tensor(targets, dtype=torch.float32), dim=-1), persistent=False)
        self.log_tau = nn.Parameter(torch.tensor(float(np.log(tau))))

    def embed(self, latent):
        return F.normalize(self.proj(latent.float()), dim=-1)

    def loss(self, latent, rows, max_pixels=4096):
        """Cross-entropy of cos(embed, targets) / tau against concept rows (IGNORE skipped); also top-1 accuracy."""
        valid = rows >= 0
        if not valid.any():
            return latent.sum() * 0.0, {}
        lat, tgt = latent[valid], rows[valid].long()
        if len(tgt) > max_pixels:
            pick = torch.randperm(len(tgt), device=tgt.device)[:max_pixels]
            lat, tgt = lat[pick], tgt[pick]
        logits = self.embed(lat) @ self.targets.T / self.log_tau.exp()
        loss = F.cross_entropy(logits, tgt)
        return loss, dict(lang=float(loss.detach()), lang_acc=float((logits.argmax(-1) == tgt).float().mean()),
                          lang_tau=float(self.log_tau.exp().detach()))