# Neural Piano main Python module

import os
import io
from pathlib import Path

import librosa
import soundfile as sf

import midirenderer

from .music2latent.inference import EncoderDecoder

from .master import master_mono_piano

def neuralpiano(input_midi_file,
                output_audio_file,
                return_audio=False
                ):
    
    filepath = os.path.abspath(__file__)
    lib_root = os.path.dirname(filepath)
    models_dir_path = lib_root + "/models/"
    
    print('=' * 70)
    print('Neural Piano')
    print('=' * 70)
    
    print('Reading and rendering MIDI file...')
    wav_data = midirenderer.render_wave_from(
        Path(models_dir_path+"SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2").read_bytes(),
        Path(input_midi_file).read_bytes()
    )
    
    print('Loading rendered MIDI...')
    with io.BytesIO(wav_data) as byte_stream:
        wv, sr = librosa.load(byte_stream, sr=48000)  # Music2Latent supports 48kHz audio as well
    
    print('Prepping model...')
    encdec = neuralpiano.EncoderDecoder()
    
    print('Encoding...')
    latent = encdec.encode(wv)
    
    print('Rendering...')
    wv_rec = encdec.decode(latent)
    
    print('Mastering...')
    stereo, diag = master_mono_piano(wv_rec)
    
    print('Saving final audio...')
    sf.write(output_audio_file, stereo.squeeze(), sr)
    
    print('=' * 70)
    print('Done!')
    print('=' * 70)
    
    if return_audio:
        return stereo