<gh_stars>1000+
import os
import csv
import pickle

import yaml
import librosa

from omnizart.utils import ensure_path_exists, LazyLoader, get_logger


# Lazy load the Spleeter pacakge for avoiding pulling large dependencies
# and boosting the import speed.
adapter = LazyLoader("adapter", globals(), "spleeter.audio.adapter")
logger = get_logger("IO")


def dump_pickle(data, save_to):
    """Dump data to the given path.

    Parameters
    ----------
    data: python objects
        Data to store. Should be python built-in types like `dict`, `list`, `str`, `int`, etc
    save_to: Path
        The full path to store the pickle file, including file name.
        Will create the directory if the given path doesn't exist.

    """
    base_dir = os.path.dirname(save_to)
    ensure_path_exists(base_dir)
    with open(save_to, "wb") as pkl_file:
        pickle.dump(data, pkl_file)


def load_pickle(pickle_file):
    """Load pickle file from the given path

    Read and returns the data from the given pickle file path

    Parameters
    ----------
    pickle_file: Path
        The full path to the pickle file for read

    Returns
    -------
    object
        Python object, could be `dict`, `list`, `str`, etc.
    """
    return pickle.load(open(pickle_file, "rb"))


def load_audio(audio_path, sampling_rate=44100, mono=True):
    """Load audio with spleeter.

    A much faster and general approach for loading audio comparing to use librosa.
    This function also allows to read .mp3 files.

    Parameters
    ----------
    audio_path: Path
        Path to the audio.
    sampling_rate: int
        Target sampling rate after loaded.
    mono: bool
        Wether to transform the audio into monophonic channel.

    Returns
    -------
    audio: 1D numpy array
        Raw data of the audio.
    fs: int
        Sampling rate of the audio. Will be the same as the given ``sampling_rate``.
    """
    try:
        audio_loader = adapter.AudioAdapter.default()
        audio, fs = audio_loader.load(audio_path, sample_rate=sampling_rate)
        if mono:
            audio = librosa.to_mono(audio.squeeze().T)

    except (ImportError, adapter.SpleeterError) as error:
        logger.warning(
            "Failed to load audio with Spleeter due to '%s'. Continue to use Librosa.", str(error)
        )
        audio, fs = load_audio_with_librosa(audio_path, sampling_rate=sampling_rate, mono=mono)
        if not mono:
            audio = audio.T

    return audio, fs


def load_audio_with_librosa(audio_path, sampling_rate=44100, mono=True):
    """Load audio from the given path with librosa.load

    Parameters
    ----------
    audio_path: Path
        Path to the audio.
    sampling_rate: int
        Target sampling rate after loaded.
    mono: bool
        Wether to transform the audio into monophonic channel.

    Returns
    -------
    audio: 1D numpy array
        Raw data of the audio.
    fs: int
        Sampling rate of the audio. Will be the same as the given ``sampling_rate``.
    """
    return librosa.load(audio_path, mono=mono, sr=sampling_rate)


def load_yaml(yaml_path):
    return yaml.load(open(yaml_path, "r"), Loader=yaml.Loader)


def write_yaml(json_obj, output_path, dump=True):
    # If dump is false, then the json_obj should be yaml string already.
    out_str = yaml.dump(json_obj) if dump else json_obj
    open(output_path, "w").write(out_str)


def write_agg_f0_results(agg_f0, output_path):
    """Write out aggregated F0 information as a CSV file.

    Parameters
    ----------
    agg_f0: list[dict]
        List of aggregated F0 information.
    output_path: Path
        Path for output the CSV file. Should contain the file name.

    See Also
    --------
    omnizart.utils.aggregate_f0_info:
        The function for generating the aggregated F0 information.
    """
    fieldnames = ["start_time", "end_time", "frequency", "pitch"]

    # Check the format is correct
    if any(list(row.keys()) != fieldnames for row in agg_f0):
        raise ValueError(f"Fields inconsistent! Expected: {fieldnames}")

    with open(output_path, "w") as out:
        writer = csv.DictWriter(out, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(agg_f0)

import os
import csv
import pickle

import yaml
import librosa

from omnizart.utils import ensure_path_exists, LazyLoader, get_logger


# Lazy load the Spleeter pacakge for avoiding pulling large dependencies
# and boosting the import speed.
adapter = LazyLoader("adapter", globals(), "spleeter.audio.adapter")
logger = get_logger("IO")


def dump_pickle(data, save_to):
    """Dump data to the given path.

    Parameters
    ----------
    data: python objects
        Data to store. Should be python built-in types like `dict`, `list`, `str`, `int`, etc
    save_to: Path
        The full path to store the pickle file, including file name.
        Will create the directory if the given path doesn't exist.

    """
    base_dir = os.path.dirname(save_to)
    ensure_path_exists(base_dir)
    with open(save_to, "wb") as pkl_file:
        pickle.dump(data, pkl_file)


def load_pickle(pickle_file):
    """Load pickle file from the given path

    Read and returns the data from the given pickle file path

    Parameters
    ----------
    pickle_file: Path
        The full path to the pickle file for read

    Returns
    -------
    object
        Python object, could be `dict`, `list`, `str`, etc.
    """
    return pickle.load(open(pickle_file, "rb"))


def load_audio(audio_path, sampling_rate=44100, mono=True):
    """Load audio with spleeter.

    A much faster and general approach for loading audio comparing to use librosa.
    This function also allows to read .mp3 files.

    Parameters
    ----------
    audio_path: Path
        Path to the audio.
    sampling_rate: int
        Target sampling rate after loaded.
    mono: bool
        Wether to transform the audio into monophonic channel.

    Returns
    -------
    audio: 1D numpy array
        Raw data of the audio.
    fs: int
        Sampling rate of the audio. Will be the same as the given ``sampling_rate``.
    """
    try:
        audio_loader = adapter.AudioAdapter.default()
        audio, fs = audio_loader.load(audio_path, sample_rate=sampling_rate)
        if mono:
            audio = librosa.to_mono(audio.squeeze().T)

    except (ImportError, adapter.SpleeterError) as error:
        logger.warning(
            "Failed to load audio with Spleeter due to '%s'. Continue to use Librosa.", str(error)
        )
        audio, fs = load_audio_with_librosa(audio_path, sampling_rate=sampling_rate, mono=mono)
        if not mono:
            audio = audio.T

    return audio, fs


def load_audio_with_librosa(audio_path, sampling_rate=44100, mono=True):
    """Load audio from the given path with librosa.load

    Parameters
    ----------
    audio_path: Path
        Path to the audio.
    sampling_rate: int
        Target sampling rate after loaded.
    mono: bool
        Wether to transform the audio into monophonic channel.

    Returns
    -------
    audio: 1D numpy array
        Raw data of the audio.
    fs: int
        Sampling rate of the audio. Will be the same as the given ``sampling_rate``.
    """
    return librosa.load(audio_path, mono=mono, sr=sampling_rate)


def load_yaml(yaml_path):
    return yaml.load(open(yaml_path, "r"), Loader=yaml.Loader)


def write_yaml(json_obj, output_path, dump=True):
    # If dump is false, then the json_obj should be yaml string already.
    out_str = yaml.dump(json_obj) if dump else json_obj
    open(output_path, "w").write(out_str)


def write_agg_f0_results(agg_f0, output_path):
    """Write out aggregated F0 information as a CSV file.

    Parameters
    ----------
    agg_f0: list[dict]
        List of aggregated F0 information.
    output_path: Path
        Path for output the CSV file. Should contain the file name.

    See Also
    --------
    omnizart.utils.aggregate_f0_info:
        The function for generating the aggregated F0 information.
    """
    fieldnames = ["start_time", "end_time", "frequency", "pitch"]

    # Check the format is correct
    if any(list(row.keys()) != fieldnames for row in agg_f0):
        raise ValueError(f"Fields inconsistent! Expected: {fieldnames}")

    with open(output_path, "w") as out:
        writer = csv.DictWriter(out, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(agg_f0)
# Lazy load the Spleeter pacakge for avoiding pulling large dependencies # and boosting the import speed. Dump data to the given path.      Parameters     ----------     data: python objects         Data to store. Should be python built-in types like `dict`, `list`, `str`, `int`, etc     save_to: Path         The full path to store the pickle file, including file name.         Will create the directory if the given path doesn't exist. Load pickle file from the given path      Read and returns the data from the given pickle file path      Parameters     ----------     pickle_file: Path         The full path to the pickle file for read      Returns     -------     object         Python object, could be `dict`, `list`, `str`, etc. Load audio with spleeter.      A much faster and general approach for loading audio comparing to use librosa.     This function also allows to read .mp3 files.      Parameters     ----------     audio_path: Path         Path to the audio.     sampling_rate: int         Target sampling rate after loaded.     mono: bool         Wether to transform the audio into monophonic channel.      Returns     -------     audio: 1D numpy array         Raw data of the audio.     fs: int         Sampling rate of the audio. Will be the same as the given ``sampling_rate``. Load audio from the given path with librosa.load      Parameters     ----------     audio_path: Path         Path to the audio.     sampling_rate: int         Target sampling rate after loaded.     mono: bool         Wether to transform the audio into monophonic channel.      Returns     -------     audio: 1D numpy array         Raw data of the audio.     fs: int         Sampling rate of the audio. Will be the same as the given ``sampling_rate``. # If dump is false, then the json_obj should be yaml string already. Write out aggregated F0 information as a CSV file.      Parameters     ----------     agg_f0: list[dict]         List of aggregated F0 information.     output_path: Path         Path for output the CSV file. Should contain the file name.      See Also     --------     omnizart.utils.aggregate_f0_info:         The function for generating the aggregated F0 information. # Check the format is correct