import numpy as np
import pandas as pd

np.random.seed(42)
epsilon = 1e-8

class Dataset:
    def __init__(self, inverse=False):
        filename = './Data/DataForThermoforming.xlsx'
        self.df = pd.read_excel(filename, sheet_name='Data')
        # remove rows by index
        self.df = self.df.drop([20, 48], axis=0)

        # normalize data
        if inverse:
            self.input_columns = ['Ply_Number', 'A1(abs)', 'B1(abs)', 'C1(abs)', 'Stress(Max) MPa']
            self.output_columns = ['Initial_Temp (degree celsius)', 'Punch_Velocity (mm/s)', 'Cooling_Time (s)']
        else:
            self.input_columns = ['Ply_Number', 'Initial_Temp (degree celsius)', 'Punch_Velocity (mm/s)', 'Cooling_Time (s)']
            self.output_columns = ['A1(abs)', 'B1(abs)', 'C1(abs)', 'Stress(Max) MPa']

        self.input_mean = self.df[self.input_columns].mean().to_numpy(dtype=np.float32)
        self.input_std = self.df[self.input_columns].std().to_numpy(dtype=np.float32) + epsilon
        self.output_mean = self.df[self.output_columns].mean().to_numpy(dtype=np.float32)
        self.output_std = self.df[self.output_columns].std().to_numpy(dtype=np.float32) + epsilon


    def get_input(self, normalize=False):
        data = self.df[self.input_columns].to_numpy(dtype=np.float32)
        if normalize:
            data = self.normalize_input(data)
        return data
    

    def get_output(self, normalize=False):
        data = self.df[self.output_columns].to_numpy(dtype=np.float32)
        if normalize:
            data = self.normalize_output(data)
        return data

    def __str__(self):
        return str(self.df.head())
    
    def normalize_input(self, input_data):
        return (input_data - self.input_mean) / self.input_std
    
    def normalize_output(self, output_data):
        return (output_data - self.output_mean) / self.output_std
    
    def denormalize_input(self, normalized_input):
        return normalized_input * self.input_std + self.input_mean
    
    def denormalize_output(self, normalized_output):
        return normalized_output * self.output_std + self.output_mean

if __name__ == "__main__":
    dataset = Dataset()
    
    # Example usage
    input_data = dataset.get_input(normalize=True)
    output_data = dataset.get_output(normalize=True)
    
    print("Input shape:", input_data.shape)
    print("Output shape:", output_data.shape)