"""MRMS data loading utilities matching the official NowcastNet contract."""

from __future__ import annotations

import struct
import zlib
from pathlib import Path

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


def _read_png_gray16(path: Path) -> np.ndarray:
    """Read the 16-bit grayscale PNG subset used by MRMS."""
    raw = path.read_bytes()
    if raw[:8] != b"\x89PNG\r\n\x1a\n":
        raise ValueError(f"Not a PNG file: {path}")
    pos = 8
    idat: list[bytes] = []
    width = height = bit_depth = color_type = None
    while pos < len(raw):
        length = struct.unpack(">I", raw[pos:pos + 4])[0]
        kind = raw[pos + 4:pos + 8]
        payload = raw[pos + 8:pos + 8 + length]
        pos += length + 12
        if kind == b"IHDR":
            width, height, bit_depth, color_type = struct.unpack(">IIBB", payload[:10])
        elif kind == b"IDAT":
            idat.append(payload)
        elif kind == b"IEND":
            break
    if width is None or height is None or bit_depth != 16 or color_type != 0:
        raise ValueError(f"Expected 16-bit grayscale PNG: {path}")

    decoded = zlib.decompress(b"".join(idat))
    row_bytes = width * 2
    previous = np.zeros(row_bytes, dtype=np.uint8)
    rows: list[np.ndarray] = []
    cursor = 0
    for _ in range(height):
        filter_type = decoded[cursor]
        row = np.frombuffer(decoded[cursor + 1:cursor + 1 + row_bytes], dtype=np.uint8).copy()
        cursor += row_bytes + 1
        for index in range(row_bytes):
            left = row[index - 2] if index >= 2 else 0
            up = previous[index]
            upper_left = previous[index - 2] if index >= 2 else 0
            if filter_type == 1:
                value = left
            elif filter_type == 2:
                value = up
            elif filter_type == 3:
                value = (int(left) + int(up)) // 2
            elif filter_type == 4:
                predictor = int(left) + int(up) - int(upper_left)
                distances = (abs(predictor - int(left)), abs(predictor - int(up)), abs(predictor - int(upper_left)))
                value = (left, up, upper_left)[int(np.argmin(distances))]
            elif filter_type == 0:
                value = 0
            else:
                raise ValueError(f"Unsupported PNG filter {filter_type}")
            row[index] = (int(row[index]) + int(value)) & 255
        rows.append(row)
        previous = row
    return np.frombuffer(b"".join(row.tobytes() for row in rows), dtype=">u2").astype(np.uint16).reshape(height, width)


def decode_mrms_event(event_dir: str | Path, image_height: int = 512, image_width: int = 512, total_length: int = 29) -> np.ndarray:
    event_dir = Path(event_dir)
    event_name = event_dir.name
    frames = [_read_png_gray16(event_dir / f"{event_name}-{index:02d}.png") for index in range(total_length)]
    data = np.stack(frames).astype(np.float32) / 10.0 - 3.0
    if data.shape[1:] != (image_height, image_width):
        raise ValueError(f"MRMS frame shape {data.shape[1:]} != {(image_height, image_width)}")
    mask = np.ones_like(data, dtype=np.float32)
    mask[data < 0] = 0
    data[data < 0] = 0
    data = np.clip(data, 0, 128)
    return np.stack([data, mask], axis=-1)


class MRMSDataset(Dataset):
    def __init__(self, data_dir: str | Path, image_height: int = 512, image_width: int = 512, total_length: int = 29, split: str = "train"):
        self.data_dir = Path(data_dir)
        self.events = sorted(path for path in self.data_dir.iterdir() if path.is_dir())
        if split == "train":
            self.events = [path for path in self.events if len(list(path.glob("*.png"))) == total_length]
        elif split != "test":
            raise ValueError("split must be train or test")
        self.image_height = image_height
        self.image_width = image_width
        self.total_length = total_length

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

    def __getitem__(self, index: int) -> dict[str, torch.Tensor | str]:
        event = self.events[index]
        frames = decode_mrms_event(event, self.image_height, self.image_width, self.total_length)
        return {"radar_frames": torch.from_numpy(frames), "event": event.name}