"""`adverse_weather_canny_webbed` data_library: the Webbed (tar-shard) version of the
adverse-weather rain/snow canny data, with the UMT5 caption embedding read straight from
the in-tar `text_emb` (lat_language) stream — no external captions store at train time.

Mirrors AdverseWeatherCannyDataset's batch contract (top-level t5_text_embeddings /
t5_text_mask / neg_*), but the per-sample positive embedding comes from the tar instead of
the hash store. The single uncond/negative embedding is loaded once from the store (it's a
fixed vector; cheap and not per-sample). Canny is computed on the fly downstream from the
RGB frames, exactly as in the local Unified path.
"""
import json

import torch

from custom.dataset.anydata_dataset import AnyDataset
from custom.dataset.registry import register_data_library
from imaginaire.utils import log


class AdverseWeatherCannyWebbedDataset(AnyDataset):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._neg = None      # cached uncond embedding (loaded once)
        self._miss = 0

    @property
    def neg(self):
        if self._neg is None:
            from custom.dataset.adverse_weather_captions import AdverseWeatherCaptions
            n = AdverseWeatherCaptions().negative()
            self._neg = (n["t5_text_embeddings"][0], n["t5_text_mask"][0])
        return self._neg

    def __getitem__(self, index):
        d = super().__getitem__(index)
        if d is None:
            return None
        lat = d.get("anydata", {}).get("lat_language", {})
        neg_emb, neg_mask = self.neg
        if lat:
            entry = next(iter(lat.values()))   # {t5_text_embeddings (512,4096), t5_text_mask, prompt}
            d["t5_text_embeddings"] = torch.as_tensor(entry["t5_text_embeddings"]).to(torch.bfloat16)
            d["t5_text_mask"] = torch.as_tensor(entry["t5_text_mask"]).long()
            d["prompt"] = str(entry.get("prompt", d.get("prompt", "")))
        else:
            # No embedding stream for this sample -> fall back to uncond (keeps shapes consistent).
            self._miss += 1
            if self._miss <= 5 or self._miss % 5000 == 0:
                log.warning(f"[adverse_weather_canny_webbed] no lat_language for sample {index} "
                            f"(miss={self._miss}); using uncond.", rank0_only=False)
            d["t5_text_embeddings"] = neg_emb.to(torch.bfloat16)
            d["t5_text_mask"] = neg_mask.long()
        d["neg_t5_text_embeddings"] = neg_emb.to(torch.bfloat16)
        d["neg_t5_text_mask"] = neg_mask.long()
        return d


register_data_library("adverse_weather_canny_webbed", AdverseWeatherCannyWebbedDataset)