"""
specialists_gpu.py — the GPU half: load the Apache/MIT specialist models and populate a
`Solids` per image, then score the built task JSON through the EXISTING vlmbench scorers.

This runs on Colab (needs torch + transformers + the model weights). It is deliberately
incremental: the detection hub (GroundingDINO) + the detection-dependent tasks are wired
and validated first (against real COCO GT) so the whole chain — detect → Solids → build →
score_vision_sample — is proven before the other loaders (SAM2, Depth, OCR, SigLIP2, RAM++)
are added. Each loader is a small function returning primitives in PIXEL space.

Model picks are pinned to the plan's Apache/MIT license ledger.
"""

from __future__ import annotations

import json
import time
from typing import Optional

import numpy as np

from .specialists import Solids, build_bbox, build_spatial, build_subject
from .tasks_vision import get_task

# ── detection hub: GroundingDINO (Apache) ────────────────────────────────────
GROUNDING_DINO_ID = "IDEA-Research/grounding-dino-base"   # Apache-2.0, local weights (NOT the API-only 1.5)

# COCO-80 (a known closed vocab for the detection *validation*; production uses the RAM++ tagger)
COCO_CLASSES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
    "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
    "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
    "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
    "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
    "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
    "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush",
]


def load_grounding_dino(device: str = "cuda"):
    """Returns (processor, model). Apache checkpoint, loads via transformers.
    No dtype kwarg — float32 is the default and `torch_dtype=` now deprecation-warns."""
    from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
    proc = AutoProcessor.from_pretrained(GROUNDING_DINO_ID)
    model = AutoModelForZeroShotObjectDetection.from_pretrained(
        GROUNDING_DINO_ID).to(device).eval()
    return proc, model


def _gdino_prompt(classes) -> str:
    # GroundingDINO wants lowercased, "." separated phrases ending in a period.
    return ". ".join(c.strip().lower() for c in classes) + "."


def detect(proc, model, image, classes, box_threshold: float = 0.30,
           text_threshold: float = 0.25, device: str = "cuda") -> list:
    """image → [{label, box:[x1,y1,x2,y2] PIXEL-ABS, score}]. Forces target_sizes so the
    boxes come back pixel-abs xyxy (else transformers returns normalized cxcywh)."""
    import torch
    inputs = proc(images=image, text=_gdino_prompt(classes), return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    W, H = image.size
    _kw = dict(text_threshold=text_threshold, target_sizes=[(H, W)])  # (height, width) → pixel-abs xyxy
    try:                                                              # transformers renamed the arg
        res = proc.post_process_grounded_object_detection(
            outputs, inputs["input_ids"], threshold=box_threshold, **_kw)[0]
    except TypeError:
        res = proc.post_process_grounded_object_detection(
            outputs, inputs["input_ids"], box_threshold=box_threshold, **_kw)[0]
    # dict.get(k, default) evaluates the default EAGERLY — touching the deprecated
    # "labels" key fires transformers' FutureWarning even when text_labels exists
    labels = res["text_labels"] if "text_labels" in res else res.get("labels")
    out = []
    for box, score, lab in zip(res["boxes"], res["scores"], labels):
        b = [float(v) for v in box.tolist()]
        out.append({"label": str(lab).strip() or "object", "box": b, "score": float(score)})
    return out


def ground_phrases(gdino, image, phrases, box_threshold: float = 0.25,
                   text_threshold: float = 0.20, max_tokens_per_chunk: int = 250,
                   device: str = "cuda") -> list:
    """Fusion-tier grounding pass: caption-derived attribute phrases → boxes.
    → [{phrase, matched_span, box:[x1,y1,x2,y2] px, score}].

    Lower thresholds than the base detection pass — fine-grained phrases score
    lower than category nouns, and the downstream containment+margin gate protects
    precision (recall matters more here: an ungrounded attribute falls back to the
    weaker caption-binding path).

    GDINO's BERT text encoder truncates at 256 TOKENS silently (the model slices
    input_ids with no warning), so phrases are chunked by the processor's OWN
    tokenizer count (≤ max_tokens_per_chunk, headroom for [CLS]/[SEP]) with an
    accounting assert — a silently dropped phrase is the failure mode. GDINO
    returns the matched text SPAN (possibly a sub-span, "earrings" from "silver
    drop earrings"): each hit is re-mapped to its source phrase by maximum token
    overlap within the chunk."""
    import torch
    proc, model = gdino
    phrases = [p.strip().lower() for p in phrases if p and p.strip()]
    if not phrases:
        return []

    tok = getattr(proc, "tokenizer", None)

    def _ntok(p):
        if tok is not None:
            return len(tok(p, add_special_tokens=False)["input_ids"]) + 1   # +1 for ". "
        return len(p.split()) + 1                       # crude fallback, ~1.3 tok/word

    chunks, cur, ntok = [], [], 0
    budget = max_tokens_per_chunk if tok is not None else max_tokens_per_chunk // 2
    for p in phrases:
        t = _ntok(p)
        if cur and ntok + t > budget:
            chunks.append(cur)
            cur, ntok = [], 0
        cur.append(p)
        ntok += t
    if cur:
        chunks.append(cur)
    assert sum(len(c) for c in chunks) == len(phrases), "phrase chunking dropped input"

    W, H = image.size
    out = []
    for chunk in chunks:
        text = ". ".join(chunk) + "."
        inputs = proc(images=image, text=text, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = model(**inputs)
        _kw = dict(text_threshold=text_threshold, target_sizes=[(H, W)])
        try:                                                          # transformers renamed the arg
            res = proc.post_process_grounded_object_detection(
                outputs, inputs["input_ids"], threshold=box_threshold, **_kw)[0]
        except TypeError:
            res = proc.post_process_grounded_object_detection(
                outputs, inputs["input_ids"], box_threshold=box_threshold, **_kw)[0]
        out.extend(_remap_spans(res, chunk))
    out.sort(key=lambda r: (r["phrase"], -r["score"]))
    return out


def _remap_spans(res, chunk) -> list:
    """GDINO post-process result → attr-box records, re-mapping each matched text
    SPAN back to its source phrase by maximum token overlap within the chunk."""
    labels = res["text_labels"] if "text_labels" in res else res.get("labels")
    out = []
    for box, score, span in zip(res["boxes"], res["scores"], labels):
        span_toks = set(str(span).lower().split())
        best, best_ov = None, 0
        for p in chunk:
            ov = len(span_toks & set(p.split()))
            if ov > best_ov or (ov == best_ov and best and len(p) > len(best)):
                if ov > 0:
                    best, best_ov = p, ov
        if best is None:
            continue
        out.append({"phrase": best, "matched_span": str(span).strip(),
                    "box": [float(v) for v in box.tolist()],
                    "score": float(score)})
    return out


# ── BATCHED specialist paths (throughput: the serial path leaves a 96GB card ──
#    ~90% idle; every model here batches across images) ────────────────────────

def detect_batch(gdino, images, classes, box_threshold: float = 0.30,
                 text_threshold: float = 0.25, device: str = "cuda") -> list:
    """Batched detection: ONE forward for N images sharing one vocabulary.
    → [boxes_list per image] (same record shape as detect())."""
    import torch
    proc, model = gdino
    images = list(images)
    text = _gdino_prompt(classes)
    inputs = proc(images=images, text=[text] * len(images),
                  return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    sizes = [(im.size[1], im.size[0]) for im in images]          # (H, W)
    _kw = dict(text_threshold=text_threshold, target_sizes=sizes)
    try:
        res = proc.post_process_grounded_object_detection(
            outputs, inputs["input_ids"], threshold=box_threshold, **_kw)
    except TypeError:
        res = proc.post_process_grounded_object_detection(
            outputs, inputs["input_ids"], box_threshold=box_threshold, **_kw)
    out = []
    for r in res:
        labels = r["text_labels"] if "text_labels" in r else r.get("labels")
        out.append([{"label": str(l).strip() or "object",
                     "box": [float(v) for v in b.tolist()], "score": float(s)}
                    for b, s, l in zip(r["boxes"], r["scores"], labels)])
    return out


def zero_shot_batch(siglip, images, labels, device: str = "cuda",
                    template: str = "a photo of a {}.") -> list:
    """Batched SigLIP2 zero-shot → per-image ranked [{label, score}]."""
    import torch
    proc, model = siglip
    texts = [template.format(l) for l in labels]
    # max_length=64 is REQUIRED: the SigLIP2 Gemma tokenizer has no model_max_length,
    # so padding="max_length" alone silently degrades to no padding (HF's own
    # zero-shot pipeline hardcodes 64 for the siglip family)
    inputs = proc(text=texts, images=list(images), return_tensors="pt",
                  padding="max_length", max_length=64, truncation=True).to(device)
    with torch.no_grad():
        logits = model(**inputs).logits_per_image                # [B, n_text]
    probs = torch.sigmoid(logits).float().cpu().numpy()
    out = []
    for row in probs:
        order = row.argsort()[::-1]
        out.append([{"label": labels[i], "score": float(row[i])} for i in order])
    return out


def depth_map_batch(dp, images) -> list:
    """Batched Depth-Anything → per-image HxW float32 nearness maps.

    The DPT image processor resizes with keep_aspect_ratio and NO padding, so a
    mixed-aspect batch produces ragged tensors and crashes. Images are therefore
    grouped by exact size (one forward per group) — byte-identical to the serial
    path, and full batching whenever a set shares a resolution (the synth set)."""
    import torch
    from PIL import Image as _I
    proc, model = dp
    images = list(images)
    groups: dict = {}
    for i, im in enumerate(images):
        groups.setdefault(im.size, []).append(i)
    out = [None] * len(images)
    for size, idxs in sorted(groups.items()):
        chunk = [images[i] for i in idxs]
        inputs = proc(images=chunk, return_tensors="pt").to(model.device)
        with torch.no_grad():
            pd = model(**inputs).predicted_depth                 # [B, h, w]
        W, H = size
        for arr, i in zip(pd, idxs):
            a = arr.squeeze().float().cpu().numpy().astype(np.float32)
            if a.shape != (H, W):
                a = np.asarray(_I.fromarray(a).resize((W, H), _I.BILINEAR),
                               dtype=np.float32)
            out[i] = a
    return out


def segment_batch(sam, images, boxes_list, device: str = "cuda") -> list:
    """Batched grounded-SAM. Variable per-image box counts are PADDED to the batch
    max (dummy [0,0,2,2] prompts) and the surplus masks dropped — SAM's processor
    needs a rectangular input_boxes tensor. Mutates boxes in place like segment()."""
    if sam is None or not any(boxes_list):
        return boxes_list
    import torch
    proc, model = sam
    keep = [i for i, bl in enumerate(boxes_list) if bl]
    imgs = [images[i] for i in keep]
    max_n = max(len(boxes_list[i]) for i in keep)
    padded = [[[float(v) for v in b["box"]] for b in boxes_list[i]]
              + [[0.0, 0.0, 2.0, 2.0]] * (max_n - len(boxes_list[i]))
              for i in keep]
    inputs = proc(imgs, input_boxes=padded, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    masks = proc.image_processor.post_process_masks(
        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(),
        inputs["reshaped_input_sizes"].cpu())                    # [n_obj, 3, H, W] per image
    scores = outputs.iou_scores.cpu().numpy()                    # [B, n_obj, 3]
    for bi, i in enumerate(keep):
        m = np.asarray(masks[bi])
        sc = scores[bi]
        for oi, b in enumerate(boxes_list[i]):                   # surplus (padded) masks ignored
            mo = m[oi]
            if mo.ndim == 3:
                best = int(sc[oi].argmax()) if oi < len(sc) else 0
                b["mask_score"] = float(sc[oi][best]) if oi < len(sc) else None
                mo = mo[best]
            b["mask"] = np.asarray(mo, dtype=bool)
    return boxes_list


def ground_phrases_batch(gdino, images, phrases_list, box_threshold: float = 0.25,
                         text_threshold: float = 0.20, max_tokens_per_chunk: int = 250,
                         device: str = "cuda") -> list:
    """Batched phrase grounding. Images whose phrase text fits ONE chunk (the
    typical case) share a single forward; oversized ones fall back to the serial
    chunked ground_phrases. → per-image attr-box record lists."""
    import torch
    proc, model = gdino
    norm = [[p.strip().lower() for p in (ph or []) if p and p.strip()]
            for ph in phrases_list]
    tok = getattr(proc, "tokenizer", None)
    out = [[] for _ in images]

    easy, hard = [], []
    for i, ph in enumerate(norm):
        if not ph:
            continue
        text = ". ".join(ph) + "."
        ntok = (len(tok(text)["input_ids"]) if tok is not None
                else 2 * len(text.split()))
        (easy if ntok <= max_tokens_per_chunk else hard).append(i)

    if easy:
        imgs = [images[i] for i in easy]
        texts = [". ".join(norm[i]) + "." for i in easy]
        inputs = proc(images=imgs, text=texts, return_tensors="pt",
                      padding=True).to(device)
        with torch.no_grad():
            outputs = model(**inputs)
        sizes = [(im.size[1], im.size[0]) for im in imgs]
        _kw = dict(text_threshold=text_threshold, target_sizes=sizes)
        try:
            res = proc.post_process_grounded_object_detection(
                outputs, inputs["input_ids"], threshold=box_threshold, **_kw)
        except TypeError:
            res = proc.post_process_grounded_object_detection(
                outputs, inputs["input_ids"], box_threshold=box_threshold, **_kw)
        for bi, i in enumerate(easy):
            recs = _remap_spans(res[bi], norm[i])
            recs.sort(key=lambda r: (r["phrase"], -r["score"]))
            out[i] = recs
    for i in hard:
        out[i] = ground_phrases(gdino, images[i], norm[i],
                                box_threshold=box_threshold,
                                text_threshold=text_threshold,
                                max_tokens_per_chunk=max_tokens_per_chunk,
                                device=device)
    return out


def solids_from_detection(image, boxes) -> Solids:
    """Minimal Solids from detection alone (feeds bbox / spatial / subject)."""
    return Solids(size=image.size, boxes=boxes)


# ── validation: detection hub vs COCO GT, through the existing scorers ────────
def validate_detection(n: int = 24, box_threshold: float = 0.30, device: str = "cuda") -> dict:
    """Run GroundingDINO on real COCO images, build the bbox JSON, and score it with the
    EXISTING vlmbench detection scorer — apples-to-apples with the VLM's bbox number."""
    from .datasets import load_gt
    from .metrics import score_vision_sample, score_vision_run

    spec = get_task("bbox_grounding")
    proc, model = load_grounding_dino(device)
    samples = load_gt(spec.gt_dataset, n=n, split=spec.gt_split, dataset="full")
    print(f"[validate_detection] {GROUNDING_DINO_ID} on {len(samples)} COCO images")

    results, t0 = [], time.perf_counter()
    for i, s in enumerate(samples):
        boxes = detect(proc, model, s.image, COCO_CLASSES, box_threshold, device=device)
        pred = build_bbox(solids_from_detection(s.image, boxes))
        mr = score_vision_sample(spec, json.dumps(pred), s.gt, mode="specialist",
                                 image_id=s.image_id, image_size=s.size)
        results.append(mr)
        if i < 3:
            print(f"  {s.image_id}: {len(boxes)} boxes, primary={mr.primary_score}")
    dt = time.perf_counter() - t0
    run = score_vision_run(results, model="grounding-dino-base", category=spec.category, mode="specialist")
    out = {"model": GROUNDING_DINO_ID, "n": len(samples),
           "primary_score_mean": run.primary_score_mean, "schema_valid_rate": run.schema_valid_rate,
           "img_per_s": round(len(samples) / max(0.001, dt), 2)}
    print(f"\n[validate_detection] mean primary={out['primary_score_mean']}  "
          f"valid={out['schema_valid_rate']}  {out['img_per_s']} img/s")
    print("Compare vs the VLM bbox_grounding effective yield (~0.16–0.30 in the vlmbench).")
    return out


# ── depth hub: Depth-Anything-V2-Small (Apache; higher = nearer) ─────────────
DEPTH_ID = "depth-anything/Depth-Anything-V2-Small-hf"   # ONLY Small is Apache

def load_depth_anything(device: str = "cuda"):
    """Returns (processor, model). Direct model call — the transformers pipeline()
    warns ("use a dataset") when invoked per-image on GPU and adds dispatch overhead."""
    from transformers import AutoImageProcessor, AutoModelForDepthEstimation
    proc = AutoImageProcessor.from_pretrained(DEPTH_ID)
    model = AutoModelForDepthEstimation.from_pretrained(DEPTH_ID).to(device).eval()
    return proc, model


def depth_map(dp, image):
    """HxW float32 relative depth; Depth-Anything convention: HIGHER = NEARER."""
    import torch
    proc, model = dp
    inputs = proc(images=image, return_tensors="pt").to(model.device)
    with torch.no_grad():
        pd = model(**inputs).predicted_depth
    arr = pd.squeeze().float().cpu().numpy().astype(np.float32)
    # resize to image size if the model returned a different resolution
    W, H = image.size
    if arr.shape != (H, W):
        from PIL import Image as _I
        arr = np.asarray(_I.fromarray(arr).resize((W, H), _I.BILINEAR), dtype=np.float32)
    return arr


# ── segmentation hub: SAM v1 (Apache), prompted by detection boxes ───────────
#   SAM2's transformers processor is currently broken (missing preprocessor_config on the -hf
#   repo); SAM v1 is equally Apache-2.0, box-promptable, and rock-solid in transformers.
SAM_ID = "facebook/sam-vit-base"

def load_sam(device: str = "cuda"):
    """Returns (processor, model) or None. SAM v1 via transformers (SamModel/SamProcessor)."""
    try:
        from transformers import SamModel, SamProcessor
        proc = SamProcessor.from_pretrained(SAM_ID)
        model = SamModel.from_pretrained(SAM_ID).to(device).eval()
        return proc, model
    except Exception as e:
        print(f"[load_sam] SAM unavailable: {type(e).__name__}: {e}")
        return None


def segment(sam, image, boxes, device: str = "cuda"):
    """Attach a boolean mask to each box dict (in place). Grounded-SAM: box → mask.
    Picks the highest-IoU of SAM's 3 mask proposals per box."""
    if sam is None or not boxes:
        return boxes
    import torch
    proc, model = sam
    input_boxes = [[[float(v) for v in b["box"]] for b in boxes]]     # [image][obj][xyxy]
    inputs = proc(image, input_boxes=input_boxes, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    masks = proc.image_processor.post_process_masks(
        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(),
        inputs["reshaped_input_sizes"].cpu())[0]                     # [obj, n_masks, H, W]
    m = np.asarray(masks)
    scores = outputs.iou_scores.cpu().numpy()[0]                      # [obj, n_masks]
    for i, b in enumerate(boxes):
        if i >= len(m):
            break
        mi = m[i]
        if mi.ndim == 3:                                             # [n_masks, H, W] → best by IoU
            best = int(scores[i].argmax()) if i < len(scores) else 0
            b["mask_score"] = float(scores[i][best]) if i < len(scores) else None
            mi = mi[best]
        b["mask"] = np.asarray(mi, dtype=bool)
    return boxes


# ── classification / style hub: SigLIP2 (Apache) zero-shot ───────────────────
SIGLIP_ID = "google/siglip2-so400m-patch14-384"
STYLE_LABELS = ["photo", "painting", "3d_render", "sketch", "anime", "other"]

def load_siglip(device: str = "cuda"):
    from transformers import AutoProcessor, AutoModel
    from transformers.utils import logging as hf_logging
    # The checkpoint's config carries CLIP-tokenizer bos/eos ids (49406/49407) that
    # newer transformers flags against the 32k text vocab. SigLIP never generates,
    # so the ids are inert — silence the config validation for just this load.
    prev = hf_logging.get_verbosity()
    hf_logging.set_verbosity_error()
    try:
        proc = AutoProcessor.from_pretrained(SIGLIP_ID)
        model = AutoModel.from_pretrained(SIGLIP_ID).to(device).eval()
    finally:
        hf_logging.set_verbosity(prev)
    return proc, model


def zero_shot(siglip, image, labels, device: str = "cuda", template: str = "a photo of a {}.") -> list:
    """SigLIP2 zero-shot → [{label, score}] sorted desc (sigmoid, not softmax)."""
    import torch
    proc, model = siglip
    texts = [template.format(l) for l in labels]
    # max_length=64 required — see zero_shot_batch (SigLIP2 tokenizer has no
    # model_max_length, so padding="max_length" alone silently doesn't pad)
    inputs = proc(text=texts, images=image, return_tensors="pt", padding="max_length",
                  max_length=64, truncation=True).to(device)
    with torch.no_grad():
        logits = model(**inputs).logits_per_image[0]
    probs = torch.sigmoid(logits).float().cpu().numpy()
    order = probs.argsort()[::-1]
    return [{"label": labels[i], "score": float(probs[i])} for i in order]


# ── OCR hub: EasyOCR (Apache, torch — no Paddle/torch CUDA conflict) ──────────
def load_ocr(device: str = "cuda"):
    try:
        import easyocr
        return easyocr.Reader(["en"], gpu=(device == "cuda"))
    except Exception as e:
        print(f"[load_ocr] EasyOCR unavailable: {type(e).__name__}: {e}")
        return None


def ocr_read(reader, image) -> dict:
    """EasyOCR → {full_text, lines:[{text, box:[quad px], conf}]}. Confidence is
    RETAINED (the fusion tier carries it; build_ocr ignores the extra key)."""
    if reader is None:
        return {"full_text": "", "lines": []}
    res = reader.readtext(np.asarray(image))     # [(quad, text, conf), ...]
    lines = [{"text": str(t), "box": [[float(x), float(y)] for x, y in quad],
              "conf": float(c)} for quad, t, c in res]
    return {"full_text": " ".join(l["text"] for l in lines), "lines": lines}


# ── generic per-task validation through the existing scorers ─────────────────
SHAPE_CLASSES = ["red circle", "green circle", "blue circle"]   # synthetic-shape GT vocab

# task → dict of which models it needs, the GDINO vocab, and (optional) a REAL-image GT
# override that replaces the synthetic GT in the task spec.
_TASK_CFG = {
    "bbox_grounding":               dict(vocab=COCO_CLASSES, gdino=True),
    "segmentation":                 dict(vocab=COCO_CLASSES, gdino=True, masks=True, gt="coco_segmentation"),
    "outline_association":          dict(vocab=COCO_CLASSES, gdino=True, masks=True, gt="coco_outline"),
    "subject_fixation":             dict(vocab=COCO_CLASSES, gdino=True, gt="coco_subject"),
    # still synthetic — need real depth (NYU/DIODE) + relations (Visual Genome); next real-GT pass
    "depth_analysis":               dict(vocab=SHAPE_CLASSES, gdino=True, depth=True, masks=True),
    "structural_spatial_awareness": dict(vocab=SHAPE_CLASSES, gdino=True, depth=True),
    "image_classification":         dict(siglip=True),                       # vocab from GT labels
    "style_structural_awareness":   dict(gdino=True, siglip=True, gray=True),  # style has no real GT
    "ocr_text":                     dict(ocr=True),
    "data_type_differentiation":    dict(ocr=True),                          # rendered-format GT is synthetic
    "data_type_utilization":        dict(ocr=True),
}


def validate_task(task: str, n: int = 24, device: str = "cuda", *, gdino=None,
                  depth_pipe=None, sam=None, siglip=None, ocr=None) -> dict:
    """Run the specialist/derive chain for one task and score it with the vlmbench scorer."""
    from .datasets import load_gt
    from .metrics import score_vision_sample, score_vision_run
    from .specialists import (Solids, build_bbox, build_spatial, build_subject,
                              build_depth_order, build_segmentation, build_outline,
                              build_classification, build_style, build_ocr,
                              build_datatype_diff, build_datatype_util)

    spec = get_task(task)
    cfg = _TASK_CFG[task]
    gt_key = cfg.get("gt", spec.gt_dataset)               # real-image GT override when available
    samples = load_gt(gt_key, n=n, split=spec.gt_split or "", dataset="full")

    # candidate label set for zero-shot classification: the classes present in this GT slice
    class_vocab = None
    if task == "image_classification":
        seen = []
        for s in samples:
            for l in (s.gt.get("labels", []) if isinstance(s.gt, dict) else []):
                if l not in seen:
                    seen.append(l)
        class_vocab = seen or ["object"]

    results = []
    for s in samples:
        sol = Solids(size=s.image.size)
        if cfg.get("gdino") and gdino is not None:
            sol.boxes = detect(gdino[0], gdino[1], s.image, cfg.get("vocab", COCO_CLASSES), device=device)
            if cfg.get("masks") and sam is not None:
                sol.boxes = segment(sam, s.image, sol.boxes, device=device)
        if cfg.get("depth") and depth_pipe is not None:
            sol.depth = depth_map(depth_pipe, s.image)
        if cfg.get("gray"):
            sol.gray = np.asarray(s.image.convert("L"), dtype=np.float32)
        if cfg.get("siglip") and siglip is not None:
            if task == "image_classification":
                sol.class_top = zero_shot(siglip, s.image, class_vocab, device=device)[:5]
            if task == "style_structural_awareness":
                sol.style = zero_shot(siglip, s.image, STYLE_LABELS, device=device)[0]["label"]
        if cfg.get("ocr") and ocr is not None:
            sol.ocr = ocr_read(ocr, s.image)

        if task == "depth_analysis":
            pred = build_depth_order(sol)
        elif task == "segmentation":
            pred = build_segmentation(sol)
        elif task == "outline_association":
            pred = build_outline(sol)
        elif task == "structural_spatial_awareness":
            pred = build_spatial(sol)
        elif task == "subject_fixation":
            pred = build_subject(sol)
        elif task == "bbox_grounding":
            pred = build_bbox(sol)
        elif task == "image_classification":
            pred = build_classification(sol)
        elif task == "style_structural_awareness":
            pred = build_style(sol)
        elif task == "ocr_text":
            pred = build_ocr(sol)
        elif task == "data_type_differentiation":
            pred = build_datatype_diff(sol)
        elif task == "data_type_utilization":
            pred = build_datatype_util(sol)[0]
        else:
            raise KeyError(task)

        mr = score_vision_sample(spec, json.dumps(pred), s.gt, mode="specialist",
                                 image_id=s.image_id, image_size=s.size)
        results.append(mr)
    run = score_vision_run(results, model="specialist", category=task, mode="specialist")
    return {"task": task, "n": len(samples), "primary": run.primary_score_mean,
            "valid": run.schema_valid_rate}


class SpecialistPipeline:
    """Simplified interface: load the Apache/MIT specialist models ONCE, then `extract(image)`
    returns every deterministic task's JSON (the production entry point for `tasks_json`).

        pipe = SpecialistPipeline()
        tasks = pipe.extract(pil_image)   # {"bbox_grounding": {...}, "segmentation": {...}, ...}
    """
    DEFAULT_VOCAB = COCO_CLASSES

    def __init__(self, device: str = "cuda", with_ocr: bool = True):
        self.device = device
        self.gdino = load_grounding_dino(device)
        self.depth = load_depth_anything(device)
        self.sam = load_sam(device)
        self.siglip = load_siglip(device)
        self.ocr = load_ocr(device) if with_ocr else None

    def solidify(self, image, vocab=None):
        """Run every specialist once → a `Solids` (primitives in pixel space)."""
        vocab = vocab or self.DEFAULT_VOCAB
        s = Solids(size=image.size)
        s.boxes = detect(self.gdino[0], self.gdino[1], image, vocab, device=self.device)
        s.boxes = segment(self.sam, image, s.boxes, device=self.device)
        s.depth = depth_map(self.depth, image)
        s.gray = np.asarray(image.convert("L"), dtype=np.float32)
        if self.siglip is not None:
            s.class_top = zero_shot(self.siglip, image, vocab, device=self.device)[:5]
            s.style = zero_shot(self.siglip, image, STYLE_LABELS, device=self.device)[0]["label"]
        if self.ocr is not None:
            s.ocr = ocr_read(self.ocr, image)
        return s

    @staticmethod
    def _build_tasks(s) -> dict:
        """Solids → {task_name: task_json} for all 11 deterministic tasks."""
        from .specialists import (build_bbox, build_segmentation, build_classification,
                                  build_ocr, build_spatial, build_depth_order, build_subject,
                                  build_outline, build_style, build_datatype_diff,
                                  build_datatype_util)
        util, _ = build_datatype_util(s)
        return {
            "bbox_grounding": build_bbox(s),
            "segmentation": build_segmentation(s),
            "outline_association": build_outline(s),
            "subject_fixation": build_subject(s),
            "depth_analysis": build_depth_order(s),
            "structural_spatial_awareness": build_spatial(s),
            "image_classification": build_classification(s),
            "style_structural_awareness": build_style(s),
            "ocr_text": build_ocr(s),
            "data_type_differentiation": build_datatype_diff(s),
            "data_type_utilization": util,
        }

    def extract(self, image, vocab=None) -> dict:
        """→ {task_name: task_json} for all 11 deterministic tasks (one solidify pass)."""
        return self._build_tasks(self.solidify(image, vocab))

    def ground(self, image, phrases) -> list:
        """Fusion grounding pass: caption phrases → attr-box records (GDINO reused)."""
        return ground_phrases(self.gdino, image, phrases, device=self.device)

    def solidify_batch(self, images, vocab=None, phrases_list=None,
                       batch: int = 16, gdino_batch: int = 2) -> list:
        """Batched solidify: SAM/depth/SigLIP run at `batch` images per forward;
        GroundingDINO runs in sub-chunks of `gdino_batch` — its deformable
        attention's activation memory EXPLODES with padded batches (measured on
        the 96GB Blackwell: B2 = 11GB, B16 = 42GB for LESS throughput), so ~2 is
        its sweet spot. EasyOCR stays serial (4% of the budget). → [Solids]
        aligned with `images`, same output contract as solidify()."""
        vocab = vocab or self.DEFAULT_VOCAB
        images = list(images)
        solids = []
        for start in range(0, len(images), batch):
            chunk = images[start:start + batch]
            p_chunk = (phrases_list[start:start + batch]
                       if phrases_list is not None else None)
            boxes_list = []
            for s2 in range(0, len(chunk), gdino_batch):
                boxes_list.extend(detect_batch(
                    self.gdino, chunk[s2:s2 + gdino_batch], vocab,
                    device=self.device))
            boxes_list = segment_batch(self.sam, chunk, boxes_list, device=self.device)
            depths = (depth_map_batch(self.depth, chunk)
                      if self.depth is not None else [None] * len(chunk))
            classes = (zero_shot_batch(self.siglip, chunk, vocab, device=self.device)
                       if self.siglip is not None else [None] * len(chunk))
            styles = (zero_shot_batch(self.siglip, chunk, STYLE_LABELS, device=self.device)
                      if self.siglip is not None else [None] * len(chunk))
            if p_chunk is not None:
                attrs = []
                for s2 in range(0, len(chunk), gdino_batch):
                    attrs.extend(ground_phrases_batch(
                        self.gdino, chunk[s2:s2 + gdino_batch],
                        p_chunk[s2:s2 + gdino_batch], device=self.device))
            else:
                attrs = [[] for _ in chunk]
            for k, im in enumerate(chunk):
                s = Solids(size=im.size)
                s.boxes = boxes_list[k]
                s.depth = depths[k]
                s.gray = np.asarray(im.convert("L"), dtype=np.float32)
                if classes[k] is not None:
                    s.class_top = classes[k][:5]
                    s.style = styles[k][0]["label"]
                if self.ocr is not None:
                    s.ocr = ocr_read(self.ocr, im)
                s.attr_boxes = attrs[k]
                solids.append(s)
        return solids

    def extract_batch(self, images, vocab=None, phrases_list=None,
                      batch: int = 16) -> list:
        """Batched extract(+digest): → [(tasks_dict, digest)] aligned with images."""
        from .fuse import solids_digest
        out = []
        for s in self.solidify_batch(images, vocab, phrases_list, batch=batch):
            out.append((self._build_tasks(s), solids_digest(s)))
        return out

    def extract_with_digest(self, image, phrases=None, vocab=None) -> tuple:
        """→ (tasks_dict, solids_digest) from ONE solidify pass (+ the phrase-
        grounding pass when `phrases` is given). The digest is the fusion tier's
        input — compact, JSON-able, carries the retained confidences."""
        from .fuse import solids_digest
        s = self.solidify(image, vocab)
        if phrases:
            s.attr_boxes = self.ground(image, phrases)
        return self._build_tasks(s), solids_digest(s)


def load_vlm(model_key: str = "qwen3vl-4b"):
    from .model_registry import get_runner
    return get_runner(model_key)


def validate_task_vlm(task: str, n: int = 24, model_key: str = "qwen3vl-4b",
                      runner=None, device: str = "cuda") -> dict:
    """Run the Qwen VLM on the SAME (real) GT as validate_task — a true apples-to-apples
    head-to-head. Reuses the existing VLMRunner + score path. Pass a pre-loaded `runner`
    to avoid reloading the model per task."""
    from .datasets import load_gt
    from .metrics import score_vision_sample, score_vision_run

    spec = get_task(task)
    gt_key = _TASK_CFG[task].get("gt", spec.gt_dataset)
    samples = load_gt(gt_key, n=n, split=spec.gt_split or "", dataset="full")
    own = runner is None
    if own:
        runner = load_vlm(model_key)
    results = []
    try:
        for s in samples:
            up = s.prompt if spec.per_sample_prompt else None
            res = runner.generate(spec, s.image, "json_mode", image_id=s.image_id,
                                  image_size=s.size, gt=s.gt, user_prompt=up)
            results.append(score_vision_sample(spec, res.raw_text, s.gt, mode="json_mode",
                                               image_id=s.image_id, image_size=s.size))
    finally:
        if own:
            close = getattr(runner, "close", None)
            if callable(close):
                close()
    run = score_vision_run(results, model=model_key, category=task, mode="json_mode")
    # effective yield = accuracy × validity (the vlmbench headline metric)
    acc = run.primary_score_mean
    return {"task": task, "vlm_primary": acc, "vlm_valid": run.schema_valid_rate,
            "vlm_yield": (acc * run.schema_valid_rate) if acc is not None else None}