import numpy as np
from sklearn import metrics
from sklearn.metrics import roc_auc_score, accuracy_score
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from Classifier.Classes.Helpers import Helpers

helper = Helpers("Model", False)

class LuekemiaNet(nn.Module):
    """ Model Class

        Model functions for the Acute Lymphoblastic Leukemia Pytorch CNN 2020.
        """
    def __init__(self):
        super(LuekemiaNet, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(2))

        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(2))

        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(2))

        self.fc = nn.Sequential(
            nn.Linear(128*2*2, 1024),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(1024, 512),
            nn.Dropout(0.2),
            nn.Linear(512, 2))

    def forward(self, x):
        """Method for Forward Prop"""
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)

        ######################
        # For model debugging #
        ######################
        #print(out.shape)

        out = out.view(x.shape[0], -1)
        out = self.fc(out)
        return out



###################
# Train model     #
###################
def learn(model, loss_function, optimizer, epochs, train_loader,
          valid_loader, scheduler, train_on_gpu=True):
    """ """

    trainset_auc = list()
    train_losses = list()
    valid_losses = list()
    val_auc = list()
    valid_auc_epoch = list()
    train_auc_epoch = list()

    for epoch in range(epochs):
        train_loss = 0.0
        valid_loss = 0.0
        model.train_model()
        for i, data in enumerate(train_loader):
            inputs, labels = data
            if train_on_gpu:
                inputs, labels = inputs.cuda(), labels.cuda()
            # first of all, set gradient of all optimized variables to zero
            optimizer.zero_grad()
            # forward pass: compute predicted outputs by passing inputs data to the model
            outputs = model(inputs)
            # compute batch loss(binary cross entropy
            loss = loss_function(outputs[:,0], labels.float())
            # Update Train loss and accuracies
            train_loss += loss.item() * inputs.size(0)
            # backward pass: compute gradient of the loss with respect to model parameters
            loss.backward()
            # perform one step of optimization
            optimizer.step()
            scheduler.step()
            y_actual = labels.data.cpu().numpy()
            prediction = outputs[:, 0].detach().cpu().numpy()
            trainset_auc.append(roc_auc_score(y_true=y_actual, y_score=prediction))

        model.eval()
        for _, data in enumerate(valid_loader):
            inputs, labels = data
            # if CUDA GPU is available
            if train_on_gpu:
                inputs, labels = inputs.cuda(), labels.cuda()
                # first of all, set gradient of all optimized variables to zero
            optimizer.zero_grad()
            # forward pass: compute predicted outputs by passing inputs data to the model
            outputs = model(inputs)
            # compute batch loss(binary cross entropy
            loss = loss_function(outputs[:,0], labels.float())
            # update average validation loss
            valid_loss += loss.item() * inputs.size(0)
            y_actual = labels.data.cpu().numpy()
            predictionx = outputs[:, 0].detach().cpu().numpy()
            val_auc.append(roc_auc_score(y_actual, predictionx))

            # calculate average losses
        train_loss = train_loss / len(train_loader.sampler)
        valid_loss = valid_loss / len(valid_loader.sampler)
        valid_auc = np.mean(val_auc)
        train_auc = np.mean(trainset_auc)
        train_auc_epoch.append(np.mean(trainset_auc))
        valid_auc_epoch.append(np.mean(val_auc))
        train_losses.append(train_loss)
        valid_losses.append(valid_loss)

        helper.logger.info(
            'Epoch: {} | Training Loss: {:.6f} | Training AUC: {:.6f}| Validation Loss: {:.6f} | Validation AUC: {:.4f}'.format(
                    epoch, train_loss, train_auc, valid_loss, valid_auc))

    torch.save(model.state_dict(), helper.config["classifier"]["model_params"]["weights"])

    plt.plot(train_losses, label='Training loss')
    plt.plot(valid_losses, label='Validation loss')
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend(frameon=False)
    plt.savefig(helper.config["classifier"]["model_params"]["plot_loss"])
    plt.show()

    plt.plot(valid_auc_epoch, label='Validation AUC/Epochs')
    plt.plot(train_auc_epoch, label='Training AUC/Epochs')
    plt.legend("")
    plt.xlabel("Epochs")
    plt.ylabel("Area Under the Curve")
    plt.legend(frameon=False)
    plt.savefig(helper.config["classifier"]["model_params"]["plot_auc"])
    plt.show()