# formulate pred generated by predictor
import sys
import json
import pdb
from copy import deepcopy
import pickle
from pprint import pprint
from utils.log_system import logger

def read_domain_slot_list(filename):
    with open(filename) as fp:
        lines = fp.readlines()
    domain_slots = []
    for line in lines:
        if line.startswith("#"):
            continue
        if len(line.strip("\n ")) == 0:
            continue
        line_arr = line.split("\t")
        ds = line_arr[0] + " " + line_arr[1]
        if line_arr[3] == "n":
            domain_slots.append(ds)
    return domain_slots


class DSTQAAnalyser():
    '''
    This class is to analyze the prediction results
    '''
    def __init__(self, config, data=None):
        self.config = config
        self.data = data
        self.DST_results = None

    def calc_pred_belief_state(self, prediction, ds_list, ontology):
        def dict2str(d):
            res = []
            for k, v in d.items():
                res.append(k + ":" + v)
            return sorted(res)

        prediction = prediction["predicted_labels"]
        turn_bs = []
        for turn in prediction:
            cur_bs = {}
            for ds in ds_list:
                if ds not in ontology: continue
                cur_bs[ds] = "none"
            for slot_value in turn:
                p = slot_value.find(":")
                slot = slot_value[:p]
                if slot not in ontology: continue
                value = slot_value[p + 1:]  # value may have ":"
                cur_bs[slot] = value
            turn_bs.append(dict2str(cur_bs))
        return turn_bs

    def calc_acc(self, true_labels, pred_labels):
        assert (len(true_labels) == len(pred_labels))
        total_turn = 0.0
        err_turn = 0.0
        total_slots = 0
        correct_slots = 0
        wrong_dialog = []
        for d in range(len(true_labels)):  # for each dialog
            err_of_dialog = 0
            assert (len(true_labels[d]) == len(pred_labels[d]))
            for t in range(len(true_labels[d])):  # for each turn
                total_turn += 1
                # print('true_labels', true_labels[d][t])
                # print('pred_labels', pred_labels[d][t])

                for x, y in zip(true_labels[d][t], pred_labels[d][t]):
                    # The label has value
                    total_slots += 1
                    if x == y:
                        correct_slots += 1

                if len(true_labels[d][t]) != len(pred_labels[d][t]):
                    # If length is not right
                    err_turn += 1
                    err_of_dialog += 1
                    continue
                for x, y in zip(true_labels[d][t], pred_labels[d][t]):
                    if x != y:
                        err_turn += 1
                        err_of_dialog += 1
                        break
            if err_of_dialog > 0:
                wrong_dialog.append(d)
        return (total_turn - err_turn) / total_turn, wrong_dialog, correct_slots / total_slots

    def get_DST_benchmark(self):
        data = self.data

        dialogs = []
        js = {}

        for result in data:
            index, model_input, prediction = result[0], json.loads(result[1]), json.loads(result[2])
            dialogs.append((model_input, prediction))

        ds_list = read_domain_slot_list(self.config.data_loader.additional.domain_slot_list_path)
        ontology = set(ds_list)
        true_labels = []
        pred_labels = []
        for dialog, prediction in dialogs:
            dialog_bs = []
            for turn in dialog["dialogue"]:
                turn_bs = []
                ds_set = set(ds_list)
                for domain, v in turn["belief_state"].items():
                    for slot, slot_value in v["semi"].items():
                        ds = domain + " " + slot
                        if ds not in ontology:
                            continue
                        if slot_value == "": slot_value = "none"
                        turn_bs.append(domain + " " + slot + ":" + slot_value)
                        ds_set.remove(domain + " " + slot)
                for ds in ds_set:
                    if ds not in ontology: continue
                    turn_bs.append(ds + ":" + "none")
                turn_bs = sorted(turn_bs)
                dialog_bs.append(turn_bs)
            true_labels.append(dialog_bs)
            pred_labels.append(self.calc_pred_belief_state(prediction, ds_list, ontology))

        # print(len(true_labels), true_labels[0])
        # print(len(pred_labels), pred_labels[0])
        acc, wrong_dialogs, slot_acc = self.calc_acc(true_labels, pred_labels)
        logger.print('results are', acc, slot_acc)
        self.DST_results = {
            'joint_acc': acc,
            'slot_acc': slot_acc,
            'wrong_dialogs': wrong_dialogs,
            'true_labels': true_labels,
            'pred_labels': pred_labels,
        }
