# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/12_training.cnn.ipynb (unless otherwise specified).

__all__ = ['Model', 'weights_init', 'Learner', 'Evaluator']

# Cell
#nbdev_comment from __future__ import annotations

# Python utils
from collections import OrderedDict
from tqdm.auto import tqdm
import functools
import operator
from pathlib import Path

import numpy as np

from sklearn.model_selection import train_test_split

from ..data.loading import load_kssl
from ..data.selection import (select_y, select_tax_order, select_X)
from ..data.transform import log_transform_y
from ..data.torch import DataLoaders, SNV_transform
from .metrics import eval_reg

# Deep Learning stack
import torch
from torch import nn
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
from torch.optim import Adam
from torch.optim.lr_scheduler import CyclicLR
from torch.nn import MSELoss

from fastcore.test import *
from fastcore.basics import store_attr
from fastcore.transform import compose

# Cell
# https://datascience.stackexchange.com/questions/40906/determining-size-of-fc-layer-after-conv-layer-in-pytorch
class Model(nn.Module):
    def __init__(self, input_dim, in_channel=1, out_channel=16, is_classifier=False,
                 dropout=0.4):
        super(Model, self).__init__()
        # Build the neural network
        self.feature_extractor = nn.Sequential(
            self.make_convpool_block(in_channel, out_channel),
            self.make_convpool_block(out_channel, out_channel*2),
            self.make_convpool_block(out_channel*2, out_channel*4),
            self.make_convpool_block(out_channel*4, out_channel*8),
            self.make_convpool_block(out_channel*8, out_channel*16))

        num_features_before_fcnn = functools.reduce(operator.mul,
                                                    self.feature_extractor(torch.rand(1,
                                                                                      in_channel,
                                                                                      input_dim)).shape)

        output_layers = [nn.Dropout(dropout),
                         nn.Linear(in_features=num_features_before_fcnn, out_features=20),
                         nn.BatchNorm1d(20),
                         nn.LeakyReLU(0.1),
                         nn.Linear(in_features=20, out_features=1)];

        if is_classifier:
            output_layers.append(nn.Sigmoid())

        self.output_block = nn.Sequential(*output_layers)

    def make_convpool_block(self, input_channels, output_channels, kernel_size=3, stride=1):
        return nn.Sequential(
            nn.Conv1d(input_channels, output_channels, kernel_size, bias=False),
            nn.BatchNorm1d(output_channels),
            nn.LeakyReLU(0.1),
            nn.AvgPool1d(3))

    def forward(self, x):
        batch_size = x.size(0)
        out = self.feature_extractor(x)
        out = out.view(batch_size, -1)  # flatten the vector
        out = self.output_block(out)
        return out

# Cell
def weights_init(m):
    if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear):
        torch.nn.init.kaiming_uniform_(m.weight)
    if isinstance(m, nn.BatchNorm1d):
        torch.nn.init.normal_(m.weight, 0.0, 0.02)
        torch.nn.init.constant_(m.bias, 0)

# Cell
class Learner():
    def __init__(self, model, criterion, opt, n_epochs=50,
                 scheduler=None, verbose=True):
        store_attr() # see https://fastpages.fast.ai/fastcore
        use_cuda = torch.cuda.is_available()
        self.device = torch.device('cuda:0' if use_cuda else 'cpu')
        self._init_losses()
        print(f'Runtime is: {self.device}')

    def fit(self, training_generator, validation_generator,
            losses_append=None, loss_valid_stop=-1, loss_train_stop=-1):
        if losses_append is None: self._init_losses()
        for epoch in tqdm(range(self.n_epochs)):
            loss_valid = []
            loss_train = []

            # Training loop
            self.model.train()
            for batches in training_generator:
                local_batch, local_labels = self._to_device(batches)
                self.opt.zero_grad() #
                y_hat, loss = self.eval(local_batch, local_labels)
                loss.backward(retain_graph=True) #
                self.opt.step() #
                loss_train.append(loss.item())


            # Validation loop
            self.model.eval()
            with torch.no_grad():
                for batches in validation_generator:
                    local_batch, local_labels = self._to_device(batches)
                    y_hat, loss = self.eval(local_batch, local_labels)
                    loss_valid.append(loss.item())

            # mean losses after n epoch
            mean_loss_valid = np.mean(np.array(loss_valid))
            mean_loss_train = np.mean(np.array(loss_train))

            if self.scheduler is not None:
                self.scheduler.step()

            if self.verbose:
                print(f'End of Epoch {epoch + 1}\n Training loss: {mean_loss_train}\n Validation loss: {mean_loss_valid}')

            self.losses['train'].append(mean_loss_train)
            self.losses['valid'].append(mean_loss_valid)

            if mean_loss_valid < loss_valid_stop:
                print('Goal achieved!')
                break
            if (mean_loss_train < loss_train_stop) and not (epoch % 10):
                print('Goal achieved!')
                break
        return self.model, self.losses

    def eval(self, X, y):
        y_hat = self.model(X)
        loss = self.criterion(y_hat, y)
        return y_hat, loss

    def predict(self, data_generator, sample=None):
        self.model.eval() if sample is None else self._MCDropoutOn()
        ys_hat, ys = [], []
        if sample is not None: raise Exception(f'MCDropout needs refactoring')
        with torch.no_grad():
            for batches in data_generator:
                local_batch, local_labels = self._to_device(batches)
                if sample is None:
                    ys_hat.append(self.model(local_batch))
                else:
                    # TO BE FIXED: MCDROPOUT
                    y_hat = np.stack([self.model(local_batch).cpu() for s in tqdm(range(sample))]).squeeze()
                ys.append(local_labels)
        return (self._npify(torch.cat(y)) for y in (ys_hat, ys))

    def _to_device(self, batches):
        if self.device.type == 'cpu': return batches
        return (batch.to(self.device) for batch in batches)

    def _init_losses(self):
        self.losses = {'train': [], 'valid': []}

    def _MCDropoutOn(self):
         # Check https://discuss.pytorch.org/t/turn-off-batch-norm-but-leave-dropout-on/14815
        # to deactivate BatchNorm during inference (maybe not critical when predicting all set)
        # and
        # for m in model.modules():
        #   if isinstance(m, nn.BatchNorm1d):
        #     m.eval()
        self.model.train()
        for m in self.model.modules():
            if isinstance(m, nn.BatchNorm1d):
                m.eval()
    def _npify(self, tensor):
        if self.device.type == 'cpu': return tensor.numpy()
        return tensor.cpu().numpy()


# Cell
class Evaluator():
    def __init__(self, model, data, depth_order,
                 seeds=range(20), device='cpu', n_epochs=5,
                 verbose=True, split_ratio=0.1):
        self.seeds = seeds
        self.model = model
        self.X, self.y = data
        self.in_shape = self.X.shape[1]
        self.depth_order = depth_order
        self.device = device
        self.split_ratio = split_ratio
        self.n_epochs = n_epochs
        self.verbose = verbose

        self.losses = []
        # self.models = []
        self.learners = []
        self.perfs = OrderedDict({'train': [], 'test': []})

    def train_multiple(self,
                       optimizer=torch.optim.Adam,
                       scheduler=torch.optim.lr_scheduler.CyclicLR,
                       criterion=nn.MSELoss(),
                       sc_kwargs={}):

        for seed in tqdm(self.seeds):
            X_train, X_test, y_train, y_test, depth_order_train, depth_order_test = self._splitter(seed)

            dls = DataLoaders(((X_train, y_train), (X_test, y_test)),
                              transform=SNV_transform())
            training_generator, validation_generator = dls.loaders()

            model = self.model(self.in_shape).to(self.device)
            opt = optimizer(model.parameters())
            model = model.apply(weights_init)

            learner = Learner(model, criterion, opt, n_epochs=self.n_epochs,
                              scheduler=scheduler(opt, **sc_kwargs), verbose=self.verbose)

            model, loss = learner.fit(training_generator, validation_generator)

            self.losses.append(loss)
            # self.models.append(model)
            self.learners.append(learner)

    def eval_on_train(self, reducer):
        perfs = []
        for i, seed in enumerate(self.seeds):
            X_train, X_test, y_train, y_test, _, _ = self._splitter(seed)

            dls = DataLoaders(((X_train, y_train), (X_test, y_test)),
                              transform=SNV_transform())
            train_generators , _ = dls.loaders()

            y_hat, y_true = self.learners[i].predict(train_generators)
            perf = eval_reg(y_true, y_hat)
            # perf = self.models[i].eval((X_train, y_train))
            perf['n'] = len(X_train)
            perfs.append(perf)
        if reducer:
            perfs = self.reduce(perfs, reducer)
        return perfs

    def eval_on_test(self, order=-1, reducer=None):
        perfs = []
        for i, seed in enumerate(self.seeds):
            X_train, X_test, y_train, y_test, depth_order_train, depth_order_test = self._splitter(seed)
            if order != - 1:
                mask = depth_order_test[:, 1] == order
                X_test, y_test = X_test[mask, :], y_test[mask]

            dls = DataLoaders(((X_train, y_train), (X_test, y_test)),
                              transform=SNV_transform())
            _, test_generators = dls.loaders()

            y_hat, y_true = self.learners[i].predict(test_generators)
            perf = eval_reg(y_true, y_hat)
            perf['n'] = len(X_test)
            perfs.append(perf)
        if reducer:
            perfs = self.reduce(perfs, reducer)
        return perfs

    def _splitter(self, seed):
        return train_test_split(self.X, self.y, self.depth_order,
                                test_size=self.split_ratio,
                                random_state=seed)

    def reduce(self, perfs, fn=np.mean):
        results = {}
        for metric in perfs[0].keys():
            result = fn(np.array([perf[metric] for perf in perfs]))
            results[metric] = result
        return results