"""Multimodal train / eval records from public Hugging Face datasets only (no synthetic media).

image : A-OKVQA, POPE (from base.pt) + ScienceQA-IMG, AI2D, VQAv2-small (yes/no -> noul, how-many -> score, other -> choice)
audio : ESC-50 (from the audio part) + Clotho-AQA (yes/no), UrbanSound8K (fold 10 held out), Speech Commands v0.02
video : TempCompass (multi-choice / yes-no / caption matching, split by video), MV-Bench-mini, UCF101-subset, Kinetics-mini

Features are cached at the grid the LM consumes (image 8x8 -> 64 tokens, video frame 4x4 -> 16 tokens); a few eval items
per task keep their raw media for end-to-end latency and demo cards. Builders take `jev` and return record lists.
"""
import glob, io, os, random, re, subprocess, tarfile, zipfile
import numpy as np, torch
from datasets import Audio, load_dataset
from huggingface_hub import hf_hub_download, snapshot_download
from mmjev import Seg

N_RAW = 12
IMG_GRID, FRAME_GRID = 8, 4


def _rec(task, mod, split, state, qs, ys):
    qs = qs if isinstance(qs, list) else [qs]
    ys = ys if isinstance(ys, list) else [ys]
    from mmjev import options_of
    tg = [[1.0 if i == y else 0.0 for i in range(len(options_of(q)[0]))] for q, y in zip(qs, ys)]
    return dict(task=task, modality=mod, split=split, state=state, qs=qs, ys=ys, targets=tg, raw=None)


def _img_feats(jev, ims, grid=IMG_GRID):
    g = jev.vision_tower_features(ims).float()
    return torch.nn.functional.avg_pool2d(g, g.shape[-1] // grid).half().cpu()


def _encode_images(jev, recs):
    """Replace PIL images in record states by cached grids; keep raw for the first N_RAW eval items of each task."""
    cnt = {}
    for r in recs:
        if r["split"] == "eval":
            cnt[r["task"]] = cnt.get(r["task"], 0) + 1
            if cnt[r["task"]] <= N_RAW:
                r["raw"] = [Seg(s.kind, s.data) for s in r["state"]]
    for s in range(0, len(recs), 32):
        chunk = recs[s:s + 32]
        f = _img_feats(jev, [r["state"][0].data.convert("RGB") for r in chunk])
        for r, x in zip(chunk, f):
            r["state"][0] = Seg("image", x)
    return recs


def _letters(opts):
    return {o: "" for o in opts} if len(set(opts)) == len(opts) else {f"{chr(65 + i)}. {o}": "" for i, o in enumerate(opts)}


# ------------------------------------------------------------------------------------------ image
def scienceqa(jev, n_tr=1500, n_ev=500):
    out = []
    for split, hf, n in (("train", "train", n_tr), ("eval", "test", n_ev)):
        ds = load_dataset("lmms-lab/ScienceQA-IMG", split=hf).shuffle(seed=0)
        k = 0
        for ex in ds:
            if k >= n:
                break
            if ex["image"] is None:
                continue
            opts = eval(ex["choices"]) if isinstance(ex["choices"], str) else list(ex["choices"])
            q = ex["question"] + (f"\nHint: {ex['hint']}" if ex.get("hint") else "")
            crit = _letters([str(o) for o in opts])
            out.append(_rec("scienceqa_img", "image", split, [Seg("image", ex["image"])],
                            {"type": "choice", "instructions": q, "criteria": crit}, int(ex["answer"]))); k += 1
    return _encode_images(jev, out)


def ai2d(jev, n_ev=600):
    ds = load_dataset("lmms-lab-encoder/ai2d", split="test").shuffle(seed=0)
    out = []
    for i, ex in enumerate(ds):
        opts = eval(ex["options"]) if isinstance(ex["options"], str) else list(ex["options"])
        split = "eval" if i < n_ev else "train"
        out.append(_rec("ai2d", "image", split, [Seg("image", ex["image"])],
                        {"type": "choice", "instructions": ex["question"], "criteria": _letters([str(o) for o in opts])},
                        int(ex["answer"])))
    return _encode_images(jev, out)


def vqav2(jev, n=6000, eval_frac=0.15, seed=0):
    rng = random.Random(seed)
    ds = load_dataset("merve/vqav2-small", split="validation").shuffle(seed=seed).select(range(n))
    qtype = lambda q: " ".join(q.lower().split()[:3])
    pool = {}
    for ex in ds:
        pool.setdefault(qtype(ex["question"]), []).append(ex["multiple_choice_answer"])
    out = []
    for i, ex in enumerate(ds):
        split = "eval" if i < n * eval_frac else "train"
        q, a = ex["question"], str(ex["multiple_choice_answer"]).strip().lower()
        st = [Seg("image", ex["image"])]
        if a in ("yes", "no"):
            out.append(_rec("vqav2_yesno", "image", split, st, {"type": "noul", "instructions": q}, int(a == "yes")))
        elif q.lower().startswith("how many") and a.isdigit() and int(a) <= 10:
            out.append(_rec("vqav2_count", "image", split, st, {"type": "score", "instructions": q,
                                                                 "criteria": [str(k) for k in range(11)]}, int(a)))
        else:
            cands = [x for x in set(pool[qtype(q)]) if x != a]
            if len(cands) < 3:
                continue
            opts = [a] + rng.sample(cands, 3); rng.shuffle(opts)
            out.append(_rec("vqav2_other", "image", split, st, {"type": "choice", "instructions": q,
                                                                 "criteria": {o: "" for o in opts}}, opts.index(a)))
    return _encode_images(jev, out)


# ------------------------------------------------------------------------------------------ audio
def _decode(a, sr=16000):
    import soundfile as sf, librosa
    y, s = sf.read(io.BytesIO(a["bytes"]), dtype="float32", always_2d=False)
    if y.ndim > 1:
        y = y.mean(1)
    return librosa.resample(y, orig_sr=s, target_sr=sr) if s != sr else y


def _encode_audio(jev, recs):
    cnt = {}
    for r in recs:
        if r["split"] == "eval":
            cnt[r["task"]] = cnt.get(r["task"], 0) + 1
            if cnt[r["task"]] <= N_RAW:
                r["raw"] = [Seg("audio", r["state"][0].data)]
    for s in range(0, len(recs), 32):
        chunk = recs[s:s + 32]
        feats = jev.audio_tower_features([r["state"][0].data[:16000 * 30] for r in chunk])
        for r, f in zip(chunk, feats):
            r["state"][0] = Seg("audio", f.cpu())
    return recs


def clotho_aqa(jev, n_tr=1500, n_ev=400):
    out = []
    for split, pat, n in (("train", "clotho_aqa_val_filtered-*", n_tr), ("eval", "clotho_aqa_test_filtered-*", n_ev)):
        files = f"hf://datasets/lmms-lab-audio/ClothoAQA/clotho_aqa/{pat}.parquet"
        ds = load_dataset("parquet", data_files=files, split="train")
        acol = next(c for c in ds.column_names if "audio" in c.lower() or c in ("context",))
        ds = ds.cast_column(acol, Audio(decode=False)).shuffle(seed=0)
        qcol = next(c for c in ds.column_names if "question" in c.lower())
        anscol = next(c for c in ds.column_names if c.lower() in ("answer", "answers", "label"))
        k = 0
        for ex in ds:
            if k >= n:
                break
            a = str(ex[anscol]).strip().lower()
            if a not in ("yes", "no"):
                continue
            out.append(_rec("clotho_aqa", "audio", split, [Seg("audio", _decode(ex[acol]))],
                            {"type": "noul", "instructions": ex[qcol]}, int(a == "yes"))); k += 1
    return _encode_audio(jev, out)


def urbansound(jev, shards=6, n_tr=1500, n_ev=400):
    files = sorted(f for f in [s.rfilename for s in __import__("huggingface_hub").HfApi().dataset_info("danavery/urbansound8K").siblings]
                   if f.endswith(".parquet"))[:shards]
    ds = load_dataset("parquet", data_files=[f"hf://datasets/danavery/urbansound8K/{f}" for f in files], split="train")
    ds = ds.cast_column("audio", Audio(decode=False)).shuffle(seed=0)
    classes = sorted(set(ds["class"]))
    crit = {c.replace("_", " "): "" for c in classes}
    out, k = [], {"train": 0, "eval": 0}
    for ex in ds:
        split = "eval" if str(ex["fold"]) == "10" else "train"
        if k[split] >= (n_ev if split == "eval" else n_tr):
            continue
        out.append(_rec("urbansound8k", "audio", split, [Seg("audio", _decode(ex["audio"]))],
                        {"type": "choice", "instructions": "Which sound is in the recording?", "criteria": crit},
                        classes.index(ex["class"]))); k[split] += 1
    return _encode_audio(jev, out)


def speech_commands(jev, n_tr=1500, n_ev=400):
    out = []
    for split, pat, n in (("train", "data/train/train-00000-of-*.parquet", n_tr), ("eval", "data/test/test-00000-of-*.parquet", n_ev)):
        ds = load_dataset("parquet", data_files=f"hf://datasets/pollen-robotics/speech-commands-v0.02/{pat}", split="train")
        ds = ds.cast_column("audio", Audio(decode=False)).shuffle(seed=0)
        words = sorted(set(ds["label"]) - {"_background_noise_", "_silence_", "_unknown_"})
        crit = {w: "" for w in words}
        k = 0
        for ex in ds:
            if k >= n:
                break
            if ex["label"] not in crit:
                continue
            out.append(_rec("speech_commands", "audio", split, [Seg("audio", _decode(ex["audio"]))],
                            {"type": "choice", "instructions": "Which word is spoken in the recording?", "criteria": crit},
                            words.index(ex["label"]))); k += 1
    return _encode_audio(jev, out)


# ------------------------------------------------------------------------------------------ video
def _frames(path, n=8):
    import cv2
    from PIL import Image
    cap = cv2.VideoCapture(path)
    tot = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    if tot <= 0:
        return None, None
    idx = np.linspace(0, tot - 1, n).round().astype(int)
    out = []
    for i in idx:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(i))
        ok, fr = cap.read()
        if ok:
            out.append(Image.fromarray(fr[:, :, ::-1]))
    cap.release()
    if len(out) < 2:
        return None, None
    eff_fps = (len(out) - 1) / max((idx[-1] - idx[0]) / fps, 1e-3)
    return out, eff_fps


def _video_rec(jev, task, split, path, qs, ys, keep_raw):
    fr, fps = _frames(path)
    if fr is None:
        return None
    r = _rec(task, "video", split, [Seg("video", _img_feats(jev, fr, FRAME_GRID), None, fps)], qs, ys)
    if keep_raw:
        r["raw"] = [Seg("video", fr, None, fps)]
    return r


def _parse_mc(text, kind):
    """TempCompass question text -> (instructions, options)."""
    lines = text.split("\n")
    if kind == "mc":
        opts = [re.sub(r"^[A-Z]\.\s*", "", l).strip() for l in lines[1:] if re.match(r"^[A-Z]\.", l)]
        return lines[0], opts
    opts = [re.sub(r"^Option \d+:\s*", "", l).strip() for l in lines[1:] if l.startswith("Option")]
    return lines[0], opts


def tempcompass(jev, eval_frac=0.25):
    root = "/content/pipe/raw/tempcompass"
    if not os.path.exists(root):
        z = hf_hub_download("lmms-eval/TempCompass", "tempcompass_videos.zip", repo_type="dataset")
        zipfile.ZipFile(z).extractall(root)
    vids = {os.path.splitext(os.path.basename(p))[0]: p for p in glob.glob(f"{root}/**/*.mp4", recursive=True)}
    per = {}
    for cfg, kind in (("multi-choice", "mc"), ("caption_matching", "cm"), ("yes_no", "yn")):
        for ex in load_dataset("lmms-eval/TempCompass", cfg, split="test"):
            vid = str(ex["video_id"])
            if vid not in vids:
                continue
            if kind == "yn":
                a = str(ex["answer"]).strip().lower()
                if a not in ("yes", "no"):
                    continue
                per.setdefault(vid, []).append(({"type": "noul", "instructions": ex["question"]}, int(a == "yes")))
            else:
                ins, opts = _parse_mc(ex["question"], kind)
                ans = re.sub(r"^([A-Z]\.|Option \d+:)\s*", "", str(ex["answer"])).strip()
                if ans not in opts or len(set(opts)) != len(opts):
                    continue
                per.setdefault(vid, []).append(({"type": "choice", "instructions": ins, "criteria": {o: "" for o in opts}},
                                                opts.index(ans)))
    keys = sorted(per); random.Random(0).shuffle(keys)
    n_ev = int(len(keys) * eval_frac)
    out = []
    for i, vid in enumerate(keys):
        split = "eval" if i < n_ev else "train"
        qa = per[vid]
        for s in range(0, len(qa), 4):                      # up to 4 questions per state / call
            chunk = qa[s:s + 4]
            r = _video_rec(jev, "tempcompass", split, vids[vid], [q for q, _ in chunk], [y for _, y in chunk],
                           keep_raw=split == "eval" and i < N_RAW)
            if r:
                out.append(r)
    return out


def mvbench_mini(jev, eval_frac=0.25):
    import csv
    root = snapshot_download("firstep-ai/MV-Bench-mini", repo_type="dataset")
    rows = list(csv.DictReader(open(f"{root}/resampled_data.csv", encoding="utf-8-sig")))
    random.Random(0).shuffle(rows)
    out = []
    for i, r in enumerate(rows):
        path = os.path.join(root, r["relpath"])
        if not os.path.exists(path):
            continue
        opts = [str(o) for o in eval(r["candidates"])]
        if r["answer"] not in opts or len(set(opts)) != len(opts):
            continue
        split = "eval" if i < len(rows) * eval_frac else "train"
        rec = _video_rec(jev, "mvbench_mini", split, path,
                         {"type": "choice", "instructions": r["question"], "criteria": {o: "" for o in opts}},
                         opts.index(r["answer"]), keep_raw=split == "eval" and i < N_RAW)
        if rec:
            out.append(rec)
    return out


def ucf101_subset(jev):
    root = "/content/pipe/raw/ucf"
    if not os.path.exists(root):
        tarfile.open(hf_hub_download("sayakpaul/ucf101-subset", "UCF101_subset.tar.gz", repo_type="dataset")).extractall(root)
    paths = glob.glob(f"{root}/**/*.avi", recursive=True)
    cls = lambda p: os.path.basename(os.path.dirname(p))
    classes = sorted(set(cls(p) for p in paths))
    crit = {re.sub(r"(?<!^)([A-Z])", r" \1", c).lower(): "" for c in classes}
    out, n_raw = [], 0
    for p in sorted(paths):
        split = "eval" if "/test/" in p else "train" if "/train/" in p else "eval"
        r = _video_rec(jev, "ucf101", split, p, {"type": "choice", "instructions": "Which action is shown in the video?",
                                                  "criteria": crit}, classes.index(cls(p)), keep_raw=split == "eval" and n_raw < N_RAW)
        if r:
            n_raw += int(r["raw"] is not None)
            out.append(r)
    return out


def kinetics_mini(jev, eval_frac=0.25):
    root = snapshot_download("nateraw/kinetics-mini", repo_type="dataset")
    paths = sorted(glob.glob(f"{root}/**/*.mp4", recursive=True))
    cls = lambda p: os.path.basename(os.path.dirname(p))
    classes = sorted(set(cls(p) for p in paths))
    crit = {c.replace("_", " "): "" for c in classes}
    random.Random(0).shuffle(paths)
    out = []
    for i, p in enumerate(paths):
        split = "eval" if i < len(paths) * eval_frac else "train"
        r = _video_rec(jev, "kinetics_mini", split, p, {"type": "choice", "instructions": "Which activity is shown in the video?",
                                                         "criteria": crit}, classes.index(cls(p)), keep_raw=split == "eval" and i < N_RAW)
        if r:
            out.append(r)
    return out


BUILDERS = {"scienceqa": scienceqa, "ai2d": ai2d, "vqav2": vqav2, "clotho_aqa": clotho_aqa, "urbansound": urbansound,
            "speech_commands": speech_commands, "tempcompass": tempcompass, "mvbench_mini": mvbench_mini,
            "ucf101": ucf101_subset, "kinetics_mini": kinetics_mini}