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

import argparse
import codecs
import os
import logging
from multiprocessing import Pool

import numpy as np
import scipy.io.wavfile as wf
from nara_wpe.utils import istft, stft
from nara_wpe.wpe import wpe_v8 as wpe


def wpe_worker(
    wav_scp,
    audio_dir="",
    output_dir="",
    channel=0,
    processing_id=None,
    processing_num=None,
):
    sampling_rate = 16000
    iterations = 5
    stft_options = dict(
        size=512,
        shift=128,
        window_length=None,
        fading=True,
        pad=True,
        symmetric_window=False,
    )
    with codecs.open(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)]
    for wav_idx in range(len(wav_lines)):
        if processing_id is None:
            processing_token = True
        else:
            if wav_idx % processing_num == processing_id:
                processing_token = True
            else:
                processing_token = False
        if processing_token:
            wav_list = wav_lines[wav_idx].split(" ")
            file_exist = True
            for wav_path in wav_list:
                file_exist = file_exist and os.path.exists(
                    wav_path.replace(audio_dir, output_dir)
                )
                if not file_exist:
                    break
            if not file_exist:
                logging.info("wait to process {} : {}".format(wav_idx, wav_list[0]))
                signal_list = []
                for f in wav_list:
                    _, data = wf.read(f)
                    data = data[:, channel - 1]
                    if data.dtype == np.int16:
                        data = np.float32(data) / 32768
                    signal_list.append(data)
                min_len = len(signal_list[0])
                max_len = len(signal_list[0])
                for i in range(1, len(signal_list)):
                    min_len = min(min_len, len(signal_list[i]))
                    max_len = max(max_len, len(signal_list[i]))
                if min_len != max_len:
                    for i in range(len(signal_list)):
                        signal_list[i] = signal_list[i][:min_len]
                y = np.stack(signal_list, axis=0)
                Y = stft(y, **stft_options).transpose(2, 0, 1)
                Z = wpe(Y, iterations=iterations, statistics_mode="full").transpose(
                    1, 2, 0
                )
                z = istft(Z, size=stft_options["size"], shift=stft_options["shift"])
                for d in range(len(signal_list)):
                    store_path = wav_list[d].replace(audio_dir, output_dir)
                    if not os.path.exists(os.path.split(store_path)[0]):
                        os.makedirs(os.path.split(store_path)[0], exist_ok=True)
                    tmpwav = np.int16(z[d, :] * 32768)
                    wf.write(store_path, sampling_rate, tmpwav)
            else:
                logging.info("file exist {} : {}".format(wav_idx, wav_list[0]))
    return None


def wpe_manager(
    wav_scp, processing_num=1, audio_dir="", output_dir="", channel=1
):
    if processing_num > 1:
        pool = Pool(processes=processing_num)
        for i in range(processing_num):
            pool.apply_async(
                wpe_worker,
                kwds={
                    "wav_scp": wav_scp,
                    "processing_id": i,
                    "processing_num": processing_num,
                    "audio_dir": audio_dir,
                    "output_dir": output_dir,
                },
            )
        pool.close()
        pool.join()
    else:
        wpe_worker(wav_scp, audio_dir=audio_dir, output_dir=output_dir, channel=channel)
    return None


if __name__ == "__main__":
    parser = argparse.ArgumentParser("run_wpe")
    parser.add_argument(
        "--wav-scp",
        type=str,
        required=True,
        help="Path pf wav scp file",
    )
    parser.add_argument(
        "--audio-dir",
        type=str,
        required=True,
        help="Directory of input audio files",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        required=True,
        help="Output directory of WPE enhanced audio files",
    )
    parser.add_argument(
        "--channel",
        type=str,
        required=True,
        help="Channel number of input audio",
    )
    parser.add_argument("--nj", type=int, default="1", help="number of process")
    args = parser.parse_args()

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

    logging.info("wavfile={}".format(args.wav_scp))
    logging.info("processingnum={}".format(args.nj))

    wpe_manager(
        wav_scp=args.wav_scp,
        processing_num=args.nj,
        audio_dir=args.audio_dir,
        output_dir=args.output_dir,
        channel=int(args.channel)
    )