import torch
import torch.distributed as dist
import numpy as np
import random
import os
import shutil

from metrics import *

def init_random_seed(seed=0):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    np.random.seed(seed)  # Numpy module.
    random.seed(seed)  # Python random module.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    
def generate_default_config():
    configs = {}
    
    # Device
    configs['use_gpu'] = torch.cuda.is_available()
    configs['device'] = torch.device('cuda' if torch.cuda.is_available() 
                                     and configs['use_gpu'] else 'cpu')
    
    # Training parameters
    configs['train_batch_size'] = 32
    configs['lr'] = 5e-5
    configs['weight_decay'] = 1e-2
    
    # Reproducibility
    configs['rand_seed'] = 1
    
    # Testing parameters
    configs['test_batch_size'] = configs['train_batch_size']
    
    # Dirs
    configs['pretained_model_dir'] = "../Pretain"
    configs['dataset_dir'] = "../Datasets"
    
    # Others
    configs['dtype'] = torch.float
    configs['eps'] = 1e-5
    
    return configs

def clear_old_logs(path):
    if not os.path.exists(path):
        os.makedirs(path)
    else:
        shutil.rmtree(path)
        os.makedirs(path)
        
def get_learning_rate(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']
        
def display(content, log_file=None):
    '''Print 'content' to stdout and file (optional)'''
    print(content)
    with open(log_file, 'a') as f:
        f.write(content + '\n')
        
class ValueMeter(object):
    """Stores the current value"""
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0

    def update(self, val):
        self.val = val

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '}'
        return fmtstr.format(**self.__dict__)
        
class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self, name, fmt=':f', val_only=False):
        self.name = name
        self.fmt = fmt
        self.val_only = val_only
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
        
    def reduce(self, dst=0):
        temp = torch.tensor([self.sum, self.count], dtype=torch.float32, device=torch.device("cuda"))
        dist.reduce(temp, dst, dist.ReduceOp.SUM)
        self.sum, self.count = temp.tolist()
        self.avg = self.sum / self.count

    def __str__(self):
        if self.val_only:
            fmtstr = '{name} {val' + self.fmt + '}'
        else:
            fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)

class ProgressMeter(object):
    def __init__(self, num_batches, meters, prefix=""):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
        self.meters = meters
        self.prefix = prefix

    def display(self, batch, log_file=None):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]
        entries += [str(meter) for meter in self.meters]
        display('  '.join(entries), log_file)

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))
        fmt = '{:' + str(num_digits) + 'd}'
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'

class AveragePrecisionMeter(object):
    """
    The APMeter measures the average precision per class.
    The APMeter is designed to operate on `NxK` Tensors `output` and `target`.
    """

    def __init__(self, norm=False):
        super(AveragePrecisionMeter, self).__init__()
        self.norm = norm
        self.reset()
        
    def reset(self):
        self.scores = []
        self.targets = []
    
    def add(self, output, target):
        if self.norm:
            output = torch.sigmoid(output)
        self.scores.append(output.detach().cpu().float())
        self.targets.append(target.detach().cpu())

    def value(self):
        """Returns the model's average precision for each class
        Return:
            ap (FloatTensor): 1xK tensor, with avg precision for each class k
        """
        self._list2tensor()
        
        if self.scores.numel() == 0:
            return 0
        ap = torch.zeros(self.scores.size(1))
        # compute average precision for each class
        for k in range(self.scores.size(1)):
            # sort scores
            scores = self.scores[:, k]
            targets = self.targets[:, k]
            # compute average precision
            ap[k] = self.average_precision(scores, targets)
        return ap
    
    def save_results(self, file_path=None, postfix='', rank=0):
        self._list2tensor()
        if rank > 0:
            save_data = torch.stack([self.scores, self.targets], dim=0)
            file_name = os.path.join(file_path, 'evaluation_tmp.{}.{}.pt'.format(postfix, rank))
            torch.save(save_data, file_name)
    
    def gather(self, file_path=None, postfix='', total_ranks=1):
        for rank in range(1, total_ranks):
            file_name = os.path.join(file_path, 'evaluation_tmp.{}.{}.pt'.format(postfix, rank))
            if os.path.isfile(file_name):
                temp = torch.load(file_name, map_location=torch.device('cpu'))
                self.scores = torch.cat([self.scores, temp[0]], dim=0)
                self.targets = torch.cat([self.targets, temp[1]], dim=0)
                os.remove(file_name)
                continue
            else:
                raise ValueError("=> no evaluation temp found at '{}'".format(file_name))

    # @staticmethod
    def average_precision(self, output, target):
        valid_ind = target > -1 # -1 means unannotated label
        output = output[valid_ind]
        target = target[valid_ind]
        if torch.sum(target) == 0: # skip label with no positive instances
            return 0.0
        
        # sort examples
        sorted, indices = torch.sort(output, dim=0, descending=True)

        # Computes ap
        pos_count = 0.
        ap = 0.
        for i in range(output.size(0)):
            if target[indices[i]] == 1:
                pos_count += 1
                ap += pos_count / (i+1)
        ap /= pos_count
        return ap

    def overall(self, threshold=0.5):
        self._list2tensor()
        
        if self.scores.numel() == 0:
            return 0
        scores = self.scores.cpu().numpy()
        targets = self.targets.cpu().numpy()
        return self.evaluation(scores, targets, threshold)

    def overall_topk(self, k, threshold=0.5):
        self._list2tensor()
        
        targets = self.targets.cpu().numpy()
        # targets[targets == -1] = 0
        n, c = self.scores.size()
        scores = np.zeros((n, c)) - 1
        index = self.scores.topk(k, 1, True, True)[1].cpu().numpy()
        tmp = self.scores.cpu().numpy()
        for i in range(n):
            for ind in index[i]:
                scores[i, ind] = 1 if tmp[i, ind] >= threshold else -1
        return self.evaluation(scores, targets)
    
    def evaluation(self, scores_, targets_, threshold=0.5):
        n, n_class = scores_.shape
        Nc, Np, Ng = np.zeros(n_class), np.zeros(n_class), np.zeros(n_class)
        
        valid_ind = targets_ > -1 # -1 means unannotated label
        for k in range(n_class):
            scores = scores_[valid_ind[:, k], k]
            targets = targets_[valid_ind[:, k], k]
            Ng[k] = np.sum(targets == 1)
            Np[k] = np.sum(scores >= threshold)
            Nc[k] = np.sum(targets * (scores >= threshold))
        Np[Np == 0] = 1
        OP = np.sum(Nc) / np.sum(Np)
        OR = np.sum(Nc) / np.sum(Ng)
        OF1 = (2 * OP * OR) / (OP + OR)

        CP = np.sum(Nc / Np) / n_class
        CR = np.sum(Nc / Ng) / n_class
        CF1 = (2 * CP * CR) / (CP + CR)
        return OP, OR, OF1, CP, CR, CF1
    
    def value_AP(self):
        self._list2tensor()
        
        if self.scores.numel() == 0:
            return 0
        scores = self.scores.cpu()
        targets = self.targets.cpu()
        
        return AveragePrecision(scores, targets)
    
    def value_COV(self):
        self._list2tensor()
        
        if self.scores.numel() == 0:
            return 0
        scores = self.scores.cpu()
        targets = self.targets.cpu()
        
        return Coverage(scores, targets)
    
    def traditional_metrics(self, threshold=0.5):
        self._list2tensor()
        
        if self.scores.numel() == 0:
            return 0
        scores = self.scores.cpu()
        targets = self.targets.cpu()
        
        metrics = {}
        metrics['AveragePrecision'] = AveragePrecision(scores, targets)
        metrics['OneError'] = OneError(scores, targets)
        metrics['Coverage'] = Coverage(scores, targets)
        metrics['RankingLoss'] = RankingLoss(scores, targets)
        
        n, c = self.scores.size()
        preds = torch.zeros((n, c))
        preds[scores >= threshold] = 1
        metrics['HammingLoss'] = HammingLoss(preds, targets)
        return metrics
        
    def _list2tensor(self):
        if isinstance(self.scores, list):
            self.scores = torch.cat(self.scores, dim=0)
            self.targets = torch.cat(self.targets, dim=0)