#!/usr/bin/env python

import IPython.display as display
import argparse
import importlib.util
import json
import magic
import os
import soundfile
import subprocess
import sys
import sys  
import torch
import torchaudio

from collections import OrderedDict
from time import time
from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
from typing import Dict, Any

sys.path.append("./acoustic-model/")
from acoustic import AcousticModel

#sys.path.insert(0, './acoustic-model/acoustic/')
#from model import AcousticModel

sys.path.append("./hifigan/")
from hifigan.generator import HifiganGenerator
from hifigan.generator import consume_prefix_in_state_dict_if_present

from gpu_utils import print_gpu_info

# ===========================================

print_gpu_info()

parser = argparse.ArgumentParser()
parser.add_argument('--acoustic_model_filename', type=str, required=True)
parser.add_argument('--hifigan_model_filename', type=str, required=True)
parser.add_argument('--source_audio_filename', type=str, required=True)
parser.add_argument('--output_audio_filename', type=str, default='output.wav')
parser.add_argument('--temporary_audio_filename', type=str, default='temp.wav') # used for truncation, etc.
parser.add_argument('--output_metadata_filename', type=str, default='metadata.json')
parser.add_argument('--truncate_seconds', type=int, default=-1) # -1 is "do not truncate", 0 is invalid
args = parser.parse_args()

print('\n'.join(f'[arg] {k} = {v}' for k, v in vars(args).items()))

# ===========================================

def load_acoustic_model(checkpoint_path):

    def modify_model(state_dict: Dict[str,Any]):
        new_state_dict = OrderedDict()
        for k in state_dict['acoustic-model'].keys():
            value = state_dict['acoustic-model'][k]
            name = k[7:] # remove `module.` because they used ddp training to save
            new_state_dict[name] = value
        return new_state_dict

    # Load the acoustic model (either hubert_soft or hubert_discrete) 
    acoustic_dict = torch.load(checkpoint_path) # this used ddp
    acoustic = AcousticModel().cuda()
    acoustic.eval()

    new_state_dict = modify_model(state_dict=acoustic_dict)
    acoustic.load_state_dict(new_state_dict)
    return acoustic

# ===========================================

def load_hifigan_model(hifigan_checkpoint_path=None):
    # NB(bt, 2022-12-07): Not sure why this "spec loading" is needed.
    spec = importlib.util.spec_from_file_location("hifigan.utils", "./hifigan/hifigan/utils.py")
    u = importlib.util.module_from_spec(spec)
    sys.modules["hifigan.utils"] = u
    spec.loader.exec_module(u)

    spec = importlib.util.spec_from_file_location("hifigan.generator", "./hifigan/hifigan/generator.py")
    g = importlib.util.module_from_spec(spec)
    sys.modules["hifigan.generator"] = g
    spec.loader.exec_module(g)

    if hifigan_checkpoint_path:
        checkpoint = torch.load(hifigan_checkpoint_path)["generator"]["model"]
        hifigan = g.HifiganGenerator()
        consume_prefix_in_state_dict_if_present(checkpoint, "module.")
        hifigan.load_state_dict(checkpoint)
        hifigan.remove_weight_norm()
        hifigan.eval()
    else:
        hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft")

    hifigan.cuda()
    return hifigan

# ===========================================


def load_wav_data(file_name):
    source, sr = torchaudio.load(f"{file_name}")
    source = torchaudio.functional.resample(source, sr, 16000)
    source = torch.mean(source, dim=0).unsqueeze(0)
    source = source.unsqueeze(0).cuda()
    return source

# ===========================================

def do_voice_conversion(hubert, acoustic, hifigan, audio_source):
    #@title Run the inference and display the result
    # Convert to the target speaker
    with torch.inference_mode():
        # Extract speech units
        start = time()
        units = hubert.units(audio_source)
        end = time()
        # Generate target spectrogram
        print(f"time for feature embedding:{end-start}")

        start = time()
        mel = acoustic.generate(units).transpose(1, 2)
        end = time()
        print(f"time for inference:{end-start}")
        # Generate audio waveform
        start = time()
        target = hifigan(mel)
        end = time()
        print(f"time for wavform {end-start}")
        return target

# ===========================================

def truncate_source_audio(audio_source_filename, output_filename, seconds):
    # Run ffmpeg's ffprobe
    command = [
        "ffmpeg",
        "-y", # Force file overwrite
        "-i",  
        audio_source_filename,
        "-t", 
        str(seconds),
        "-c",
        "copy",
        output_filename,
    ]

    print(f"Truncating source audio to {seconds} seconds", flush=True)
    pipe = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    out, err = pipe.communicate()


def save_output_audio(audio_data, output_audio_filename):
  soundfile.write(args.output_audio_filename, audio_data.squeeze().cpu(), samplerate=16000)

def save_metadata(audio_filename, output_metadata_filename):
    # Run ffmpeg's ffprobe
    command = [
        "ffprobe",
        "-loglevel",  "quiet",
        "-print_format", "json",
        "-show_format",
        "-show_streams",
        audio_filename
    ]

    pipe = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    out, err = pipe.communicate()
    ffmpeg_metadata = json.loads(out)

    # Collect metadata
    print(ffmpeg_metadata)
    mime_type = magic.from_file(audio_filename, mime=True)
    file_size_bytes = os.path.getsize(audio_filename)

    duration_millis = int(float(ffmpeg_metadata['format']['duration']) * 1000)

    metadata = {
        'duration_millis': duration_millis,
        'mimetype': mime_type,
        'file_size_bytes': file_size_bytes,
    }

    with open(output_metadata_filename, 'w') as json_file:
        json.dump(metadata, json_file)

intermediate_audio_filename = args.source_audio_filename

if args.truncate_seconds not in [-1, 0]:
    intermediate_audio_filename = args.temporary_audio_filename
    truncate_source_audio(args.source_audio_filename, intermediate_audio_filename, args.truncate_seconds)

hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()

acoustic = load_acoustic_model(args.acoustic_model_filename)
hifigan = load_hifigan_model(args.hifigan_model_filename)

source_audio = load_wav_data(intermediate_audio_filename)
target_audio = do_voice_conversion(hubert, acoustic, hifigan, source_audio)

save_output_audio(target_audio, args.output_audio_filename)
save_metadata(args.output_audio_filename, args.output_metadata_filename)