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
from funasr.utils.types import str2triple_str
import logging
from typing import List, Union, Tuple, Sequence
from funasr.bin.sv_inference import inference_modelscope
import soundfile
import torch


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

    def run(self):
        parser = argparse.ArgumentParser("")
        # Task-independent options
        parser.add_argument("--njobs", 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("--log_level", type=str, default="INFO")
        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.njobs))
        if args.verbose:
            print("Split {} tasks into {} sub-tasks with chunk_size {}".format(len(task_list), args.njobs, chunk_size))
        subtask_list = [(i, task_list[i * chunk_size: (i + 1) * chunk_size], shared_param, args)
                        for i in range(args.njobs)]
        result_list = self.pool_run(subtask_list, args)
        self.post(result_list, args)

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

    def post(self, results_list: list, args: argparse.Namespace):
        raise NotImplementedError("Please implement the post function.")

    def pool_run(self, tasks: list, args: argparse.Namespace):
        results = []
        if args.debug:
            one_result = self.process(tasks[0])
            results.append(one_result)
        else:
            pool = Pool(args.njobs)
            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 MyRunner(MultiProcessRunner):
    def prepare(self, parser: argparse.ArgumentParser):
        parser.add_argument(
            "--gpu_inference",
            type=bool,
            default=False
        )
        parser.add_argument(
            "--data_path_and_name_and_type",
            type=str2triple_str,
            required=True,
            action="append"
        )
        parser.add_argument(
            "--gpu_devices",
            type=lambda devices: devices.split(","),
            default=None,
        )
        args = parser.parse_args()

        logging.basicConfig(
            level=args.log_level,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )

        if args.gpu_inference and (args.gpu_devices is None or len(args.gpu_devices) == 0):
            logging.warning("gpu_inference is set to True, but gpu_devices is not given, use CPU instead.")
            args.gpu_inference = False

        if args.gpu_inference:
            args.njobs = args.njobs * len(args.gpu_devices)

        speech_dict = {}
        ref_speech_dict = {}
        for _path, _name, _type in args.data_path_and_name_and_type:
            if _name == "speech":
                speech_dict = self.read_data_path(_path)
            elif _name == "ref_speech":
                ref_speech_dict = self.read_data_path(_path)

        task_list, args.njobs = self.get_key_list(args.data_path_and_name_and_type, args.njobs)

        return task_list, [speech_dict, ref_speech_dict], args

    def read_data_path(self, file_path):
        results = {}
        for line in open(file_path, "r"):
            key, path = line.strip().split(" ", 1)
            results[key] = path

        return results

    def get_key_list(
            self,
            data_path_and_name_and_type: Sequence[Tuple[str, str, str]],
            njobs: int
    ):
        first_data = data_path_and_name_and_type[0]
        content = open(first_data[0], "r").readlines()
        line_number = len(content)
        njobs = min(njobs, line_number)
        logging.warning("njobs is reduced to {}, since only {} lines exist in {}".format(
            njobs, line_number, first_data[0],
        ))
        key_list = [line.strip().split(" ", 1)[0] for line in content]
        return key_list, njobs

    def post(self, results_list: list, args: argparse.Namespace):
        for results in results_list:
            for key, value in results:
                logging.info("{} {}".format(key, value))


def process(task_args):
    task_id, key_list, [speech_dict, ref_speech_dict], args = task_args
    if args.gpu_inference:
        device = args.gpu_devices[task_id % len(args.gpu_devices)]
        torch.cuda.set_device("cuda:".format(device))
    inference_func = inference_modelscope(
        output_dir=None,
        batch_size=1,
        dtype="float32",
        ngpu=1 if args.gpu_inference else 0,
        seed=0,
        num_workers=0,
        log_level=logging.INFO,
        key_file=None,
        sv_train_config="sv.yaml",
        sv_model_file="sv.pb",
        model_tag=None,
        allow_variable_data_keys=True,
        streaming=False,
        embedding_node="resnet1_dense",
        sv_threshold=0.9465,
    )
    results = {}
    for key in key_list:
        speech = soundfile.read(speech_dict[key])[0]
        ref_speech = soundfile.read(ref_speech_dict[key])[0]
        ret = inference_func(None, (speech, ref_speech))
        results[key] = ret["value"]

    return results


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