"""

models.py

=======================

Houses all of the PyTorch models to access and the corresponding Scikit-Learn like model trainer.

"""

from pathflowai.unet import UNet

# from pathflowai.unet2 import NestedUNet

# from pathflowai.unet4 import UNetSmall as UNet2

from pathflowai.fast_scnn import get_fast_scnn

import torch

import torchvision

from torchvision import models

from torchvision.models import segmentation as segmodels

from torch import nn

from torch.nn import functional as F

import pandas as pd, numpy as np

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt

import seaborn as sns

from pathflowai.schedulers import *

import pysnooper

from torch.autograd import Variable

import copy

from sklearn.metrics import roc_curve, confusion_matrix, classification_report, r2_score

sns.set()

from pathflowai.losses import GeneralizedDiceLoss, FocalLoss

from apex import amp

from torch.nn import functional as F

import time, os

class MLP(nn.Module):

    """Multi-layer perceptron model.

    Parameters

    ----------

    n_input:int

        Number input dimensions.

    hidden_topology:list

        List of hidden topology

    dropout_p:float

        Amount dropout.

    n_outputs:int

        Number outputs.

    binary:bool

        Binary output with sigmoid transform.

    softmax:bool

        Whether to apply softmax on output.

    """

    def __init__(self, n_input, hidden_topology, dropout_p, n_outputs=1, binary=True, softmax=False):

        super(MLP,self).__init__()

        self.topology = [n_input]+hidden_topology+[n_outputs]

        layers = [nn.Linear(self.topology[i],self.topology[i+1]) for i in range(len(self.topology)-2)]

        for layer in layers:

            torch.nn.init.xavier_uniform_(layer.weight)

        self.layers = [nn.Sequential(layer,nn.LeakyReLU(),nn.Dropout(p=dropout_p)) for layer in layers]

        self.output_layer = nn.Linear(self.topology[-2],self.topology[-1])

        torch.nn.init.xavier_uniform_(self.output_layer.weight)

        if binary:

            output_transform = nn.Sigmoid()

        elif softmax:

            output_transform = nn.Softmax()

        else:

            output_transform = nn.Dropout(p=0.)

        self.layers.append(nn.Sequential(self.output_layer,output_transform))

        self.mlp = nn.Sequential(*self.layers)

    def forward(self,x):

        return self.mlp(x)

class FixedSegmentationModule(nn.Module):

    """Special model modification for segmentation tasks. Gets output from some of the models' forward loops.

    Parameters

    ----------

    segnet:nn.Module

        Segmentation network

    """

    def __init__(self, segnet):

        super(FixedSegmentationModule, self).__init__()

        self.segnet=segnet

    def forward(self, x):

        """Forward pass.

        Parameters

        ----------

        x:Tensor

            Input

        Returns

        -------

        Tensor

            Output from model.

        """

        return self.segnet(x)['out']

def generate_model(pretrain,architecture,num_classes, add_sigmoid=True, n_hidden=100, segmentation=False):

    """Generate a nn.Module for use.

    Parameters

    ----------

    pretrain:bool

        Pretrain using ImageNet?

    architecture:str

        See model_training for list of all architectures you can train with.

    num_classes:int

        Number of classes to predict.

    add_sigmoid:type

        Add sigmoid non-linearity at end.

    n_hidden:int

        Number of hidden fully connected layers.

    segmentation:bool

        Whether segment task?

    Returns

    -------

    nn.Module

        Pytorch model.

    """

    # to add: https://github.com/zhanghang1989/PyTorch-Encoding/blob/master/encoding/models/model_zoo.py

    #architecture = 'resnet' + str(num_layers)

    model = None

    if architecture =='unet':

        model = UNet(n_channels=3, n_classes=num_classes)

    elif architecture =='unet2':

        print('Deprecated for now, defaulting to UNET.')

        model = UNet(n_channels=3, n_classes=num_classes)#UNet2(3,num_classes)

    elif architecture == 'fast_scnn':

        model = get_fast_scnn(num_classes)

    elif architecture == 'nested_unet':

        print('Nested UNET is deprecated for now, defaulting to UNET.')

        model = UNet(n_channels=3, n_classes=num_classes)#NestedUNet(3, num_classes)

    elif architecture.startswith('efficientnet'):

        from efficientnet_pytorch import EfficientNet

        if pretrain:

            model = EfficientNet.from_pretrained(architecture, override_params=dict(num_classes=num_classes))

        else:

            model = EfficientNet.from_name(architecture, override_params=dict(num_classes=num_classes))

        print(model)

    elif architecture.startswith('sqnxt'):

        from pytorchcv.model_provider import get_model as ptcv_get_model

        model = ptcv_get_model(architecture, pretrained=pretrain)

        num_ftrs=int(128*int(architecture.split('_')[-1][1]))

        model.output=MLP(num_ftrs, [1000], dropout_p=0., n_outputs=num_classes, binary=add_sigmoid, softmax=False).mlp

    else:

        #for pretrained on imagenet

        model_names = [m for m in dir(models) if not m.startswith('__')]

        segmentation_model_names = [m for m in dir(segmodels) if not m.startswith('__')]

        if architecture in model_names:

            model = getattr(models, architecture)(pretrained=pretrain)

        if segmentation:

            if architecture in segmentation_model_names:

                model = getattr(segmodels, architecture)(pretrained=pretrain)

            else:

                model = UNet(n_channels=3, n_classes=num_classes)

            if architecture.startswith('deeplab'):

                model.classifier[4] = nn.Conv2d(256, num_classes, kernel_size=(1, 1), stride=(1, 1))

                model = FixedSegmentationModule(model)

            elif architecture.startswith('fcn'):

                model.classifier[4] = nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))

                model = FixedSegmentationModule(model)

        elif architecture.startswith('resnet') or architecture.startswith('inception'):

            num_ftrs = model.fc.in_features

            #linear_layer = nn.Linear(num_ftrs, num_classes)

            #torch.nn.init.xavier_uniform(linear_layer.weight)

            model.fc = MLP(num_ftrs, [1000], dropout_p=0., n_outputs=num_classes, binary=add_sigmoid, softmax=False).mlp#nn.Sequential(*([linear_layer]+([nn.Sigmoid()] if (add_sigmoid) else [])))

        elif architecture.startswith('alexnet') or architecture.startswith('vgg') or architecture.startswith('densenet'):

            num_ftrs = model.classifier[6].in_features

            #linear_layer = nn.Linear(num_ftrs, num_classes)

            #torch.nn.init.xavier_uniform(linear_layer.weight)

            model.classifier[6] = MLP(num_ftrs, [1000], dropout_p=0., n_outputs=num_classes, binary=add_sigmoid, softmax=False).mlp#nn.Sequential(*([linear_layer]+([nn.Sigmoid()] if (add_sigmoid) else [])))

    return model

#@pysnooper.snoop("dice_loss.log")

def dice_loss(logits, true, eps=1e-7):

    """https://github.com/kevinzakka/pytorch-goodies

    Computes the Sørensen–Dice loss.

    Note that PyTorch optimizers minimize a loss. In this

    case, we would like to maximize the dice loss so we

    return the negated dice loss.

    Args:

        true: a tensor of shape [B, 1, H, W].

        logits: a tensor of shape [B, C, H, W]. Corresponds to

            the raw output or logits of the model.

        eps: added to the denominator for numerical stability.

    Returns:

        dice_loss: the Sørensen–Dice loss.

    """

    #true=true.long()

    num_classes = logits.shape[1]

    if num_classes == 1:

        true_1_hot = torch.eye(num_classes + 1)[true.squeeze(1)]

        true_1_hot = true_1_hot.permute(0, 3, 1, 2).float()

        true_1_hot_f = true_1_hot[:, 0:1, :, :]

        true_1_hot_s = true_1_hot[:, 1:2, :, :]

        true_1_hot = torch.cat([true_1_hot_s, true_1_hot_f], dim=1)

        pos_prob = torch.sigmoid(logits)

        neg_prob = 1 - pos_prob

        probas = torch.cat([pos_prob, neg_prob], dim=1)

    else:

        true_1_hot = torch.eye(num_classes)[true.squeeze(1)]

        #print(true_1_hot.size())

        true_1_hot = true_1_hot.permute(0, 3, 1, 2).float()

        probas = F.softmax(logits, dim=1)

    true_1_hot = true_1_hot.type(logits.type())

    dims = (0,) + tuple(range(2, true.ndimension()))

    intersection = torch.sum(probas * true_1_hot, dims)

    cardinality = torch.sum(probas + true_1_hot, dims)

    dice_loss = (2. * intersection / (cardinality + eps)).mean()

    return (1 - dice_loss)

class ModelTrainer:

    """Trainer for the neural network model that wraps it into a scikit-learn like interface.

    Parameters

    ----------

    model:nn.Module

        Deep learning pytorch model.

    n_epoch:int

        Number training epochs.

    validation_dataloader:DataLoader

        Dataloader of validation dataset.

    optimizer_opts:dict

        Options for optimizer.

    scheduler_opts:dict

        Options for learning rate scheduler.

    loss_fn:str

        String to call a particular loss function for model.

    reduction:str

        Mean or sum reduction of loss.

    num_train_batches:int

        Number of training batches for epoch.

    """

    def __init__(self, model, n_epoch=300, validation_dataloader=None, optimizer_opts=dict(name='adam',lr=1e-3,weight_decay=1e-4), scheduler_opts=dict(scheduler='warm_restarts',lr_scheduler_decay=0.5,T_max=10,eta_min=5e-8,T_mult=2), loss_fn='ce', reduction='mean', num_train_batches=None, seg_out_class=-1, apex_opt_level="O2", checkpointing=False):

        self.model = model

        optimizers = {'adam':torch.optim.Adam, 'sgd':torch.optim.SGD}

        loss_functions = {'bce':nn.BCEWithLogitsLoss(reduction=reduction), 'ce':nn.CrossEntropyLoss(reduction=reduction), 'mse':nn.MSELoss(reduction=reduction), 'nll':nn.NLLLoss(reduction=reduction), 'dice':dice_loss, 'focal':FocalLoss(num_class=2), 'gdl':GeneralizedDiceLoss(add_softmax=True)}

        loss_functions['dice+ce']=(lambda y_pred, y_true: dice_loss(y_pred,y_true)+loss_functions['ce'](y_pred,y_true))

        if 'name' not in list(optimizer_opts.keys()):

            optimizer_opts['name']='adam'

        self.optimizer = optimizers[optimizer_opts.pop('name')](self.model.parameters(),**optimizer_opts)

        if torch.cuda.is_available():

            self.model, self.optimizer = amp.initialize(self.model, self.optimizer, opt_level=apex_opt_level)

            self.cuda=True

        else:

            self.cuda=False

        self.scheduler = Scheduler(optimizer=self.optimizer,opts=scheduler_opts)

        self.n_epoch = n_epoch

        self.validation_dataloader = validation_dataloader

        self.loss_fn = loss_functions[loss_fn]

        self.loss_fn_name = loss_fn

        self.bce=(self.loss_fn_name=='bce' or self.validation_dataloader.dataset.mt_bce)

        self.sigmoid = nn.Sigmoid()

        self.original_loss_fn = copy.deepcopy(loss_functions[loss_fn])

        self.num_train_batches = num_train_batches

        self.val_loss_fn = copy.deepcopy(loss_functions[loss_fn])

        self.seg_out_class=seg_out_class

        self.checkpointing=checkpointing

        self.checkpoint_dir='./checkpoints'

        if self.checkpointing:

            os.makedirs(self.checkpoint_dir,exist_ok=True)

    def save_model(self, model=None, epoch=0):

        torch.save((model if isinstance(model,type(None)) else self.model).state_dict(),os.path.join(self.checkpoint_dir,f'checkpoint.{epoch}.pth'))

    def calc_loss(self, y_pred, y_true):

        """Calculates loss supplied in init statement and modified by reweighting.

        Parameters

        ----------

        y_pred:tensor

            Predictions.

        y_true:tensor

            True values.

        Returns

        -------

        loss

        """

        return self.loss_fn(y_pred, y_true)

    def calc_val_loss(self, y_pred, y_true):

        """Calculates loss supplied in init statement on validation set.

        Parameters

        ----------

        y_pred:tensor

            Predictions.

        y_true:tensor

            True values.

        Returns

        -------

        val_loss

        """

        return self.val_loss_fn(y_pred, y_true)

    def reset_loss_fn(self):

        """Resets loss to original specified loss."""

        self.loss_fn = self.original_loss_fn

    def add_class_balance_loss(self, dataset, custom_weights=''):

        """Updates loss function to handle class imbalance by weighting inverse to class appearance.

        Parameters

        ----------

        dataset:DynamicImageDataset

            Dataset to balance by.

        """

        self.class_weights = dataset.get_class_weights() if not custom_weights else np.array(list(map(float,custom_weights.split(','))))

        if custom_weights:

            self.class_weights=self.class_weights/sum(self.class_weights)

        print('Weights:',self.class_weights)

        self.original_loss_fn = copy.deepcopy(self.loss_fn)

        weight=torch.tensor(self.class_weights,dtype=torch.float)

        if torch.cuda.is_available():

            weight=weight.cuda()

        if self.loss_fn_name=='ce':

            self.loss_fn = nn.CrossEntropyLoss(weight=weight)

        elif self.loss_fn_name=='nll':

            self.loss_fn = nn.NLLLoss(weight=weight)

        else: # modify below for multi-target

            self.loss_fn = lambda y_pred,y_true: sum([self.class_weights[i]*self.original_loss_fn(y_pred[y_true==i],y_true[y_true==i]) if sum(y_true==i) else 0. for i in range(2)])

    def calc_best_confusion(self, y_pred, y_true):

        """Calculate confusion matrix on validation set for classification/segmentation tasks, optimize threshold where positive.

        Parameters

        ----------

        y_pred:array

            Predictions.

        y_true:array

            Ground truth.

        Returns

        -------

        float

            Optimized threshold to use on test set.

        dataframe

            Confusion matrix.

        """

        fpr, tpr, thresholds = roc_curve(y_true, y_pred)

        threshold=thresholds[np.argmin(np.sum((np.array([0,1])-np.vstack((fpr, tpr)).T)**2,axis=1)**.5)]

        y_pred = (y_pred>threshold).astype(int)

        return threshold, pd.DataFrame(confusion_matrix(y_true,y_pred),index=['F','T'],columns=['-','+']).iloc[::-1,::-1].T

    def loss_backward(self,loss):

        """Backprop using mixed precision for added speed boost.

        Parameters

        ----------

        loss:loss

            Torch loss calculated.

        """

        if self.cuda:

            with amp.scale_loss(loss,self.optimizer) as scaled_loss:

                scaled_loss.backward()

        else:

            loss.backward()

    # @pysnooper.snoop('train_loop.log')

    def train_loop(self, epoch, train_dataloader):

        """One training epoch, calculate predictions, loss, backpropagate.

        Parameters

        ----------

        epoch:int

            Current epoch.

        train_dataloader:DataLoader

            Training data.

        Returns

        -------

        float

            Training loss for epoch

        """

        self.model.train(True)

        running_loss = 0.

        n_batch = len(train_dataloader.dataset)//train_dataloader.batch_size if self.num_train_batches == None else self.num_train_batches

        for i, batch in enumerate(train_dataloader):

            starttime=time.time()

            if i == n_batch:

                break

            X = Variable(batch[0], requires_grad=True)

            y_true = Variable(batch[1])

            if not train_dataloader.dataset.segmentation and self.loss_fn_name=='ce' and y_true.shape[1]>1:

                y_true=y_true.argmax(1).long()

            if train_dataloader.dataset.segmentation and self.loss_fn_name!='dice':

                y_true=y_true.squeeze(1)

            if torch.cuda.is_available():

                X = X.cuda()

                y_true=y_true.cuda()

            y_pred = self.model(X)

            #sizes=(y_pred.size(),y_true.size())

            #print(y_true)

            loss = self.calc_loss(y_pred,y_true)

            train_loss=loss.item()

            running_loss += train_loss

            self.optimizer.zero_grad()

            self.loss_backward(loss)#loss.backward()

            self.optimizer.step()

            endtime=time.time()

            print("Epoch {}[{}/{}] Time:{}, Train Loss:{}".format(epoch,i,n_batch,round(endtime-starttime,3),train_loss))

        self.scheduler.step()

        running_loss/=n_batch

        return running_loss

    def val_loop(self, epoch, val_dataloader, print_val_confusion=True, save_predictions=True):

        """Calculate loss over validation set.

        Parameters

        ----------

        epoch:int

            Current epoch.

        val_dataloader:DataLoader

            Validation iterator.

        print_val_confusion:bool

            Calculate confusion matrix and plot.

        save_predictions:int

            Print validation results.

        Returns

        -------

        float

            Validation loss for epoch.

        """

        self.model.train(False)

        n_batch = len(val_dataloader.dataset)//val_dataloader.batch_size

        running_loss = 0.

        Y = {'pred':[],'true':[]}

        with torch.no_grad():

            for i, batch in enumerate(val_dataloader):

                X = Variable(batch[0],requires_grad=False)

                y_true = Variable(batch[1])

                if not val_dataloader.dataset.segmentation and self.loss_fn_name=='ce' and y_true.shape[1]>1:

                    y_true=y_true.argmax(1).long()

                if val_dataloader.dataset.segmentation and self.loss_fn_name!='dice':

                    y_true=y_true.squeeze(1)

                if torch.cuda.is_available():

                    X = X.cuda()

                    y_true=y_true.cuda()

                y_pred = self.model(X)

                if save_predictions:

                    if val_dataloader.dataset.segmentation:

                        Y['true'].append(torch.flatten(y_true if not val_dataloader.dataset.gdl else y_true).detach().cpu().numpy().astype(int).flatten()) # .argmax(axis=1)

                        Y['pred'].append((y_pred.detach().cpu().numpy().argmax(axis=1)).astype(int).flatten())

                    else:

                        Y['true'].append(y_true.detach().cpu().numpy().astype(int).flatten())

                        y_pred_numpy=((y_pred if not self.bce else self.sigmoid(y_pred)).detach().cpu().numpy()).astype(float)

                        if len(y_pred_numpy)>1 and y_pred_numpy.shape[1]>1 and not val_dataloader.dataset.mt_bce:

                            y_pred_numpy=y_pred_numpy.argmax(axis=1)

                        Y['pred'].append(y_pred_numpy.flatten())

                loss = self.calc_val_loss(y_pred,y_true)

                val_loss=loss.item()

                running_loss += val_loss

                print("Epoch {}[{}/{}] Val Loss:{}".format(epoch,i,n_batch,val_loss))

        if print_val_confusion and save_predictions:

            y_pred,y_true = np.hstack(Y['pred']),np.hstack(Y['true'])

            if not val_dataloader.dataset.segmentation:

                if self.loss_fn_name in ['bce','mse'] and not val_dataloader.dataset.mt_bce:

                    threshold, best_confusion = self.calc_best_confusion(y_pred,y_true)

                    print("Epoch {} Val Confusion, Threshold {}:".format(epoch,threshold))

                    print(best_confusion)

                    y_true = y_true.astype(int)

                    y_pred = (y_pred>=threshold).astype(int)

                elif val_dataloader.dataset.mt_bce:

                    n_targets = len(val_dataloader.dataset.targets)

                    y_pred=y_pred[y_true>0]

                    y_true=y_true[y_true>0]

                    y_true=y_true[np.isnan(y_pred)==False]

                    y_pred=y_pred[np.isnan(y_pred)==False]

                    if 0 and n_targets > 1:

                        n_row=len(y_true)/n_targets

                        y_pred=y_pred.reshape(int(n_row),n_targets)

                        y_true=y_true.reshape(int(n_row),n_targets)

                    print("Epoch {} Val Regression, R2 Score {}".format(epoch, str(r2_score(y_true, y_pred))))

            else:

                print(classification_report(y_true,y_pred))

        running_loss/=n_batch

        return running_loss

    #@pysnooper.snoop("test_loop.log")

    def test_loop(self, test_dataloader):

        """Calculate final predictions on loss.

        Parameters

        ----------

        test_dataloader:DataLoader

            Test dataset.

        Returns

        -------

        array

            Predictions or embeddings.

        """

        #self.model.train(False) KEEP DROPOUT? and BATCH NORM??

        y_pred = []

        running_loss = 0.

        with torch.no_grad():

            for i, (X,y_test) in enumerate(test_dataloader):

                #X = Variable(batch[0],requires_grad=False)

                if torch.cuda.is_available():

                    X = X.cuda()

                if test_dataloader.dataset.segmentation:

                    prediction=self.model(X).detach().cpu().numpy()

                    if self.seg_out_class>=0:

                        prediction=prediction[:,self.seg_out_class,...]

                    else:

                        prediction=prediction.argmax(axis=1).astype(int)

                    pred_size=prediction.shape#size()

                    #pred_mean=prediction[0].mean(axis=0)

                    y_pred.append(prediction)

                else:

                    prediction=self.model(X)

                    if self.loss_fn_name != 'mse' and ((len(test_dataloader.dataset.targets)-1) or self.bce):

                        prediction=self.sigmoid(prediction)

                    elif test_dataloader.dataset.classify_annotations:

                        prediction=F.softmax(prediction,dim=1)

                    y_pred.append(prediction.detach().cpu().numpy())

        y_pred = np.concatenate(y_pred,axis=0)#torch.cat(y_pred,0)

        return y_pred

    def fit(self, train_dataloader, verbose=False, print_every=10, save_model=True, plot_training_curves=False, plot_save_file=None, print_val_confusion=True, save_val_predictions=True):

        """Fits the segmentation or classification model to the patches, saving the model with the lowest validation score.

        Parameters

        ----------

        train_dataloader:DataLoader

            Training dataset.

        verbose:bool

            Print training and validation loss?

        print_every:int

            Number of epochs until print?

        save_model:bool

            Whether to save model when reaching lowest validation loss.

        plot_training_curves:bool

            Plot training curves over epochs.

        plot_save_file:str

            File to save training curves.

        print_val_confusion:bool

            Print validation confusion matrix.

        save_val_predictions:bool

            Print validation results.

        Returns

        -------

        self

            Trainer.

        float

            Minimum val loss.

        int

            Best validation epoch with lowest loss.

        """

        # choose model with best f1

        self.train_losses = []

        self.val_losses = []

        for epoch in range(self.n_epoch):

            start_time=time.time()

            train_loss = self.train_loop(epoch,train_dataloader)

            current_time=time.time()

            train_time=current_time-start_time

            self.train_losses.append(train_loss)

            val_loss = self.val_loop(epoch,self.validation_dataloader, print_val_confusion=print_val_confusion, save_predictions=save_val_predictions)

            val_time=time.time()-current_time

            self.val_losses.append(val_loss)

            if verbose and not (epoch % print_every):

                if plot_training_curves:

                    self.plot_train_val_curves(plot_save_file)

                print("Epoch {}: Train Loss {}, Val Loss {}, Train Time {}, Val Time {}".format(epoch,train_loss,val_loss,train_time,val_time))

            if val_loss <= min(self.val_losses) and save_model:

                min_val_loss = val_loss

                best_epoch = epoch

                best_model = copy.deepcopy(self.model)

                if self.checkpointing:

                    self.save_model(best_model,epoch)

        if save_model:

            self.model = best_model

        return self, min_val_loss, best_epoch

    def plot_train_val_curves(self, save_file=None):

        """Plots training and validation curves.

        Parameters

        ----------

        save_file:str

            File to save to.

        """

        plt.figure()

        sns.lineplot('epoch','value',hue='variable',

                     data=pd.DataFrame(np.vstack((np.arange(len(self.train_losses)),self.train_losses,self.val_losses)).T,

                                       columns=['epoch','train','val']).melt(id_vars=['epoch'],value_vars=['train','val']))

        if save_file is not None:

            plt.savefig(save_file, dpi=300)

    def predict(self, test_dataloader):

        """Make classification segmentation predictions on testing data.

        Parameters

        ----------

        test_dataloader:DataLoader

            Test data.

        Returns

        -------

        array

            Predictions.

        """

        y_pred = self.test_loop(test_dataloader)

        return y_pred

    def fit_predict(self, train_dataloader, test_dataloader):

        """Fit model to training data and make classification segmentation predictions on testing data.

        Parameters

        ----------

        train_dataloader:DataLoader

            Train data.

        test_dataloader:DataLoader

            Test data.

        Returns

        -------

        array

            Predictions.

        """

        return self.fit(train_dataloader)[0].predict(test_dataloader)

    def return_model(self):

        """Returns pytorch model.

        """

        return self.model