import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pd

class NrvDataset(Dataset):
    def __init__(self, dataframe, data_config, sequence_length=96, predict_sequence_length=96):
        self.data_config = data_config
        
        self.nrv = torch.tensor(dataframe['nrv'].to_numpy(), dtype=torch.float32)
        self.load_forecast = torch.tensor(dataframe['load_forecast'].to_numpy(), dtype=torch.float32)

        self.sequence_length = sequence_length
        self.predict_sequence_length = predict_sequence_length

    def __len__(self):
        return len(self.nrv) - self.sequence_length - self.predict_sequence_length
    
    def __getitem__(self, idx):
        features = []

        if self.data_config.NRV_HISTORY:
            nrv = self.nrv[idx:idx+self.sequence_length]
            features.append(nrv.view(-1))

        if self.data_config.LOAD_HISTORY:
            load_history = self.load_forecast[idx:idx+self.sequence_length]
            features.append(load_history.view(-1))

        if self.data_config.LOAD_FORECAST:
            load_forecast = self.load_forecast[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]
            features.append(load_forecast.view(-1))

        if not features:
            raise ValueError("No features are configured to be included in the dataset.")

        # Concatenate along dimension 0 to create a one-dimensional feature vector
        all_features = torch.cat(features, dim=0)
        
        # Target sequence, flattened if necessary
        nrv_target = self.nrv[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length].view(-1)

        return all_features, nrv_target

    def random_day_autoregressive(self, idx: int):
        features = []

        # we already have the NRV history with the newly predicted values, so we don't need to include the last 96 values
        
        if self.data_config.LOAD_HISTORY:
            load_history = self.load_forecast[idx:idx+self.sequence_length]
            features.append(load_history.view(-1))

        if self.data_config.LOAD_FORECAST:
            load_forecast = self.load_forecast[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]
            features.append(load_forecast.view(-1))

        target = self.nrv[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]

        if len(features) == 0:
            return None, target

        all_features = torch.cat(features, dim=0)
        return all_features, target