"""`adverse_weather_canny` data_library: AnyData Unified loader + precomputed
UMT5 caption embeddings for the locally-downloaded GAIA adverse_weather subset.

Each scene's `metadata.json` `language.prompt` (surfaced as `data_dict['prompt']`
by AnyDataset.supplement_sample) is mapped to its precomputed UMT5-XXL embedding
under `captions_t5_all/` via `AdverseWeatherCaptions`, and attached to the sample
as top-level tensors:

    t5_text_embeddings      (512, 4096)   bf16
    t5_text_mask            (512,)        i64
    neg_t5_text_embeddings  (512, 4096)   bf16   (uncond, for CFG dropout)
    neg_t5_text_mask        (512,)        i64

`any4d_collate` stacks these to (B, 512, 4096) / (B, 512), and
`a4d_vae_wan.encode_text_as_needed` consumes them via its fast-path — so no UMT5
text encoder is loaded at train time (set the experiment's text_encoder_path='').

Captions that aren't found in the store fall back to the uncond embedding so the
keys are always present (keeps batch shapes consistent); misses are logged
sparsely. Set ADVERSE_WEATHER_DATA_ROOT to the data root (default /home/ubuntu/data;
inside the training container it is /workspace/data).

Registered as data_library 'adverse_weather_canny' (see custom/dataset/registry.py);
the experiment imports this module for its registration side-effect.
"""

import json
import os

from custom.dataset.anydata_dataset import AnyDataset
from custom.dataset.adverse_weather_captions import AdverseWeatherCaptions, DEFAULT_DATA_ROOT
from custom.dataset.registry import register_data_library
from imaginaire.utils import log

# Scene dirs live under <data_root>/<_SCENES_SUBDIR>/<chunk>.parquet/<scene>/metadata.json
_SCENES_SUBDIR = 'cv_unified_adverse_weather/videos/parquet_all_weather'


class AdverseWeatherCannyDataset(AnyDataset):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._caps = None
        self._hit = 0
        self._miss = 0
        self._scenes_root = os.path.join(DEFAULT_DATA_ROOT, _SCENES_SUBDIR)
        self._scene_prompt_cache = {}
        # Weather-transfer validation: if the YAML's `custom:` block sets
        # `weather_transfer_prompts` (path to a JSON list of rain/snow captions that
        # already exist in the t5 store), override each sample's caption with one of
        # these (cycled deterministically by index). Paired with the clear-scene val
        # split, this asks the model to render rain/snow onto non-adverse structure.
        self._transfer_prompts = None
        tp = self.custom.get('weather_transfer_prompts') if isinstance(self.custom, dict) else None
        if tp:
            try:
                with open(tp) as f:
                    self._transfer_prompts = json.load(f)
                log.info(f'[adverse_weather_canny] weather-transfer val: '
                         f'{len(self._transfer_prompts)} prompts from {tp}', rank0_only=True)
            except Exception as e:
                log.warning(f'[adverse_weather_canny] failed to load '
                            f'weather_transfer_prompts {tp!r}: {e!r}')

    @property
    def caps(self):
        # Lazily build the text->hash map + embedding store (per dataloader worker).
        if self._caps is None:
            self._caps = AdverseWeatherCaptions()
        return self._caps

    def _scene_prompt(self, d):
        """Read the caption straight from the scene's metadata.json (language.prompt).

        We load labels=[rgb] only (no lowdim), so the prompt isn't surfaced via the
        'language' modality; instead derive the scene dir from metadata['path']
        ('<chunk>.parquet/<scene>/{label}/<cam>') and read its metadata.json.
        """
        md = d.get('anydata', {}).get('metadata', {})
        path = md.get('path')
        if not isinstance(path, str):
            return None
        scene = path.split('/{label}')[0] if '/{label}' in path else path.rsplit('/rgb/', 1)[0]
        if scene in self._scene_prompt_cache:
            return self._scene_prompt_cache[scene]
        prompt = None
        try:
            mj = os.path.join(self._scenes_root, scene, 'metadata.json')
            prompt = json.load(open(mj)).get('language', {}).get('prompt')
        except Exception:
            prompt = None
        if len(self._scene_prompt_cache) < 50000:
            self._scene_prompt_cache[scene] = prompt
        return prompt

    def _attach_text(self, d, force_prompt=None):
        prompt = force_prompt if isinstance(force_prompt, str) else self._scene_prompt(d)
        if isinstance(prompt, str):
            d['prompt'] = prompt  # surface the real caption (was the YAML fallback)
        emb = self.caps.get(prompt) if isinstance(prompt, str) else None

        if emb is None:
            self._miss += 1
            if self._miss <= 5 or self._miss % 5000 == 0:
                log.warning(
                    f'[adverse_weather_canny] caption not in t5 store '
                    f'(miss={self._miss}, hit={self._hit}): {str(prompt)[:90]!r}',
                    rank0_only=False)
            neg = self.caps.negative()
            d['t5_text_embeddings'] = neg['t5_text_embeddings'][0]
            d['t5_text_mask'] = neg['t5_text_mask'][0]
            d['neg_t5_text_embeddings'] = neg['t5_text_embeddings'][0]
            d['neg_t5_text_mask'] = neg['t5_text_mask'][0]
            return d

        self._hit += 1
        d['t5_text_embeddings'] = emb['t5_text_embeddings'][0]
        d['t5_text_mask'] = emb['t5_text_mask'][0]
        d['neg_t5_text_embeddings'] = emb['t5_text_embeddings_neg'][0]
        d['neg_t5_text_mask'] = emb['t5_text_mask_neg'][0]
        return d

    def __getitem__(self, index):
        d = super().__getitem__(index)
        try:
            force = None
            if self._transfer_prompts:
                force = self._transfer_prompts[index % len(self._transfer_prompts)]
            d = self._attach_text(d, force_prompt=force)
        except Exception as e:
            log.warning(f'[adverse_weather_canny] _attach_text failed: {e!r}',
                        rank0_only=False)
        return d


register_data_library('adverse_weather_canny', AdverseWeatherCannyDataset)