#!/usr/bin/env python
# coding: utf8

""" Spectrogram specific data augmentation. """

# pyright: reportMissingImports=false
# pylint: disable=import-error
import numpy as np
import tensorflow as tf
from tensorflow.signal import hann_window, stft

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


def compute_spectrogram_tf(
    waveform: tf.Tensor,
    frame_length: int = 2048,
    frame_step: int = 512,
    spec_exponent: float = 1.0,
    window_exponent: float = 1.0,
) -> tf.Tensor:
    """
    Compute magnitude / power spectrogram from waveform as a
    `n_samples x n_channels` tensor.

    Parameters:
        waveform (tensorflow.Tensor):
            Input waveform as `(times x number of channels)` tensor.
        frame_length (int):
            Length of a STFT frame to use.
        frame_step (int):
            HOP between successive frames.
        spec_exponent (float):
            Exponent of the spectrogram (usually 1 for magnitude
            spectrogram, or 2 for power spectrogram).
        window_exponent (float):
            Exponent applied to the Hann windowing function (may be
            useful for making perfect STFT/iSTFT reconstruction).

    Returns:
        tensorflow.Tensor:
            Computed magnitude / power spectrogram as a
            `(T x F x n_channels)` tensor.
    """
    stft_tensor: tf.Tensor = tf.transpose(
        stft(
            tf.transpose(waveform),
            frame_length,
            frame_step,
            window_fn=lambda f, dtype: hann_window(
                f, periodic=True, dtype=waveform.dtype
            )
            ** window_exponent,
        ),
        perm=[1, 2, 0],
    )
    # result = tf.abs(stft_tensor) ** spec_exponent
    # return result

    return tf.abs(stft_tensor) ** spec_exponent

#
# def time_stretch(
#     spectrogram: tf.Tensor,
#     factor: float = 1.0,
#     method: tf.image.ResizeMethod = tf.image.ResizeMethod.BILINEAR,
# ) -> tf.Tensor:
#     """
#     Time stretch a spectrogram preserving shape in tensorflow. Note that
#     this is an approximation in the frequency domain.
#
#     Parameters:
#         spectrogram (tensorflow.Tensor):
#             Input spectrogram to be time stretched as tensor.
#         factor (float):
#             (Optional) Time stretch factor, must be > 0, default to `1`.
#         method (tensorflow.image.ResizeMethod):
#             (Optional) Interpolation method, default to `BILINEAR`.
#
#     Returns:
#         tensorflow.Tensor:
#             Time stretched spectrogram as tensor with same shape.
#     """
#     T = tf.shape(spectrogram)[0]
#     T_ts = tf.cast(tf.cast(T, tf.float32) * factor, tf.int32)[0]
#     F = tf.shape(spectrogram)[1]
#     ts_spec = tf.image.resize_images(
#         spectrogram, [T_ts, F], method=method, align_corners=True
#     )
#     return tf.image.resize_image_with_crop_or_pad(ts_spec, T, F)
#
#
# def random_time_stretch(
#     spectrogram: tf.Tensor, factor_min: float = 0.9, factor_max: float = 1.1, **kwargs
# ) -> tf.Tensor:
#     """
#     Time stretch a spectrogram preserving shape with random ratio in
#     tensorflow. Applies time_stretch to spectrogram with a random ratio
#     drawn uniformly in `[factor_min, factor_max]`.
#
#     Parameters:
#         spectrogram (tensorflow.Tensor):
#             Input spectrogram to be time stretched as tensor.
#         factor_min (float):
#             (Optional) Min time stretch factor, default to `0.9`.
#         factor_max (float):
#             (Optional) Max time stretch factor, default to `1.1`.
#
#     Returns:
#         tensorflow.Tensor:
#             Randomly time stretched spectrogram as tensor with same shape.
#     """
#     factor = (
#         tf.random_uniform(shape=(1,), seed=0) * (factor_max - factor_min) + factor_min
#     )
#     return time_stretch(spectrogram, factor=factor, **kwargs)
#
#
# def pitch_shift(
#     spectrogram: tf.Tensor,
#     semitone_shift: float = 0.0,
#     method: tf.image.ResizeMethod = tf.image.ResizeMethod.BILINEAR,
# ) -> tf.Tensor:
#     """
#     Pitch shift a spectrogram preserving shape in tensorflow. Note that
#     this is an approximation in the frequency domain.
#
#     Parameters:
#         spectrogram (tensorflow.Tensor):
#             Input spectrogram to be pitch shifted as tensor.
#         semitone_shift (float):
#             (Optional) Pitch shift in semitone, default to `0.0`.
#         method (tensorflow.image.ResizeMethod):
#             (Optional) Interpolation method, default to `BILINEAR`.
#
#     Returns:
#         tensorflow.Tensor:
#             Pitch shifted spectrogram (same shape as spectrogram).
#     """
#     factor = 2 ** (semitone_shift / 12.0)
#     T = tf.shape(spectrogram)[0]
#     F = tf.shape(spectrogram)[1]
#     F_ps = tf.cast(tf.cast(F, tf.float32) * factor, tf.int32)[0]
#     ps_spec = tf.image.resize_images(
#         spectrogram, [T, F_ps], method=method, align_corners=True
#     )
#     paddings = [[0, 0], [0, tf.maximum(0, F - F_ps)], [0, 0]]
#     return tf.pad(ps_spec[:, :F, :], paddings, "CONSTANT")
#
#
# def random_pitch_shift(
#     spectrogram: tf.Tensor, shift_min: float = -1.0, shift_max: float = 1.0, **kwargs
# ) -> tf.Tensor:
#     """
#     Pitch shift a spectrogram preserving shape with random ratio in
#     tensorflow. Applies pitch_shift to spectrogram with a random shift
#     amount (expressed in semitones) drawn uniformly in
#     `[shift_min, shift_max]`.
#
#     Parameters:
#         spectrogram (tensorflow.Tensor):
#             Input spectrogram to be pitch shifted as tensor.
#         shift_min (float):
#             (Optional) Min pitch shift in semitone, default to -1.
#         shift_max (float):
#             (Optional) Max pitch shift in semitone, default to 1.
#
#     Returns:
#         tensorflow.Tensor:
#             Randomly pitch shifted spectrogram (same shape as spectrogram).
#     """
#     semitone_shift = (
#         tf.random_uniform(shape=(1,), seed=0) * (shift_max - shift_min) + shift_min
#     )
#     return pitch_shift(spectrogram, semitone_shift=semitone_shift, **kwargs)
