"""JSONL data format (same as the report). One record per line:
{"id", "images": [path,...], "state": str|obj, "questions": {qid: {type, instructions, criteria}},
 "teacher": {qid: {option_id: prob}}  (optional; teacher distribution, e.g. Clef via Workers AI or local)
 "gold":    {qid: option_id}          (optional; ground-truth label, e.g. from synthetic generator)
 "meta":    {...}}
Image paths are resolved relative to the JSONL file.
For noul questions option ids are "true"/"false". A teacher value may be a bare float = p(true)
(that's what the Workers AI / SystemOne response gives), it is expanded automatically.

Box questions (Jev 1.14 `box`, see glint/box.py): questions[qid] = {"type": "box", "instructions": ...};
gold[qid] = {"top", "bottom", "left", "right"} as fractions (0-1) of image height/width from the top-left,
optionally with "present": bool ({"present": false} alone is fine). Box gold comes from ground truth, never
from the teacher; teacher entries for box questions are ignored. present=false -> no box loss for that question.
Presence itself is a separate noul question, by convention `<qid>_visible`; if it has no gold, it is derived
from `present` (`derive_presence_gold`)."""
from __future__ import annotations
import json, random
from pathlib import Path
from PIL import Image
from .box import is_box_question

def load_jsonl(path):
    path = Path(path)
    rows = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                ex = json.loads(line)
                ex["_base"] = str(path.parent)
                rows.append(ex)
    return rows

def normalize_teacher(ex: dict) -> dict:
    out = {}
    for qid, v in (ex.get("teacher") or {}).items():
        q = ex["questions"].get(qid)
        if q is None or is_box_question(q):
            continue
        if isinstance(v, (int, float)):            # noul: p(true)
            v = {"true": float(v), "false": 1.0 - float(v)}
        elif isinstance(v, dict) and "probabilities" in v:   # SystemOne choice/score answer
            v = v["probabilities"]
        elif isinstance(v, dict) and "noul" in v:
            v = {"true": float(v["noul"]), "false": 1.0 - float(v["noul"])}
        out[qid] = {str(k): float(p) for k, p in v.items()}
    return out

def open_images(ex: dict):
    base = Path(ex.get("_base", "."))
    return [Image.open(p if Path(p).is_absolute() else base / p).convert("RGB") for p in ex.get("images", [])]

def to_record(ex: dict, augment: bool = False) -> dict:
    # encode_record は choice の選択肢を option_id でソートするので選択肢順シャッフルは無意味。
    # 増強は質問(フィールド)順のシャッフルのみ。
    items = list(ex["questions"].items())
    if augment:
        random.shuffle(items)
    return {"id": ex.get("id", "unknown"), "state": ex.get("state", ""), "images": open_images(ex),
            "questions": dict(items), **({"prompt_style": ex["prompt_style"]} if ex.get("prompt_style") else {})}

def teacher_questions(questions: dict) -> dict:
    """Questions a Clef teacher can answer (drops `box` questions, which Clef does not support)."""
    return {qid: q for qid, q in questions.items() if not is_box_question(q)}

def normalize_box_gold(ex: dict) -> dict:
    """{qid: [top, bottom, left, right] or None (= not present / no gold box)} for every box question with gold.
    Accepts {"top",..} or {"box": {"top",..}} with optional "present"; swaps inverted edges, clamps to [0, 1]."""
    out = {}
    for qid, q in ex["questions"].items():
        g = (ex.get("gold") or {}).get(qid)
        if not is_box_question(q) or g is None:
            continue
        if not isinstance(g, dict):
            raise ValueError(f"{ex.get('id')}/{qid}: box gold must be an object, got {g!r}")
        b = g.get("box", g)
        if g.get("present", True) is False or not all(k in b for k in ("top", "bottom", "left", "right")):
            out[qid] = None
            continue
        t, bo, l, r = (min(max(float(b[k]), 0.0), 1.0) for k in ("top", "bottom", "left", "right"))
        out[qid] = [min(t, bo), max(t, bo), min(l, r), max(l, r)]
    return out

def derive_presence_gold(ex: dict) -> dict:
    """gold, plus `<qid>_visible` = "true"/"false" from a box gold's `present` when that noul question exists
    and has no gold of its own. Returns the original dict when there is nothing to add."""
    gold = ex.get("gold") or {}
    extra = {}
    for qid, q in ex["questions"].items():
        g, vis = gold.get(qid), f"{qid}_visible"
        if is_box_question(q) and isinstance(g, dict) and vis in ex["questions"] and vis not in gold:
            extra[vis] = "false" if g.get("present", True) is False else "true"
    return {**gold, **extra} if extra else gold