# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/11_audio.ipynb.

# %% auto 0
__all__ = ['AudioMetadata', 'Spectrogram', 'Wav', 'play_audio']

# %% ../../nbs/11_audio.ipynb 2
import copy
import tempfile
from pathlib import Path
from dataclasses import dataclass
from typing import Self, assert_never

import librosa
import soundfile
import numpy as np

from IPython.display import Audio, display

# %% ../../nbs/11_audio.ipynb 5
@dataclass
class AudioMetadata:
    """
    Metadata for structured audio data.

    It uses librosa default values.
    """
    source_file: str | Path | None = None
    sample_rate: int = 22050
    duration: float | None = None
    n_fft: int = 2048
    hop_length: int = 512

# %% ../../nbs/11_audio.ipynb 6
class Spectrogram:
    """Class to handle spectrogram data in 2D format."""

    def __init__(
        self,
        data: np.ndarray,
        stft: np.ndarray | None = None,
        metadata: AudioMetadata | None = None
    ) -> None:
        if data.ndim != 2:
            raise ValueError(f"Spectrogram must be 2D array, got {data.ndim}D")

        self.data = data
        self.stft = stft
        self.metadata = metadata or AudioMetadata()

    @property
    def freq_bins(self) -> int:
        """Get number of frequency bins."""
        return self.data.shape[0]

    @property
    def time_frames(self) -> int:
        """Get number of time frames."""
        return self.data.shape[1]

    def copy(self) -> 'Spectrogram':
        """
        Create a deep copy of the Spectrogram object.

        Returns:
            New Spectrogram object with copied data and metadata.
        """
        return Spectrogram(
            data=self.data.copy(),
            stft=self.stft.copy() if self.stft is not None else None,
            metadata=copy.deepcopy(self.metadata)
        )

# %% ../../nbs/11_audio.ipynb 7
class Wav:
    def __init__(
        self,
        *,
        waveform: np.ndarray,
        metadata: AudioMetadata,
    ) -> None:
        if waveform.ndim != 1:
            raise ValueError(f"Waveform must be 1D array, got {waveform.ndim}D")

        self.waveform = waveform
        self.metadata = metadata

        if self.metadata.duration is None:
            self.metadata.duration = len(waveform) / metadata.sample_rate

    def to_spectrogram(self) -> Spectrogram:
        # Short-time Fourier transform
        stft = librosa.stft(
            self.waveform,
            n_fft=self.metadata.n_fft,
            hop_length=self.metadata.hop_length
        )

        magnitude = librosa.amplitude_to_db(np.abs(stft), ref=np.max)

        return Spectrogram(
            data=magnitude,
            stft=stft,
            metadata=self.metadata
        )

    @classmethod
    def from_spectrogram(cls, spectrogram: Spectrogram) -> Self:
        if spectrogram.stft is None:
            raise ValueError("Cannot reconstruct audio: STFT data is missing")

        # Use phase information for accurate reconstruction
        reconstructed_waveform = librosa.istft(
            spectrogram.stft,
            hop_length=spectrogram.metadata.hop_length
        )

        new_metadata = AudioMetadata(
            source_file=f"{spectrogram.metadata.source_file}_reconstructed" if spectrogram.metadata.source_file else "reconstructed",
            sample_rate=spectrogram.metadata.sample_rate,
            duration=len(reconstructed_waveform) / spectrogram.metadata.sample_rate,
            n_fft=spectrogram.metadata.n_fft,
            hop_length=spectrogram.metadata.hop_length
        )

        return cls(waveform=reconstructed_waveform, metadata=new_metadata)

    @classmethod
    def load(
        cls,
        data: str | Path | np.ndarray,
        metadata: AudioMetadata | None = None,
    ) -> Self:
        if metadata is None:
            metadata = AudioMetadata()

        if isinstance(data, (str, Path)):
            waveform, actual_sr = librosa.load(data, sr=metadata.sample_rate)

            metadata.sample_rate = int(actual_sr)
            metadata.source_file = str(data)
        elif isinstance(data, np.ndarray):
            waveform = data
        else:
            assert_never()

        return cls(waveform=waveform, metadata=metadata)

    def copy(self) -> 'Wav':
        """
        Create a deep copy of the Wav object.

        Returns:
            New Wav object with copied waveform and metadata.
        """
        return Wav(
            waveform=self.waveform.copy(),
            metadata=copy.deepcopy(self.metadata)
        )

# %% ../../nbs/11_audio.ipynb 8
def play_audio(
    audio_data: np.ndarray | str | Path | Spectrogram | Wav,
    sample_rate: int | None = None
) -> None:
    """
    Universal function to play audio from various input types.
    """
    if isinstance(audio_data, (str, Path)):
        # File path - use directly
        display(Audio(str(audio_data)))

    elif isinstance(audio_data, Wav):
        # Wav object - use waveform and sample rate
        display(Audio(audio_data.waveform, rate=audio_data.metadata.sample_rate))

    elif isinstance(audio_data, Spectrogram):
        if audio_data.stft is not None:
            # Spectrogram object - convert to wav first
            wav_obj = Wav.from_spectrogram(audio_data)
            display(Audio(wav_obj.waveform, rate=wav_obj.metadata.sample_rate))
        else:
            # No STFT -> estimate phase with Griffin-Lim from magnitude (data is in dB)
            # This yields an approximate reconstruction suitable for playback.
            mag = librosa.db_to_amplitude(audio_data.data)  # convert dB -> linear magnitude
            n_fft = audio_data.metadata.n_fft
            hop_length = audio_data.metadata.hop_length
            
            # Use improved Griffin-Lim parameters for better reconstruction
            waveform = librosa.griffinlim(
                mag, 
                n_iter=128,
                hop_length=hop_length, 
                win_length=n_fft,
                length=None  # Let it determine the length from spectrogram
            )
            display(Audio(waveform, rate=audio_data.metadata.sample_rate))

    elif isinstance(audio_data, np.ndarray):
        # Numpy array - need sample rate
        if sample_rate is None:
            raise ValueError("sample_rate is required when playing numpy arrays")
        display(Audio(audio_data, rate=sample_rate))

    else:
        raise TypeError(f"Unsupported audio data type: {type(audio_data)}")
