import torch
import scipy.io as sio
import torch.nn as nn
import numpy as np
from pymoo.core.problem import Problem
# from layer_config_forward_oracle import MultiLayerPerceptron_forward
from layer_config_forward import MultiLayerPerceptron_forward

class AB_BO_surrogate_uncertainty(Problem):
    def __init__(self, n_iter):
        super().__init__(n_var=44,
                         n_obj=4,
                         xl=0,
                         xu=1)

        # Load the forward model
        input_size = 44
        hidden_size_mu = [100, 50, 100]
        self.num_classes = 8
        # Load the model
        self.mu_models = []
        for n_net in range(1, 11):
            # load mu models
            mu_model = MultiLayerPerceptron_forward(input_size, hidden_size_mu, self.num_classes, n_net)
            model_name = 'Models/iter_%d/mu_net_%d.ckpt' % (n_iter, n_net)
            mu_model.load_state_dict(torch.load(model_name, map_location=torch.device('cpu')))
            mu_model.eval()
            self.mu_models.append(mu_model)

    def _evaluate(self, design, out, *args, **kwargs):
        designs = torch.tensor(design).float()
        batchsize = designs.shape[0]
        reproduced_Performance_ensembel = torch.empty(10, batchsize, self.num_classes)

        reproduced_Performance_mu = torch.empty(batchsize, self.num_classes)
        uncertainty_epistemic = torch.empty((batchsize, self.num_classes))

        for net_n in range(10):
            reproduced_Performance_ensembel[net_n, :, :] = self.mu_models[net_n](designs)

        reproduced_Performance_mu = (1 / 10) * torch.sum(reproduced_Performance_ensembel, 0)
        uncertainty_epistemic = (1 / 10) * torch.sum(
            reproduced_Performance_ensembel ** 2 - reproduced_Performance_mu.repeat(10, 1, 1) ** 2,
            0)
        reproduced_Performance_mu[:, 1] = -reproduced_Performance_mu[:, 1]
        out["F"] = torch.cat(
            (reproduced_Performance_mu[:, :2].detach(), -uncertainty_epistemic[:, :2].detach()), 1)
        # out["E_A"] = uncertainty_epistemic[:,:2].detach()
