import os
import glob
import math
from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from PIL import Image

# canonical rain-type ordering shared by training/eval modules
RAIN_VARIANTS = ("raindrop", "rainstreak", "rainstreak_raindrop")
VARIANT_TO_IDX = {name: idx for idx, name in enumerate(RAIN_VARIANTS)}

# ----------------------------
# Dataset & dataloaders
# ----------------------------

class PairedRainDataset(Dataset):
    def __init__(
        self,
        root: str = "./dataset",
        split: str = "train_set",
        load_variant: str = "any",
        image_size: Optional[int] = None,
    ):
        """
        Args:
            root: path to 'dataset'
            split: 'train_set' or 'test_set'
            load_variant:
              - "any": pick a random variant per sample
              - "raindrop" | "rainstreak" | "rainstreak_raindrop": fixed variant
              - "all": include ALL three variants per scene (triples)
            image_size: if not None, resize GT and rainy images to [image_size, image_size]
        """
        super().__init__()
        assert split in ("train_set", "test_set")
        assert load_variant in ("any", "raindrop", "rainstreak", "rainstreak_raindrop", "all")
        self.root = root
        self.split = split
        self.load_variant = load_variant
        self.image_size = image_size

        self.gt_dir   = os.path.join(root, split, "gt")
        self.rd_dir   = os.path.join(root, split, "raindrop")
        self.rs_dir   = os.path.join(root, split, "rainstreak")
        self.both_dir = os.path.join(root, split, "rainstreak_raindrop")

        self.ids = sorted([os.path.basename(p) for p in glob.glob(os.path.join(self.gt_dir, "*.png"))])

        if len(self.ids) == 0:
            raise FileNotFoundError(
                f"No PNG files found in {self.gt_dir}. "
                f"Resolved path: {os.path.abspath(self.gt_dir)}\n"
                f"Check that your dataset layout is:\n"
                f"{self.root}/{self.split}/gt/*.png and images actually exist."
            )

        self.variants = RAIN_VARIANTS

        if self.load_variant == "all":
            self.samples = [(name, v) for name in self.ids for v in self.variants]
        else:
            self.samples = [(name, None) for name in self.ids]

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

    def _load_img(self, path: str) -> torch.Tensor:
        with Image.open(path).convert("RGB") as im:
            x = torch.from_numpy(np.array(im)).float() / 255.0
        x = x.permute(2, 0, 1)  # [C,H,W]
        if self.image_size is not None:
            x = F.interpolate(
                x.unsqueeze(0),
                size=(self.image_size, self.image_size),
                mode="bilinear",
                align_corners=False,
            ).squeeze(0)
        return x

    def __getitem__(self, idx: int):
        name, preset_variant = self.samples[idx]

        # load gt
        gt = self._load_img(os.path.join(self.gt_dir, name))

        # decide variant
        if self.load_variant == "any":
            variant = np.random.choice(self.variants)
        elif self.load_variant == "all":
            variant = preset_variant
        else:
            variant = self.load_variant

        if variant == "raindrop":
            ypath = os.path.join(self.rd_dir, name)
        elif variant == "rainstreak":
            ypath = os.path.join(self.rs_dir, name)
        else:
            ypath = os.path.join(self.both_dir, name)

        y = self._load_img(ypath)

        return {"y": y, "gt": gt, "variant": variant}


def get_dataloaders(
    data_root: str,
    batch_size: int = 4,
    image_size: Optional[int] = None,
    train_variant: str = "all",
    test_variant: str = "rainstreak_raindrop",
    train_num_workers: int = 4,
    test_num_workers: int = 2,
    pin_memory: bool = True,
):
    """
    Build train/test dataloaders.

    image_size:
        - if None: keep original resolution
        - if int: resize images to [image_size, image_size]
    """
    # Train: include ALL three variants per scene
    train_ds = PairedRainDataset(
        data_root,
        "train_set",
        load_variant=train_variant,
        image_size=image_size,
    )
    # Test: rainstreak_raindrop only
    test_ds  = PairedRainDataset(
        data_root,
        "test_set",
        load_variant=test_variant,
        image_size=image_size,
    )

    train_loader = DataLoader(
        train_ds, batch_size=batch_size, shuffle=True,
        num_workers=train_num_workers, pin_memory=pin_memory,
    )
    test_loader  = DataLoader(
        test_ds,  batch_size=1, shuffle=False,
        num_workers=test_num_workers, pin_memory=pin_memory,
    )
    return train_loader, test_loader