from __future__ import print_function
import numpy as np
import os
import kaldiio
from multiprocessing import Pool
import argparse
from tqdm import tqdm
import math


class MultiProcessRunner:
    def __init__(self, fn):
        self.args = None
        self.process = fn

    def run(self):
        parser = argparse.ArgumentParser("")
        # Task-independent options
        parser.add_argument("--nj", type=int, default=16)
        parser.add_argument("--debug", action="store_true", default=False)
        parser.add_argument("--no_pbar", action="store_true", default=False)
        parser.add_argument("--verbose", action="store_ture", default=False)

        task_list, args = self.prepare(parser)
        result_list = self.pool_run(task_list, args)
        self.post(result_list, args)

    def prepare(self, parser):
        raise NotImplementedError("Please implement the prepare function.")

    def post(self, result_list, args):
        raise NotImplementedError("Please implement the post function.")

    def pool_run(self, tasks, args):
        results = []
        if args.debug:
            one_result = self.process(tasks[0])
            results.append(one_result)
        else:
            pool = Pool(args.nj)
            for one_result in tqdm(pool.imap(self.process, tasks), total=len(tasks), ascii=True, disable=args.no_pbar):
                results.append(one_result)
            pool.close()

        return results


class MultiProcessRunnerV2:
    def __init__(self, fn):
        self.args = None
        self.process = fn

    def run(self):
        parser = argparse.ArgumentParser("")
        # Task-independent options
        parser.add_argument("--nj", type=int, default=16)
        parser.add_argument("--debug", action="store_true", default=False)
        parser.add_argument("--no_pbar", action="store_true", default=False)
        parser.add_argument("--verbose", action="store_true", default=False)

        task_list, args = self.prepare(parser)
        chunk_size = int(math.ceil(float(len(task_list)) / args.nj))
        if args.verbose:
            print("Split {} tasks into {} sub-tasks with chunk_size {}".format(len(task_list), args.nj, chunk_size))
        subtask_list = [task_list[i*chunk_size: (i+1)*chunk_size] for i in range(args.nj)]
        result_list = self.pool_run(subtask_list, args)
        self.post(result_list, args)

    def prepare(self, parser):
        raise NotImplementedError("Please implement the prepare function.")

    def post(self, result_list, args):
        raise NotImplementedError("Please implement the post function.")

    def pool_run(self, tasks, args):
        results = []
        if args.debug:
            one_result = self.process(tasks[0])
            results.append(one_result)
        else:
            pool = Pool(args.nj)
            for one_result in tqdm(pool.imap(self.process, tasks), total=len(tasks), ascii=True, disable=args.no_pbar):
                results.append(one_result)
            pool.close()

        return results


class MultiProcessRunnerV3(MultiProcessRunnerV2):
    def run(self):
        parser = argparse.ArgumentParser("")
        # Task-independent options
        parser.add_argument("--nj", type=int, default=16)
        parser.add_argument("--debug", action="store_true", default=False)
        parser.add_argument("--no_pbar", action="store_true", default=False)
        parser.add_argument("--verbose", action="store_true", default=False)
        parser.add_argument("--sr", type=int, default=16000)

        task_list, shared_param, args = self.prepare(parser)
        chunk_size = int(math.ceil(float(len(task_list)) / args.nj))
        if args.verbose:
            print("Split {} tasks into {} sub-tasks with chunk_size {}".format(len(task_list), args.nj, chunk_size))
        subtask_list = [(i, task_list[i * chunk_size: (i + 1) * chunk_size], shared_param, args)
                        for i in range(args.nj)]
        result_list = self.pool_run(subtask_list, args)
        self.post(result_list, args)



class MyRunner(MultiProcessRunnerV3):
    def prepare(self, parser):
        assert isinstance(parser, argparse.ArgumentParser)
        parser.add_argument("enroll_dir", type=str)
        parser.add_argument("trial_in", type=str)
        parser.add_argument("trial_out", type=str)
        args = parser.parse_args()

        if not os.path.exists(os.path.dirname(args.trial_out)):
            os.makedirs(os.path.dirname(args.trial_out))

        flist_path = os.path.join(args.enroll_dir, "spk2xvec.flist")
        spk2xvec = {}
        for _path in open(flist_path, "r"):
            for key, value in kaldiio.load_ark(_path.strip()):
                if "-enroll" in key:
                    key = key.replace("-enroll", "")
                spk2xvec[key] = value

        flist_path = os.path.join(args.enroll_dir, "utt2xvec.flist")
        utt2xvec = {}
        for _path in open(flist_path, 'r'):
            for key, value in kaldiio.load_ark(_path.strip()):
                utt2xvec[key] = value

        task_list = [one_line.strip().split(" ") for one_line in open(args.trial_in, "rt")]
        return task_list, [spk2xvec, utt2xvec], args

    def post(self, results_list, args):
        with open(args.trial_out, "wt") as fs:
            for results in results_list:
                for one_item in results:
                    fs.write(one_item+"\n")


def process(task_args):
    task_id, task_list, [spk2xvec, utt2xvec], args = task_args
    results = []
    for spk, utt, _ in task_list:
        xvec = utt2xvec[utt]
        normed_x = xvec / np.linalg.norm(xvec)
        normed_y = spk2xvec[spk] / np.linalg.norm(spk2xvec[spk])
        score = np.sum(normed_x * normed_y)
        results.append("{} {} {:.5f}".format(spk, utt, score))

    return results


if __name__ == '__main__':
    my_runner = MyRunner(process)
    my_runner.run()