import os.path as osp
import torch
import pandas as pd
import numpy as np

class CateAttrEvaluator(object):
    """ Evaluate DeepFashion topk recall.

    Args:
        category_topk (list(int)) topk category recall
        attr_topk (list(int)): topk attribute recall
        data_root (str): root directory of dataset

    Usage:
        ```python
        evaluator = CateAttrEvaluator()
        for batch in loader:
            out = model(batch)
            evaluator.add(out, batch)
        res = evaluator.evaluate()
        for topk, accuracy in res['category_accuracy_topk'].items():
            print('metrics/category_top{}'.format(topk), accuracy)
        for topk, accuracy in res['attr_group_recall'].items():
            for attr_type in range(1, 6):
                print('metrics/attr_top{}_type_{}_{}_recall'.format(
                    topk, attr_type, const.attrtype2name[attr_type]), accuracy[attr_type - 1]
                )
            print('metrics/attr_top{}_all_recall'.format(topk), res['attr_recall'][topk])
        ```

    """

    def __init__(self,
            category_topk=(1, 3, 5),
            attr_topk=(3, 5),
            data_root="data/df"):
        self.category_topk = category_topk
        self.attr_topk = attr_topk
        self.reset()
        with open(osp.join(data_root, 'annotations/list_attr_cloth.txt')) as f:
            ret = []
            f.readline()
            f.readline()
            for line in f:
                line = line.split(' ')
                while line[-1].strip().isdigit() is False:
                    line = line[:-1]
                ret.append([
                    ' '.join(line[0:-1]).strip(),
                    int(line[-1])
                ])
        attr_type = pd.DataFrame(ret, columns=['attr_name', 'type'])
        attr_type['attr_index'] = ['attr_' + str(i) for i in range(1000)]
        attr_type.set_index('attr_index', inplace=True)
        self.attr_type = attr_type
        self.attrtype2name = {1: 'texture', 2: 'fabric', 3: 'shape', 4: 'part', 5: 'style'}

    def reset(self):
        self.category_accuracy = []
        self.attr_group_gt = np.array([0.] * 5)
        self.attr_group_tp = np.zeros((5, len(self.attr_topk)))
        self.attr_all_gt = 0
        self.attr_all_tp = np.zeros((len(self.attr_topk),))

    def category_topk_accuracy(self, output, target):
        """ Compute topk category recall

        Notes:
            N: number of samples
            K: number of categorys

        Args:
            output (torch.Tensor(N, K)): prediction cateogory score
            target (torch.Tensor(N, ))
        """
        if isinstance(output, np.ndarray):
            output = torch.from_numpy(output)
        if isinstance(target, np.ndarray):
            target = torch.from_numpy(target)
        with torch.no_grad():
            maxk = max(self.category_topk)
            batch_size = target.shape[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 self.category_topk:
                correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
                res.append(correct_k.mul_(100 / batch_size))
            for i in range(len(res)):
                res[i] = res[i].cpu().numpy()[0] / 100

            self.category_accuracy.append(res)

    def attr_count(self, output, sample):
        attr_group_gt = np.array([0.] * 5)
        attr_group_tp = np.zeros((5, len(self.attr_topk)))
        attr_all_tp = np.zeros((len(self.attr_topk),))

        if isinstance(sample['attr'], torch.Tensor):
            target = sample['attr'].cpu().numpy()
        elif isinstance(sample['attr'], np.ndarray):
            target = sample['attr']
        else:
            raise TypeError("Error type of sampe['attr']")
        target = np.split(target, target.shape[0])
        target = [x[0, :] for x in target]

        if isinstance(output['attr'], torch.Tensor):
            pred = output['attr'].cpu().detach().numpy()
        elif isinstance(output['attr'], np.ndarray):
            pred = output['attr']
        else:
            raise TypeError("Error type of output['attr']")
        pred = np.split(pred, pred.shape[0])
        pred = [x[0, 1, :] for x in pred]

        for batch_idx in range(len(target)):
            result_df = pd.DataFrame([target[batch_idx], pred[batch_idx]],
                                     index=['target', 'pred'], columns=['attr_' + str(i) for i in range(1000)])
            result_df = result_df.transpose()
            result_df = result_df.join(self.attr_type[['type']])
            ret = []
            for i in range(1, 6):
                ret.append(result_df[result_df['type'] == i]['target'].sum())
            attr_group_gt += np.array(ret)
            ret = []
            result_df = result_df.sort_values('pred', ascending=False)
            attr_all_tp += np.array([
                result_df.head(k)['target'].sum() for k in self.attr_topk
            ])
            for i in range(1, 6):
                sort_df = result_df[result_df['type'] == i]
                ret.append([
                    sort_df.head(k)['target'].sum() for k in self.attr_topk
                ])
            attr_group_tp += np.array(ret)

        self.attr_group_gt += attr_group_gt
        self.attr_group_tp += attr_group_tp

        self.attr_all_gt += attr_group_gt.sum()
        self.attr_all_tp += attr_all_tp

    def add(self, output, sample):
        self.category_topk_accuracy(output['cate'], sample['category_id'])
        self.attr_count(output, sample)

    def evaluate(self):
        category_accuracy = np.array(self.category_accuracy).mean(axis=0)
        category_accuracy_topk = {}
        for i, top_n in enumerate(self.category_topk):
            category_accuracy_topk[top_n] = category_accuracy[i]

        attr_group_recall = {}
        attr_recall = {}
        for i, top_n in enumerate(self.attr_topk):
            attr_group_recall[top_n] = self.attr_group_tp[..., i] / self.attr_group_gt
            attr_recall[top_n] = self.attr_all_tp[i] / self.attr_all_gt

        return {
            'category_accuracy_topk': category_accuracy_topk,
            'attr_group_recall': attr_group_recall,
            'attr_recall': attr_recall,
        }


class CateCalculator(object):
    """Calculate Category prediction top-k recall rate

    Usage:
        ```python
        cate_calculator = CateCalculator()
        for data in loader:
            result = model(data)
            cate_calculator.collect_result(result, data)
        cate_calculator.show_result()
        ```
    """
    def __init__(self, category_num, topns=[1, 3], show_cate_name=False, cate_name_file=None):
        self.collector = dict()
        self.num_cate = category_num
        self.topns = topns
        # true positive
        for topi in topns:
            self.collector['top%s' % str(topi)] = dict()
            tp = np.zeros(self.num_cate)
            fn = np.zeros(self.num_cate)
            self.collector['top%s' % str(topi)]['tp'] = tp
            self.collector['top%s' % str(topi)]['fn'] = fn

        """ topn recall rate """
        self.recall = dict()
        
        # collect target per category
        self.target_per_cate = np.zeros(self.num_cate)

        # num of total predictions
        self.total = 0
        # topn category prediction
        self.topns = topns

        self.show_cate_name = show_cate_name
        if self.show_cate_name:
            assert cate_name_file is not None
            self.cate_dict = {}
            cate_names = open(cate_name_file).readlines()
            for i, cate_name in enumerate(cate_names[2:]):
                self.cate_dict[i] = cate_name.split()[0]

    def collect(self, indexes, target, topk):
        """calculate and collect recall rate
        Args:
            indexes(list): predicted top-k indexes  
            target(list): ground-truth
            topk(str): top-k, e.g., top1, top3, top5
        """
        for i, cnt in enumerate(self.collector[topk]['tp']): # true-positive
            if i in indexes and i in target:
                self.collector[topk]['tp'][i] += 1

        for i, cnt in enumerate(self.collector[topk]['fn']): # false negative
            if i not in indexes and i in target:
                self.collector[topk]['fn'][i] += 1

    def collect_result(self, pred, target):
        """collect_result.

        Args:
            pred: Tensor of shape [N, category_num]
            target: Tensor of shape [N, 1]
        """
        if isinstance(pred, torch.Tensor):
            data = pred.data.cpu().numpy()
        elif isinstance(pred, np.ndarray):
            data = pred
        else:
            raise TypeError('type {} cannot be calculated.'.format(type(pred)))

        for i in range(pred.shape[0]):
            self.total += 1
            indexes = np.argsort(data[i])[::-1]
            for k in self.topns:
                idx = indexes[:k]
                self.collect(idx, target[i], 'top%d'%k)

    def compute_one_recall(self, tp, fn):
        """ Mean recall for each category. """
        empty = 0
        recall = np.zeros(tp.shape)
        for i, num in enumerate(tp):
            # ground truth number = true_positive(tp) + false_negative(fn)
            if tp[i] + fn[i] == 0:
                empty += 1
                continue
            else:
                recall[i] = float(tp[i]) / float(tp[i] + fn[i])
        sorted_recall = sorted(recall)[::-1]
        return 100 * sum(sorted_recall) / (len(sorted_recall) - empty)

    def compute_recall(self):
        for key, top in self.collector.items():
            self.recall[key] = self.compute_one_recall(top['tp'], top['fn'])


    def show_result(self, batch_idx=None):
        print('----------- Category Prediction ----------')
        if batch_idx is not None:
            print('Batch[%d]' % batch_idx)
        else:
            print('Total')

        self.compute_recall()
        print('[ Recall Rate ]')
        for k in self.topns:
            print('top%d = %.2f' % (k, self.recall['top%d' % k]))

    def show_per_cate_result(self):
        for key, top in self.collector.items():
            tp = top['tp']
            fn = top['fn']
            recall = np.zeros(tp.shape)
            print("------------- %s Per-Category Recall -------------" % key)
            for i, num in enumerate(tp):
                if tp[i] + fn[i] == 0:
                    print('%s(count: 0)' % self.cate_dict[i])
                    continue
                else:
                    recall[i] = float(tp[i]) / float(tp[i] + fn[i])
                    print('%s(count: %d) %.2f'% (self.cate_dict[i], (tp[i]+fn[i]), recall[i]))

class AttrCalculator(object):

    def __init__(self,
                 num_of_cls,
                 topns=[3, 5, 10],
                 show_attr_name=False,
                 attr_name_file=None):
        """Create the empty array to count true positive(tp),
            true negative(tn), false positive(fp) and false negative(fn).

        Args:
            num_of_cls:
            class_num(int): number of classes in the dataset.
            topns(list of int): default calculate top3, top5 and top10
                accuracy.
            show_attr_name(bool): print predicted attribute name, for demo
                usage.
            attr_name_file(str): file of attribute name, used for mapping
                attribute index to attribute names.

        Usage:
            ```python
            for data in loader:
                result = model(data)
                attr_calculator.collect_result(result, data)
            attr_calculator.show_result()
            attr_calculator.show_per_attr_result()
            ```
        """
        self.collector = dict()
        self.total = 0  # the number of total predictions
        num_classes = num_of_cls
        self.topns = topns
        for i in self.topns:
            tp, tn, fp, fn = np.zeros(num_classes), np.zeros(
                num_classes), np.zeros(num_classes), np.zeros(num_classes)
            pos = np.zeros(num_classes)
            self.collector['top%s' % str(i)] = dict()
            self.collector['top%s' % str(i)]['tp'] = tp
            self.collector['top%s' % str(i)]['tn'] = tn
            self.collector['top%s' % str(i)]['fp'] = fp
            self.collector['top%s' % str(i)]['fn'] = fn
            self.collector['top%s' % str(i)]['pos'] = pos
        """ precision = true_positive/(true_positive+false_positive)"""
        self.precision = dict()
        """ accuracy = (true_positive+true_negative)/total_precision"""
        self.accuracy = dict()
        """ topn recall rate """
        self.recall = dict()
        self.topn = 50

        self.show_attr_name = show_attr_name
        if self.show_attr_name:
            assert attr_name_file is not None
            # map the index of attribute to attribute name
            self.attr_dict = {}
            attr_names = open(attr_name_file).readlines()
            for i, attr_name in enumerate(attr_names[2:]):
                self.attr_dict[i] = attr_name.split()[0]

    def get_dict(self, fn):
        rf = open(fn).readlines()
        dic = dict()
        for i, line in enumerate(rf):
            id_num = int(line.strip('\n'))
            dic[i] = id_num

        return dic

    def collect(self, indexes, target, top):
        for i, t in enumerate(target):
            if t == 1:
                if i in indexes:
                    top['pos'][i] += 1
                    top['tp'][i] += 1
                else:
                    top['fn'][i] += 1
            if t == 0:
                if i in indexes:
                    top['pos'][i] += 1
                    top['fp'][i] += 1
                else:
                    top['tn'][i] += 1

    def index_to_attribute_name(self, index):
        for pred_i in index:
            pred_attr_name = self.attr_dict[pred_i]
            print(pred_attr_name)

    def collect_result(self, pred, target):
        """collect_result.

        Args:
            pred: Tensor of shape [N, attribute_num]
            target: Tensor of shape [N, attribute_num]
        """
        if isinstance(pred, torch.Tensor):
            data = pred.data.cpu().numpy()
        elif isinstance(pred, np.ndarray):
            data = pred
        else:
            raise TypeError('type {} cannot be calculated.'.format(type(pred)))

        for i in range(pred.shape[0]):
            self.total += 1
            indexes = np.argsort(data[i])[::-1]
            for k in self.topns:
                idx = indexes[:k]
                self.collect(idx, target[i], self.collector['top%d'%k])

    def compute_one_recall(self, tp, fn):
        empty = 0
        recall = np.zeros(tp.shape)
        for i, num in enumerate(tp):
            # ground truth number = true_positive(tp) + false_negative(fn)
            if tp[i] + fn[i] == 0:
                empty += 1
                continue
            else:
                recall[i] = float(tp[i]) / float(tp[i] + fn[i])
        sorted_recall = sorted(recall)[::-1]
        return 100 * sum(sorted_recall[:self.topn]) / min(
            self.topn,
            len(sorted_recall) - empty)

    def compute_recall(self):
        for key, top in self.collector.items():
            self.recall[key] = self.compute_one_recall(top['tp'], top['fn'])

    def compute_one_precision(self, tp, fp, pos):
        empty = 0
        precision = np.zeros(tp.shape)
        for i, num in enumerate(tp):
            if pos[i] == 0:
                empty += 1
                continue
            else:
                precision[i] = float(tp[i]) / float(pos[i])
        sorted_precision = sorted(precision)[::-1]
        return 100 * sum(sorted_precision[:self.topn]) / min(
            self.topn,
            len(sorted_precision) - empty)

    def compute_precision(self):
        for key, top in self.collector.items():
            self.precision[key] = self.compute_one_precision(
                top['tp'], top['fp'], top['pos'])

    def compute_one_accuracy(self, tp, tn):
        accuracy = np.zeros(tp.shape)
        for i, num in enumerate(tp):
            accuracy[i] = float(tp[i] + tn[i]) / float(self.total)
        return 100 * float(np.sum(accuracy)) / len(accuracy)

    def compute_accuracy(self):
        for key, top in self.collector.items():
            self.accuracy[key] = self.compute_one_accuracy(
                top['tp'], top['tn'])

    def show_result(self, batch_idx=None):
        print('----------- Attribute Prediction ----------')
        if batch_idx is not None:
            print('Batch[%d]' % batch_idx)
        else:
            print('Total')

        self.compute_recall()
        print('[ Recall Rate ]')
        for k in self.topns:
            print('top%d = %.2f' % (k, self.recall['top%d'%k]))

        self.compute_accuracy()
        print('[ Accuracy ]')
        for k in self.topns:
            print('top%d = %.2f' % (k, self.accuracy['top%d'%k]))

        self.compute_precision()
        print('[ Precision ]')
        for k in self.topns:
            print('top%d = %.2f' % (k, self.precision['top%d'%k]))

    def show_per_attr_result(self):
        for key, top in self.collector.items():
            tp = top['tp']
            tn = top['tn']
            accuracy = np.zeros(tp.shape)
            recall = np.zeros(tp.shape)
            print("------------- %s Per-Attribute Accuracy -------------" % key)
            for i, num in enumerate(tp):
                accuracy[i] = float(tp[i] + tn[i]) / float(self.total)
                # i: attribute index
                print('%s %.2f'% (self.attr_dict[i], accuracy[i]))
            print("------------- %s Per-Attribute Recall -------------" % key)
            for i, num in enumerate(tp):
                recall[i] = float(tp[i]) / float(tp[i] + tn[i])
                # i: attribute index
                print('%s %.2f'% (self.attr_dict[i], recall[i]))
