# Neural Piano main Python module
"""
Neural Piano - MIDI to high-quality piano audio rendering pipeline.

This module exposes a single high-level function `render_midi` which converts a
MIDI file into a polished, mastered piano audio file using a multi-stage
pipeline:

1. **MIDI rendering**: Uses a SoundFont (SF2) and `midirenderer` to synthesize
   an initial raw waveform from the input MIDI file.
2. **Silence trimming (optional)**: Removes leading/trailing silence using
   `librosa.effects.trim`.
3. **Latent encoding / neural rendering**: Encodes the rendered waveform into a
   latent representation and decodes it with a neural model (via
   `music2latent.inference.EncoderDecoder`) to produce a refined audio tensor.
4. **Denoising (optional)**: Applies a learned denoiser to reduce artifacts.
5. **Bass enhancement (optional)**: Applies a low-frequency enhancement stage.
6. **Mastering (optional)**: Applies final gain and mastering tailored for
   mono piano output.
7. **Saving**: Writes the final waveform to disk using `soundfile.write`.

Notes
-----
- The function expects a `models` directory in the current working directory
  (os.getcwd()) containing the specified SoundFont file (default:
  'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'). You may override the SF2
  filename with `sf2_name` or provide a custom model path for the neural model
  with `custom_model_path`.
- The neural encoder/decoder is instantiated from
  `music2latent.inference.EncoderDecoder`. Model selection flags are provided
  to support legacy or multi-instrument variants.
- The function uses `librosa` for audio loading and trimming and `soundfile`
  for writing the final WAV file.
- GPU usage is supported via the `device` parameter (e.g., 'cuda' or 'cpu').
- This module does not perform input validation beyond what the underlying
  libraries perform; callers should ensure file paths exist and that the
  environment has the required dependencies installed.

Dependencies
------------
- torch
- librosa
- soundfile
- midirenderer
- music2latent (local package)
- denoise, bass, master modules (local)

"""

import os
import io
from pathlib import Path
import importlib

import torch

import librosa
import soundfile as sf

import midirenderer

from .music2latent.inference import EncoderDecoder

from .denoise import denoise_audio

from .bass import enhance_bass

from .master import master_mono_piano


def render_midi(input_midi_file,
                output_audio_file='neuralpiano_output.wav',
                sample_rate=48000,
                denoising_steps=10,
                max_batch_size=None,
                show_progress=True,
                use_v1_piano_model=False,
                load_multi_instrumental_model=False,
                custom_model_path=None,
                sf2_name='SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2',
                trim_silence=True,
                trim_top_db=60,
                trim_frame_length=2048,
                trim_hop_length=512,
                denoise=True,
                enhance_bass=True,
                low_gain_db=8.0,
                master=True,
                gain_db=10.0,
                device='cuda',
                return_audio=False,
                verbose=True,
                verbose_diag=False
               ):
    
    """
    Render a MIDI file to a polished piano audio file using a neural pipeline.

    This function orchestrates a full rendering pipeline:
    - Render MIDI -> raw waveform using a SoundFont (SF2) via `midirenderer`.
    - Optionally trim leading/trailing silence using `librosa.effects.trim`.
    - Encode the waveform into a latent representation and decode it with a
      neural model (EncoderDecoder) to produce a refined audio tensor.
    - Optionally denoise, enhance bass, and master the audio.
    - Save the final audio to disk (WAV) and optionally return the final audio
      numpy array.

    Parameters
    ----------
    input_midi_file : str or pathlib.Path
        Path to the input MIDI file to render.
    output_audio_file : str, optional
        Path where the final WAV file will be written. Default is
        'neuralpiano_output.wav'.
    sample_rate : int, optional
        Target sample rate for audio loading and processing. Default is 48000.
    denoising_steps : int, optional
        Number of denoising steps passed to the neural decoder. Higher values
        may produce smoother output at the cost of compute. Default is 10.
    use_v1_piano_model : bool, optional
        If True, load the legacy v1 piano model variant in EncoderDecoder.
        Default is False.
    load_multi_instrumental_model : bool, optional
        If True, load a multi-instrument model variant in EncoderDecoder.
        Default is False.
    custom_model_path : str or None, optional
        Path to a custom neural model checkpoint to load for inference. If
        None, the EncoderDecoder will use its default model selection logic.
        Default is None.
    sf2_name : str, optional
        Filename of the SoundFont (SF2) located in the `models` directory in
        the current working directory. Default is
        'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'.
    trim_silence : bool, optional
        If True, trim leading and trailing silence from the initial rendered
        waveform before encoding. Default is True.
    trim_top_db : float, optional
        The `top_db` parameter passed to `librosa.effects.trim` controlling the
        silence threshold in decibels. Default is 60.
    trim_frame_length : int, optional
        Frame length parameter for `librosa.effects.trim`. Default is 2048.
    trim_hop_length : int, optional
        Hop length parameter for `librosa.effects.trim`. Default is 512.
    denoise : bool, optional
        If True, run the denoising stage after neural decoding. Default is True.
    enhance_bass : bool, optional
        If True, run the bass enhancement stage after denoising. Default is True.
    low_gain_db : float, optional
        Gain in decibels applied by the bass enhancement stage. Default is 8.0.
    master : bool, optional
        If True, run the mastering stage (final gain and processing). Default
        is True.
    gain_db : float, optional
        Mastering gain in decibels applied in the mastering stage. Default is
        10.0.
    device : str or torch.device, optional
        Device identifier for GPU/CPU processing (e.g., 'cuda' or 'cpu'). This
        is forwarded to denoising, bass enhancement, and mastering functions.
        Default is 'cuda'.
    return_audio : bool, optional
        If True, return the final audio as a NumPy array in addition to
        writing it to `output_audio_file`. Default is False.
    verbose : bool, optional
        If True, print high-level progress messages. Default is True.
    verbose_diag : bool, optional
        If True, print diagnostic information (shapes, intermediate values).
        Default is False.

    Returns
    -------
    numpy.ndarray or None
        If `return_audio` is True, returns the final audio as a NumPy array
        shaped (samples,) or (channels, samples) depending on processing.
        Otherwise returns None.

    Raises
    ------
    FileNotFoundError
        If the specified SF2 file or MIDI file cannot be read from the expected
        `models` directory or the provided paths.
    RuntimeError
        If the neural model fails to initialize or inference fails; underlying
        libraries (torch, midirenderer, librosa, soundfile) may raise their
        own exceptions which are propagated.
    ValueError
        If audio shapes or types are incompatible with downstream processing.

    Example
    -------
    >>> # Basic usage: render a MIDI file to WAV on CPU and return the audio array
    >>> render_midi('example.mid',
    ...             output_audio_file='out.wav',
    ...             sample_rate=44100,
    ...             device='cpu',
    ...             return_audio=True)

    Implementation details
    ----------------------
    - The function expects a `models` directory in the current working
      directory. The SoundFont path is constructed as:
      `os.path.join(os.getcwd(), "models", sf2_name)`.
    - The neural model is created by calling `EncoderDecoder(...)`. The
      `custom_model_path`, `use_v1_piano_model`, and
      `load_multi_instrumental_model` flags are forwarded to that class.
    - The initial MIDI rendering is performed by `midirenderer.render_wave_from`
      which accepts raw bytes for the SF2 and MIDI files.
    - Audio tensors returned by the neural decoder are PyTorch tensors on the
      device; they are moved to CPU and converted to NumPy before writing with
      `soundfile.write`.

    See also
    --------
    - :mod:`midirenderer`
    - :mod:`librosa`
    - :mod:`soundfile`
    - :mod:`torch`
    - :mod:`music2latent.inference.EncoderDecoder`
    - Local modules: :mod:`.denoise`, :mod:`.bass`, :mod:`.master`
    """
    
    home_root = os.getcwd()
    models_dir = os.path.join(home_root, "models")
    sf2_path = os.path.join(models_dir, sf2_name)

    def _pv(msg):
        if verbose:
            print(msg)

    _pv('=' * 70)
    _pv('Neural Piano')
    _pv('=' * 70)

    _pv('Prepping model...')
    encdec = EncoderDecoder(load_multi_instrumental_model=load_multi_instrumental_model,
                            use_v1_piano_model=use_v1_piano_model,
                            load_path_inference=custom_model_path
                            )

    if verbose_diag:
        _pv(encdec.gen)
        _pv('=' * 70)

    _pv('Reading and rendering MIDI file...')
    wav_data = midirenderer.render_wave_from(
        Path(sf2_path).read_bytes(),
        Path(input_midi_file).read_bytes()
    )

    if verbose_diag:
        _pv(len(wav_data))
        _pv('=' * 70)

    _pv('Loading rendered MIDI...')
    with io.BytesIO(wav_data) as byte_stream:
        wv, sr = librosa.load(byte_stream, sr=sample_rate)

    if verbose_diag:
        _pv(sr)
        _pv(wv.shape)
        _pv('=' * 70)

    if trim_silence:
        _pv('Trimming leading and trailing silence from rendered waveform...')
        wv_trimmed, trim_interval = librosa.effects.trim(
            wv,
            top_db=trim_top_db,
            frame_length=trim_frame_length,
            hop_length=trim_hop_length
        )
        start_sample, end_sample = trim_interval
        orig_dur = len(wv) / sr
        trimmed_dur = len(wv_trimmed) / sr
        if verbose:
            _pv(f'  Trimmed samples: start={start_sample}, end={end_sample}')
            _pv(f'  Duration before={orig_dur:.3f}s, after={trimmed_dur:.3f}s')
        wv = wv_trimmed
    else:
        _pv('Silence trimming disabled; using full rendered waveform.')

    if verbose_diag:
        _pv(wv.shape)
        _pv('=' * 70)

    _pv('Encoding...')
    latent = encdec.encode(wv,
                           max_batch_size=max_batch_size,
                           show_progress=show_progress
                           )

    if verbose_diag:
        _pv(latent.shape)
        _pv('=' * 70)

    _pv('Rendering...')
    audio = encdec.decode(latent,
                          denoising_steps=denoising_steps,
                          max_batch_size=max_batch_size,
                          show_progress=show_progress
                         )

    if verbose_diag:
        _pv(audio.shape)
        _pv('=' * 70)

    if denoise:
        _pv('Denoising...')
        audio, den_diag = denoise_audio(audio,
                                        sr=sr,
                                        device=torch.device(device)
                                        )

        if verbose_diag:
            _pv(den_diag)
            _pv('=' * 70)


    if enhance_bass:
        _pv('Enhancing bass...')
        audio, bass_diag = enhance_bass(audio,
                                        sr=sr,
                                        low_gain_db=low_gain_db,
                                        device=torch.device(device)
                                       )

        if verbose_diag:
            _pv(bass_diag)
            _pv('=' * 70)

    if master:
        _pv('Mastering...')
        audio, mas_diag = master_mono_piano(audio,
                                            gain_db=gain_db,
                                            device=torch.device(device)
                                           )

        if verbose_diag:
            _pv(mas_diag)
            _pv('=' * 70)

    if verbose_diag:
        _pv(audio.shape)
        _pv('=' * 70)

    _pv('Saving final audio...')
    final_audio = audio.cpu().numpy().squeeze().T

    if verbose_diag:
        _pv(final_audio.shape)
        _pv('=' * 70)

    sf.write(output_audio_file, final_audio, samplerate=sr)

    _pv('=' * 70)
    _pv('Done!')
    _pv('=' * 70)

    if return_audio:
        return final_audio