# bora_dataset.py -- dataset for the Bora latent shards (latents_256/ on Ami1ne/Bora)
#
# mode="clip":  whole clip, [16,10,32,32]
# mode="block": block_len latent frames, [16,block_len,32,32]; block starts step by block_len - overlap (0, 3, 6 for 4/1)
# Latents are already normalized. Decode with vae.decode(z * ds.latents_std + ds.latents_mean).
#
# Manifest: latents_256_manifest_labeled.csv (one row per unique clip, written by cell 6 of bora_load_latents.ipynb)
#   global_idx, latent_file, latent_row, clip_path, source_file, clip_idx, source, category, n_copies, caption_ids
#   caption_ids is a JSON list of indices into captions.json; index 0 is the empty caption "" (for classifier-free guidance).
# Falls back to latents_256_manifest.csv / latents_256_manifests/part_*.csv (no captions, duplicates kept) if the labeled one is absent.
#
# Validation: a whole source video is held out when crc32(source_file) % val_mod == 0 (val_mod=100 -> about 1% of the videos).
#
# Captions per clip are ordered most reliable first; an item takes caption 0 with probability primary_p, else a uniform pick.
#
# text_embeds: optional path to a safetensors file with "emb" [N_captions, L, D] float16 and "mask" [N_captions, L] bool
# (written by cell 9). The item then carries text_emb [L,D] and text_mask [L] for its chosen caption.
import json
import zlib
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from safetensors import safe_open
from torch.utils.data import Dataset

LABELED = "latents_256_manifest_labeled.csv"
MERGED = "latents_256_manifest.csv"
PARTS_DIR = "latents_256_manifests"
CAPTIONS = "captions.json"


def source_key(df):
    """A string per row that identifies the source video (the split must keep every clip of a video on one side)."""
    if "source_file" in df.columns:
        return df["source_file"]
    # clip_path looks like clips_5s/shard_00000/0000000__pexels-25560546__c0000.mp4 -> pexels-25560546
    return df["clip_path"].str.rsplit("/", n=1).str[-1].str.split("__").str[1]


def is_val(source, val_mod):
    return np.fromiter((zlib.crc32(s.encode()) % val_mod == 0 for s in source), dtype=bool, count=len(source))


class BoraLatents(Dataset):
    def __init__(self, root, split="train", mode="block", block_len=4, overlap=1, random_block=True,
                 dtype=torch.float32, val_mod=100, manifest=None, captions_file=None, text_embeds=None,
                 random_caption=None, primary_p=0.5):
        assert mode in ("clip", "block") and split in ("train", "val") and 0 <= overlap < block_len
        root = Path(root)
        self.root, self.mode, self.dtype = root, mode, dtype
        self.block_len, self.random_block = block_len, random_block
        self.random_caption = random_block if random_caption is None else random_caption
        self.primary_p = primary_p  # chance of taking the first (most reliable) caption; otherwise uniform over all of them

        # ---- manifest ----
        if manifest is not None:
            manifest = Path(manifest)
        elif (root / LABELED).exists():
            manifest = root / LABELED
        elif (root / MERGED).exists():
            manifest = root / MERGED
        if manifest is not None:
            df = pd.read_csv(manifest, dtype=str, keep_default_na=False)
            self.manifest_source = manifest.name
        else:  # per-shard ledger
            parts = sorted((root / PARTS_DIR).glob("part_*.csv"))
            assert parts, f"no manifest found under {root}"
            df = pd.concat([pd.read_csv(p, dtype=str, keep_default_na=False) for p in parts], ignore_index=True)
            self.manifest_source = "part_*.csv"
        df["_g"] = df["global_idx"].astype(int)
        df["latent_row"] = df["latent_row"].astype(int)
        df = df.sort_values("_g").reset_index(drop=True)
        self.manifest = df
        self.labeled = "caption_ids" in df.columns

        # ---- captions ----
        self.captions = None
        if self.labeled:
            cf = Path(captions_file) if captions_file else (manifest.parent / CAPTIONS)
            assert cf.exists(), f"{cf} not found (it is written next to the labeled manifest)"
            self.captions = json.load(open(cf, encoding="utf-8"))
            assert self.captions[0] == "", "captions.json must start with the empty caption"

        # ---- split by source video ----
        val = is_val(source_key(df), val_mod)
        sub = (df[val] if split == "val" else df[~val]).reset_index(drop=True)
        assert len(sub), f"empty {split} split"
        self.n_videos = int(source_key(sub).nunique())

        # ---- shards ----
        self._handles = {}  # opened lazily, so every DataLoader worker owns its own handles
        first = sub["latent_file"].iloc[0]
        self.latents_mean, self.latents_std = self._read_stats(first)
        _, self.C, self.T, self.H, self.W = self._open(first).get_slice("latents").get_shape()
        self.starts = list(range(0, self.T - block_len + 1, block_len - overlap))

        if mode == "block" and not random_block:  # eval: every block of every clip, in order
            sub = sub.loc[sub.index.repeat(len(self.starts))].reset_index(drop=True)
            self._bidx = [i % len(self.starts) for i in range(len(sub))]
        self._file = sub["latent_file"].tolist()
        self._row = sub["latent_row"].tolist()
        self._gidx = sub["_g"].tolist()
        self._cids = [json.loads(s) for s in sub["caption_ids"]] if self.labeled else None

        # ---- precomputed text embeddings ----
        self.text_embeds = Path(text_embeds) if text_embeds else None
        self._emb = None
        if self.text_embeds is not None:
            assert self.labeled, "text_embeds needs the labeled manifest"
            with safe_open(str(self.text_embeds), framework="pt") as f:
                n, self.L, self.D = f.get_slice("emb").get_shape()
            assert n == len(self.captions), f"{self.text_embeds} has {n} rows, captions.json has {len(self.captions)}"

    def _open(self, rel):
        h = self._handles.get(rel)
        if h is None:
            h = self._handles[rel] = safe_open(str(self.root / rel), framework="pt")
        return h

    def _read_stats(self, rel):
        meta = self._open(rel).metadata()
        shape = (1, -1, 1, 1, 1)
        return (torch.tensor(json.loads(meta["latents_mean"])).view(shape),
                torch.tensor(json.loads(meta["latents_std"])).view(shape))

    def __len__(self):
        return len(self._file)

    def __getitem__(self, i):
        row = self._row[i]
        z = self._open(self._file[i]).get_slice("latents")[row:row + 1][0]  # [16,10,32,32] float16, only this row is read
        item = {"global_idx": self._gidx[i]}
        if self.mode == "block":
            b = int(torch.randint(len(self.starts), (1,))) if self.random_block else self._bidx[i]
            s = self.starts[b]
            z = z[:, s:s + self.block_len]
            item["block_index"], item["start"] = b, s
        item["latents"] = z.to(self.dtype)
        if self._cids is not None:
            cids = self._cids[i]
            if self.random_caption and len(cids) > 1 and torch.rand(()) >= self.primary_p:
                cid = cids[int(torch.randint(len(cids), (1,)))]
            else:
                cid = cids[0]
            item["caption_id"] = cid
            item["text"] = self.captions[cid]
            if self.text_embeds is not None:
                if self._emb is None:
                    self._emb = safe_open(str(self.text_embeds), framework="pt")
                item["text_emb"] = self._emb.get_slice("emb")[cid:cid + 1][0]    # [L,D] float16
                item["text_mask"] = self._emb.get_slice("mask")[cid:cid + 1][0]  # [L] bool
        return item