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.dataframe = dataframe
        # reset dataframe index
        self.dataframe.reset_index(drop=True, inplace=True)

        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.total_load = torch.tensor(dataframe['total_load'].to_numpy(), dtype=torch.float32)
        self.pv_gen_forecast = torch.tensor(dataframe['pv_forecast'].to_numpy(), dtype=torch.float32)
        self.wind_gen_forecast = torch.tensor(dataframe['wind_forecast'].to_numpy(), dtype=torch.float32)

        self.sequence_length = sequence_length
        self.predict_sequence_length = predict_sequence_length

        self.samples_to_skip = self.skip_samples()
        total_indices = set(range(len(self.nrv) - self.sequence_length - self.predict_sequence_length))
        self.valid_indices = sorted(list(total_indices - set(self.samples_to_skip)))

    def skip_samples(self):
        nan_rows = self.dataframe[self.dataframe.isnull().any(axis=1)]
        nan_indices = nan_rows.index
        skip_indices = [list(range(idx-self.sequence_length-self.predict_sequence_length, idx+1)) for idx in nan_indices]

        skip_indices = [item for sublist in skip_indices for item in sublist]
        skip_indices = list(set(skip_indices))
        skip_indices.sort()
        return skip_indices

    def __len__(self):
        return len(self.valid_indices)
    
    def __getitem__(self, idx):
        actual_idx = self.valid_indices[idx]
        features = []

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

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

        if self.data_config.PV_HISTORY:
            pv_history = self.pv_gen_forecast[actual_idx:actual_idx+self.sequence_length]
            features.append(pv_history.view(-1))

        if self.data_config.WIND_HISTORY:
            wind_history = self.wind_gen_forecast[actual_idx:actual_idx+self.sequence_length]
            features.append(wind_history.view(-1))

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

        if self.data_config.PV_FORECAST:
            pv_forecast = self.pv_gen_forecast[actual_idx+self.sequence_length:actual_idx+self.sequence_length+self.predict_sequence_length]
            features.append(pv_forecast.view(-1))

        if self.data_config.WIND_FORECAST:
            wind_forecast = self.wind_gen_forecast[actual_idx+self.sequence_length:actual_idx+self.sequence_length+self.predict_sequence_length]
            features.append(wind_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[actual_idx+self.sequence_length:actual_idx+self.sequence_length+self.predict_sequence_length].view(-1)

        # check if nan values are present
        if torch.isnan(all_features).any():
            print(f"Found nan values in the features of sample {idx}.")
            print(f"Actual index: {actual_idx}")
            raise ValueError("There are nan values in the features.")

        return all_features, nrv_target

    def random_day_autoregressive(self, idx: int):
        idx = self.valid_indices[idx]
        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.total_load[idx:idx+self.sequence_length]
            features.append(load_history.view(-1))

        if self.data_config.PV_HISTORY:
            pv_history = self.pv_gen_forecast[idx:idx+self.sequence_length]
            features.append(pv_history.view(-1))

        if self.data_config.WIND_HISTORY:
            wind_history = self.wind_gen_forecast[idx:idx+self.sequence_length]
            features.append(wind_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 self.data_config.PV_FORECAST:
            pv_forecast = self.pv_gen_forecast[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]
            features.append(pv_forecast.view(-1))

        if self.data_config.WIND_FORECAST:
            wind_forecast = self.wind_gen_forecast[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]
            features.append(wind_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