class MfModel(ABC):
    """
    Wrapping class around a multifidelity model to define it, train it and evaluate it.

    Attributes:
        dim (int): problem dimension.
        model_dict (dict): dictionary containing all necessary model parameters.
        outdir (str): path to the model folder.
        seed (int): seed to enforce reproducibility.
        x_lf_DOE (np.ndarray): the lf_DOE the model was trained with.
        x_hf_DOE (np.ndarray): the hf_DOE the model was trained with.
        y_lf_DOE (np.ndarray): the lf_DOE the model was trained with.
        y_hf_DOE (np.ndarray): the hf_DOE the model was trained with.
    """
    def __init__(self, dim: int, model_dict: dict, outdir: str, seed: int):
        self.dim = dim
        self.model_dict = model_dict
        self.outdir = outdir
        self.seed = seed
        self.requires_nested_doe: bool
        self.x_lf_DOE: np.ndarray | None = None
        self.x_hf_DOE: np.ndarray | None = None
        self.y_lf_DOE: np.ndarray | None = None
        self.y_hf_DOE: np.ndarray | None = None
        # seed numpy
        np.random.seed(seed)

    @abstractmethod
    def train(self):
        """
        Trains a model with low and high fidelity data.
        """

    @abstractmethod
    def evaluate(self, x: np.ndarray):
        """
        Returns a model prediction for a given input x.
        """

    def evaluate_std(self, x: np.ndarray):
        raise Exception("Not implemented")

    def get_DOE(self) -> np.ndarray:
        """
        Returns the DOE the model was trained with.
        """
        assert self.x_lf_DOE is not None and self.x_hf_DOE is not None
        if self.requires_nested_doe:
            return np.copy(self.x_lf_DOE)
        else:
            return np.vstack((self.x_lf_DOE, self.x_hf_DOE))

    def set_DOE(self, x_lf: np.ndarray, x_hf: np.ndarray, y_lf: np.ndarray, y_hf: np.ndarray):
        """
        Sets the hf and lf DOEs the model was trained with.
        """
        if self.x_lf_DOE is None:
            self.x_lf_DOE = x_lf
            self.x_hf_DOE = x_hf
            self.y_lf_DOE = y_lf
            self.y_hf_DOE = y_hf
        else:
            assert (
                self.x_lf_DOE is not None
                and self.x_hf_DOE is not None
                and self.y_lf_DOE is not None
                and self.y_hf_DOE is not None
            )
            self.x_lf_DOE = np.vstack((self.x_lf_DOE, x_lf))
            self.x_hf_DOE = np.vstack((self.x_hf_DOE, x_hf))
            if self.y_lf_DOE.ndim == 2:
                self.y_lf_DOE = np.vstack((self.y_lf_DOE, y_lf))
                self.y_hf_DOE = np.vstack((self.y_hf_DOE, y_hf))
            else:
                self.y_lf_DOE = np.hstack((self.y_lf_DOE, y_lf))
                self.y_hf_DOE = np.hstack((self.y_hf_DOE, y_hf))

    def get_y_star(self) -> float | np.ndarray:
        """
        Returns the current best lf fitness (SOO).
        """
        # single objective, y_star is simply the min value
        assert self.y_hf_DOE is not None
        return min(self.y_hf_DOE)

class MfDNN(MfModel):
    """
    Wrapping class around a torch multifidelity deep neural network model.
    """
    def __init__(self, dim: int, model_dict: dict, outdir: str, seed: int):
        super().__init__(dim, model_dict, outdir, seed)
        check_dir(outdir)
        self.requires_nested_doe = model_dict.get("nested_doe", False)
        self.name = model_dict.get("name", "MFDNN")
        # seed torch
        torch.manual_seed(seed)
        # select device
        self.device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # the first layer of both network is created here
        self.model_dict["NNL"]["layer_sizes_NNL"].insert(0, dim)
        self.n_obj = self.model_dict["NNL"]["layer_sizes_NNL"][-1]
        self.model_dict["NNH"]["layer_sizes_NNH1"].insert(0, dim + self.n_obj)
        self.model_dict["NNH"]["layer_sizes_NNH2"].insert(0, dim + self.n_obj)
        # low fidelity network
        self.NNL = NNL(self.model_dict["NNL"]["layer_sizes_NNL"])
        # high fidelity network
        self.NNH = NNH(
            self.model_dict["NNH"]["layer_sizes_NNH1"],
            self.model_dict["NNH"]["layer_sizes_NNH2"]
        )
        # to device
        self.NNL.to(self.device)
        self.NNH.to(self.device)
        # init weights
        self.NNL.apply(weights_init)
        self.NNH.apply(weights_init)

    def train(self):
        # pretrain low fidelity model
        weight_decay_NNL = self.model_dict["NNL"]["optimizer"].get("weight_decay", 0)
        if self.model_dict["pretraining"]:
            loss_target = self.model_dict["NNL"].get("loss_target", 1e-3)
            niter = self.model_dict["NNL"].get("niter", 10000)
            # define optimizer
            lr = self.model_dict["NNL"]["optimizer"].get("lr", 1e-3)
            NNL_optimizer = torch.optim.Adam(
                self.NNL.parameters(),
                lr=lr,
                weight_decay=weight_decay_NNL
            )
            # define scheduler
            if "scheduler" in self.model_dict["NNL"]:
                NNL_scheduler = torch.optim.lr_scheduler.MultiStepLR(
                    NNL_optimizer, **self.model_dict["NNL"]["scheduler"]
                )
            else:
                NNL_scheduler = None
            # pretrain
            assert self.x_lf_DOE is not None and self.y_lf_DOE is not None
            NNL_pretrain(
                self.NNL,
                NNL_optimizer,
                self.x_lf_DOE, self.y_lf_DOE,
                loss_target,
                niter,
                self.device,
                scheduler=NNL_scheduler
            )
        # coupled training
        loss_target = self.model_dict["NNH"].get("loss_target", 1e-5)
        niter = self.model_dict["NNH"].get("niter", 15000)
        # define optimizer
        lr = self.model_dict["NNH"]["optimizer"].get("lr", 1e-4)
        weight_decay_NNH1 = (
            self.model_dict["NNH"]["optimizer"].get("weight_decay_NNH1", 0)
        )
        weight_decay_NNH2 = (
            self.model_dict["NNH"]["optimizer"].get("weight_decay_NNH2", 0)
        )
        MfDNN_optimizer = torch.optim.Adam(
            [{'params': self.NNH.NNH1.parameters(), 'weight_decay': weight_decay_NNH1},
             {'params': self.NNH.NNH2.parameters(), 'weight_decay': weight_decay_NNH2},
             {'params': self.NNH.alpha},
             {'params': self.NNL.parameters(), 'weight_decay': weight_decay_NNL}], lr=lr
        )
        # define scheduler
        if "scheduler" in self.model_dict["NNH"]:
            MfDNN_scheduler = torch.optim.lr_scheduler.MultiStepLR(
                MfDNN_optimizer,
                **self.model_dict["NNH"]["scheduler"]
            )
        else:
            MfDNN_scheduler = None
        # train
        assert self.x_lf_DOE is not None and self.y_lf_DOE is not None
        assert self.x_hf_DOE is not None and self.y_hf_DOE is not None
        MfDNN_train(
            self.NNL,
            self.NNH,
            MfDNN_optimizer,
            self.x_lf_DOE, self.y_lf_DOE,
            self.x_hf_DOE, self.y_hf_DOE,
            loss_target,
            niter,
            self.device,
            self.outdir,
            scheduler=MfDNN_scheduler
        )

    def evaluate(self, x: np.ndarray) -> np.ndarray | list[np.ndarray]:
        x_torch = torch.from_numpy(x).float()
        x_torch = x_torch.to(self.device)
        y_lo_hi = self.NNL.eval()(x_torch)
        y = self.NNH.eval()(torch.cat((x_torch, y_lo_hi), dim=1))
        y_np = y.cpu().detach().numpy()
        return y_np if self.n_obj == 1 else [y_np[:, c_][:, None] for c_ in range(self.n_obj)]

class MfSMT(MfModel):
    """
    Wrapping class around an smt multifidelity cokriging model.
    """
    def __init__(self, dim: int, model_dict: dict, outdir: str, seed: int):
        super().__init__(dim, model_dict, outdir, seed)
        self.requires_nested_doe = model_dict.get("nested_doe", True)
        self.name = model_dict.get("name", "AR1")
        self.eval_noise = model_dict.get("eval_noise", False)
        self.model = MFK(theta0=dim * [1.0], eval_noise=self.eval_noise, print_global=False)

    def train(self):
        self.model.set_training_values(self.x_lf_DOE, self.y_lf_DOE, name=0)
        self.model.set_training_values(self.x_hf_DOE, self.y_hf_DOE)
        self.model.train()

    def evaluate(self, x: np.ndarray) -> np.ndarray:
        y = self.model.predict_values(x)
        return y

    def evaluate_std(self, x: np.ndarray) -> np.ndarray:
        return np.sqrt(self.model.predict_variances(x))

class MultiObjectiveModel(MfModel):
    """
    Multi-objective surrogate model.
    """
    def __init__(self, list_of_models: list[MfSMT | MfKPLS | SfSMT | MfLGP]):
        self.models: list[MfSMT | MfKPLS | SfSMT | MfLGP] = list_of_models
        self.requires_nested_doe = self.models[0].requires_nested_doe
        self.name = self.models[0].name

    def train(self):
        for model in self.models:
            model.train()

    def evaluate(self, x: np.ndarray) -> list[np.ndarray]:
        return [model.evaluate(x) for model in self.models]

    def evaluate_std(self, x: np.ndarray) -> list[np.ndarray]:
        return [model.evaluate_std(x) for model in self.models]

    def get_DOE(self) -> np.ndarray:
        return self.models[0].get_DOE()

    def set_DOE(
            self,
            x_lf: np.ndarray, x_hf: np.ndarray,
            y_lf: np.ndarray | list[np.ndarray], y_hf: np.ndarray | list[np.ndarray]
    ):
        for model, yy_lf, yy_hf in zip(self.models, y_lf, y_hf):
            model.set_DOE(x_lf, x_hf, yy_lf, yy_hf)

    def get_y_star(self) -> float | np.ndarray:
        """
        Returns the current best lf fitness (SO) or pareto front (MO).
        """
        raise Exception("Not implemented")

class EDProblem(Problem):
    """
    Euclidean Distance problem.
    """
    def __init__(self, DOE: np.ndarray, n_var: int, bound: list):
        super().__init__(n_var=n_var, n_obj=1, xl=bound[0], xu=bound[-1])
        self.DOE = DOE

    def _evaluate(self, x: np.ndarray, out: np.ndarray, *args, **kwargs):
        out["F"] = -ED_acquisition_function(x, self.DOE)

class AcquisitionFunctionProblem(Problem):
    """
    Generic class for Bayesian acquisition function optimization problems.
    """
    def __init__(
            self,
            function: Callable,
            model: MfLGP | MfKPLS | MfSMT | SfSMT | MultiObjectiveModel,
            n_var: int,
            bound: list,
            min: bool = True
    ):
        Problem.__init__(
            self, n_var=n_var, n_ieq_constr=0, xl=bound[0], xu=bound[1]
        )
        self.model = model
        self.function = function
        self.min = min

    def _evaluate(self, X: np.ndarray, out: np.ndarray, *args, **kwargs):
        out["F"] = self.function(X, self.model) if self.min else -self.function(X, self.model)

class RegCritProblem(Problem):
    """
    Regularized infill criterion problem:
    see R. Grapin et al. (2022): https://doi.org/10.2514/6.2022-4053
    """
    def __init__(
            self, function: Callable, model: MultiObjectiveModel,
            n_var: int, bound: list, gamma: float
    ):
        super().__init__(n_var=n_var, n_obj=1, xl=bound[0], xu=bound[-1])
        self.model = model
        self.gamma = gamma
        self.function = function

    def _evaluate(self, x: np.ndarray, out: np.ndarray):
        u1 = self.model.models[0].evaluate(x)
        u2 = self.model.models[1].evaluate(x)
        out["F"] = -(self.gamma * self.function(x, self.model)
                     - np.sum(np.column_stack([u1, u2]), axis=1))