import json
import os
import pickle
import re

import h5py
import numpy as np
from scipy.io import wavfile
from scipy.io.wavfile import write as write_wav
from tensorflow import keras

from generators.generator import InverSynthGenerator, SoundGenerator, VSTGenerator
from generators.parameters import ParameterSet

"""
This module generates comparisons - takes the original sound + params,
then generates a file with the predicted parameters
"""


def compare(
    model: keras.Model,
    generator: SoundGenerator,
    parameters: ParameterSet,
    orig_file: str,
    output_dir: str,
    orig_params,
    length: float,
    sample_rate: int,
    extra: dict = {},
):
    # (copy original file if given)
    base_filename = orig_file.replace(".wav", "")
    base_filename = re.sub(r".*/", "", base_filename)
    copy_file: str = f"{output_dir}/{base_filename}_copy.wav"
    regen_file: str = f"{output_dir}/{base_filename}_duplicate.wav"
    reconstruct_file: str = f"{output_dir}/{base_filename}_reconstruct.wav"
    print(f"Creating copy as {copy_file}")

    # Load the wave file
    fs, data = wavfile.read(orig_file)
    # Copy original file to make sure
    write_wav(copy_file, sample_rate, data)

    # Decode original params, and regenerate output (make sure its correct)
    orig = parameters.encoding_to_settings(orig_params)
    generator.generate(orig, regen_file, length, sample_rate, extra)

    # Run the wavefile into the model for prediction
    X = [data]
    Xd = np.expand_dims(np.vstack(X), axis=2)
    # Get encoded parameters out of model
    result = model.predict(Xd)[0]

    # Decode prediction, and reconstruct output
    predicted = parameters.encoding_to_settings(result)
    generator.generate(predicted, reconstruct_file, length, sample_rate, extra)


def run_comparison(
    model: keras.Model,
    generator: SoundGenerator,
    run_name: str,
    indices=None,
    num_samples=10,
    data_dir="./test_datasets",
    output_dir="./comparison",
    length=1.0,
    sample_rate=16384,
    shuffle=True,
    extra={},
):
    # Figure out data file and params file from run name
    data_file = f"{data_dir}/{run_name}_data.hdf5"
    parameters_file = f"{data_dir}/{run_name}_params.pckl"
    print(f"Reading parameters from {parameters_file}")
    parameters = pickle.load(open(parameters_file, "rb"))

    output_dir = f"{output_dir}/{run_name}/"
    os.makedirs(output_dir, exist_ok=True)

    database = h5py.File(data_file, "r")

    if not indices:
        ids = np.array(range(len(database["files"])))
        if shuffle:
            np.random.shuffle(ids)
        indices = ids[0:num_samples]

    # filename
    for i in indices:
        print("Looking at index: {}".format(i))
        filename = database["files"][i]
        labels = database["labels"][i]
        compare(
            model=model,
            generator=generator,
            parameters=parameters,
            orig_file=filename,
            output_dir=output_dir,
            orig_params=labels,
            length=length,
            sample_rate=sample_rate,
            extra=extra,
        )
    # Generate


if __name__ == "__main__":

    note_length = 0.8
    sample_rate = 16384

    lokomotiv = True
    fm = True

    if lokomotiv:
        from generators.vst_generator import *

        run_name = "lokomotiv_full"
        model_file = "output/lokomotiv_full_e2e_best.h5"
        plugin = "/Library/Audio/Plug-Ins/VST/Lokomotiv.vst"
        config_file = "plugin_config/lokomotiv.json"
        generator = VSTGenerator(vst=plugin, sample_rate=sample_rate)
        with open(config_file, "r") as f:
            config = json.load(f)

        model = keras.models.load_model(model_file)
        run_comparison(
            model,
            generator,
            run_name,
            num_samples=100,
            extra={"note_length": note_length, "config": config},
        )

    if fm:
        from generators.fm_generator import *

        run_name = "inversynth_full"
        model_file = "output/inversynth_full_e2e_best.h5"
        generator = InverSynthGenerator()
        model = keras.models.load_model(model_file)
        run_comparison(model, generator, run_name, num_samples=100)