# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/05_deepLearningNewsvendor.ipynb.

# %% auto 0
__all__ = ['NewsvendorData', 'DeepLearningNewsvendor']

# %% ../../nbs/05_deepLearningNewsvendor.ipynb 5
from ._base import BaseNewsvendor, DataDrivenMixin
from ..utils.validation import check_cu_co
from ..metrics import average_costs

import torch
from torch.utils.data import Dataset
import torch.nn as nn  # Lazy import
from torch.utils.data import DataLoader


from sklearn.utils.validation import check_is_fitted
import numpy as np

# class NewsvendorData:
#     def __init__(self, x, y):
#         # Dynamically assign torch Dataset as base class
#         self.__class__ = type(self.__class__.__name__, (Dataset,), dict(self.__class__.__dict__))

#         self.x = torch.from_numpy(x).float()
#         self.y = torch.from_numpy(y).float()
#         self.n_samples = y.shape[0]

#     def __getitem__(self, index):
#         return self.x[index], self.y[index]

#     def __len__(self):
#         return self.n_samples



class NewsvendorData:
    def __init__(self, x, y):
        self.x = torch.from_numpy(x).float()
        self.y = torch.from_numpy(y).float()
        self.n_samples = y.shape[0]

    def __getitem__(self, index):
        return self.x[index], self.y[index]

    def __len__(self):
        return self.n_samples


class DeepLearningNewsvendor(BaseNewsvendor, DataDrivenMixin):

    ACTIVATION_FUNCTIONS = {
        'ReLU': lambda: __import__('torch').nn.ReLU(),
        'Sigmoid': lambda: __import__('torch').nn.Sigmoid(),
        'Tanh': lambda: __import__('torch').nn.Tanh(),
        'LeakyReLU': lambda: __import__('torch').nn.LeakyReLU(),
        'ELU': lambda: __import__('torch').nn.ELU(),
        'SELU': lambda: __import__('torch').nn.SELU(),
        'GELU': lambda: __import__('torch').nn.GELU(),
    }

    def __init__(self, cu=None, co=None, neurons=[100, 50], activations=['ReLU', 'ReLU'], optimizer='adam', 
                 epochs=100, batch_size=64, learning_rate=0.003, drop_prob=0.0, random_state=None, verbose=1, print_freq=10):
        self.neurons = neurons
        self.activations = activations
        self.optimizer = optimizer
        self.epochs = epochs
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.drop_prob = drop_prob
        self.random_state = random_state
        self.verbose = verbose
        self.print_freq = print_freq
        super().__init__(cu=cu, co=co)

    @staticmethod
    def max_or_zero(data):
        return torch.max(data, torch.zeros_like(data))

    @staticmethod
    def pinball_loss(cu, co, demand, order_quantity):
        if len(demand.shape) == 1:
            demand = demand.unsqueeze(1)

        cu = torch.tensor(cu, dtype=torch.float32).to(demand.device)
        co = torch.tensor(co, dtype=torch.float32).to(demand.device)

        underage = cu * DeepLearningNewsvendor.max_or_zero(demand - order_quantity)
        overage = co * DeepLearningNewsvendor.max_or_zero(order_quantity - demand)
        return underage + overage

    def _create_model(self):


        if self.random_state is not None:
            torch.manual_seed(self.random_state)

        layers = []
        for i, (neurons, activation) in enumerate(zip(self.neurons, self.activations)):
            if i == 0:
                layers.append(nn.Linear(self.n_features_, neurons))
            else:
                layers.append(nn.Linear(self.neurons[i - 1], neurons))
            layers.append(DeepLearningNewsvendor.ACTIVATION_FUNCTIONS[activation]())
            if self.drop_prob > 0:
                layers.append(nn.Dropout(self.drop_prob))
        layers.append(nn.Linear(self.neurons[-1], self.n_outputs_))
        return nn.Sequential(*layers)

    def fit(self, X, y, validation_data=None):


        X, y = self._validate_data(X, y, multi_output=True)
        self.X_ = X
        self.y_ = y
        self.n_features_ = X.shape[1]
        self.n_outputs_ = y.shape[1]
        self.cu_, self.co_ = check_cu_co(self.cu, self.co, self.n_outputs_)

        dataset = NewsvendorData(X, y)
        train_loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)

        model = self._create_model()
        optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate) if self.optimizer == 'adam' else torch.optim.SGD(model.parameters(), lr=self.learning_rate)

        for epoch in range(self.epochs):
            model.train()
            epoch_loss = 0
            for batch_x, batch_y in train_loader:
                optimizer.zero_grad()
                outputs = model(batch_x)
                loss = self.pinball_loss(self.cu_, self.co_, batch_y, outputs).mean()
                loss.backward()
                optimizer.step()
                epoch_loss += loss.item()

            if self.verbose > 0 and (epoch + 1) % self.print_freq == 0:
                print(f"Epoch {epoch+1}/{self.epochs}, Loss: {epoch_loss:.4f}")
        self.model_ = model

    def predict(self, X):

        X = self._validate_data(X)

        dataset = NewsvendorData(X, np.zeros((X.shape[0], self.n_outputs_)))
        test_loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False)

        self.model_.eval()
        preds = []
        for batch_x, _ in test_loader:
            preds.extend(self.model_(batch_x).detach().numpy())
        return np.array(preds)
