#!/usr/bin/env python3
"""Generates precomputed benchmark audio examples for the TEC HF Space demo."""

import json
import os
import sys
import numpy as np
import tensorflow.compat.v1 as tf

REPO_ROOT = '/usr/local/google/home/quanw/Code/github/tec'
if REPO_ROOT not in sys.path:
  sys.path.insert(0, REPO_ROOT)

from tec import data_prep  # noqa: E402
from tec import inference  # noqa: E402

TFRECORD_PATH = (
    '/usr/local/google/home/quanw/Data/tec_datasets/tfrecords/'
    'single_test_clean.tfrecord'
)
CHECKPOINT_ROOT = (
    '/usr/local/google/home/quanw/Data/tec_datasets/checkpoints'
)
EXAMPLES_DIR = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), 'examples'
)

TARGET_SAMPLES = [
    {
        'id': 1,
        'utt_id': '4446_2275_000048_000001__LJ012-0176',
        'lj_id': 'LJ012-0176',
        'clean_ref': '"We\'ve tortured each other enough for tonight.',
        'interfering_text': (
            'To obtain possession of this, Thurtell with his two associates'
            ' resolved to kill him.'
        ),
    },
    {
        'id': 2,
        'utt_id': '1995_1837_000024_000001__LJ030-0183',
        'lj_id': 'LJ030-0183',
        'clean_ref': 'The chair was empty; but he knew.',
        'interfering_text': (
            "As he told the driver, quote, Let's get out of here; we are hit,"
            ' end quote,'
        ),
    },
    {
        'id': 3,
        'utt_id': '260_123286_000017_000001__LJ040-0208',
        'lj_id': 'LJ040-0208',
        'clean_ref': 'Have we followed that road?',
        'interfering_text': (
            'Carro reported that when questioned about his mother Lee said,'
            ' quote,'
        ),
    },
    {
        'id': 4,
        'utt_id': '8224_274384_000017_000000__LJ043-0134',
        'lj_id': 'LJ043-0134',
        'clean_ref': (
            'The parliament and the Scots laid their proposals before the king.'
        ),
        'interfering_text': (
            'The Commission has concluded that on April ten, nineteen'
            ' sixty-three,'
        ),
    },
]

MODELS = [
    ('tec_single', 'TecSingleInterfering'),
    ('tec_multi', 'TecMultiInterfering'),
    ('aec_single', 'AecSingleInterfering'),
]


def _find_tfrecord_records():
  """Finds the 4 benchmark records inside single_test_clean.tfrecord."""
  by_clean = {}
  for idx, raw in enumerate(tf.io.tf_record_iterator(TFRECORD_PATH)):
    ex = tf.train.Example()
    ex.ParseFromString(raw)
    feat = ex.features.feature
    clean_txt = feat['clean_transcript'].bytes_list.value[0].decode('utf-8')
    int_txt = feat['interfering_transcript'].bytes_list.value[0].decode('utf-8')
    utt_id = feat['utt_id'].bytes_list.value[0].decode('utf-8')
    mixed_wav = np.array(
        feat['mixed_waveform'].float_list.value, dtype=np.float32)
    int_wav = np.array(
        feat['interfering_waveform'].float_list.value, dtype=np.float32)
    clean_wav = np.array(
        feat['clean_waveform'].float_list.value, dtype=np.float32)
    by_clean[clean_txt.strip()] = {
        'tfrecord_index': idx,
        'utt_id': utt_id,
        'clean_txt': clean_txt,
        'int_txt': int_txt,
        'mixed_wav': mixed_wav,
        'int_wav': int_wav,
        'clean_wav': clean_wav,
    }
  return by_clean


def main():
  os.makedirs(EXAMPLES_DIR, exist_ok=True)
  by_clean = _find_tfrecord_records()
  metadata = []

  for item in TARGET_SAMPLES:
    sid = item['id']
    rec = by_clean.get(item['clean_ref'])
    mixed_path = os.path.join(EXAMPLES_DIR, f'sample_{sid}_mixed.wav')
    if rec is not None:
      data_prep.write_wav_file(mixed_path, rec['mixed_wav'], sample_rate=24000)
    mixed_wav, sr = data_prep.read_wav_file(mixed_path, target_sample_rate=24000)

    if rec is not None:
      int_wav = rec['int_wav']
      utt_id = rec['utt_id']
      tfr_idx = rec['tfrecord_index']
    else:
      lj_wav_path = (
          '/usr/local/google/home/quanw/Data/tec_datasets/ljspeech/LJSpeech-1.1/wavs/'
          f"{item['lj_id']}.wav"
      )
      int_wav, _ = data_prep.read_wav_file(lj_wav_path, target_sample_rate=24000)
      if len(int_wav) < len(mixed_wav):
        int_wav = np.pad(int_wav, (0, len(mixed_wav) - len(int_wav)))
      else:
        int_wav = int_wav[:len(mixed_wav)]
      utt_id = item['utt_id']
      tfr_idx = -1

    int_path = os.path.join(EXAMPLES_DIR, f'sample_{sid}_interfering.wav')
    data_prep.write_wav_file(int_path, int_wav, sample_rate=24000)

    entry = {
        'id': sid,
        'utt_id': utt_id,
        'tfrecord_index': tfr_idx,
        'clean_ref': item['clean_ref'],
        'interfering_text': item['interfering_text'],
        'mixed_wav': f'examples/sample_{sid}_mixed.wav',
        'interfering_wav': f'examples/sample_{sid}_interfering.wav',
        'outputs': {},
    }

    for suffix, model_name in MODELS:
      ckpt = os.path.join(CHECKPOINT_ROOT, model_name, 'best.ckpt')
      out_rel = f'examples/sample_{sid}_{suffix}.wav'
      out_abs = os.path.join(EXAMPLES_DIR, f'sample_{sid}_{suffix}.wav')
      res = inference.run_inference(
          mixed_waveform=mixed_wav,
          model_name=model_name,
          interfering_text=item['interfering_text'],
          interfering_waveform=int_wav,
          checkpoint_path=ckpt,
      )
      data_prep.write_wav_file(
          out_abs, res['predicted_waveform'], sample_rate=sr)
      entry['outputs'][suffix] = out_rel
      if suffix == 'tec_single':
        compat_abs = os.path.join(EXAMPLES_DIR, f'sample_{sid}_tec.wav')
        data_prep.write_wav_file(
            compat_abs, res['predicted_waveform'], sample_rate=sr)
      print(f'Wrote {out_rel} ({len(res["predicted_waveform"])} samples)')

    metadata.append(entry)

  meta_path = os.path.join(EXAMPLES_DIR, 'metadata.json')
  with open(meta_path, 'w', encoding='utf-8') as f:
    json.dump(metadata, f, indent=2)
  print(f'Saved metadata to {meta_path}')


if __name__ == '__main__':
  main()