# from eval_file import *

import argparse
from collections import defaultdict
from itertools import chain
from math import pow
from pathlib import Path

# from common_utils.common_io import load_bio_file_into_sents
# from common_utils.common_log import create_logger
# -*- coding: utf-8 -*-

# -*- coding: utf-8 -*-

import json
import pickle as pkl


def read_from_file(ifn):
    with open(ifn, "r") as f:
        text = f.read()
    return text


def write_to_file(text, ofn):
    with open(ofn, "w") as f:
        f.write(text)
    return True


def pkl_load(ifn):
    with open(ifn, "rb") as f:
        pdata = pkl.load(f)
    return pdata


def pkl_dump(pdata, ofn):
    with open(ofn, "wb") as f:
        pkl.dump(pdata, f)
    return True


def json_load(ifn):
    with open(ifn, "r") as f:
        jdata = json.load(f)
    return jdata


def json_dump(jdata, ofn):
    with open(ofn, "w") as f:
        json.dump(jdata, f)
    return True


def load_bio_file_into_sents(bio_file, word_sep=" ", do_lower=False):
    bio_text = read_from_file(bio_file)
    bio_text = bio_text.strip()
    if do_lower:
        bio_text = bio_text.lower()

    new_sents = []
    sents = bio_text.split("\n\n")

    for sent in sents:
        new_sent = []
        words = sent.split("\n")
        for word in words:
            new_word = word.split(word_sep)
            new_sent.append(new_word)
        new_sents.append(new_sent)

    return new_sents


def output_bio(bio_data, output_file, sep=" "):
    with open(output_file, "w") as f:
        for sent in bio_data:
            for word in sent:
                line = sep.join(word)
                f.write(line)
                f.write("\n")
            f.write("\n")


class PRF:
    def __init__(self):
        self.true = 0
        self.false = 0

    def add_true_case(self):
        self.true += 1

    def add_false_case(self):
        self.false += 1

    def get_true_false_counts(self):
        return self.true, self.false

    def __str__(self):
        return str(self.__dict__)


class BioEval:
    def __init__(self):
        self.acc = PRF()
        # prediction
        self.all_strict = PRF()
        self.all_relax = PRF()
        self.cat_strict = defaultdict(PRF)
        self.cat_relax = defaultdict(PRF)
        # gold standard
        self.gs_all = 0
        self.gs_cat = defaultdict(int)
        self.performance = dict()
        self.counts = dict()
        self.beta = 1
        self.label_not_for_eval = {'o'}

    def reset(self):
        self.acc = PRF()
        self.all_strict = PRF()
        self.all_relax = PRF()
        self.cat_strict = defaultdict(PRF)
        self.cat_relax = defaultdict(PRF)
        self.gs_all = 0
        self.gs_cat = defaultdict(int)
        self.performance = dict()
        self.counts = dict()

    def set_beta_for_f_score(self, beta):
        print("Using beta={} for calculating F-score".format(beta))
        self.beta = beta

    # def set_logger(self, logger):
    #     self.logger = logger

    def add_labels_not_for_eval(self, *labels):
        for each in labels:
            self.label_not_for_eval.add(each.lower())

    def __calc_prf(self, tp, fp, tp_tn):
        """
        Using this function to calculate F-beta score, beta=1 is f_score-score, set beta=2 favor recall, and set beta=0.5 favor precision.
        Using set_beta_for_f_score function to change beta value.
        """
        tp_fp = tp + fp
        pre = 1.0 * tp / tp_fp if tp_fp > 0 else 0.0
        rec = 1.0 * tp / tp_tn if tp_tn > 0 else 0.0
        beta2 = pow(self.beta, 2)
        f_beta = (1 + beta2) * pre * rec / (beta2 * pre + rec) if (pre + rec) > 0 else 0.0
        return pre, rec, f_beta

    def __measure_performance(self):
        self.performance['overall'] = dict()

        acc_true_num, acc_false_num = self.acc.get_true_false_counts()
        total_acc_num = acc_true_num + acc_false_num
        # calc acc
        overall_acc = round(1.0 * acc_true_num / total_acc_num, 4) if total_acc_num > 0 else 0.0
        self.performance['overall']['acc'] = overall_acc

        strict_true_counts, strict_false_counts = self.all_strict.get_true_false_counts()
        strict_pre, strict_rec, strict_f_score = self.__calc_prf(strict_true_counts, strict_false_counts, self.gs_all)
        self.performance['overall']['strict'] = dict()
        self.performance['overall']['strict']['precision'] = strict_pre
        self.performance['overall']['strict']['recall'] = strict_rec
        self.performance['overall']['strict']['f_score'] = strict_f_score

        relax_true_counts, relax_false_counts = self.all_relax.get_true_false_counts()
        relax_pre, relax_rec, relax_f_score = self.__calc_prf(relax_true_counts, relax_false_counts, self.gs_all)
        self.performance['overall']['relax'] = dict()
        self.performance['overall']['relax']['precision'] = relax_pre
        self.performance['overall']['relax']['recall'] = relax_rec
        self.performance['overall']['relax']['f_score'] = relax_f_score

        self.performance['category'] = dict()
        self.performance['category']['strict'] = dict()
        for k, v in self.cat_strict.items():
            self.performance['category']['strict'][k] = dict()
            stc, sfc = v.get_true_false_counts()
            p, r, f = self.__calc_prf(stc, sfc, self.gs_cat[k])
            self.performance['category']['strict'][k]['precision'] = p
            self.performance['category']['strict'][k]['recall'] = r
            self.performance['category']['strict'][k]['f_score'] = f

        self.performance['category']['relax'] = dict()
        for k, v in self.cat_relax.items():
            self.performance['category']['relax'][k] = dict()
            rtc, rfc = v.get_true_false_counts()
            p, r, f = self.__calc_prf(rtc, rfc, self.gs_cat[k])
            self.performance['category']['relax'][k]['precision'] = p
            self.performance['category']['relax'][k]['recall'] = r
            self.performance['category']['relax'][k]['f_score'] = f

    def __measure_counts(self):
        # gold standard
        self.counts['expect'] = dict()
        self.counts['expect']['overall'] = self.gs_all
        for k, v in self.gs_cat.items():
            self.counts['expect'][k] = v
        # prediction
        self.counts['prediction'] = {'strict': dict(), 'relax': dict()}
        # strict
        strict_true_counts, strict_false_counts = self.all_strict.get_true_false_counts()
        self.counts['prediction']['strict']['overall'] = dict()
        self.counts['prediction']['strict']['overall']['total'] = strict_true_counts + strict_false_counts
        self.counts['prediction']['strict']['overall']['true'] = strict_true_counts
        self.counts['prediction']['strict']['overall']['false'] = strict_false_counts
        for k, v in self.cat_strict.items():
            t, f = v.get_true_false_counts()
            self.counts['prediction']['strict'][k] = dict()
            self.counts['prediction']['strict'][k]['total'] = t + f
            self.counts['prediction']['strict'][k]['true'] = t
            self.counts['prediction']['strict'][k]['false'] = f
        # relax
        relax_true_counts, relax_false_counts = self.all_relax.get_true_false_counts()
        self.counts['prediction']['relax']['overall'] = dict()
        self.counts['prediction']['relax']['overall']['total'] = relax_true_counts + relax_false_counts
        self.counts['prediction']['relax']['overall']['true'] = relax_true_counts
        self.counts['prediction']['relax']['overall']['false'] = relax_false_counts
        for k, v in self.cat_relax.items():
            t, f = v.get_true_false_counts()
            self.counts['prediction']['relax'][k] = dict()
            self.counts['prediction']['relax'][k]['total'] = t + f
            self.counts['prediction']['relax'][k]['true'] = t
            self.counts['prediction']['relax'][k]['false'] = f

    @staticmethod
    def __strict_match(gs, pred, s_idx, e_idx, en_type):
        if e_idx < len(gs) and gs[e_idx] == f"i-{en_type}":
            # check token after end in GS is not continued entity token
            return False
        elif gs[s_idx] != f"b-{en_type}" or pred[s_idx] != f"b-{en_type}":
            # force first token to be B-
            return False
        # check every token in span is the same
        for idx in range(s_idx, e_idx):
            if gs[idx] != pred[idx]:
                return False
        return True

    @staticmethod
    def __relax_match(gs, pred, s_idx, e_idx, en_type):
        # we adopt the partial match strategy which is very loose compare to right-left or approximate match
        for idx in range(s_idx, e_idx):
            gs_cate = gs[idx].split("-")[-1]
            pred_bound, pred_cate = pred[idx].split("-")
            if gs_cate == pred_cate == en_type:
                return True
        return False

    @staticmethod
    def __check_evaluated_already(gs_dict, cate, start_idx, end_idx):
        for k, v in gs_dict.items():
            c, s, e = k
            if not (e < start_idx or s > end_idx) and c == cate:
                if v == 0:
                    return True
                else:
                    gs_dict[k] -= 1
                    return False
        return False

    def __process_bio(self, gs_bio, pred_bio):
        # measure acc
        for w_idx, (gs_word, pred_word) in enumerate(zip(gs_bio, pred_bio)):
            # measure acc
            if gs_word == pred_word:
                self.acc.add_true_case()
            else:
                self.acc.add_false_case()

        # process gold standard
        llen = len(gs_bio)
        gs_dict = defaultdict(int)
        cur_idx = 0
        while cur_idx < llen:
            if gs_bio[cur_idx].strip() in self.label_not_for_eval:
                cur_idx += 1
            else:
                start_idx = cur_idx
                end_idx = start_idx + 1
                _, cate = gs_bio[start_idx].strip().split('-')
                while end_idx < llen and gs_bio[end_idx].strip() == f"i-{cate}":
                    end_idx += 1
                self.gs_all += 1
                self.gs_cat[cate] += 1
                gs_dict[(cate, start_idx, end_idx)] += 1
                cur_idx = end_idx
        # process predictions
        cur_idx = 0
        while cur_idx < llen:
            if pred_bio[cur_idx].strip() in self.label_not_for_eval:
                cur_idx += 1
            else:
                start_idx = cur_idx
                end_idx = start_idx + 1
                _, cate = pred_bio[start_idx].strip().split("-")
                while end_idx < llen and pred_bio[end_idx].strip() == f"i-{cate}":
                    end_idx += 1
                if self.__strict_match(gs_bio, pred_bio, start_idx, end_idx, cate):
                    self.all_strict.add_true_case()
                    self.cat_strict[cate].add_true_case()
                    self.all_relax.add_true_case()
                    self.cat_relax[cate].add_true_case()
                elif self.__relax_match(gs_bio, pred_bio, start_idx, end_idx, cate):
                    if self.__check_evaluated_already(gs_dict, cate, start_idx, end_idx):
                        cur_idx = end_idx
                        continue
                    self.all_strict.add_false_case()
                    self.cat_strict[cate].add_false_case()
                    self.all_relax.add_true_case()
                    self.cat_relax[cate].add_true_case()
                else:
                    self.all_strict.add_false_case()
                    self.cat_strict[cate].add_false_case()
                    self.all_relax.add_false_case()
                    self.cat_relax[cate].add_false_case()
                cur_idx = end_idx

    def eval_file(self, gs_file, pred_file):
        print("processing gold standard file: {} and prediciton file: {}".format(gs_file, pred_file))
        pred_bio_sents = load_bio_file_into_sents(pred_file, do_lower=True)
        gs_bio_sents = load_bio_file_into_sents(gs_file, do_lower=True)
        # process bio data
        # check two data have same amount of sents
        assert len(gs_bio_sents) == len(pred_bio_sents), \
            "gold standard and prediction have different dimension: gs: {}; pred: {}".format(len(gs_bio_sents), len(pred_bio_sents))
        # measure performance
        for s_idx, (gs_sent, pred_sent) in enumerate(zip(gs_bio_sents, pred_bio_sents)):
            # check two sents have same No. of words
            assert len(gs_sent) == len(pred_sent), \
                "In {}th sentence, the words counts are different; gs: {}; pred: {}".format(s_idx, gs_sent, pred_sent)
            gs_sent = list(map(lambda x: x[-1], gs_sent))
            pred_sent = list(map(lambda x: x[-1], pred_sent))
            self.__process_bio(gs_sent, pred_sent)
        # get the evaluation matrix
        self.__measure_performance()
        self.__measure_counts()

    def eval_mem(self, gs, pred, do_flat=False):
        # flat sents to sent; we assume input sequences only have 1 dimension (only labels)
        if do_flat:
            print('Sentences have been flatten to 1 dim.')
            gs = list(chain(*gs))
            pred = list(chain(*pred))
            gs = list(map(lambda x: x.lower(), gs))
            pred = list(map(lambda x: x.lower(), pred))
            self.__process_bio(gs, pred)
        else:
            for sidx, (gs_s, pred_s) in enumerate(zip(gs, pred)):
                gs_s = list(map(lambda x: x.lower(), gs_s))
                pred_s = list(map(lambda x: x.lower(), pred_s))
                self.__process_bio(gs_s, pred_s)

        self.__measure_performance()
        self.__measure_counts()

    def evaluate_annotations(self, gs, pred, do_lower=False):
        for gs_sent, pred_sent in zip(gs, pred):
            if do_lower:
              gs_sent = list(map(lambda x: x.lower(), gs_sent))
              pred_sent = list(map(lambda x: x.lower(), pred_sent))
            self.__process_bio(gs_sent, pred_sent)

        self.__measure_performance()
        self.__measure_counts()

    def get_performance(self):
        return self.performance

    def get_counts(self):
        return self.counts

    def save_evaluation(self, file):
        with open(file, "w") as f:
            json.dump(self.performance, f)

    def show_evaluation(self, digits=4):
        if len(self.performance) == 0:
            raise RuntimeError('call eval_mem() first to get the performance attribute')

        cate = self.performance['category']['strict'].keys()

        headers = ['precision', 'recall', 'f1']
        width = max(max([len(c) for c in cate]), len('overall'), digits)
        head_fmt = '{:>{width}s} ' + ' {:>9}' * len(headers)

        report = head_fmt.format(u'', *headers, width=width)
        report += '\n\nstrict\n'

        row_fmt = '{:>{width}s} ' + ' {:>9.{digits}f}' * 3 + '\n'
        for c in cate:
            precision = self.performance['category']['strict'][c]['precision']
            recall = self.performance['category']['strict'][c]['recall']
            f1 = self.performance['category']['strict'][c]['f_score']
            report += row_fmt.format(c, *[precision, recall, f1], width=width, digits=digits)

        report += '\nrelax\n'

        for c in cate:
            precision = self.performance['category']['relax'][c]['precision']
            recall = self.performance['category']['relax'][c]['recall']
            f1 = self.performance['category']['relax'][c]['f_score']
            report += row_fmt.format(c, *[precision, recall, f1], width=width, digits=digits)

        report += '\n\noverall\n'
        report += 'acc: ' + str(self.performance['overall']['acc'])
        report += '\nstrict\n'
        report += row_fmt.format('', *[self.performance['overall']['strict']['precision'],
                                       self.performance['overall']['strict']['recall'],
                                       self.performance['overall']['strict']['f_score']], width=width, digits=digits)

        report += '\nrelax\n'
        report += row_fmt.format('', *[self.performance['overall']['relax']['precision'],
                                       self.performance['overall']['relax']['recall'],
                                       self.performance['overall']['relax']['f_score']], width=width, digits=digits)
        return report