import argparse
import glob
import multiprocessing
import re
from functools import partial
from pathlib import Path

import librosa
import numpy

from become_yukarin import AcousticConverter
from become_yukarin import SuperResolution
from become_yukarin.config.config import create_from_json as create_config

parser = argparse.ArgumentParser()
parser.add_argument('-cmd', '--conv_model_directory', type=Path, default=Path('/mnt/dwango/hiroshiba/become-yukarin/'))
parser.add_argument('-srmd', '--sr_model_directory', type=Path, default=Path('/mnt/dwango/hiroshiba/become-yukarin/'))
parser.add_argument('-iwd', '--input_directory', type=Path, default=Path('./input'))
parser.add_argument('-owd', '--output_directory', type=Path, default=Path('./output'))
parser.add_argument('-it', '--iteration', type=int)
parser.add_argument('-g', '--gpu', type=int)
args = parser.parse_args()

conv_model = args.conv_model_directory  # type: Path
sr_model = args.sr_model_directory  # type: Path
input_directory = args.input_directory
output_directory = args.output_directory
it = args.iteration
gpu = args.gpu

def extract_number(f):
    s = re.findall("\d+", str(f))
    return int(s[-1]) if s else -1

def load_acoustic_converter():
  conv_config = create_config(conv_model / 'config.json')
  if it is not None:
      conv_model_path = conv_model / 'predictor_{}.npz'.format(it)
  else:
      conv_model_paths = conv_model.glob('predictor_*.npz')
      conv_model_path = list(sorted(conv_model_paths, key=extract_number))[-1]
  print(conv_model_path)
  return AcousticConverter(conv_config, conv_model_path, gpu=gpu)

def load_super_resolution():
  sr_config = create_config(sr_model / 'config.json')
  sr_model_paths = base_model.glob('predictor*.npz')
  sr_model_path = list(sorted(sr_model_paths, key=extract_number))[-1]
  print(sr_model_path)
  return SuperResolution(sr_config, sr_model_path, gpu=gpu)

def process(p: Path, acoustic_converter: AcousticConverter, super_resolution: SuperResolution, output: Path):
  try:
    print(str(p))
    input = acoustic_converter(p)
    wave = super_resolution(input.spectrogram, acoustic_feature=input, sampling_rate=param.voice_param.sample_rate)
    librosa.output.write_wav(str(output / p.stem) + '.wav', wave.wave, wave.sampling_rate, norm=True)
  except:
    import traceback
    print('error!', str(p))
    print(traceback.format_exc())

def run():
  input_paths = list(input_directory.glob('*.wav'))

  acoustic_converter = load_acoustic_converter()
  super_resolution = load_super_resolution()

  output = output_directory.absolute()
  output.mkdir(exist_ok=True)

  process_partial = partial(process, acoustic_converter=acoustic_converter, super_resolution=super_resolution, output=output)
  if gpu is None:
      pool = multiprocessing.Pool()
      pool.map(process_partial, input_paths)
  else:
      list(map(process_partial, input_paths))

run()