import torch, torchvision
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torchvision.datasets as datasets
import torch.utils.data as data
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
import torchvision.models as models
import matplotlib.pyplot as plt
import time, os, copy, numpy as np
from sklearn.metrics import confusion_matrix
from tqdm import tqdm
# from livelossplot import PlotLosses
import sys

# from util.masking import Mask_attr

from pdb import set_trace as bp

def train_model(model, dataloaders, dataset_sizes, criterion, optimizer, scheduler, num_epochs, device_id):
    device = torch.device(device_id if torch.cuda.is_available() else "cpu")
    since = time.time()
    # liveloss = PlotLosses()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for i, (inputs, labels) in enumerate(dataloaders[phase]):
                # bp()
                inputs = inputs.to(device)
                labels = labels.to(device)
                # weights = torch.tensor([torch.count_nonzero(labels==i)  for i in range(21)]).to(device)
                # weights = weights/weights.sum()

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    
                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                print("\rIteration: {}/{}, Loss: {}.".format(i+1, len(dataloaders[phase]), loss.item() * inputs.size(0)), end="")

#                 print( (i+1)*100. / len(dataloaders[phase]), "% Complete" )
                sys.stdout.flush()
                
                
            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            if phase == 'train':
                avg_loss = epoch_loss
                t_acc = epoch_acc
            else:
                val_loss = epoch_loss
                val_acc = epoch_acc
            
#             print('{} Loss: {:.4f} Acc: {:.4f}'.format(
#                 phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                
        

        print(f'log loss: {avg_loss}, val_log loss: {val_loss}, accuracy: {t_acc}, val_accuracy: {val_acc}')
                
        # liveloss.draw()
        print('Train Loss: {:.4f} Acc: {:.4f}'.format(avg_loss, t_acc))
        print(  'Val Loss: {:.4f} Acc: {:.4f}'.format(val_loss, val_acc))
        print('Best Val Accuracy: {}'.format(best_acc))
        print()
        sys.stdout.flush()
        sys.stderr.flush()
    
    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

def infer_model(model, dataloaders, dataset_sizes, device_id):
    device = torch.device(device_id if torch.cuda.is_available() else "cpu")
    since = time.time()
    # liveloss = PlotLosses()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    print('...Inference...')
    all_trues, all_preds = [], []
    images_correct, labels_correct, annts_correct = [], [], []
    # Each epoch has a training and validation phase
    for phase in ['val']:
        
        model.eval()   # Set model to evaluate mode

        running_loss = 0.0
        running_corrects = 0

        # Iterate over data.
        for i, (inputs, labels, annts) in enumerate(dataloaders[phase]):
            # bp()
            inputs = inputs.to(device)
            labels = labels.to(device)
            # weights = torch.tensor([torch.count_nonzero(labels==i)  for i in range(21)]).to(device)
            # weights = weights/weights.sum()

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            # statistics
            running_corrects += torch.sum(preds == labels.data)
            print("\rIteration: {}/{}.".format(i+1, len(dataloaders[phase])), end="")
            all_trues = all_trues + labels.data.detach().cpu().tolist()
            all_preds = all_preds + preds.detach().cpu().tolist()

            images_correct.append(inputs[preds == labels.data])
            labels_correct.append(labels.data[preds == labels.data])
            annts_correct.append(np.array(annts, dtype=object)[preds.detach().cpu() == labels.data.detach().cpu()])

#                 print( (i+1)*100. / len(dataloaders[phase]), "% Complete" )
            sys.stdout.flush()
            
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            
            val_acc = epoch_acc
            
#             print('{} Loss: {:.4f} Acc: {:.4f}'.format(
#                 phase, epoch_loss, epoch_acc))

        
        print(f'val_accuracy: {val_acc}')
        matrix = confusion_matrix(all_trues, all_preds)
        print(f'Classwise acc. scores: {matrix.diagonal()/matrix.sum(axis=1)}')
        print(f'Classwise total gnd truths: {np.unique(all_trues, return_counts=True)[1]}')
        
        images_correct = torch.cat(images_correct, dim=0)
        labels_correct = torch.cat(labels_correct, dim=0)
        annts_correct = np.concatenate(annts_correct, axis=0)
        print('saving the tensors.. may take large memory.')
        torch.save([images_correct, labels_correct, annts_correct], 'VOC_filtered/correct_preds_and_annotations.pt')
        
        # liveloss.draw()
        # print('Train Loss: {:.4f} Acc: {:.4f}'.format(avg_loss, t_acc))
        # print(  'Val Loss: {:.4f} Acc: {:.4f}'.format(val_loss, val_acc))
        # print('Best Val Accuracy: {}'.format(best_acc))
        # print()
        sys.stdout.flush()
        sys.stderr.flush()
    
    # time_elapsed = time.time() - since
    # print('Training complete in {:.0f}m {:.0f}s'.format(
    #     time_elapsed // 60, time_elapsed % 60))
    # print('Best val Acc: {:4f}'.format(best_acc))

    # # load best model weights
    # model.load_state_dict(best_model_wts)
    return None

class MaskedDataset(Dataset):
            """Face Landmarks dataset."""

            def __init__(self, images, labels):
                """
                Arguments:
                    csv_file (string): Path to the csv file with annotations.
                    root_dir (string): Directory with all the images.
                    transform (callable, optional): Optional transform to be applied
                        on a sample.
                """
                self.images = images
                self.labels = labels

            def __len__(self):
                return len(self.labels)

            def __getitem__(self, idx):
                if torch.is_tensor(idx):
                    idx = idx.tolist()

                return self.images[idx], self.labels[idx]

def train_model_masking(model, model_ft_pretr, dataloaders, dataset_sizes, criterion, optimizer, scheduler, batch_size, mask_attr, num_epochs, device_id):
    device = torch.device(device_id if torch.cuda.is_available() else "cpu")

    if mask_attr:
        
        train_images, train_labels = np.load('train_images_80_100.npy'), np.load('train_labels_80_100.npy')
        train_dataset = MaskedDataset(torch.tensor(train_images), torch.tensor(train_labels))
        
        dataloaders['train'] = torch.utils.data.DataLoader(dataset=train_dataset,
                                            batch_size=batch_size,
                                            shuffle=False,
                                            pin_memory=True)
        #...........................................

    else:
        print('masking not activated')

    since = time.time()
    # liveloss = PlotLosses()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for i,(inputs, labels) in enumerate(dataloaders[phase]):
                # bp()
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                print("\rIteration: {}/{}, Loss: {}.".format(i+1, len(dataloaders[phase]), loss.item() * inputs.size(0)), end="")

#                 print( (i+1)*100. / len(dataloaders[phase]), "% Complete" )
                sys.stdout.flush()
                
                
            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            if phase == 'train':
                avg_loss = epoch_loss
                t_acc = epoch_acc
            else:
                val_loss = epoch_loss
                val_acc = epoch_acc
            
#             print('{} Loss: {:.4f} Acc: {:.4f}'.format(
#                 phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                
        # liveloss.update({
        #     'log loss': avg_loss,
        #     'val_log loss': val_loss,
        #     'accuracy': t_acc,
        #     'val_accuracy': val_acc
        # })

        print(f'log loss: {avg_loss}, val_log loss: {val_loss}, accuracy: {t_acc}, val_accuracy: {val_acc}')
                
        # liveloss.draw()
        print('Train Loss: {:.4f} Acc: {:.4f}'.format(avg_loss, t_acc))
        print(  'Val Loss: {:.4f} Acc: {:.4f}'.format(val_loss, val_acc))
        print('Best Val Accuracy: {}'.format(best_acc))
        print()
    
    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model