import time
import shutil
import os.path
import sys

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim
import torch.utils.data
import torch.utils.model_zoo as model_zoo
import torchvision.transforms as transforms
import torchvision.datasets
import torchvision.models

from .utils import AverageMeter, check_acc
from ..models.densenet import DenseNet
from .sampler import BatchLoader

if torch.cuda.is_available():
  dtype = {'float': torch.cuda.FloatTensor, 'long': torch.cuda.LongTensor, 'byte': torch.cuda.ByteTensor} 
else:
  dtype = {'float': torch.FloatTensor, 'long': torch.LongTensor, 'byte': torch.ByteTensor} 


class Solver(object):
    """Solver
    Args:
        model: 
        data:
        optimizer: e.g., torch.optim.Adam(model.parameters())
        loss_fn: loss function; e.g., torch.nn.CrossEntropy()
        resume: file path to checkpoint
    """
    def __init__(self, model, data, optimizer, loss_fn, resume=None):
        self.model = model
        self.data = data
        self.optimizer = optimizer
        self.loss_fn = loss_fn
        
        # keep track of loss and accuracy during training
        self.losses_train = []
        self.losses_val = []
        self.acc_train = []
        self.acc_val = []
        self.best_acc_val = 0
        self.epoch_counter = 0
        
        if resume:
            if os.path.isfile(resume):
                checkpoint = torch.load(resume)
                self.model.load_state_dict(checkpoint['model_state'])
                self.optimizer = checkpoint['optimizer']
                self.best_acc_val = checkpoint['best_acc_val']
                self.epoch_counter = checkpoint['epoch']
                self.losses_train = checkpoint['losses_train']
                self.losses_val = checkpoint['losses_val']
                self.acc_train = checkpoint['acc_train']
                self.acc_val = checkpoint['acc_val']
            else:
                print("==> No checkpoint found at '{}'".format(resume))
        
    def _reset_avg_meter(self):
        """reset loss_epoch, top1, top5, batch_time at the beginning of each epoch
        """
        self.loss_epoch = AverageMeter()
        self.top1 = AverageMeter()
        self.top5 = AverageMeter()
        self.batch_time = AverageMeter()
        
    
    def run_one_epoch(self, epoch, batch_size=100, num_samples=None, print_every=100, 
                      training=True, balanced_sample=False, topk=5):
        """run one epoch for training or validating
        Args:
            epoch: int; epoch_counter; used for printing only
            batch_size: int, default: 100
            num_samples: int, default: None. 
                How many samples to use in case we don't want train a whole epoch
            print_every: int, default: 100
            training: bool, default:True. If true, train; else validate
            balanced_sample: default: False. Used for unbalanced dataset
        """
        if 'train_loader' in self.data:
            # This is for image related tasks
            dataloader = self.data['train_loader'] if training else self.data['val_loader']
            # This is very important! dataloader.batch_size is controlled by dataloader.batch_sampler.batch_size
            # not the other way around. This is (probably) due to the fact that dataloader was created by setting batch_size
            dataloader.batch_sampler.batch_size = batch_size
            N = len(dataloader.dataset.imgs)
            num_chunks = (N + batch_size - 1) // batch_size
        elif 'X_train' in self.data:
            X, y = (self.data['X_train'], self.data['y_train']) if training else (self.data['X_val'], self.data['y_val'])
            N = X.size(0)
            if num_samples:
                if num_samples < N and num_samples > 0:
                    N = num_samples
                    
            if balanced_sample and isinstance(y, dtype['long']):
                dataloader = BatchLoader((X[:N], y[:N]), batch_size)
                num_chunks = len(dataloader)
            else:
                shuffle_idx = torch.randperm(N)
                X = torch.index_select(X, 0, shuffle_idx)
                y = torch.index_select(y, 0, shuffle_idx)
                num_chunks = (N + batch_size - 1) // batch_size
                X_chunks = X.chunk(num_chunks)
                y_chunks = y.chunk(num_chunks)
                dataloader = zip(X_chunks, y_chunks)
        else:
            raise ValueError('data must contain either X_train or train_loader')
        
        if training:
            print("Training:")
        else:
            print("Validating:")
            
        self._reset_avg_meter()
        end_time = time.time()
        for i, (X, y) in enumerate(dataloader):
            X = Variable(X)
            y = Variable(y)
            
            y_pred = self.model(X)
            loss = self.loss_fn(y_pred, y)
            
            if training:
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
            
            self.loss_epoch.update(loss.item(), y.size(0))
            # For classification tasks, y.data is torch.LongTensor
            # For regression tasks, y.data is torch.FloatTensor
            is_classification = isinstance(y.data, dtype['long'])
            if is_classification:
                res = check_acc(y_pred, y, (1, topk))
                self.top1.update(res[0].item())
                self.top5.update(res[1].item())
            else:
                # top1 is approximately the 'inverse' of loss
                self.top1.update(1. / (loss.item() + 1.), y.size(0))
            self.batch_time.update(time.time() - end_time)
            end_time = time.time()
            
            if training:
                self.losses_train.append(self.loss_epoch.avg)
                self.acc_train.append(self.top1.avg)
            else:
                self.losses_val.append(self.loss_epoch.avg)
                self.acc_val.append(self.top1.avg)
                
            if print_every:
                if (i + 1) % print_every == 0:
                    print('Epoch {0}: iteration {1}/{2}\t'
                          'loss: {losses.val:.3f}, avg: {losses.avg:.3f}\t'
                          'Prec@1: {prec1.val:.3f}, avg: {prec1.avg:.3f}\t'
                          'Prec@5: {prec5.val:.3f}, avg: {prec5.avg:.3f}\t'
                          'batch time: {batch_time.val:.3f} avg: {batch_time.avg:.3f}'.format(
                              epoch + 1, i + 1, num_chunks, losses=self.loss_epoch, prec1=self.top1, 
                              prec5=self.top5, batch_time=self.batch_time))
                    sys.stdout.flush()
            
        return self.top1.avg
    
    def train_eval(self, num_iter=100, batch_size=100, X=None, y=None, X_val=None, y_val=None,
                   X_test=None, y_test=None, eval_test=False, balanced_sample=False, allow_duplicate=False,
                   max_redundancy=1000, seed=None):
        if X is None or y is None:
            X, y = self.data['X_train'], self.data['y_train']
        # Currently only for classification tasks, y is torch.LongTensor 
        assert isinstance(y, dtype['long'])
        if X_val is None or y_val is None:
            X_val, y_val = self.data['X_val'], self.data['y_val']
        if eval_test and (X_test is None or y_test is None):
            X_test, y_test = self.data['X_test'], self.data['y_test']
        
        dataloader_train = BatchLoader((X, y), batch_size, balanced=balanced_sample, 
            num_iter=num_iter, allow_duplicate=allow_duplicate, max_redundancy=max_redundancy, 
            shuffle=True, seed=seed)
        dataloader_val = BatchLoader((X_val, y_val), batch_size, balanced=balanced_sample, 
            num_iter=num_iter, allow_duplicate=allow_duplicate, max_redundancy=max_redundancy, 
            shuffle=True, seed=seed)
        if X_test is not None:
            dataloader_test = BatchLoader((X_test, y_test), batch_size, balanced=balanced_sample, 
                num_iter=num_iter, allow_duplicate=allow_duplicate, max_redundancy=max_redundancy, 
                shuffle=True, seed=seed)
        else:
            dataloader_test = [None]*num_iter

        loss_train_meter = AverageMeter()
        loss_train = {'avg':[], 'batch':[]}
        acc_train_meter = AverageMeter()
        acc_train = {'avg':[], 'batch':[]}
        loss_val_meter = AverageMeter()
        loss_val = {'avg':[], 'batch':[]}
        acc_val_meter = AverageMeter()
        acc_val = {'avg':[], 'batch':[]}
        loss_test_meter = AverageMeter()
        loss_test = {'avg':[], 'batch':[]}
        acc_test_meter = AverageMeter()
        acc_test = {'avg':[], 'batch':[]}

        def forward(X, y, loss_meter, losses, acc_meter, acc, training=False):
            X = Variable(X)
            y = Variable(y)
            y_pred = self.model(X)
            loss = self.loss_fn(y_pred, y)
            loss_meter.update(loss.item(), y.size(0))
            losses['avg'].append(loss_meter.avg)
            losses['batch'].append(loss.item())
            res = check_acc(y_pred, y, (1,))
            acc_meter.update(res[0].item(), y.size(0))
            acc['avg'].append(acc_meter.avg)
            acc['batch'].append(res[0].item())

            if training:
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
            
            return y_pred, loss

        for (X, y), (X_val, y_val), test_data in zip(dataloader_train, 
                dataloader_val, dataloader_test):        
            forward(X, y, loss_train_meter, loss_train, acc_train_meter, acc_train, 
                training=True)
            forward(X_val, y_val, loss_val_meter, loss_val, acc_val_meter, acc_val, 
                training=False)
            if test_data is not None:
                X_test, y_test = test_data
                forward(X_test, y_test, loss_test_meter, loss_test, acc_test_meter, 
                    acc_test, training=False)
        
        if eval_test:
            return loss_train, acc_train, loss_val, acc_val, loss_test, acc_test
        else:
            return loss_train, acc_train, loss_val, acc_val

    
    def train(self, num_epoch = 10, batch_size=100, num_samples=None, print_every=100, 
              use_validation = True, save_checkpoint=True, file_prefix='', balanced_sample=False, topk=5):
        """train
        Args:
            num_epoch: int, default: 100 
            batch_size: int, default: 100
            num_samples: int, default: None
            print_every: int, default: 100
            use_validation: bool, default: True. If True, run_one_epoch for both training and validating
            save_checkpoint: bool, default: True. If True, save checkpoint with name (file_prefix + 'checkpoint%d.pth' % self.epoch_counter) and best model (file_prefix + 'model_best.pth').
            file_prefix: str, default:''
            balanced_sample: bool; used for sampling balanced batches from unbalanced dataset
        """
        for i in range(self.epoch_counter, self.epoch_counter + num_epoch):
            accuracy = self.run_one_epoch(i, batch_size, num_samples, print_every,
                                          balanced_sample=balanced_sample, topk=topk)
            # In case we don't want validation set. Very rare
            if use_validation:
                accuracy = self.run_one_epoch(i, batch_size, num_samples, print_every, 
                                              training=False, balanced_sample=balanced_sample, topk=topk)
            
            if accuracy > self.best_acc_val:
                self.best_acc_val = accuracy
                if save_checkpoint:
                    state = {'model_state': self.model.state_dict(), 
                            'optimizer': self.optimizer,
                            'best_acc_val': self.best_acc_val,
                            'epoch': i + 1,
                            'losses_train': self.losses_train,
                            'losses_val': self.losses_val,
                            'acc_train': self.acc_train,
                            'acc_val': self.acc_val}
                    filename = file_prefix + 'checkpoint%d.pth' % (i + 1)
                    torch.save(state, filename)
                    shutil.copyfile(filename, file_prefix + 'model_best.pth')
    
    def predict(self, batch_size=100, save_file=True, file_prefix='', X=None, y=None, topk=5, verbose=False):
        """predict
        Args:
            batch_size: int, default: 100; can be larger for large memory
            save_file: bool, default: True; if true, save file
            file_prefix: save file name: file_prefix + 'y_test.pth'
            X: default: None. If not None, use X instead of self.data['X_test']
            y: default: None. Similary to X
        """
        if X is None:
            if 'X_test' in self.data:
                X = self.data['X_test']
            elif 'test_loader' in self.data:
                X = self.data['test_loader']
                dataloader = X
            else:
                raise ValueError('If X is None, then self.data '
                                 'must contain either X_test or test_loader')
            
        if y is None and 'y_test' in self.data:
                y = self.data['y_test']
        
        is_truth_avail = isinstance(y, dtype['long']) or isinstance(y, dtype['float'])
        
        if isinstance(X, dtype['float']):
            N = X.size(0)
            num_chunks = (N + batch_size - 1) // batch_size
            X_chunks = X.chunk(num_chunks)
            dataloader = X_chunks
        
        if is_truth_avail:
            N = y.size(0)
            num_chunks = (N + batch_size - 1) // batch_size
            y_chunks = y.chunk(num_chunks)
        else:
            y_chunks = [None] * num_chunks
        
        self._reset_avg_meter()
        end_time = time.time()
        y_pred = []
        for X, y in zip(X_chunks, y_chunks):  
            X = Variable(X)
            y = Variable(y)
            
            y_pred_tmp = self.model(X) # sometimes model output a tuple
            
            if is_truth_avail:
                loss = self.loss_fn(y_pred_tmp, y)
                self.loss_epoch.update(loss.item(), y.size(0))
                if isinstance(y.data, dtype['long']):
                    res = check_acc(y_pred_tmp, y, (1, topk))
                    self.top1.update(res[0].item())
                    self.top5.update(res[1].item())
                else:
                    self.top1.update(1. / (loss.item() + 1.), y.size(0))
            self.batch_time.update(time.time() - end_time)
            end_time = time.time()
            if isinstance(y_pred_tmp, tuple):
                y_pred_tmp = y_pred_tmp[0]
            y_pred.append(y_pred_tmp)
        
        if is_truth_avail and verbose:
            print('Test set: loss: {losses.avg:.3f}\t'
                  'AP@1: {prec1.avg:.3f}\t'
                  'AP@5: {prec5.avg:.3f}\t'
                  'batch time: {batch_time.avg:.3f}'.format(
                      losses=self.loss_epoch, prec1=self.top1, 
                      prec5=self.top5, batch_time=self.batch_time))
            sys.stdout.flush()
        y_pred = torch.cat(y_pred, 0)
        if save_file:
            torch.save({'y_pred': y_pred}, file_prefix + 'y_pred.pth')
        return y_pred


if __name__ == '__main__':

    mnist_train = torchvision.datasets.MNIST('/projects/academic/jamesjar/tianlema/dl-datasets/mnist',
                                             transform=transforms.Compose([transforms.ToTensor(), 
                                                                           transforms.Normalize((0.1307,), (0.3081,))]))
    train_loader = torch.utils.data.DataLoader(mnist_train, batch_size=200)

    mnist_test = torchvision.datasets.MNIST('/projects/academic/jamesjar/tianlema/dl-datasets/mnist',
                                            transform=transforms.Compose([transforms.ToTensor(), 
                                                                          transforms.Normalize((0.1307,), (0.3081,))]), 
                                             train=False)
    test_loader = torch.utils.data.DataLoader(mnist_test, batch_size=200)

    train = list(train_loader)
    train = list(zip(*train))
    X_train = torch.cat(train[0], 0)
    y_train = torch.cat(train[1], 0)

    X_val = X_train[50000:]
    y_val = y_train[50000:]
    X_train = X_train[:50000]
    y_train = y_train[:50000]

    test = list(test_loader)
    test = list(zip(*test))
    X_test = torch.cat(test[0], 0)
    y_test = torch.cat(test[1], 0)

    data = {'X_train': X_train, 'y_train': y_train, 'X_val': X_val, 'y_val': y_val, 
            'X_test': X_test, 'y_test': y_test}



    
    model = DenseNet(input_param=(1, 64), block_layers=(6, 4), num_classes=10, 
                     growth_rate=32, bn_size=2, dropout_rate=0, transition_pool_param=(3, 1, 1))



    loss_fn = nn.CrossEntropyLoss()



    optimizer = torch.optim.Adam(model.parameters(), weight_decay=1e-4)


    
    solver = Solver(model, data, optimizer, loss_fn)
    solver.train(num_epoch=2, file_prefix='mnist-')
    solver.predict(file_prefix='mnist-')