import os
from funasr.utils.job_runner import MultiProcessRunnerV3
import numpy as np
from funasr.utils.misc import load_scp_as_list, load_scp_as_dict
from collections import OrderedDict
from tqdm import tqdm
from scipy.ndimage import median_filter


class MyRunner(MultiProcessRunnerV3):
    def prepare(self, parser):
        parser.add_argument("label_txt", type=str)
        parser.add_argument("map_scp", type=str)
        parser.add_argument("out_rttm", type=str)
        parser.add_argument("--n_spk", type=int, default=4)
        parser.add_argument("--chunk_len", type=int, default=1600)
        parser.add_argument("--shift_len", type=int, default=400)
        parser.add_argument("--ignore_len", type=int, default=5)
        parser.add_argument("--smooth_size", type=int, default=7)
        parser.add_argument("--vote_prob", type=float, default=0.5)
        args = parser.parse_args()

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

        utt2labels = load_scp_as_list(args.label_txt, 'list')
        utt2labels = sorted(utt2labels, key=lambda x: x[0])
        meeting2map = load_scp_as_dict(args.map_scp)
        meeting2labels = OrderedDict()
        for utt_id, chunk_label in utt2labels:
            mid = utt_id.split("-")[0]
            if mid not in meeting2labels:
                meeting2labels[mid] = []
            meeting2labels[mid].append(chunk_label)
        task_list = [(mid, labels, meeting2map[mid]) for mid, labels in meeting2labels.items()]

        return task_list, None, args

    def post(self, result_list, args):
        with open(args.out_rttm, "wt") as fd:
            for results in result_list:
                fd.writelines(results)


def int2vec(x, vec_dim=8, dtype=np.int):
    b = ('{:0' + str(vec_dim) + 'b}').format(x)
    # little-endian order: lower bit first
    return (np.array(list(b)[::-1]) == '1').astype(dtype)


def seq2arr(seq, vec_dim=8):
    return np.row_stack([int2vec(int(x), vec_dim) for x in seq])


def sample2ms(sample, sr=16000):
    return int(float(sample) / sr * 100)


def calc_multi_labels(chunk_label_list, chunk_len, shift_len, n_spk, vote_prob=0.5):
    n_chunk = len(chunk_label_list)
    last_chunk_valid_frame = len(chunk_label_list[-1]) - (chunk_len - shift_len)
    n_frame = (n_chunk - 2) * shift_len + chunk_len + last_chunk_valid_frame
    multi_labels = np.zeros((n_frame, n_spk), dtype=float)
    weight = np.zeros((n_frame, 1), dtype=float)
    for i in range(n_chunk):
        raw_label = chunk_label_list[i]
        for k in range(len(raw_label)):
            if raw_label[k] == '<unk>':
                raw_label[k] = raw_label[k-1] if k > 0 else '0'
        chunk_multi_label = seq2arr(raw_label, n_spk)
        chunk_len = chunk_multi_label.shape[0]
        multi_labels[i*shift_len:i*shift_len+chunk_len, :] += chunk_multi_label
        weight[i*shift_len:i*shift_len+chunk_len, :] += 1
    multi_labels = multi_labels / weight  # normalizing vote
    multi_labels = (multi_labels > vote_prob).astype(int)  # voting results
    return multi_labels


def calc_spk_turns(label_arr, spk_list):
    turn_list = []
    length = label_arr.shape[0]
    n_spk = label_arr.shape[1]
    for k in range(n_spk):
        if spk_list[k] == "None":
            continue
        in_utt = False
        start = 0
        for i in range(length):
            if label_arr[i, k] == 1 and in_utt is False:
                start = i
                in_utt = True
            if label_arr[i, k] == 0 and in_utt is True:
                turn_list.append([spk_list[k], start, i - start])
                in_utt = False
        if in_utt:
            turn_list.append([spk_list[k], start, length - start])
    return turn_list


def smooth_multi_labels(multi_label, win_len):
    multi_label = median_filter(multi_label, (win_len, 1), mode="constant", cval=0.0).astype(int)
    return multi_label


def process(task_args):
    _, task_list, _, args = task_args
    spk_list = ["spk{}".format(i+1) for i in range(args.n_spk)]
    template = "SPEAKER {} 1 {:.2f} {:.2f} <NA> <NA> {} <NA> <NA>\n"
    results = []
    for mid, chunk_label_list, map_file_path in tqdm(task_list, total=len(task_list), ascii=True, disable=args.no_pbar):
        utt2map = load_scp_as_list(map_file_path, 'list')
        multi_labels = calc_multi_labels(chunk_label_list, args.chunk_len, args.shift_len, args.n_spk, args.vote_prob)
        multi_labels = smooth_multi_labels(multi_labels, args.smooth_size)
        org_len = sample2ms(int(utt2map[-1][1][1]), args.sr)
        org_multi_labels = np.zeros((org_len, args.n_spk))
        for seg_id, [org_st, org_ed, st, ed] in utt2map:
            org_st, org_dur = sample2ms(int(org_st), args.sr), sample2ms(int(org_ed) - int(org_st), args.sr)
            st, dur = sample2ms(int(st), args.sr), sample2ms(int(ed) - int(st), args.sr)
            ll = min(org_multi_labels[org_st: org_st+org_dur, :].shape[0], multi_labels[st: st+dur, :].shape[0])
            org_multi_labels[org_st: org_st+ll, :] = multi_labels[st: st+ll, :]
        spk_turns = calc_spk_turns(org_multi_labels, spk_list)
        spk_turns = sorted(spk_turns, key=lambda x: x[1])
        for spk, st, dur in spk_turns:
            # TODO: handle the leak of segments at the change points
            if dur > args.ignore_len:
                results.append(template.format(mid, float(st)/100, float(dur)/100, spk))
    return results


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