#!/usr/bin/env python3

# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita)
# Licensed under the MIT license.
#
# This script generates simulated multi-talker mixtures for diarization
#
# common/make_mixture.py \
#     mixture.scp \
#     data/mixture \
#     wav/mixture


import argparse
import os
from funasr.modules.eend_ola.utils import kaldi_data
import numpy as np
import math
import soundfile as sf
import json

parser = argparse.ArgumentParser()
parser.add_argument('script',
                    help='list of json')
parser.add_argument('out_data_dir',
                    help='output data dir of mixture')
parser.add_argument('out_wav_dir',
                    help='output mixture wav files are stored here')
parser.add_argument('--rate', type=int, default=16000,
                    help='sampling rate')
args = parser.parse_args()

# open output data files
segments_f = open(args.out_data_dir + '/segments', 'w')
utt2spk_f = open(args.out_data_dir + '/utt2spk', 'w')
wav_scp_f = open(args.out_data_dir + '/wav.scp', 'w')

# "-R" forces the default random seed for reproducibility
resample_cmd = "sox -R -t wav - -t wav - rate {}".format(args.rate)

for line in open(args.script):
    recid, jsonstr = line.strip().split(None, 1)
    indata = json.loads(jsonstr)
    wavfn = indata['recid']
    # recid now include out_wav_dir
    recid = os.path.join(args.out_wav_dir, wavfn).replace('/','_')
    noise = indata['noise']
    noise_snr = indata['snr']
    mixture = []
    for speaker in indata['speakers']:
        spkid = speaker['spkid']
        utts = speaker['utts']
        intervals = speaker['intervals']
        rir = speaker['rir']
        data = []
        pos = 0
        for interval, utt in zip(intervals, utts):
            # append silence interval data
            silence = np.zeros(int(interval * args.rate))
            data.append(silence)
            # utterance is reverberated using room impulse response
            preprocess = "wav-reverberate --print-args=false " \
                         " --impulse-response={} - -".format(rir)
            if isinstance(utt, list):
                rec, st, et = utt
                st = np.rint(st * args.rate).astype(int)
                et = np.rint(et * args.rate).astype(int)
            else:
                rec = utt
                st = 0
                et = None
            if rir is not None:
                wav_rxfilename = kaldi_data.process_wav(rec, preprocess)
            else:
                wav_rxfilename = rec
            wav_rxfilename = kaldi_data.process_wav(
                    wav_rxfilename, resample_cmd)
            speech, _ = kaldi_data.load_wav(wav_rxfilename, st, et)
            data.append(speech)
            # calculate start/end position in samples
            startpos = pos + len(silence)
            endpos = startpos + len(speech)
            # write segments and utt2spk
            uttid = '{}_{}_{:07d}_{:07d}'.format(
                    spkid, recid, int(startpos / args.rate * 100),
                    int(endpos / args.rate * 100))
            print(uttid, recid,
                  startpos / args.rate, endpos / args.rate, file=segments_f)
            print(uttid, spkid, file=utt2spk_f)
            # update position for next utterance
            pos = endpos
        data = np.concatenate(data)
        mixture.append(data)

    # fitting to the maximum-length speaker data, then mix all speakers
    maxlen = max(len(x) for x in mixture)
    mixture = [np.pad(x, (0, maxlen - len(x)), 'constant') for x in mixture]
    mixture = np.sum(mixture, axis=0)
    # noise is repeated or cutted for fitting to the mixture data length
    noise_resampled = kaldi_data.process_wav(noise, resample_cmd)
    noise_data, _ = kaldi_data.load_wav(noise_resampled)
    if maxlen > len(noise_data):
        noise_data = np.pad(noise_data, (0, maxlen - len(noise_data)), 'wrap')
    else:
        noise_data = noise_data[:maxlen]
    # noise power is scaled according to selected SNR, then mixed
    signal_power = np.sum(mixture**2) / len(mixture)
    noise_power = np.sum(noise_data**2) / len(noise_data)
    scale = math.sqrt(
                math.pow(10, - noise_snr / 10) * signal_power / noise_power)
    mixture += noise_data * scale
    # output the wav file and write wav.scp
    outfname = '{}.wav'.format(wavfn)
    outpath = os.path.join(args.out_wav_dir, outfname)
    sf.write(outpath, mixture, args.rate)
    print(recid, os.path.abspath(outpath), file=wav_scp_f)

wav_scp_f.close()
segments_f.close()
utt2spk_f.close()