import numpy as np
import os
import time
import torch
from torch.utils.data import Dataset, DataLoader
import random

DEBUG = 1


class SpatialSlicesDataset(Dataset):
    def __init__(self, spatial_slices):
        self.spatial_slices = spatial_slices
        self.num, self.x, self.y, _ = self.spatial_slices.shape


    def __getitem__(self, idx):
        spatial_slice = self.spatial_slices[idx]
        coords = torch.stack([
            (torch.arange(self.x).float() / (self.x - 1) * 2 - 1).unsqueeze(1).repeat(1, self.y),
            (torch.arange(self.y).float() / (self.y - 1) * 2 - 1).unsqueeze(0).repeat(self.x, 1)
            # (torch.arange(self.x).float() / (self.x - 1) * 2 - 1).expand(-1, self.y),
            # (torch.arange(self.y).float() / (self.y - 1) * 2 - 1).expand(self.x, -1)
        ], dim=-1).to(torch.float32)
        targets = spatial_slice.to(torch.float32)
        return coords, targets
    

    def __len__(self):
        return self.num
    

    def get_batch(self, batch_size, sample_num):
        n_indices = torch.randint(0, self.num-1, (batch_size,))
        timesteps = (n_indices.float() / 25).int()

        selected_slices = self.spatial_slices[n_indices]

        x_indices = torch.randint(0, self.x, (batch_size, sample_num))
        y_indices = torch.randint(0, self.y, (batch_size, sample_num))

        # - make sure both corrds and targets are float type
        # targets = torch.stack([selected_slices[i][x_indices[i], y_indices[i]].float() for i in range(batch_size)])
        targets = selected_slices[torch.arange(batch_size)[:, None], x_indices, y_indices].float()

        coords = torch.stack([
            (x_indices.float() / (self.x - 1) * 2 - 1),
            (y_indices.float() / (self.y - 1) * 2 - 1)
        ], dim=-1)

        return coords, targets, timesteps
    

    def get_slice(self, index):
        return self.__getitem__(index)
    

class SpatialFieldsDataset(Dataset):
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.field_idx = []
        
        for filename in os.listdir(data_dir):
            # if filename.endswith(".pt"):  
            self.field_idx.append(filename[:4])

        #%TODO
        self.t = 1001
        self.x = 512
        self.y = 512


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


    def __getitem__(self, idx):
        #%TODO n_idx is not continuous integer, n_idx+=160 if we pick 50 fields in all
        n_idx, t_idx = divmod(idx, self.t)
        n_idx = n_idx*160
        spatial_slice = torch.load(os.path.join(self.data_dir, f'/{n_idx}/{n_idx}_{t_idx}.pt'))
        
        coords = torch.stack([
            (torch.arange(self.x).float() / (self.x - 1) * 2 - 1).expand(-1, self.y),
            (torch.arange(self.y).float() / (self.y - 1) * 2 - 1).expand(self.x, -1)
        ], dim=-1)
        
        # - reshaping the spatial slice to match the coordinates shape
        targets = spatial_slice
        
        return coords, targets                

    def get_batch(self, batch_size, sample_num):
        coords = torch.empty((batch_size, sample_num, 2))
        targets = torch.empty((batch_size, sample_num, 2))

        n_indices = random.choices(self.field_idx, k=batch_size)
        t_indices = torch.randint(0, self.t, (batch_size,))

        # - load spatial slices in batch
        spatial_slices = [torch.load(os.path.join(self.data_dir, f'{n_idx}/{n_idx}_{t_idx}.pt')) for n_idx, t_idx in zip(n_indices, t_indices)]
        spatial_slices_tensor = torch.stack(spatial_slices)

        # - randomly select x and y coordinates within the field for each sample in the batch
        x_indices = torch.randint(0, self.x, (batch_size, sample_num))
        y_indices = torch.randint(0, self.y, (batch_size, sample_num))

        # - use advanced indexing to select targets
        targets = spatial_slices_tensor[torch.arange(batch_size)[:, None], x_indices, y_indices]

        coords = torch.stack([
            (x_indices.float() / (self.x - 1) * 2 - 1).unsqueeze(-1), 
            (y_indices.float() / (self.y - 1) * 2 - 1).unsqueeze(-1)
        ], dim=-1)
        
        return coords, targets

# # - test SpatialSlicesDataset
# dataset = torch.load(os.path.join('input/7999_train.pt'))
# data_loader = SpatialSlicesDataset(dataset)
# batch = data_loader.get_batch(8, 500)
# if DEBUG:
#     print(batch[0].shape)
#     print(batch[1].shape)
#     print(f'coords range: {batch[0].min()},{batch[0].max()}')
#     print(f'target range: {batch[1].min()},{batch[1].max()}')
#     print(f'timestep range: {batch[2].min()},{batch[2].max()}')

# # - test SpatialSlicesDataset   
# data_loader = SpatialFieldsDataset('/media/wangx84/Elements/fuild_flow_eth') 
# batch = data_loader.get_batch(8, 100)
# if DEBUG:
#     print(batch[0].shape)
#     print(batch[1].shape)
#     print(f'coords range: {batch[0].min()},{batch[0].max()}')
#     print(f'target range: {batch[1].min()},{batch[1].max()}')




为什么在spatial_slices = [torch.load(os.path.join(self.data_dir, f'{n_idx}/{n_idx}_{t_idx}.pt')) for n_idx, t_idx in zip(n_indices, t_indices)]这一行，n_idx=spat?