import torch
import numpy as np
from copy import deepcopy
from typing import List, Optional, Tuple
from torch.utils.data import DataLoader
from supervised.utils import ids, keys, typeddicts
from supervised import saving, data, networks


VERBOSE = False  # Default: whether the code output should be verbose
NR_EPOCHS = 50  # Default max number of epochs for training in case of no early stopping
EARLY_STOPPING_PATIENCE = 5  # Number of epochs to be patient before early stopping the training


def train_epoch_multi_readout(model: networks.classes.MultitaskLearner, data_loader: DataLoader,
                              optimizer: torch.optim.Optimizer, track_pmdd: Optional[bool] = False,
                              training_tracker: Optional[typeddicts.TrackingOutcome] = None,
                              train_params: Optional[typeddicts.TrainParameters] = None,
                              pmdd_dataset: Optional[data.load.DatasetType] = None,
                              validation_dl: Optional[typeddicts.DL] = None
                              ) -> Tuple[List[float], typeddicts.TrackingOutcome]:

    # List to store the train loss for each batch
    loss_list = []
    for batch_idx, ((data_batch, _), (task, targets)) in enumerate(data_loader):
        optimizer.zero_grad()
        out = model.forward(data_batch)
        if len(model.rbs) != 1:
            out = torch.gather(out, 0, torch.add(task.view(1, -1), train_params[keys.K_FIRST_TASK_ID])).reshape(-1)

        loss = networks.evaluate.mse_loss(out, targets)
        loss.backward()
        optimizer.step()
        loss_list += [float(loss.detach())]
        if track_pmdd:
            if (batch_idx + 1) in networks.evaluate.PMDD_TRACK_BATCHES:
                training_tracker = networks.evaluate.track_network(model=model, pmdd_dataset=pmdd_dataset,
                                                                   training_tracker=training_tracker,
                                                                   train_loss=loss_list[-1],
                                                                   validation_dl=validation_dl,
                                                                   train_params=train_params)

    return loss_list, training_tracker


def train_epoch_context_learner(model: networks.classes.MultitaskLearner, data_loader: DataLoader,
                                optimizer: torch.optim.Optimizer, track_pmdd: Optional[bool] = False,
                                training_tracker: Optional[typeddicts.TrackingOutcome] = None,
                                train_params: Optional[typeddicts.TrainParameters] = None,
                                pmdd_dataset: Optional[data.load.DatasetType] = None,
                                validation_dl: Optional[typeddicts.DL] = None
                                ) -> Tuple[List[float], typeddicts.TrackingOutcome]:

    # List to store the train loss for each batch
    loss_list = []

    # Train one batch at a time
    for batch_idx, ((data_batch, _), (tasks, targets)) in enumerate(data_loader):
        optimizer.zero_grad()
        out = model.forward(data_batch, tasks + train_params[keys.K_FIRST_TASK_ID]).reshape(-1)
        loss = networks.evaluate.mse_loss(out, targets)
        loss_list += [float(loss.detach())]
        loss.backward()
        optimizer.step()

        # Track pmdd loss, validation performance & train loss
        if track_pmdd:
            if (batch_idx + 1) in networks.evaluate.PMDD_TRACK_BATCHES:
                training_tracker = networks.evaluate.track_network(model=model, pmdd_dataset=pmdd_dataset,
                                                                   training_tracker=training_tracker,
                                                                   train_loss=loss_list[-1],
                                                                   validation_dl=validation_dl,
                                                                   train_params=train_params)

    return loss_list, training_tracker


def train_epoch(model: networks.classes.MultitaskLearner, data_loader: DataLoader, optimizer: torch.optim.Optimizer,
                prog_params: typeddicts.ProgramParameters, train_params: typeddicts.TrainParameters,
                save_params: typeddicts.SavingParameters, pmdd_dataset: data.load.DatasetType,
                validation_dl: Optional[typeddicts.DL] = None, track_pmdd: Optional[bool] = False) -> List[float]:
    # Initialize the tracker for pmdd loss, validation performance & train loss to a database
    if track_pmdd:
        assert validation_dl is not None
        training_tracker = networks.evaluate.init_track_network(model=model, train_params=train_params,
                                                                pmdd_dataset=pmdd_dataset,
                                                                validation_dl=validation_dl)
    else:
        training_tracker = None

    # Train for one epoch
    if prog_params[keys.K_MODEL_ID] == ids.ID_MULTI_READOUT:
        loss_list, training_tracker = train_epoch_multi_readout(
            model=model, data_loader=data_loader, optimizer=optimizer, track_pmdd=track_pmdd, train_params=train_params,
            training_tracker=training_tracker, pmdd_dataset=pmdd_dataset, validation_dl=validation_dl)
    else:
        loss_list, training_tracker = train_epoch_context_learner(
            model=model, data_loader=data_loader, optimizer=optimizer, track_pmdd=track_pmdd, train_params=train_params,
            pmdd_dataset=pmdd_dataset, validation_dl=validation_dl)

    # Save the tracked pmdd loss, validation performance & train loss to a database
    if track_pmdd:
        saving.save.save_first_epoch_batches_pmdd(prog_params=prog_params,
                                                  train_params=train_params,
                                                  save_params=save_params,
                                                  tracked_pmdd_batches=training_tracker)

    return loss_list


def train_model(model: networks.classes.MultitaskLearner, train_data: data.datasets.KTaskNClassMDatasetData,
                validation_data: data.datasets.KTaskNClassMDatasetData, optimizer: torch.optim.Optimizer,
                prog_params: typeddicts.ProgramParameters, train_params: typeddicts.TrainParameters,
                save_params: typeddicts.SavingParameters, verbose: Optional[bool] = VERBOSE
                ) -> Tuple[networks.classes.MultitaskLearner, typeddicts.PerformanceOutcome]:
    # Prepare to evaluate best time to stop training
    validation_data_loader, nr_samples = data.load.get_dataloader(validation_data)
    validation_dl: typeddicts.DL = {keys.K_NUMBER_SAMPLES: nr_samples,
                                    keys.K_VALIDATION_DATALOADER: validation_data_loader}
    best_model = deepcopy(model)
    best_validation_performance = -1.
    best_validation_loss = -1.
    stagnation_counter = 0
    training_tracker = None
    if save_params[keys.K_SAVE_PMDD_LOSS]:
        if save_params[keys.K_PMDD_LOSS_TRACK_DATASET] in [ids.ID_EMNIST, ids.ID_K49, ids.ID_CIFAR100]:
            track_dataset = data.load.get_dataset(save_params[keys.K_PMDD_LOSS_TRACK_DATASET], False).data.reshape(
                [-1, data.datasets.n_input_dimension(save_params[keys.K_PMDD_LOSS_TRACK_DATASET])])
        else:
            raise NotImplementedError(f"{save_params[keys.K_PMDD_LOSS_TRACK_DATASET]} track pmdd loss")

        if save_params[keys.K_SAVE_PMDD_LOSS] is not None:
            training_tracker = {keys.K_TRACKED_TRAIN_LOSS: [], keys.K_TRACKED_VALIDATION_PERFORMANCE: [],
                                keys.K_TRACKED_PMDD_LOSS: []}
            pmdd_loss = networks.evaluate.get_pmdd_loss(dataset=track_dataset,
                                                        weights=np.transpose(model.ws[0].detach().numpy()))
            training_tracker[keys.K_TRACKED_PMDD_LOSS] += [pmdd_loss]
            print(f"Epoch 0 had pmdd L2 {pmdd_loss}")
    else:
        track_dataset = None

    # Train epochs until performance stops to improve
    for epoch in range(1, NR_EPOCHS + 1):
        data_loader = DataLoader(train_data, batch_size=train_params[keys.K_BATCH_SIZE], shuffle=True)

        # Train for one epoch
        if epoch == 1:
            train_loss = train_epoch(model=model, data_loader=data_loader, optimizer=optimizer, prog_params=prog_params,
                                     train_params=train_params, save_params=save_params, pmdd_dataset=track_dataset,
                                     validation_dl=validation_dl, track_pmdd=save_params[keys.K_SAVE_PMDD_LOSS])
        else:
            train_loss = train_epoch(model=model, data_loader=data_loader, optimizer=optimizer, prog_params=prog_params,
                                     train_params=train_params, save_params=save_params, pmdd_dataset=track_dataset)

        # Evaluate model performance on validation dataset
        validation_performance, validation_loss = networks.evaluate.evaluate_performance(
            model_id=prog_params[keys.K_MODEL_ID],
            model=model,
            dataset=validation_data_loader,
            nr_samples=nr_samples,
            task_id0=train_params[keys.K_FIRST_TASK_ID])

        if verbose:
            print(f"Epoch {epoch} had training loss {train_loss} and validation performance {validation_performance}%")

        if save_params[keys.K_SAVE_PMDD_LOSS]:
            training_tracker[keys.K_TRACKED_TRAIN_LOSS] += [train_loss]
            training_tracker[keys.K_TRACKED_VALIDATION_PERFORMANCE] += [validation_performance]
            pmdd_loss = networks.evaluate.get_pmdd_loss(dataset=track_dataset,
                                                        weights=np.transpose(model.ws[0].detach().numpy()))
            training_tracker[keys.K_TRACKED_PMDD_LOSS] += [pmdd_loss]
            print(f"Epoch {epoch} had "
                  f"training loss {sum(train_loss)}, "
                  f"validation performance {validation_performance}% and "
                  f"pmdd L2 {pmdd_loss}")

        # Check whether performance is still improving and stop training otherwise
        if validation_performance > best_validation_performance:
            stagnation_counter = 0
            best_validation_performance = validation_performance
            best_validation_loss = validation_loss
            best_model = deepcopy(model)
        else:
            stagnation_counter += 1
            if stagnation_counter >= EARLY_STOPPING_PATIENCE:
                if verbose:
                    print(f"Early stopping training at epoch {epoch} "
                          f"with validation performance of {best_validation_performance}%")
                break

    # Save the pmdd, train and validation losses during training epochs (if the saving parameters require it)
    saving.save.save_train_epochs_pmdd(prog_params=prog_params, train_params=train_params, save_params=save_params,
                                       tracked_pmdd_epochs=training_tracker)

    # Evaluate training performance
    train_performance, train_loss = networks.evaluate.evaluate_performance(model_id=prog_params[keys.K_MODEL_ID],
                                                                           model=best_model,
                                                                           dataset=train_data,
                                                                           task_id0=train_params[keys.K_FIRST_TASK_ID])

    outcome: typeddicts.PerformanceOutcome = {keys.K_TRAIN_PERFORMANCE: train_performance,
                                              keys.K_TRAIN_LOSS: train_loss,
                                              keys.K_VALIDATION_PERFORMANCE: best_validation_performance,
                                              keys.K_VALIDATION_LOSS: best_validation_loss}

    return best_model, outcome


def get_optimizer(training_type: str, model: networks.classes.MultitaskLearner, train_params: typeddicts.TrainParameters
                  ) -> torch.optim.Optimizer:
    parameters = None
    optimizer = None

    # Train both shared and contextual parameters
    if ids.ID_TRAIN in training_type:
        lr = train_params[keys.K_SHARED_PARAM_LEARNING_RATE]

        if ids.ID_SINGLE_TASK in training_type:
            optimizer = torch.optim.Adam(model.parameters(), lr)

        elif ids.ID_CONTEXT_B_SHARED_W in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws},
                                          {"params": model.bs, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        elif ids.ID_CONTEXT_G_SHARED_BW in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws}, {"params": model.bs},
                                          {"params": model.gs, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        elif ids.ID_CONTEXT_G_SHARED_XW in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws}, {"params": model.xs},
                                          {"params": model.gs, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        elif ids.ID_CONTEXT_G_SHARED_BXW in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws}, {"params": model.xs}, {"params": model.bs},
                                          {"params": model.gs, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        elif ids.ID_CONTEXT_BG_SHARED_W in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws},
                                          {"params": model.bs, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]},
                                          {"params": model.gs, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        elif ids.ID_CONTEXT_M_SHARED_BW in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws}, {"params": model.bs},
                                          {"params": model.ms, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        elif ids.ID_CONTEXT_M_SHARED_XW in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws}, {"params": model.xs},
                                          {"params": model.ms, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        elif ids.ID_CONTEXT_M_SHARED_BXW in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws}, {"params": model.xs}, {"params": model.bs},
                                          {"params": model.ms, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        elif ids.ID_CONTEXT_BM_SHARED_W in training_type:
            optimizer = torch.optim.Adam([{"params": model.ws},
                                          {"params": model.bs, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]},
                                          {"params": model.ms, "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                         lr)

        # Train multiple labels with individual readout neurons each all at the same time
        elif ids.ID_MULTI_READOUT in training_type:
            optimizer = torch.optim.Adam([{"params": [w for w in model.ws] + [b for b in model.bs]},
                                          {"params": [r for r in model.rs] + [rb for rb in model.rbs],
                                           "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}], lr)

    # Train only contextual parameters
    elif ids.ID_TRANSFER in training_type:
        lr = train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]
        if "deepen" in training_type:
            if ids.ID_CONTEXT_B_SHARED_W in training_type:
                optimizer = torch.optim.Adam([{"params": model.ws[1:]},
                                              {"params": model.bs,
                                               "lr": train_params[keys.K_CONTEXT_PARAM_LEARNING_RATE]}],
                                             train_params[keys.K_SHARED_PARAM_LEARNING_RATE])
            else:
                raise NotImplementedError

        elif ids.ID_CONTEXT_B_SHARED_W in training_type:
            parameters = model.bs

        elif ids.ID_CONTEXT_BG_SHARED_W in training_type:
            optimizer = torch.optim.Adam([{"params": model.bs, "lr": lr},
                                          {"params": model.gs, "lr": lr}], lr)

        elif ids.ID_CONTEXT_G_SHARED_BW in training_type or ids.ID_CONTEXT_G_SHARED_XW in training_type \
                or ids.ID_CONTEXT_G_SHARED_BXW in training_type:
            parameters = model.gs

        # Train only readout neurons on top of a fixed trunk
        elif ids.ID_MULTI_READOUT in training_type:
            parameters = [r for r in model.rs] + [rb for rb in model.rbs]

    else:
        raise ValueError(training_type)

    if optimizer is None:
        optimizer = torch.optim.Adam(parameters, lr)

    return optimizer
