from itertools import product
from time import time

import torch
from torch import optim, nn
from torch.utils.data import DataLoader, TensorDataset

from ..builders import models_builders
from ..builders.models_builders import MLPConfig
from ..visualisation import stats_visualisation

OPT_MAP = {
    "SGD": (optim.SGD, ("lr", "momentum", "dampening", "weight_decay", "nesterov")),
    "NAG": (optim.SGD, ("lr", "momentum", "dampening", "weight_decay", "nesterov")),
    "ADAM": (optim.Adam, ("lr", "betas", "eps", "weight_decay", "amsgrad")),
    "RMSPROP": (optim.RMSprop, ("lr", "alpha")),
    "ADAGRAD": (optim.Adagrad, ("lr", "lr_decay")),
}

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

PATIENCE = 10


def build_and_train(x_train, x_test, y_train, y_test, search=False, **kwargs):
    if search:
        initial_search = {
            "lr": tuple(10 ** (-x) for x in range(2, 6)),
            "optimizer": (
                "sgd",
                "nag",
                "adagrad",
                "adam",
                "rmsprop"
            ),
            "momentum": (0.9, 0.99)
        }
        configs, stats = MLPGridSearch().search(
            kwargs["inputs"], kwargs["hiddens"], kwargs["outputs"],
            x_train, x_test, y_train, y_test,
            batch_size=100, epochs=20,
            **initial_search
        )

        configs_errors = [(configs[i], stats[i]) for i in range(len(configs))]
        configs_errors.sort(key=lambda x: min(x[1]))

        best_errors = [min(e) for e in stats]
        best_error = min(best_errors)
        best_index = best_errors.index(best_error)
        best_config = configs[best_index]

        stats_visualisation.plot_stats_mlps(configs, stats, start=7)

        return best_config

    else:
        config = models_builders.MLPConfig(**kwargs)

        try:
            config.load_model()
            stats = config.load_stats()
        except Exception as e:
            model = config.build_model()

            stats = MLPTrainer().train(model, x_train, x_test, y_train, y_test, model_config=config)
            config.save_model(model)
            config.save_stats(stats)

        # stats_visualisation.plot_stats_mlp(config, stats, config.epochs // 2)

        return config


class Trainer:
    def get_optimizer(self, model, config):
        opt_name = config.optimizer.upper()
        opt_class, _opt_args = OPT_MAP[opt_name]
        opt_args = dict((k, v) for k, v in config.items() if k in _opt_args)
        return opt_class(model.parameters(), **opt_args)


class MLPTrainer(Trainer):
    def train(self,
              model,
              x_train,
              x_test,
              y_train,
              y_test,
              model_config=None,
              train_only_classifier=False):

        model = model.to(DEVICE)

        loss_class = getattr(nn, model_config.loss)
        loss = loss_class()

        """ Optimizer setup """

        opt_name = model_config.optimizer.upper()

        opt_class, _opt_args = OPT_MAP.get(opt_name)
        opt_args = dict((k, v) for k, v in model_config.dict.items() if k in _opt_args)

        optimizer = opt_class(model.parameters(), **opt_args)

        """ Go """

        test_accuracy_history = []
        test_loss_history = []

        train_accuracy_history = []
        train_loss_history = []

        train_set = TensorDataset(x_train, y_train)

        train_loader = DataLoader(
            train_set,
            batch_size=model_config.batch_size,
            shuffle=True
        )

        x_test = x_test.to(DEVICE)
        y_test = y_test.to(DEVICE)

        hidden_disabled = False

        if train_only_classifier:
            lock_hiddens(model)
            hidden_disabled = True

        best_epoch_err = None
        best_epoch_state = None
        bad_errs_count = 0

        for epoch in range(model_config.epochs):
            start = time()

            xs = None
            ys = None

            for idx, (xs, ys) in enumerate(train_loader):
                xs = xs.to(DEVICE)
                ys = ys.to(DEVICE)

                optimizer.zero_grad()

                predicted = model(xs)

                loss_value = loss(predicted, ys)
                loss_value.backward()

                optimizer.step()

                # if epoch > 50 and hidden_disabled:
                #     unlock_hiddens(model)
                #     hidden_disabled = False

            train_predictions = model.forward(xs).data
            train_loss_history.append(loss(train_predictions, ys).item())
            accuracy_train = (train_predictions.argmax(dim=1) == ys).float().mean().item()
            train_accuracy_history.append(accuracy_train)

            test_predictions = model.forward(x_test).data

            error = loss(test_predictions, y_test).item()
            test_loss_history.append(error)
            accuracy_test = (test_predictions.argmax(dim=1) == y_test).float().mean().item()
            test_accuracy_history.append(accuracy_test)

            print(f"Epoch: {epoch} >>> "
                  f"TEST ACC: {accuracy_test:.4f} :: TEST ERR: {test_loss_history[-1]:.4f} <<< | "
                  f"Train acc: {accuracy_train:.4f} | "
                  f"Train err: {train_loss_history[-1]:.4f} | :: "
                  f"Spent {time() - start:.4f} sec")

            if best_epoch_err is None or best_epoch_err > error:
                best_epoch_err = error
                best_epoch_state = model.state_dict()
                bad_errs_count = 0
                print(f"New Best err {best_epoch_err:.6f}; Best ep {epoch}")
            elif best_epoch_err < error:
                bad_errs_count += 1

                if bad_errs_count > PATIENCE:
                    model.load_state_dict(best_epoch_state)
                    break

        return {
            "test_errors": test_loss_history,
            "train_errors": train_loss_history,
            "test_accuracy": test_accuracy_history,
            "train_accuracy": train_accuracy_history
        }


class MLPGridSearch:
    def search(
            self,
            inputs, hiddens, outputs,
            x_train, x_test, y_train, y_test,
            batch_size, epochs,
            lr, optimizer, momentum):

        configs = []
        reports = []

        grid_search_config = list(product(optimizer, lr, momentum))
        print(f"Search in {len(grid_search_config)} combinations")

        for idx, variant in enumerate(grid_search_config):
            params = {
                "inputs": inputs,
                "hiddens": hiddens,
                "outputs": outputs,
                "batch_size": batch_size,
                "optimizer": variant[0],
                "lr": variant[1],
                "momentum": variant[2],
                "epochs": epochs
            }

            config = MLPConfig(**params)
            configs.append(config)

            try:
                config.load_model()
                stats = config.load_stats()
            except Exception as e:
                print(e)
                print(f"\nCurrent index {idx + 1}: {config.name}")
                model = config.build_model()
                start = time()
                stats = MLPTrainer().train(
                    model,
                    x_train,
                    x_test,
                    y_train,
                    y_test,
                    model_config=config
                )
                print(f"Spent: {time() - start}")
                config.save_model(model)
                config.save_stats(stats)

            reports.append(stats["test_errors"])

        return configs, reports


def lock_hiddens(model):
    print("Disable layers")
    print(model.layers[:-1])
    for idx, layer in enumerate(model.layers[:-1]):
        params = layer.parameters()
        for p in params:
            p.requires_grad = False


def unlock_hiddens(model):
    print("Enable layers")
    print(model.layers[:-1])
    for idx, layer in enumerate(model.layers[:-1]):
        params = layer.parameters()
        for p in params:
            p.requires_grad = True


def validate_mlp(data, config, **kwargs):
    _, _, x_label_train, x_label_test, y_label_train, y_label_test = data

    x_label_test = x_label_test.to(DEVICE)
    y_label_test = y_label_test.to(DEVICE)

    model = config.load_model().to(DEVICE)

    loss_class = getattr(nn, config.loss)
    loss = loss_class()

    test_predictions = model.forward(x_label_test).data
    error = loss(test_predictions, y_label_test).item()
    accuracy_test = (test_predictions.argmax(dim=1) == y_label_test).float().mean().item()

    print(f"\nFULL TEST DATASET: ACC {accuracy_test:.6f} :: ERR {error:.6f}")

from itertools import product
from time import time

import torch
from torch import optim, nn
from torch.utils.data import DataLoader, TensorDataset

from ..builders import models_builders
from ..builders.models_builders import MLPConfig
from ..visualisation import stats_visualisation

OPT_MAP = {
    "SGD": (optim.SGD, ("lr", "momentum", "dampening", "weight_decay", "nesterov")),
    "NAG": (optim.SGD, ("lr", "momentum", "dampening", "weight_decay", "nesterov")),
    "ADAM": (optim.Adam, ("lr", "betas", "eps", "weight_decay", "amsgrad")),
    "RMSPROP": (optim.RMSprop, ("lr", "alpha")),
    "ADAGRAD": (optim.Adagrad, ("lr", "lr_decay")),
}

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

PATIENCE = 10


def build_and_train(x_train, x_test, y_train, y_test, search=False, **kwargs):
    if search:
        initial_search = {
            "lr": tuple(10 ** (-x) for x in range(2, 6)),
            "optimizer": (
                "sgd",
                "nag",
                "adagrad",
                "adam",
                "rmsprop"
            ),
            "momentum": (0.9, 0.99)
        }
        configs, stats = MLPGridSearch().search(
            kwargs["inputs"], kwargs["hiddens"], kwargs["outputs"],
            x_train, x_test, y_train, y_test,
            batch_size=100, epochs=20,
            **initial_search
        )

        configs_errors = [(configs[i], stats[i]) for i in range(len(configs))]
        configs_errors.sort(key=lambda x: min(x[1]))

        best_errors = [min(e) for e in stats]
        best_error = min(best_errors)
        best_index = best_errors.index(best_error)
        best_config = configs[best_index]

        stats_visualisation.plot_stats_mlps(configs, stats, start=7)

        return best_config

    else:
        config = models_builders.MLPConfig(**kwargs)

        try:
            config.load_model()
            stats = config.load_stats()
        except Exception as e:
            model = config.build_model()

            stats = MLPTrainer().train(model, x_train, x_test, y_train, y_test, model_config=config)
            config.save_model(model)
            config.save_stats(stats)

        # stats_visualisation.plot_stats_mlp(config, stats, config.epochs // 2)

        return config


class Trainer:
    def get_optimizer(self, model, config):
        opt_name = config.optimizer.upper()
        opt_class, _opt_args = OPT_MAP[opt_name]
        opt_args = dict((k, v) for k, v in config.items() if k in _opt_args)
        return opt_class(model.parameters(), **opt_args)


class MLPTrainer(Trainer):
    def train(self,
              model,
              x_train,
              x_test,
              y_train,
              y_test,
              model_config=None,
              train_only_classifier=False):

        model = model.to(DEVICE)

        loss_class = getattr(nn, model_config.loss)
        loss = loss_class()

        """ Optimizer setup """

        opt_name = model_config.optimizer.upper()

        opt_class, _opt_args = OPT_MAP.get(opt_name)
        opt_args = dict((k, v) for k, v in model_config.dict.items() if k in _opt_args)

        optimizer = opt_class(model.parameters(), **opt_args)

        """ Go """

        test_accuracy_history = []
        test_loss_history = []

        train_accuracy_history = []
        train_loss_history = []

        train_set = TensorDataset(x_train, y_train)

        train_loader = DataLoader(
            train_set,
            batch_size=model_config.batch_size,
            shuffle=True
        )

        x_test = x_test.to(DEVICE)
        y_test = y_test.to(DEVICE)

        hidden_disabled = False

        if train_only_classifier:
            lock_hiddens(model)
            hidden_disabled = True

        best_epoch_err = None
        best_epoch_state = None
        bad_errs_count = 0

        for epoch in range(model_config.epochs):
            start = time()

            xs = None
            ys = None

            for idx, (xs, ys) in enumerate(train_loader):
                xs = xs.to(DEVICE)
                ys = ys.to(DEVICE)

                optimizer.zero_grad()

                predicted = model(xs)

                loss_value = loss(predicted, ys)
                loss_value.backward()

                optimizer.step()

                # if epoch > 50 and hidden_disabled:
                #     unlock_hiddens(model)
                #     hidden_disabled = False

            train_predictions = model.forward(xs).data
            train_loss_history.append(loss(train_predictions, ys).item())
            accuracy_train = (train_predictions.argmax(dim=1) == ys).float().mean().item()
            train_accuracy_history.append(accuracy_train)

            test_predictions = model.forward(x_test).data

            error = loss(test_predictions, y_test).item()
            test_loss_history.append(error)
            accuracy_test = (test_predictions.argmax(dim=1) == y_test).float().mean().item()
            test_accuracy_history.append(accuracy_test)

            print(f"Epoch: {epoch} >>> "
                  f"TEST ACC: {accuracy_test:.4f} :: TEST ERR: {test_loss_history[-1]:.4f} <<< | "
                  f"Train acc: {accuracy_train:.4f} | "
                  f"Train err: {train_loss_history[-1]:.4f} | :: "
                  f"Spent {time() - start:.4f} sec")

            if best_epoch_err is None or best_epoch_err > error:
                best_epoch_err = error
                best_epoch_state = model.state_dict()
                bad_errs_count = 0
                print(f"New Best err {best_epoch_err:.6f}; Best ep {epoch}")
            elif best_epoch_err < error:
                bad_errs_count += 1

                if bad_errs_count > PATIENCE:
                    model.load_state_dict(best_epoch_state)
                    break

        return {
            "test_errors": test_loss_history,
            "train_errors": train_loss_history,
            "test_accuracy": test_accuracy_history,
            "train_accuracy": train_accuracy_history
        }


class MLPGridSearch:
    def search(
            self,
            inputs, hiddens, outputs,
            x_train, x_test, y_train, y_test,
            batch_size, epochs,
            lr, optimizer, momentum):

        configs = []
        reports = []

        grid_search_config = list(product(optimizer, lr, momentum))
        print(f"Search in {len(grid_search_config)} combinations")

        for idx, variant in enumerate(grid_search_config):
            params = {
                "inputs": inputs,
                "hiddens": hiddens,
                "outputs": outputs,
                "batch_size": batch_size,
                "optimizer": variant[0],
                "lr": variant[1],
                "momentum": variant[2],
                "epochs": epochs
            }

            config = MLPConfig(**params)
            configs.append(config)

            try:
                config.load_model()
                stats = config.load_stats()
            except Exception as e:
                print(e)
                print(f"\nCurrent index {idx + 1}: {config.name}")
                model = config.build_model()
                start = time()
                stats = MLPTrainer().train(
                    model,
                    x_train,
                    x_test,
                    y_train,
                    y_test,
                    model_config=config
                )
                print(f"Spent: {time() - start}")
                config.save_model(model)
                config.save_stats(stats)

            reports.append(stats["test_errors"])

        return configs, reports


def lock_hiddens(model):
    print("Disable layers")
    print(model.layers[:-1])
    for idx, layer in enumerate(model.layers[:-1]):
        params = layer.parameters()
        for p in params:
            p.requires_grad = False


def unlock_hiddens(model):
    print("Enable layers")
    print(model.layers[:-1])
    for idx, layer in enumerate(model.layers[:-1]):
        params = layer.parameters()
        for p in params:
            p.requires_grad = True


def validate_mlp(data, config, **kwargs):
    _, _, x_label_train, x_label_test, y_label_train, y_label_test = data

    x_label_test = x_label_test.to(DEVICE)
    y_label_test = y_label_test.to(DEVICE)

    model = config.load_model().to(DEVICE)

    loss_class = getattr(nn, config.loss)
    loss = loss_class()

    test_predictions = model.forward(x_label_test).data
    error = loss(test_predictions, y_label_test).item()
    accuracy_test = (test_predictions.argmax(dim=1) == y_label_test).float().mean().item()

    print(f"\nFULL TEST DATASET: ACC {accuracy_test:.6f} :: ERR {error:.6f}")

# stats_visualisation.plot_stats_mlp(config, stats, config.epochs // 2) Optimizer setup Go # if epoch > 50 and hidden_disabled: #     unlock_hiddens(model) #     hidden_disabled = False