"""
PyTorch DataLoader mappings for Mashariki Weather Dataset.
"""
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

from .dataset import Windows

class MasharikiTorchDataset(Dataset):
    """
    A unified spatiotemporal PyTorch Dataset handling local Numpy windows 
    while natively masking network latency. Natively encapsulates the `Windows`
    dataclass extracting from PyTorch DataLoader logic completely bypassing 
    memory-hungry Xarray tensor conversions.
    """
    def __init__(self, windows: Windows):
        self.windows = windows
        
        # We implicitly capture the NumPy references without physically allocating memory
        # since PyTorch handles NumPy strided tensors flawlessly.
        self.x = torch.as_tensor(windows.x, dtype=torch.float32)
        self.y = torch.as_tensor(windows.y, dtype=torch.float32)
        self.mask_x = torch.as_tensor(windows.mask_x, dtype=torch.bool)
        self.mask_y = torch.as_tensor(windows.mask_y, dtype=torch.bool)
        
        self.has_gridded = windows.gridded_x is not None
        if self.has_gridded:
            self.gridded_x = torch.as_tensor(windows.gridded_x, dtype=torch.float32)
            self.has_gridded_y = windows.gridded_y is not None
            if self.has_gridded_y:
                self.gridded_y = torch.as_tensor(windows.gridded_y, dtype=torch.float32)

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

    def __getitem__(self, idx):
        # Natively fetches the zero-copy array pointers instantly (0s time complexity)
        batch = {
            "x": self.x[idx],
            "y": self.y[idx],
            "mask_x": self.mask_x[idx],
            "mask_y": self.mask_y[idx]
        }
        
        if self.has_gridded:
            batch["gridded_x"] = self.gridded_x[idx]
            if self.has_gridded_y:
                batch["gridded_y"] = self.gridded_y[idx]
                
        return batch

def create_dataloader(windows: Windows, batch_size: int = 32, shuffle: bool = True, num_workers: int = 4) -> DataLoader:
    """
    Factory function to automatically bootstrap a PyTorch DataLoader out of a generated Mashariki `Windows` structure.
    
    Args:
        windows: The Mashariki `Windows` dataclass generated by `ds.get_windows()`.
        batch_size: Target batch size for the Model architecture.
        shuffle: Whether to shuffle the timesteps (Recommended True for training).
        num_workers: Number of subprocesses to offload the batch assembly (Masks any potential Python GIL latency).
    """
    dataset = MasharikiTorchDataset(windows)
    
    # Offloading dataloader threads to aggressively map the NumPy pointers asynchronously 
    return DataLoader(
        dataset, 
        batch_size=batch_size, 
        shuffle=shuffle, 
        num_workers=num_workers,
        pin_memory=True # Automates GPU VRAM memory transfers directly from CPU
    )
