"""
dataset.py — RadioMapSeer loader for radio coverage prediction.

Input  (2 channels): city building map + transmitter location
Target (1 channel) : radio coverage / path-loss map ("gain")

Expected layout under data_root:
    png/buildings_complete/{map}.png       buildings vs background
    png/antennas/{map}_{tx}.png            transmitter marker
    gain/DPM/{map}_{tx}.png                coverage heatmap (target)

Naming: map id 0..699, tx id 0..79  ->  "{map}_{tx}.png"
"""

from pathlib import Path

import numpy as np
import torch
from torch.utils.data import Dataset
from PIL import Image

# dataset.py lives in src/, so the project's data/ folder is one level up.
DEFAULT_DATA_ROOT = Path(r"D:\dataset\RadioMapSeer")


def _load_gray(path: Path, img_size: int | None = None) -> np.ndarray:
    """Load a PNG as float32 in [0, 1]; optionally resize to img_size."""
    img = Image.open(path).convert("L")
    if img_size is not None and img.size != (img_size, img_size):
        img = img.resize((img_size, img_size), Image.BILINEAR)
    return np.asarray(img, dtype=np.float32) / 255.0


class RadioMapSeerDataset(Dataset):
    """Returns (input[2,H,W], target[1,H,W]) tensors in [0, 1]."""

    SPLITS = {
        "train": range(0, 500),
        "val":   range(500, 600),
        "test":  range(600, 700),
    }

    def __init__(self, data_root=DEFAULT_DATA_ROOT, split="train",
                 sim="DPM", img_size=256, max_samples=None):
        self.data_root = Path(data_root)
        self.img_size = img_size

        self.gain_dir  = self.data_root / "gain" / sim
        self.build_dir = self.data_root / "png" / "buildings_complete"
        self.ant_dir   = self.data_root / "png" / "antennas"

        for d in (self.gain_dir, self.build_dir, self.ant_dir):
            if not d.is_dir():
                raise FileNotFoundError(f"Expected folder not found: {d}")

        if split not in self.SPLITS:
            raise ValueError(f"split must be train/val/test, got '{split}'")
        allowed = set(self.SPLITS[split])

        # Build the sample index by scanning the gain folder for "{map}_{tx}.png".
        self.samples = []
        for f in sorted(self.gain_dir.glob("*.png")):
            try:
                map_id, tx_id = (int(x) for x in f.stem.split("_"))
            except ValueError:
                continue
            if map_id in allowed:
                self.samples.append((map_id, tx_id))

        if not self.samples:
            raise RuntimeError(f"No samples found in {self.gain_dir} — check naming.")

        if max_samples is not None:
            self.samples = self.samples[:max_samples]

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

    def __getitem__(self, idx):
        map_id, tx_id = self.samples[idx]

        buildings = _load_gray(self.build_dir / f"{map_id}.png", self.img_size)
        antenna   = _load_gray(self.ant_dir   / f"{map_id}_{tx_id}.png", self.img_size)
        gain      = _load_gray(self.gain_dir  / f"{map_id}_{tx_id}.png", self.img_size)

        x = np.stack([buildings, antenna], axis=0)   # [2, H, W]
        y = gain[None, :, :]                          # [1, H, W]
        return torch.from_numpy(x), torch.from_numpy(y)


if __name__ == "__main__":
    # Self-test: confirms paths, naming, shapes, and value ranges.
    print("Data root:", DEFAULT_DATA_ROOT)
    ds = RadioMapSeerDataset(split="train", img_size=256)
    print(f"Train samples: {len(ds)}")

    x, y = ds[0]
    print("Input shape :", tuple(x.shape), "| min/max:", float(x.min()), float(x.max()))
    print("Target shape:", tuple(y.shape), "| min/max:", float(y.min()), float(y.max()))

    for s in ("val", "test"):
        print(f"{s} samples:", len(RadioMapSeerDataset(split=s, img_size=256)))