# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


import argparse
from collections import defaultdict
from itertools import chain
from pathlib import Path

import numpy as np
import torchaudio
import torchaudio.sox_effects as ta_sox
import yaml
from tqdm import tqdm

from examples.speech_to_text.data_utils import load_tsv_to_dicts
from examples.speech_synthesis.preprocessing.speaker_embedder import SpkrEmbedder


def extract_embedding(audio_path, embedder):
    wav, sr = torchaudio.load(audio_path)  # 2D
    if sr != embedder.RATE:
        wav, sr = ta_sox.apply_effects_tensor(
            wav, sr, [["rate", str(embedder.RATE)]]
        )
    try:
        emb = embedder([wav[0].cuda().float()]).cpu().numpy()
    except RuntimeError:
        emb = None
    return emb


def process(args):
    print("Fetching data...")
    raw_manifest_root = Path(args.raw_manifest_root).absolute()
    samples = [load_tsv_to_dicts(raw_manifest_root / (s + ".tsv"))
               for s in args.splits]
    samples = list(chain(*samples))
    with open(args.config, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    with open(f"{config['audio_root']}/{config['speaker_set_filename']}") as f:
        speaker_to_id = {r.strip(): i for i, r in enumerate(f)}

    embedder = SpkrEmbedder(args.ckpt).cuda()
    speaker_to_cnt = defaultdict(float)
    speaker_to_emb = defaultdict(float)
    for sample in tqdm(samples, desc="extract emb"):
        emb = extract_embedding(sample["audio"], embedder)
        if emb is not None:
            speaker_to_cnt[sample["speaker"]] += 1
            speaker_to_emb[sample["speaker"]] += emb
    if len(speaker_to_emb) != len(speaker_to_id):
        missed = set(speaker_to_id) - set(speaker_to_emb.keys())
        print(
            f"WARNING: missing embeddings for {len(missed)} speaker:\n{missed}"
        )
    speaker_emb_mat = np.zeros((len(speaker_to_id), len(emb)), float)
    for speaker in speaker_to_emb:
        idx = speaker_to_id[speaker]
        emb = speaker_to_emb[speaker]
        cnt = speaker_to_cnt[speaker]
        speaker_emb_mat[idx, :] = emb / cnt
    speaker_emb_name = "speaker_emb.npy"
    speaker_emb_path = f"{config['audio_root']}/{speaker_emb_name}"
    np.save(speaker_emb_path, speaker_emb_mat)
    config["speaker_emb_filename"] = speaker_emb_name

    with open(args.new_config, "w") as f:
        yaml.dump(config, f)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--raw-manifest-root", "-m", required=True, type=str)
    parser.add_argument("--splits", "-s", type=str, nargs="+",
                        default=["train"])
    parser.add_argument("--config", "-c", required=True, type=str)
    parser.add_argument("--new-config", "-n", required=True, type=str)
    parser.add_argument("--ckpt", required=True, type=str,
                        help="speaker embedder checkpoint")
    args = parser.parse_args()

    process(args)


if __name__ == "__main__":
    main()