"""
Generic video clip dataset wrapper.

Pulls fixed-length clips of length L from any underlying frame-indexed
trajectory dataset. Used for video training (temporal Kalman filter loss
ablations) and video eval.

Each item is a dict with stacked-frame tensors:
    image       : (L, 3, H, W)
    depth       : (L, 1, H, W)
    mask        : (L, 1, H, W)
    sequence_id : str
    frame_ids   : list[int]
"""
from __future__ import annotations

import json
import os
import re
from typing import Sequence

import cv2
import numpy as np
import torch
from omegaconf import DictConfig
from torchvision.transforms import Compose

from ppd.utils.logger import Log


def _read_rgb(path: str) -> np.ndarray:
    rgb = cv2.imread(path)
    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
    return (rgb / 255.0).astype(np.float32)


def _read_depth_npy(path: str) -> np.ndarray:
    d = np.load(path).astype(np.float32)
    if d.ndim == 3 and d.shape[-1] == 1:
        d = d[..., 0]
    return d


class TartanAirVideoClip:
    """Sample fixed-length clips from extracted TartanAir trajectories.

    Expected layout (from the V1 zips we downloaded under
    /mnt/sig/datasets/train/tartanair/extracted/):

        <data_root>/<scene>/<difficulty>/<P###>/image_left/NNNNNN_left.png
        <data_root>/<scene>/<difficulty>/<P###>/depth_left/NNNNNN_left_depth.npy
    """

    def __init__(
        self,
        data_root: str,
        clip_length: int = 8,
        stride: int = 1,
        scenes: Sequence[str] | None = None,
        transforms: list | None = None,
        max_depth: float = 80.0,
    ):
        self.cfg = DictConfig(
            dict(data_root=data_root, clip_length=clip_length, stride=stride, max_depth=max_depth)
        )
        self.dataset_name = "tartanair_video"
        self.transform = Compose(transforms or [])
        self._build(scenes)

    def _build(self, scenes: Sequence[str] | None) -> None:
        root = self.cfg.data_root
        L = self.cfg.clip_length
        S = self.cfg.stride
        self.clips: list[tuple[str, list[str], list[str]]] = []
        if not os.path.isdir(root):
            Log.warn(f"TartanAir video root missing: {root}")
            return
        for scene in sorted(os.listdir(root)):
            if scenes is not None and scene not in scenes:
                continue
            scene_dir = os.path.join(root, scene)
            if not os.path.isdir(scene_dir):
                continue
            for difficulty in ("Easy", "Hard"):
                diff_dir = os.path.join(scene_dir, difficulty)
                if not os.path.isdir(diff_dir):
                    continue
                for traj in sorted(os.listdir(diff_dir)):
                    img_dir = os.path.join(diff_dir, traj, "image_left")
                    dpt_dir = os.path.join(diff_dir, traj, "depth_left")
                    if not (os.path.isdir(img_dir) and os.path.isdir(dpt_dir)):
                        continue
                    frames = sorted(
                        f for f in os.listdir(img_dir) if f.endswith("_left.png")
                    )
                    if len(frames) < L * S:
                        continue
                    for start in range(0, len(frames) - L * S + 1, max(L // 2, 1)):
                        idx = [start + i * S for i in range(L)]
                        rgb_paths = [os.path.join(img_dir, frames[i]) for i in idx]
                        dpt_paths = [
                            os.path.join(
                                dpt_dir, frames[i].replace("_left.png", "_left_depth.npy")
                            )
                            for i in idx
                        ]
                        seq_id = f"{scene}/{difficulty}/{traj}"
                        self.clips.append((seq_id, rgb_paths, dpt_paths))
        Log.info(f"TartanAir video: {len(self.clips)} clips")

    def __len__(self) -> int:
        return len(self.clips)

    def __getitem__(self, idx: int) -> dict:
        seq_id, rgb_paths, dpt_paths = self.clips[idx]
        images = []
        depths = []
        masks = []
        for r, d in zip(rgb_paths, dpt_paths):
            rgb = _read_rgb(r)
            depth = _read_depth_npy(d)
            mask = np.logical_and(depth > 0.1, ~np.isnan(depth)) & ~np.isinf(depth)
            mask = np.logical_and(mask, depth < self.cfg.max_depth)
            sample = {"image": rgb, "depth": depth, "mask": mask.astype(np.uint8)}
            sample = self.transform(sample)
            images.append(sample["image"])
            depths.append(sample["depth"])
            masks.append(sample["mask"])
        return {
            "image": np.stack(images, axis=0),
            "depth": np.stack(depths, axis=0),
            "mask": np.stack(masks, axis=0),
            "dataset_name": self.dataset_name,
            "sequence_id": seq_id,
            "frame_ids": list(range(len(images))),
        }


class BonnRGBDVideoClip:
    """ Bonn dynamic RGB-D dataset clip loader.

    Each unzipped sequence has:
        rgb/<timestamp>.png
        depth/<timestamp>.png  (16-bit, mm)
        rgb.txt, depth.txt     (timestamps)
        associated.txt         (rgb-depth pairing, optional)

    For simplicity, we pair frames by index after sorting.
    """

    def __init__(
        self,
        data_root: str,
        clip_length: int = 8,
        stride: int = 1,
        sequences: Sequence[str] | None = None,
        transforms: list | None = None,
    ):
        self.cfg = DictConfig(
            dict(data_root=data_root, clip_length=clip_length, stride=stride)
        )
        self.dataset_name = "bonn_rgbd_video"
        self.transform = Compose(transforms or [])
        self._build(sequences)

    def _build(self, sequences) -> None:
        root = self.cfg.data_root
        L = self.cfg.clip_length
        S = self.cfg.stride
        self.clips: list[tuple[str, list[str], list[str]]] = []
        if not os.path.isdir(root):
            Log.warn(f"Bonn root missing: {root}")
            return
        # bonn sequences live in subdirectories
        for d in sorted(os.listdir(root)):
            if sequences is not None and d not in sequences:
                continue
            seq_dir = os.path.join(root, d)
            if not os.path.isdir(seq_dir):
                continue
            rgb_dir = os.path.join(seq_dir, "rgb")
            dpt_dir = os.path.join(seq_dir, "depth")
            if not (os.path.isdir(rgb_dir) and os.path.isdir(dpt_dir)):
                continue
            rgb_files = sorted(f for f in os.listdir(rgb_dir) if f.endswith(".png"))
            dpt_files = sorted(f for f in os.listdir(dpt_dir) if f.endswith(".png"))
            n = min(len(rgb_files), len(dpt_files))
            if n < L * S:
                continue
            for start in range(0, n - L * S + 1, max(L // 2, 1)):
                idx = [start + i * S for i in range(L)]
                rgb_paths = [os.path.join(rgb_dir, rgb_files[i]) for i in idx]
                dpt_paths = [os.path.join(dpt_dir, dpt_files[i]) for i in idx]
                self.clips.append((d, rgb_paths, dpt_paths))
        Log.info(f"Bonn video: {len(self.clips)} clips from {len(set(c[0] for c in self.clips))} sequences")

    def __len__(self) -> int:
        return len(self.clips)

    def __getitem__(self, idx: int) -> dict:
        seq_id, rgb_paths, dpt_paths = self.clips[idx]
        images, depths, masks = [], [], []
        for r, d in zip(rgb_paths, dpt_paths):
            rgb = _read_rgb(r)
            # 16-bit PNG, mm scale → meters /5000 by Bonn convention
            depth = cv2.imread(d, cv2.IMREAD_ANYDEPTH).astype(np.float32) / 5000.0
            mask = np.logical_and(depth > 0.01, depth < 10.0)
            sample = {"image": rgb, "depth": depth, "mask": mask.astype(np.uint8)}
            sample = self.transform(sample)
            images.append(sample["image"])
            depths.append(sample["depth"])
            masks.append(sample["mask"])
        return {
            "image": np.stack(images, axis=0),
            "depth": np.stack(depths, axis=0),
            "mask": np.stack(masks, axis=0),
            "dataset_name": self.dataset_name,
            "sequence_id": seq_id,
            "frame_ids": list(range(len(images))),
        }