# 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', 'RNNsgdMetaAgent', 'LagLLama', 'LagLlamasgdAgent']

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

# 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, feature_map=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
        self.feature_map = feature_map
        
        # convert to torch float32
        self.x=self.x.float()
        self.y=self.y.float()

        if self.feature_map is not None:
            self.n_samples=y.shape[0]*y.shape[1]
            self.index_mapping = dict()
            index_used=0
            for i in range(y.shape[0]):
                for j in range(y.shape[1]):
                    self.index_mapping[index_used] = (i,j)
                    index_used+=1
        else:
            self.n_samples=y.shape[0]

    def __getitem__(self, index):
        
        if self.feature_map is not None:
            
            coordinates = self.index_mapping[index]

            time = coordinates[0]
            product = coordinates[1]

            y = self.y[time, product]

            # print("time: ", time, "product: ", product, "y: ", y)

            # print(self.feature_map)
            # print(self.feature_map[:, product])
            # print(self.x[time])

            x = self.x[time, self.feature_map[:, product].astype(bool)]

            if self.mask is not None:
                mask = self.mask[time, product]
        
        else:
            x = self.x[index]
            y = self.y[index]
            if self.mask is not None:
                mask = self.mask[index]
        
        if self.feature_map is not None:
            if self.mask is not None:
                return x, y, mask, product
            else:
                return x, y, product
        else:
            if self.mask is not None:
                return x, y, mask
            else:
                return x, y

    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, size_in, size_out, drop_prob):
        super(MLP_BLOCK, self).__init__()
        self.MLP_block = nn.Sequential(
            nn.Linear(size_in, size_out),
            nn.ReLU(),
            nn.Dropout(p=drop_prob),
            # add batch norm
            nn.BatchNorm1d(size_out)
        )
    
    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,
                num_RNN_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), int(hidden_size/2), num_RNN_layers, batch_first = True)
        self._h1_dropout = nn.Dropout(p=drop_prob)

        self.mlp_blocks = nn.ModuleList()
        for i in range(num_hidden_layers):
            if i == 0:
                self.mlp_blocks.append(MLP_BLOCK(size_in = int(hidden_size/2), size_out = hidden_size, drop_prob=drop_prob))
            else:
                self.mlp_blocks.append(MLP_BLOCK(size_in = int(hidden_size/(2**(i-1))), size_out = int(hidden_size/(2**i)), drop_prob=drop_prob))

        if num_hidden_layers==0:
            self.output_layer = nn.Linear(int(hidden_size/2), output_size)
        else:
            self.output_layer = nn.Linear(int(hidden_size/(2**i)), 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, num_RNN_layers=None, drop_prob=0.0, l2_reg=0.0, num_time_series_features=None, lag_window=None, relu_output=False):
        if self.model_type=="Linear":
            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, num_RNN_layers=num_RNN_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()

        total_loss = 0
        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))
                
                total_loss += loss.item()

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

        print("training loss: ", total_loss)   
        self.model.eval()

        return self.model

    def predict(self, X):

        # if len(X.shape)==1:
        #     X = np.expand_dims(X, axis=0)
        
        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):

        # TODO: check that this works in termns of shapes

        if len(demand.shape)==1:
            demand = demand.unsqueeze(1)
        
        assert demand.shape == order_quantity.shape

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

        cu = cu.to(demand.device)
        co = co.to(demand.device)

        underage=cu*SGDBase.max_or_zero(demand-order_quantity)
        overage=co*SGDBase.max_or_zero(order_quantity-demand)

        loss=underage+overage
        
        if mask is not None:
            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, 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 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_RNN_layers = 2,
                    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, num_RNN_layers=num_RNN_layers, drop_prob=drop_prob, l2_reg=l2_reg, num_time_series_features=num_time_series_features, lag_window=lag_window, 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 RNNsgdMetaAgent(SGDBase):
    def __init__(self,
                    feature_map,
                    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_RNN_layers = 2,
                    num_hidden_layers=3,
                    drop_prob=0.0,
                    l2_reg=0.0,
                    ): 
        
        self.feature_map = feature_map
        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"

        print("feature map: ", self.feature_map)
    
        if final_activation=="identity":
            self.final_activation = False
        elif final_activation=="relu":
            self.final_activation = True

        input_size = np.sum(self.feature_map[:,0])

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

    def fit_epoch(self, features_train, demand_train, mask=None):
        self.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 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.feature_map)

        self.model.to(device)
        
        train_loader=DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=False) # make shuffle true

        self.model.train()

        total_loss = 0
        for i, (output) in tqdm(enumerate(train_loader)):

                if len(output)==3:
                    feat, labels, product = output
                    masks=None
                else:
                    feat, labels, masks, product = output
                    masks=masks.to(device)   

                feat=feat.to(device)
                labels=labels.to(device)
                outputs=self.model(feat)

                cu_selected = cu[product]
                co_selected = co[product]
                
                loss_per_product = self.pinball_loss(cu_selected, co_selected, labels, outputs, masks)
                loss = torch.mean(loss_per_product)
            
                # show = len(labels)
                # for i in range (show):
                #     print("labels: ", labels[i], "feat: ", feat[i])
                #     print("outputs: ", outputs[i])
                #     print("loss: ", loss_per_product[i])
                #     print(" ")

                # print("#####")
                
                # sleep(5)
                
                total_loss += loss.item()

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

        print("training loss: ", total_loss)   
        self.model.eval()

        return self.model

    def draw_action(self, X):
        return self.predict(X)
    
    def predict(self, X):
        self.model.eval()
        for i in range(self.feature_map.shape[1]):
            if len(X.shape)==1:
                X = np.expand_dims(X, axis=0)
            X_i = X[:,self.feature_map[:,i].astype(bool)]
            with torch.no_grad():
                X_i=torch.from_numpy(X_i)
                X_i=X_i.float()
                output=self.model(X_i)
                output=output.numpy()
            
            if i == 0:
                outputs = output
            else:
                outputs = np.hstack((outputs, output))

        # print(outputs)

        return outputs

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

        if len(demand.shape)==1:
            demand = demand.unsqueeze(1)
        
        assert demand.shape == order_quantity.shape

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

        cu = cu.unsqueeze(1)
        co = co.unsqueeze(1)

        cu = cu.to(demand.device)
        co = co.to(demand.device)

        underage_quantity = SGDBase.max_or_zero(demand-order_quantity)
        overage_quantity = SGDBase.max_or_zero(order_quantity-demand)

        assert cu.shape == underage_quantity.shape
        assert co.shape == overage_quantity.shape

        underage=cu*underage_quantity
        overage=co*overage_quantity

        loss=underage+overage
 
        if mask is not None:
            loss = loss*mask
        return loss


# %% ../../../nbs/agents/benchmark_agents/03_ERM_agents.ipynb 7
class LagLLama(nn.Module):
    def __init__(self,
                input_size,
                hidden_size,
                output_size,
                num_hidden_layers=2,
                num_RNN_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), int(hidden_size/2), num_RNN_layers, batch_first = True)
        self._h1_dropout = nn.Dropout(p=drop_prob)

        self.mlp_blocks = nn.ModuleList()
        for i in range(num_hidden_layers):
            if i == 0:
                self.mlp_blocks.append(MLP_BLOCK(size_in = int(hidden_size/2), size_out = hidden_size, drop_prob=drop_prob))
            else:
                self.mlp_blocks.append(MLP_BLOCK(size_in = int(hidden_size/(2**(i-1))), size_out = int(hidden_size/(2**i)), drop_prob=drop_prob))

        if num_hidden_layers==0:
            self.output_layer = nn.Linear(int(hidden_size/2), output_size)
        else:
            self.output_layer = nn.Linear(int(hidden_size/(2**i)), 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 LagLlamasgdAgent(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_RNN_layers = 2,
                    num_hidden_layers=3,
                    drop_prob=0.0,
                    l2_reg=0.0,
                    ): 
        
        self.name=agent_name
        self.model_type="LagLlama"
        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, num_RNN_layers=num_RNN_layers, drop_prob=drop_prob, l2_reg=l2_reg, num_time_series_features=num_time_series_features, lag_window=lag_window, 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)
