import torch
from torch import nn
import torch.nn.functional as F
import numpy as np
from tools import optimization_oracle_tools
import random

# ICEO framework
class iceo_loss_func(nn.Module):
    def __init__(self, params):
        super().__init__()
        self.c_list = params['c_list']
        self.rho = params['regularization_rho']
        self.dim_sol = params['dim_sol']

    def forward(self, w_pred, z_real):
        input = w_pred - z_real
        c_tensor = torch.as_tensor(self.c_list, dtype=torch.float32)
        loss = ( c_tensor* torch.square(input) + (0.5*self.rho)*torch.square(torch.norm(w_pred, dim=1))).mean()
        return loss
    
class nf_loss_func(nn.Module):
    def __init__(self, params):
        super().__init__()
        self.c_list = params['c_list']
        self.rho = params['regularization_rho']
        self.dim_sol = params['dim_sol']

    def forward(self, w_pred, z_real):
        input = w_pred - z_real
        c_tensor = torch.as_tensor(self.c_list, dtype=torch.float32)
        loss = ( c_tensor* torch.square(input)).mean()
        return loss

class two_layers_pred(nn.Module):
    def __init__(self, input_size, output_size, num_nodes):
        super(two_layers_pred, self).__init__()
        self.fc1 = nn.Linear(input_size, num_nodes)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout(p=0.01)
        self.bn1 = nn.BatchNorm1d(num_nodes)
        self.fc2 = nn.Linear(num_nodes, num_nodes)
        self.fc3 = nn.Linear(num_nodes, output_size)

        nn.init.normal_(self.fc1.weight, mean=0.0, std=np.sqrt(1)/np.sqrt(num_nodes))
        nn.init.normal_(self.fc2.weight, mean=0.0, std=np.sqrt(1)/np.sqrt(num_nodes))
        nn.init.normal_(self.fc2.weight, mean=0.0, std=np.sqrt(1)/np.sqrt(output_size))
    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        # out = self.sigmoid(out)
        # out = self.bn1(out)
        out = self.fc2(out)
        out = self.relu(out)
        # out = self.dropout(out)
        out = self.fc3(out)
        out = F.softmax(out, dim=1)
        return out

class linear_pred(nn.Module):
    def __init__(self, input_size, output_size, num_nodes):
        super(linear_pred, self).__init__()
        self.fc1 = nn.Linear(input_size, num_nodes)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout(p=0.01)
        self.bn1 = nn.BatchNorm1d(num_nodes)
        self.fc2 = nn.Linear(num_nodes, num_nodes)
        self.fc3 = nn.Linear(num_nodes, output_size)
        self.fc4 = nn.Linear(input_size, output_size)

        nn.init.normal_(self.fc1.weight, mean=0.0, std=np.sqrt(1)/np.sqrt(num_nodes))
        nn.init.normal_(self.fc2.weight, mean=0.0, std=np.sqrt(1)/np.sqrt(num_nodes))
        nn.init.normal_(self.fc3.weight, mean=0.0, std=np.sqrt(1)/np.sqrt(output_size))
        nn.init.normal_(self.fc4.weight, mean=0.0, std=np.sqrt(1)/np.sqrt(output_size))

    def forward(self, x):
        out = self.fc1(x)
        # out = self.relu(out)
        # out = self.sigmoid(out)
        # out = self.bn1(out)
        # out = self.fc2(out)
        # out = self.relu(out)
        # out = self.dropout(out)
        out = self.fc3(out)
        out = self.relu(out)
        out = F.softmax(out, dim=1)
        return out

class two_layers_true(nn.Module):
    def __init__(self, input_size, output_size, num_nodes):
        super(two_layers_true, self).__init__()
        self.fc1 = nn.Linear(input_size, num_nodes)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(num_nodes, num_nodes)
        self.fc3 = nn.Linear(num_nodes, output_size)

        # nn.init.normal_(self.fc1.weight, mean=0.0, std=0.5)
        # nn.init.normal_(self.fc2.weight, mean=0.0, std=0.5)
        # nn.init.normal_(self.fc3.weight, mean=0.0, std=0.5)
        nn.init.normal_(self.fc1.weight, mean=0.0, std=0.5)
        nn.init.normal_(self.fc2.weight, mean=0.0, std=0.5)
        nn.init.normal_(self.fc3.weight, mean=0.1, std=0.5)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc3(out)
        out = F.softmax(out, dim=1)
        return out


class multi_layers_true(nn.Module):
    def __init__(self, input_size, output_size, num_nodes):
        super(multi_layers_true, self).__init__()
        self.fc1 = nn.Linear(input_size, num_nodes)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.fc2 = nn.Linear(num_nodes, num_nodes)
        self.fc4 = nn.Linear(num_nodes, num_nodes)
        self.fc5 = nn.Linear(num_nodes, num_nodes)
        self.fc3 = nn.Linear(num_nodes, output_size)
        self.fc = nn.Linear(input_size, output_size)
        

        def init_weights(self):
            random.seed(123) 
            for layer in [self.fc1, self.fc3]:
                n_inputs, n_outputs = layer.in_features, layer.out_features
                variance = 2.0 / (n_inputs + n_outputs)
                std_dev = torch.sqrt(torch.FloatTensor([variance]))
                layer.weight.data.normal_(0.0, std_dev)



    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        # out = self.fc4(out)
        # out = self.relu(out)
        out = self.fc5(out)
        out = self.relu(out)
        out = self.fc3(out)
        out = F.softmax(out, dim=1)
        return out




class ICEO_model(nn.Module):
    def __init__(self, params, predict_model, oracle_model):
        super(ICEO_model, self).__init__()
        # two NNs
        self.predict_model = predict_model
        self.oracle_model = oracle_model
        # self.loss_func = params['loss_func']
        self.dim_features = params['dim_features']
        self.dim_sol = params['dim_sol']
        self.dim_scenario = params['dim_scenario']
        # self.opt_oracle = params['opt_oracle']
        self.optimizer = params.get('optimizer', None)
        self.optimizer_config = params.get('optimizer_config', None)
        self.minibatch_size = params.get('minibatch_size', 32)
    def forward(self, x):
        x = self.predict_model(x)
        out1 = x
        # print(x)
        x = self.oracle_model(x)
        return x, out1
    

