# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/agents/benchmark_agents/03_ERM_agents.ipynb.

# %% auto 0
__all__ = ['FakePolicy', 'NewsvendorData', 'LinearModel', 'MLP', 'MLP_BLOCK', 'RNN', 'SGDBase', 'LERMsgdAgent', 'MLPsgdAgent',
           'RNNsgdAgent']

# %% ../../../nbs/agents/benchmark_agents/03_ERM_agents.ipynb 4
# General libraries:
import numpy as np
from scipy.stats import norm
from tqdm import trange

# Torch
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from ..processors.processors import GetTimeSeriesAndStaticFeatures

# Mushroom libraries
from mushroom_rl.core import Agent

# %% ../../../nbs/agents/benchmark_agents/03_ERM_agents.ipynb 6
class FakePolicy():
    def reset(*args, **kwargs):
        pass

class NewsvendorData(Dataset):

    def __init__(self, x, y, mask=None):
        # create torch tensors
        self.x=torch.from_numpy(x)
        self.y=torch.from_numpy(y)
        if mask is not None:
            self.mask = torch.from_numpy(mask)
            self.mask = self.mask.bool()
            if len(self.mask.shape) == 1:
                self.mask = self.mask.unsqueeze(1)
        else:
            self.mask = None
        
        # convert to torch float32
        self.x=self.x.float()
        self.y=self.y.float()

        self.n_samples=y.shape[0]

    def __getitem__(self, index):
        if self.mask is not None:
            return self.x[index], self.y[index], self.mask[index]
        else:
            return self.x[index], self.y[index]

    def __len__(self):
        return self.n_samples

class LinearModel(nn.Module):
    def __init__(self, input_size, output_size, relu_output=False):
        super().__init__()
        self.l1=nn.Linear(input_size, output_size)
        if relu_output:
            self.final_activation = nn.ReLU()
        else:
            self.final_activation = nn.Identity()

    def forward(self, x):
        out=self.l1(x)
        out=self.final_activation(out)
        return out

class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_hidden_layers=3, drop_prob=0.0, relu_output=False):
        super().__init__()

        # List of layers
        layers = []

        # Input layer
        layers.append(nn.Linear(input_size, hidden_size))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(p=drop_prob))

        # Hidden layers
        for _ in range(num_hidden_layers-1): 
            layers.append(nn.Linear(hidden_size, hidden_size))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(p=drop_prob))

        # Output layer
        layers.append(nn.Linear(hidden_size, output_size))
        if relu_output:
            layers.append(nn.ReLU()) # output is non-negative
        else:
            layers.append(nn.Identity())

        # Combine layers
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

class MLP_BLOCK(nn.Module):
    def __init__(self, hidden_size, drop_prob):
        super(MLP_BLOCK, self).__init__()
        self.MLP_block = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(p=drop_prob),
        )
    
    def forward(self, x):
        return self.MLP_block(x)

class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_hidden_layers=2, drop_prob=0.0, num_time_series_features=None, lag_window=None, relu_output=False):
        super().__init__()

        self.input_processor = GetTimeSeriesAndStaticFeatures(num_time_series_features,lag_window)
        # self._n_output = output_shape[0]
        # self._n_features = n_features

        self._h1_recurrent = nn.GRU(int(input_size/lag_window), hidden_size, 1, batch_first = True)
        self._h1_dropout = nn.Dropout(p=drop_prob)

        self.mlp_blocks = nn.ModuleList()
        for i in range(num_hidden_layers):
            self.mlp_blocks.append(MLP_BLOCK(hidden_size, drop_prob))

        self.output_layer = nn.Linear(hidden_size, output_size)
        
        if relu_output:
            self.final_activation = nn.ReLU()  # output is non-negative
        else:
            self.final_activation = nn.Identity()

    def forward(self, x):

        # if shape is 1 then expand
        if len(x.shape)==1:
            x=x.unsqueeze(0)
        x = self.input_processor(x)[0] # no static data returned as no inventory

        # RNN block
        x, _ = self._h1_recurrent(x)
        x = self._h1_dropout(x)
        x = x[:, -1, :] # only take the last time step

        # MLP block
        for mlp_block in self.mlp_blocks:
            x = mlp_block(x)  
        
        x = self.output_layer(x) 
        x = self.final_activation(x)

        # print(x.shape)
    
        return x    

class SGDBase(Agent):

    def __init__(self, input_size, hidden_size=64, output_size=1, learning_rate=0.01, num_hidden_layers=3, drop_prob=0.0, l2_reg=0.0, num_time_series_features=None, lag_window=None, relu_output=False):
        if self.model_type=="Linear":
            #super(LinearModel, self).__init__(input_size, output_size)
            self.model=LinearModel(input_size, output_size, relu_output=relu_output)
        elif self.model_type=="MLP":
            self.model=MLP(input_size, hidden_size, output_size, num_hidden_layers=num_hidden_layers, drop_prob=drop_prob, relu_output=relu_output)
        elif self.model_type=="RNN":
            self.model=RNN(input_size, hidden_size, output_size, num_hidden_layers=num_hidden_layers, drop_prob=drop_prob, num_time_series_features=num_time_series_features, lag_window=lag_window, relu_output=relu_output)
            #super(MLP, self).__init__(input_size, hidden_size, output_size)
        
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate, weight_decay=l2_reg)
        self.criterion = nn.MSELoss()


    def fit(self, X_train, y_train, mask, cu, co, batch_size=64, learning_rate=0.01, device="cpu"):
        
        if y_train.ndim == 1:
            y_train = y.reshape(-1, 1)
        
        dataset_train=NewsvendorData(X_train, y_train, mask)

        self.model.to(device)

        train_loader=DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True)

        self.model.train()

        for i, (output) in enumerate(train_loader):

                if len(output)==2:
                    feat, labels = output
                    masks=None
                else:
                    feat, labels, masks = output
                    masks=masks.to(device)
                
                feat=feat.to(device)
                labels=labels.to(device)
                outputs=self.model(feat)

                loss = torch.mean(SGDBase.pinball_loss(cu, co, labels, outputs, masks))

                #backward
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

        self.model.eval()

        return self.model

    def predict(self, X):
        # print(X.shape)
        self.model.eval()
        with torch.no_grad():
            X=torch.from_numpy(X)
            X=X.float()
            output=self.model(X)
            output=output.numpy()
        
        if len(output.shape)==2:
            output=output.reshape(-1)
        
        return output

    def train(self):
        self.model.train()
    
    def eval(self):
        self.model.eval()   
    
    ### Helper functions

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

    @staticmethod
    def pinball_loss(cu, co, demand, order_quantity, mask):

        cu = torch.tensor(cu, dtype=torch.float32)
        co = torch.tensor(co, dtype=torch.float32)

        underage=cu*SGDBase.max_or_zero(demand-order_quantity)
        overage=co*SGDBase.max_or_zero(order_quantity-demand)
        loss=underage+overage
        loss = loss*mask
        return loss

class LERMsgdAgent(SGDBase):
    def __init__(self,
                    input_size,
                    output_size,
                    cu,
                    co,
                    final_activation="identity",
                    batch_size=128,
                    learning_rate=0.01,
                    device="cpu",
                    agent_name = "LERM"
                    ):
        self.name=agent_name
        self.model_type="Linear"
        self.cu = cu
        self.co = co
        self.device = device
        self.batch_size=batch_size
        self.learning_rate=learning_rate
        self.device=device


        self.policy=FakePolicy()
        self._postprocessors = list()
        self._preprocessors = list()
        self.train_directly=True
        self.train_mode = "epochs"

        if final_activation=="identity":
            self.final_activation = False
        elif final_activation=="relu":
            self.final_activation = True

        super().__init__(input_size=input_size, hidden_size=None, output_size=output_size, learning_rate=learning_rate, relu_output=self.final_activation)

    def fit_epoch(self, features_train, demand_train, mask=None):
        super().fit(features_train, demand_train, mask=mask, cu=self.cu, co=self.co, batch_size=self.batch_size, learning_rate=self.learning_rate, device=self.device)

    def draw_action(self, X):
        return super().predict(X)

class MLPsgdAgent(SGDBase):
    def __init__(self,
                    input_size,
                    output_size,
                    cu,
                    co,
                    final_activation="identity",
                    hidden_size=64,
                    batch_size=128,
                    learning_rate=0.01,
                    device="cpu",
                    agent_name = "DLNV",
                    num_hidden_layers=3,
                    drop_prob=0.0,
                    l2_reg=0.0,
                    ): 
        
        self.name=agent_name
        self.model_type="MLP"
        self.cu = cu
        self.co = co
        self.device = device
        self.hidden_size = hidden_size
        self.batch_size=batch_size
        self.learning_rate=learning_rate
        self.device=device

        self.policy=FakePolicy()
        self._postprocessors = list()
        self._preprocessors = list()
        self.train_directly=True
        self.train_mode = "epochs"

        if final_activation=="identity":
            self.final_activation = False
        elif final_activation=="relu":
            self.final_activation = True

        super().__init__(input_size=input_size, hidden_size=hidden_size, output_size=output_size, learning_rate=learning_rate, num_hidden_layers=num_hidden_layers, drop_prob=drop_prob, l2_reg=l2_reg, final_activation=self.final_activation)

    def fit_epoch(self, features_train, demand_train, mask=None):
        super().fit(features_train, demand_train, mask=mask, cu=self.cu, co=self.co, batch_size=self.batch_size, learning_rate=self.learning_rate, device=self.device)

    def draw_action(self, X):
        return super().predict(X)

class RNNsgdAgent(SGDBase):
    def __init__(self,
                    input_size,
                    output_size,
                    cu,
                    co,
                    num_time_series_features,
                    lag_window,
                    final_activation="identity",
                    hidden_size=64,
                    batch_size=128,
                    learning_rate=0.01,
                    device="cpu",
                    agent_name = "DLNV_RNN",
                    num_hidden_layers=3,
                    drop_prob=0.0,
                    l2_reg=0.0,
                    ): 
        
        self.name=agent_name
        self.model_type="RNN"
        self.cu = cu
        self.co = co
        self.device = device
        self.hidden_size = hidden_size
        self.batch_size=batch_size
        self.learning_rate=learning_rate
        self.device=device

        self.policy=FakePolicy()
        self._postprocessors = list()
        self._preprocessors = list()
        self.train_directly=True
        self.train_mode = "epochs"
    
        if final_activation=="identity":
            self.final_activation = False
        elif final_activation=="relu":
            self.final_activation = True

        super().__init__(input_size=input_size, hidden_size=hidden_size, output_size=output_size, learning_rate=learning_rate, num_hidden_layers=num_hidden_layers, drop_prob=drop_prob, l2_reg=l2_reg, num_time_series_features=num_time_series_features, lag_window=lag_window, final_activation=self.final_activation)

    def fit_epoch(self, features_train, demand_train, mask=None):
        super().fit(features_train, demand_train, mask=mask, cu=self.cu, co=self.co, batch_size=self.batch_size, learning_rate=self.learning_rate, device=self.device)

    def draw_action(self, X):
        return super().predict(X)

