import numpy as np
import os
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

from naslib.utils.utils import AverageMeterGroup
from naslib.predictors.utils.encodings import encode
from naslib.predictors import Predictor

# NOTE: faster on CPU
device = torch.device("cpu")
print("device:", device)


def accuracy_mse(prediction, target, scale=100.0):
    prediction = prediction.detach() * scale
    target = (target) * scale
    return F.mse_loss(prediction, target)


class FeedforwardNet(nn.Module):
    def __init__(
        self,
        input_dims: int = 5,
        num_layers: int = 3,
        layer_width: list = [10, 10, 10],
        output_dims: int = 1,
        activation="relu",
    ):
        super(FeedforwardNet, self).__init__()
        assert (
            len(layer_width) == num_layers
        ), "number of widths should be \
        equal to the number of layers"

        self.activation = eval("F." + activation)

        all_units = [input_dims] + layer_width
        self.layers = nn.ModuleList(
            [nn.Linear(all_units[i], all_units[i + 1]) for i in range(num_layers)]
        )

        self.out = nn.Linear(all_units[-1], 1)

        # make the init similar to the tf.keras version
        for l in self.layers:
            torch.nn.init.xavier_uniform_(l.weight)
            torch.nn.init.zeros_(l.bias)
        torch.nn.init.xavier_uniform_(self.out.weight)
        torch.nn.init.zeros_(self.out.bias)

    def forward(self, x):
        for layer in self.layers:
            x = self.activation(layer(x))
        return self.out(x)

    def basis_funcs(self, x):
        for layer in self.layers:
            x = self.activation(layer(x))
        return x


class MLPPredictor(Predictor):
    def __init__(
        self,
        encoding_type="adjacency_one_hot",
        ss_type="nasbench201",
        hpo_wrapper=False,
        hparams_from_file=False
    ):
        self.encoding_type = encoding_type
        self.ss_type = ss_type
        self.hpo_wrapper = hpo_wrapper
        self.default_hyperparams = {
            "num_layers": 20,
            "layer_width": 20,
            "batch_size": 32,
            "lr": 0.001,
            "regularization": 0.2,
        }
        self.hyperparams = None
        self.hparams_from_file = hparams_from_file

    def get_model(self, **kwargs):
        predictor = FeedforwardNet(**kwargs)
        return predictor

    def fit(self, xtrain, ytrain, train_info=None, epochs=500, loss="mae", verbose=0):

        if self.hparams_from_file and self.hparams_from_file not in ['False', 'None'] \
        and os.path.exists(self.hparams_from_file):
            self.hyperparams = json.load(open(self.hparams_from_file, 'rb'))['mlp']
            print('loaded hyperparams from', self.hparams_from_file)
        elif self.hyperparams is None:
            self.hyperparams = self.default_hyperparams.copy()

        num_layers = self.hyperparams["num_layers"]
        layer_width = self.hyperparams["layer_width"]
        batch_size = self.hyperparams["batch_size"]
        lr = self.hyperparams["lr"]
        regularization = self.hyperparams["regularization"]

        self.mean = np.mean(ytrain)
        self.std = np.std(ytrain)
        if self.encoding_type is not None:
            _xtrain = np.array(
                [
                    encode(arch, encoding_type=self.encoding_type, ss_type=self.ss_type)
                    for arch in xtrain
                ]
            )
        else:
            _xtrain = xtrain
        _ytrain = np.array(ytrain)

        X_tensor = torch.FloatTensor(_xtrain).to(device)
        y_tensor = torch.FloatTensor(_ytrain).to(device)
        train_data = TensorDataset(X_tensor, y_tensor)
        data_loader = DataLoader(
            train_data,
            batch_size=batch_size,
            shuffle=True,
            drop_last=False,
            pin_memory=False,
        )

        self.model = self.get_model(
            input_dims=_xtrain.shape[1],
            num_layers=num_layers,
            layer_width=num_layers * [layer_width],
        )
        self.model.to(device)
        optimizer = optim.Adam(self.model.parameters(), lr=lr, betas=(0.9, 0.99))

        if loss == "mse":
            criterion = nn.MSELoss().to(device)
        elif loss == "mae":
            criterion = nn.L1Loss().to(device)

        self.model.train()

        for e in range(epochs):
            meters = AverageMeterGroup()
            for b, batch in enumerate(data_loader):
                optimizer.zero_grad()
                input = batch[0].to(device)
                target = batch[1].to(device)
                prediction = self.model(input).view(-1)

                loss_fn = criterion(prediction, target)
                # add L1 regularization
                params = torch.cat(
                    [
                        x[1].view(-1)
                        for x in self.model.named_parameters()
                        if x[0] == "out.weight"
                    ]
                )
                loss_fn += regularization * torch.norm(params, 1)
                loss_fn.backward()
                optimizer.step()

                mse = accuracy_mse(prediction, target)
                meters.update(
                    {"loss": loss_fn.item(), "mse": mse.item()}, n=target.size(0)
                )

            if verbose and e % 100 == 0:
                print("Epoch {}, {}, {}".format(e, meters["loss"], meters["mse"]))

        train_pred = np.squeeze(self.query(xtrain))
        train_error = np.mean(abs(train_pred - ytrain))
        return train_error

    def query(self, xtest, info=None, eval_batch_size=None):
        if self.encoding_type is not None:
            xtest = np.array(
                [
                    encode(arch, encoding_type=self.encoding_type, ss_type=self.ss_type)
                    for arch in xtest
                ]
            )
        X_tensor = torch.FloatTensor(xtest).to(device)
        test_data = TensorDataset(X_tensor)

        eval_batch_size = len(xtest) if eval_batch_size is None else eval_batch_size
        test_data_loader = DataLoader(
            test_data, batch_size=eval_batch_size, pin_memory=False
        )

        self.model.eval()
        pred = []
        with torch.no_grad():
            for _, batch in enumerate(test_data_loader):
                prediction = self.model(batch[0].to(device)).view(-1)
                pred.append(prediction.cpu().numpy())

        pred = np.concatenate(pred)
        return np.squeeze(pred)

    def set_random_hyperparams(self):

        if self.hyperparams is None:
            params = self.default_hyperparams.copy()

        else:
            params = {
                "num_layers": int(np.random.choice(range(5, 25))),
                "layer_width": int(np.random.choice(range(5, 25))),
                "batch_size": 32,
                "lr": np.random.choice([0.1, 0.01, 0.005, 0.001, 0.0001]),
                "regularization": 0.2,
            }

        self.hyperparams = params
        return params

[((15, 9, 15, 28), 'torch.device', 'torch.device', ({(15, 22, 15, 27): '"""cpu"""'}, {}), "('cpu')", False, 'import torch\n'), ((22, 11, 22, 41), 'torch.nn.functional.mse_loss', 'F.mse_loss', ({(22, 22, 22, 32): 'prediction', (22, 34, 22, 40): 'target'}, {}), '(prediction, target)', True, 'import torch.nn.functional as F\n'), ((47, 19, 47, 46), 'torch.nn.Linear', 'nn.Linear', ({(47, 29, 47, 42): 'all_units[-1]', (47, 44, 47, 45): '1'}, {}), '(all_units[-1], 1)', True, 'import torch.nn as nn\n'), ((53, 8, 53, 54), 'torch.nn.init.xavier_uniform_', 'torch.nn.init.xavier_uniform_', ({(53, 38, 53, 53): 'self.out.weight'}, {}), '(self.out.weight)', False, 'import torch\n'), ((54, 8, 54, 43), 'torch.nn.init.zeros_', 'torch.nn.init.zeros_', ({(54, 29, 54, 42): 'self.out.bias'}, {}), '(self.out.bias)', False, 'import torch\n'), ((107, 20, 107, 35), 'numpy.mean', 'np.mean', ({(107, 28, 107, 34): 'ytrain'}, {}), '(ytrain)', True, 'import numpy as np\n'), ((108, 19, 108, 33), 'numpy.std', 'np.std', ({(108, 26, 108, 32): 'ytrain'}, {}), '(ytrain)', True, 'import numpy as np\n'), ((118, 18, 118, 34), 'numpy.array', 'np.array', ({(118, 27, 118, 33): 'ytrain'}, {}), '(ytrain)', True, 'import numpy as np\n'), ((122, 21, 122, 54), 'torch.utils.data.TensorDataset', 'TensorDataset', ({(122, 35, 122, 43): 'X_tensor', (122, 45, 122, 53): 'y_tensor'}, {}), '(X_tensor, y_tensor)', False, 'from torch.utils.data import DataLoader, TensorDataset\n'), ((123, 22, 129, 9), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, TensorDataset\n'), ((188, 20, 188, 43), 'torch.utils.data.TensorDataset', 'TensorDataset', ({(188, 34, 188, 42): 'X_tensor'}, {}), '(X_tensor)', False, 'from torch.utils.data import DataLoader, TensorDataset\n'), ((191, 27, 193, 9), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, TensorDataset\n'), ((202, 15, 202, 35), 'numpy.concatenate', 'np.concatenate', ({(202, 30, 202, 34): 'pred'}, {}), '(pred)', True, 'import numpy as np\n'), ((203, 15, 203, 31), 'numpy.squeeze', 'np.squeeze', ({(203, 26, 203, 30): 'pred'}, {}), '(pred)', True, 'import numpy as np\n'), ((51, 12, 51, 51), 'torch.nn.init.xavier_uniform_', 'torch.nn.init.xavier_uniform_', ({(51, 42, 51, 50): 'l.weight'}, {}), '(l.weight)', False, 'import torch\n'), ((52, 12, 52, 40), 'torch.nn.init.zeros_', 'torch.nn.init.zeros_', ({(52, 33, 52, 39): 'l.bias'}, {}), '(l.bias)', False, 'import torch\n'), ((95, 12, 95, 50), 'os.path.exists', 'os.path.exists', ({(95, 27, 95, 49): 'self.hparams_from_file'}, {}), '(self.hparams_from_file)', False, 'import os\n'), ((147, 21, 147, 40), 'naslib.utils.utils.AverageMeterGroup', 'AverageMeterGroup', ({}, {}), '()', False, 'from naslib.utils.utils import AverageMeterGroup\n'), ((197, 13, 197, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((44, 13, 44, 54), 'torch.nn.Linear', 'nn.Linear', ({(44, 23, 44, 35): 'all_units[i]', (44, 37, 44, 53): 'all_units[i + 1]'}, {}), '(all_units[i], all_units[i + 1])', True, 'import torch.nn as nn\n'), ((120, 19, 120, 45), 'torch.FloatTensor', 'torch.FloatTensor', ({(120, 37, 120, 44): '_xtrain'}, {}), '(_xtrain)', False, 'import torch\n'), ((121, 19, 121, 45), 'torch.FloatTensor', 'torch.FloatTensor', ({(121, 37, 121, 44): '_ytrain'}, {}), '(_ytrain)', False, 'import torch\n'), ((187, 19, 187, 43), 'torch.FloatTensor', 'torch.FloatTensor', ({(187, 37, 187, 42): 'xtest'}, {}), '(xtest)', False, 'import torch\n'), ((215, 22, 215, 73), 'numpy.random.choice', 'np.random.choice', ({(215, 39, 215, 72): '[0.1, 0.01, 0.005, 0.001, 0.0001]'}, {}), '([0.1, 0.01, 0.005, 0.001, 0.0001])', True, 'import numpy as np\n'), ((112, 20, 112, 88), 'naslib.predictors.utils.encodings.encode', 'encode', (), '', False, 'from naslib.predictors.utils.encodings import encode\n'), ((140, 24, 140, 36), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((163, 44, 163, 65), 'torch.norm', 'torch.norm', ({(163, 55, 163, 61): 'params', (163, 63, 163, 64): '(1)'}, {}), '(params, 1)', False, 'import torch\n'), ((183, 20, 183, 88), 'naslib.predictors.utils.encodings.encode', 'encode', (), '', False, 'from naslib.predictors.utils.encodings import encode\n'), ((142, 24, 142, 35), 'torch.nn.L1Loss', 'nn.L1Loss', ({}, {}), '()', True, 'import torch.nn as nn\n')]