import torch
import shutil
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
from sklearn.utils.multiclass import unique_labels
import logging
import datetime 


class ImbalancedDatasetSampler(torch.utils.data.sampler.Sampler):

    def __init__(self, dataset, indices=None, num_samples=None):
                
        # if indices is not provided, 
        # all elements in the dataset will be considered
        self.indices = list(range(len(dataset))) \
            if indices is None else indices
            
        # if num_samples is not provided, 
        # draw `len(indices)` samples in each iteration
        self.num_samples = len(self.indices) \
            if num_samples is None else num_samples
            
        # distribution of classes in the dataset 
        label_to_count = [0] * len(np.unique(dataset.targets))
        for idx in self.indices:
            label = self._get_label(dataset, idx)
            label_to_count[label] += 1
            
        beta = 0.9999
        effective_num = 1.0 - np.power(beta, label_to_count)
        per_cls_weights = (1.0 - beta) / np.array(effective_num)

        # weight for each sample
        weights = [per_cls_weights[self._get_label(dataset, idx)]
                   for idx in self.indices]
        self.weights = torch.DoubleTensor(weights)
        
    def _get_label(self, dataset, idx):
        return dataset.targets[idx]
                
    def __iter__(self):
        return iter(torch.multinomial(self.weights, self.num_samples, replacement=True).tolist())

    def __len__(self):
        return self.num_samples

def calc_confusion_mat(val_loader, model, args):
    
    model.eval()
    all_preds = []
    all_targets = []
    with torch.no_grad():
        for i, (input, target) in enumerate(val_loader):
            if args.gpu is not None:
                input = input.cuda(args.gpu, non_blocking=True)
            target = target.cuda(args.gpu, non_blocking=True)

            # compute output
            output = model(input)
            _, pred = torch.max(output, 1)
            all_preds.extend(pred.cpu().numpy())
            all_targets.extend(target.cpu().numpy())
    cf = confusion_matrix(all_targets, all_preds).astype(float)

    cls_cnt = cf.sum(axis=1)
    cls_hit = np.diag(cf)

    cls_acc = cls_hit / cls_cnt

    print('Class Accuracy : ')
    print(cls_acc)
    classes = [str(x) for x in args.cls_num_list]
    plot_confusion_matrix(all_targets, all_preds, classes)
    plt.savefig(os.path.join(args.root_log, args.store_name, 'confusion_matrix.png'))

def plot_confusion_matrix(y_true, y_pred, classes,
                          normalize=False,
                          title=None,
                          cmap=plt.cm.Blues):
    
    if not title:
        if normalize:
            title = 'Normalized confusion matrix'
        else:
            title = 'Confusion matrix, without normalization'

    # Compute confusion matrix
    cm = confusion_matrix(y_true, y_pred)
    
    fig, ax = plt.subplots()
    im = ax.imshow(cm, interpolation='nearest', cmap=cmap)
    ax.figure.colorbar(im, ax=ax)
    # We want to show all ticks...
    ax.set(xticks=np.arange(cm.shape[1]),
           yticks=np.arange(cm.shape[0]),
           # ... and label them with the respective list entries
           xticklabels=classes, yticklabels=classes,
           title=title,
           ylabel='True label',
           xlabel='Predicted label')

    # Rotate the tick labels and set their alignment.
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right",
             rotation_mode="anchor")

    # Loop over data dimensions and create text annotations.
    fmt = '.2f' if normalize else 'd'
    thresh = cm.max() / 2.
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j, i, format(cm[i, j], fmt),
                    ha="center", va="center",
                    color="white" if cm[i, j] > thresh else "black")
    fig.tight_layout()
    return ax

def prepare_folders(args):
    
    folders_util = [args.root_log, args.root_model,
                    os.path.join(args.root_log, args.store_name),
                    os.path.join(args.root_model, args.store_name)]
    for folder in folders_util:
        if not os.path.exists(folder):
            print('creating folder ' + folder)
            os.mkdir(folder)

def prepare_folders_eval(args):
    
    folders_util = [args.root_eval, 
                    os.path.join(args.root_eval, args.store_name)]
    for folder in folders_util:
        if not os.path.exists(folder):
            print('creating folder ' + folder)
            os.mkdir(folder)

def save_checkpoint(logdir, state, is_best):
    
    filename = '%s/ckpt.pth.tar' % logdir
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, filename.replace('pth.tar', 'best.pth.tar'))


class AverageMeter(object):
    
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        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 __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)


def accuracy(output, target, topk=(1,)):
    
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


class Confidence_Diagram():
    def __init__(self,n_classes):
        # self.logdir = logdir
        self.num_class = n_classes
        self.correct = np.zeros((10,n_classes))
        self.count = np.zeros((10,n_classes))
        self.confidence = np.zeros((10,n_classes))

        # self.file_name = file_name

    def aggregate_stats(self,prob,predicted,targets):
        # prob [batch,w,h]
        # predicted [batch,w,h]
        # targets [batch, w,h]
        
        bins = prob // 0.1 #[batch,w,h] confidence
        bins = bins.masked_fill(bins == 10.,9)
        # for i in range(targets.shape[1]):
        #     mask = predicted == targets[:,i,:,:]

        for i in np.linspace(0,9,10):
            mask_bin = bins == i
            for j in range(16):
                mask_class = targets == j
                self.correct[int(i),int(j)] += (predicted[mask_bin*mask_class] == targets[mask_bin*mask_class]).sum()
                self.count[int(i),int(j)] += (mask_bin*mask_class).sum()
                self.confidence[int(i),int(j)] +=  prob[mask_bin*mask_class].sum()

    def compute_ece(self):  
        self.accuracy = self.correct/self.count
        self.confidence = self.confidence/self.count
        # import ipdb;ipdb.set_trace()
        ratio = self.count/np.expand_dims(np.nansum(self.count,0),0)
        self.ece_cls = np.nansum(ratio*abs(self.accuracy - self.confidence),0)
        self.ece = np.sum(self.ece_cls)/self.num_class
        
    def save(self,logdir):
        with open(os.path.join(logdir,'calibration.csv'), mode='w') as calibration:
            writer = csv.writer(calibration, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
            # import ipdb;ipdb.set_trace()
            writer.writerow(['mean accuracy per bin:']+(np.nansum(self.accuracy,1)/self.num_class).tolist())
            writer.writerow(['ece per class:'] + self.ece_cls.tolist())
            writer.writerow(['mean ece:',self.ece])

    def print(self):
        print(np.nansum(self.accuracy,1)/self.num_class)
        print(self.ece_cls)
        print(self.ece)

def get_logger(logdir):
    logger = logging.getLogger("imbalance")
    ts = str(datetime.datetime.now()).split(".")[0].replace(" ", "_")
    ts = ts.replace(":", "_").replace("-", "_")
    file_path = os.path.join(logdir, "run_{}.log".format(ts))
    hdlr = logging.FileHandler(file_path)
    formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
    hdlr.setFormatter(formatter)
    logger.addHandler(hdlr)
    logger.setLevel(logging.INFO)
    return logger
import torch
import shutil
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
from sklearn.utils.multiclass import unique_labels
import logging
import datetime 


class ImbalancedDatasetSampler(torch.utils.data.sampler.Sampler):

    def __init__(self, dataset, indices=None, num_samples=None):
                
        # if indices is not provided, 
        # all elements in the dataset will be considered
        self.indices = list(range(len(dataset))) \
            if indices is None else indices
            
        # if num_samples is not provided, 
        # draw `len(indices)` samples in each iteration
        self.num_samples = len(self.indices) \
            if num_samples is None else num_samples
            
        # distribution of classes in the dataset 
        label_to_count = [0] * len(np.unique(dataset.targets))
        for idx in self.indices:
            label = self._get_label(dataset, idx)
            label_to_count[label] += 1
            
        beta = 0.9999
        effective_num = 1.0 - np.power(beta, label_to_count)
        per_cls_weights = (1.0 - beta) / np.array(effective_num)

        # weight for each sample
        weights = [per_cls_weights[self._get_label(dataset, idx)]
                   for idx in self.indices]
        self.weights = torch.DoubleTensor(weights)
        
    def _get_label(self, dataset, idx):
        return dataset.targets[idx]
                
    def __iter__(self):
        return iter(torch.multinomial(self.weights, self.num_samples, replacement=True).tolist())

    def __len__(self):
        return self.num_samples

def calc_confusion_mat(val_loader, model, args):
    
    model.eval()
    all_preds = []
    all_targets = []
    with torch.no_grad():
        for i, (input, target) in enumerate(val_loader):
            if args.gpu is not None:
                input = input.cuda(args.gpu, non_blocking=True)
            target = target.cuda(args.gpu, non_blocking=True)

            # compute output
            output = model(input)
            _, pred = torch.max(output, 1)
            all_preds.extend(pred.cpu().numpy())
            all_targets.extend(target.cpu().numpy())
    cf = confusion_matrix(all_targets, all_preds).astype(float)

    cls_cnt = cf.sum(axis=1)
    cls_hit = np.diag(cf)

    cls_acc = cls_hit / cls_cnt

    print('Class Accuracy : ')
    print(cls_acc)
    classes = [str(x) for x in args.cls_num_list]
    plot_confusion_matrix(all_targets, all_preds, classes)
    plt.savefig(os.path.join(args.root_log, args.store_name, 'confusion_matrix.png'))

def plot_confusion_matrix(y_true, y_pred, classes,
                          normalize=False,
                          title=None,
                          cmap=plt.cm.Blues):
    
    if not title:
        if normalize:
            title = 'Normalized confusion matrix'
        else:
            title = 'Confusion matrix, without normalization'

    # Compute confusion matrix
    cm = confusion_matrix(y_true, y_pred)
    
    fig, ax = plt.subplots()
    im = ax.imshow(cm, interpolation='nearest', cmap=cmap)
    ax.figure.colorbar(im, ax=ax)
    # We want to show all ticks...
    ax.set(xticks=np.arange(cm.shape[1]),
           yticks=np.arange(cm.shape[0]),
           # ... and label them with the respective list entries
           xticklabels=classes, yticklabels=classes,
           title=title,
           ylabel='True label',
           xlabel='Predicted label')

    # Rotate the tick labels and set their alignment.
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right",
             rotation_mode="anchor")

    # Loop over data dimensions and create text annotations.
    fmt = '.2f' if normalize else 'd'
    thresh = cm.max() / 2.
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j, i, format(cm[i, j], fmt),
                    ha="center", va="center",
                    color="white" if cm[i, j] > thresh else "black")
    fig.tight_layout()
    return ax

def prepare_folders(args):
    
    folders_util = [args.root_log, args.root_model,
                    os.path.join(args.root_log, args.store_name),
                    os.path.join(args.root_model, args.store_name)]
    for folder in folders_util:
        if not os.path.exists(folder):
            print('creating folder ' + folder)
            os.mkdir(folder)

def prepare_folders_eval(args):
    
    folders_util = [args.root_eval, 
                    os.path.join(args.root_eval, args.store_name)]
    for folder in folders_util:
        if not os.path.exists(folder):
            print('creating folder ' + folder)
            os.mkdir(folder)

def save_checkpoint(logdir, state, is_best):
    
    filename = '%s/ckpt.pth.tar' % logdir
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, filename.replace('pth.tar', 'best.pth.tar'))


class AverageMeter(object):
    
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        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 __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)


def accuracy(output, target, topk=(1,)):
    
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


class Confidence_Diagram():
    def __init__(self,n_classes):
        # self.logdir = logdir
        self.num_class = n_classes
        self.correct = np.zeros((10,n_classes))
        self.count = np.zeros((10,n_classes))
        self.confidence = np.zeros((10,n_classes))

        # self.file_name = file_name

    def aggregate_stats(self,prob,predicted,targets):
        # prob [batch,w,h]
        # predicted [batch,w,h]
        # targets [batch, w,h]
        
        bins = prob // 0.1 #[batch,w,h] confidence
        bins = bins.masked_fill(bins == 10.,9)
        # for i in range(targets.shape[1]):
        #     mask = predicted == targets[:,i,:,:]

        for i in np.linspace(0,9,10):
            mask_bin = bins == i
            for j in range(16):
                mask_class = targets == j
                self.correct[int(i),int(j)] += (predicted[mask_bin*mask_class] == targets[mask_bin*mask_class]).sum()
                self.count[int(i),int(j)] += (mask_bin*mask_class).sum()
                self.confidence[int(i),int(j)] +=  prob[mask_bin*mask_class].sum()

    def compute_ece(self):  
        self.accuracy = self.correct/self.count
        self.confidence = self.confidence/self.count
        # import ipdb;ipdb.set_trace()
        ratio = self.count/np.expand_dims(np.nansum(self.count,0),0)
        self.ece_cls = np.nansum(ratio*abs(self.accuracy - self.confidence),0)
        self.ece = np.sum(self.ece_cls)/self.num_class
        
    def save(self,logdir):
        with open(os.path.join(logdir,'calibration.csv'), mode='w') as calibration:
            writer = csv.writer(calibration, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
            # import ipdb;ipdb.set_trace()
            writer.writerow(['mean accuracy per bin:']+(np.nansum(self.accuracy,1)/self.num_class).tolist())
            writer.writerow(['ece per class:'] + self.ece_cls.tolist())
            writer.writerow(['mean ece:',self.ece])

    def print(self):
        print(np.nansum(self.accuracy,1)/self.num_class)
        print(self.ece_cls)
        print(self.ece)

def get_logger(logdir):
    logger = logging.getLogger("imbalance")
    ts = str(datetime.datetime.now()).split(".")[0].replace(" ", "_")
    ts = ts.replace(":", "_").replace("-", "_")
    file_path = os.path.join(logdir, "run_{}.log".format(ts))
    hdlr = logging.FileHandler(file_path)
    formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
    hdlr.setFormatter(formatter)
    logger.addHandler(hdlr)
    logger.setLevel(logging.INFO)
    return logger
# if indices is not provided, # all elements in the dataset will be considered # if num_samples is not provided, # draw `len(indices)` samples in each iteration # distribution of classes in the dataset # weight for each sample # compute output # Compute confusion matrix # We want to show all ticks... # ... and label them with the respective list entries # Rotate the tick labels and set their alignment. # Loop over data dimensions and create text annotations. # self.logdir = logdir # self.file_name = file_name # prob [batch,w,h] # predicted [batch,w,h] # targets [batch, w,h] #[batch,w,h] confidence # for i in range(targets.shape[1]): #     mask = predicted == targets[:,i,:,:] # import ipdb;ipdb.set_trace() # import ipdb;ipdb.set_trace()