#!/usr/bin/env python
# -- coding: UTF-8

import argparse
import codecs
import glob
import logging
import os
from nara_wpe.utils import stft, istft
import numpy as np
import scipy.io.wavfile as wf
from tqdm import tqdm

from test_gss import *


def get_parser():
    parser = argparse.ArgumentParser("Doing GSS based enhancement.")
    parser.add_argument(
        "--wav-scp",
        type=str,
        required=True,
        help="Wav scp file for enhancement.",
    )
    parser.add_argument(
        "--segments",
        type=str,
        required=True,
        help="Wav scp file for enhancement.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        required=True,
        help="Output directory of GSS enhanced data.",
    )

    return parser


def wfread(f):
    fs, data = wf.read(f)
    if data.dtype == np.int16:
        data = np.float32(data) / 32768
    return data, fs


def wfwrite(z, fs, store_path):
    tmpwav = np.int16(z * 32768)
    wf.write(store_path, fs, tmpwav)


def main():
    args = get_parser().parse_args()

    # logging info
    logfmt = "%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"
    logging.basicConfig(level=logging.INFO, format=logfmt)

    stft_window, stft_shift = 512, 256
    gss = GSS(iterations=20, iterations_post=1)
    bf = Beamformer("mvdrSouden_ban", "mask_mul")

    with codecs.open(args.wav_scp, "r") as handle:
        lines_content = handle.readlines()
    wav_lines = [*map(lambda x: x[:-1] if x[-1] in ["\n"] else x, lines_content)]

    cnt = 0
    
    session2spk2dur = {}
    with codecs.open(args.segments, "r") as handle:
        for line in handle.readlines():
            uttid, spkid, stime, etime = line.strip().split(" ")
            sessionid = spkid.split("-")[0]
            if sessionid not in session2spk2dur.keys():
                session2spk2dur[sessionid] = {}
            if spkid not in session2spk2dur[sessionid].keys():
                session2spk2dur[sessionid][spkid] = []
            session2spk2dur[sessionid][spkid].append((float(stime), float(etime)))
    # import pdb;pdb.set_trace()
    
    for wav_idx in tqdm(range(len(wav_lines)), leave=True, desc="0"):
        # get wav files from scp file
        file_list = wav_lines[wav_idx].split(" ")
        sessionid, wav_list = file_list[0], file_list[1:]

        signal_list = []
        time_activity = []
        cnt += 1
        logging.info("Processing {}: {}".format(cnt, wav_list[0]))

        # read all wavs
        for wav in wav_list:
            data, fs = wfread(wav)
            signal_list.append(data)
        try:
            obstft = np.stack(signal_list, axis=0)
        except:
            minlen = min([len(s) for s in signal_list])
            obstft = np.stack([s[:minlen] for s in signal_list])
        wavlen = obstft.shape[1]
        obstft = stft(obstft, stft_window, stft_shift)

        # get activated timestamps and frequencies
        speaker_list = []
        for spk, dur in session2spk2dur[sessionid].items():
            speaker_list.append(spk.split("-")[-1])
            time_activity.append(get_time_activity(dur, wavlen, fs))
        time_activity.append([True] * wavlen)
        frequency_activity = get_frequency_activity(
            time_activity, stft_window, stft_shift
        )
        # import pdb;pdb.set_trace()

        # generate mask
        masks = gss(obstft, frequency_activity)
        masks_bak = masks

        for i in range(masks.shape[0] - 1):
            target_mask = masks[i]
            distortion_mask = np.sum(np.delete(masks, i, axis=0), axis=0)
            xhat = bf(obstft, target_mask=target_mask, distortion_mask=distortion_mask)
            xhat = istft(xhat, stft_window, stft_shift)
            audio_dir = "/".join(wav_list[0].split("/")[:-1])
            store_path = (
                wav_list[0]
                .replace(audio_dir, args.output_dir)
                .replace(".wav", "-{}.wav".format(speaker_list[i]))
            )
            if not os.path.exists(os.path.split(store_path)[0]):
                os.makedirs(os.path.split(store_path)[0], exist_ok=True)

            logging.info("Save wav file {}.".format(store_path))
            wfwrite(xhat, fs, store_path)
            masks = masks_bak


if __name__ == "__main__":
    main()