import torch
import torchaudio
from typing import Optional, Literal, List, Dict, Callable
import torchaudio.transforms as T
import inspect

__all__ = ["register_audio_func_loader", "_filter_kwargs", "get_label", "load_audio", "PRIM_AUD_FUNC_REGISTRY"]

PRIM_AUD_FUNC_REGISTRY: Dict[str, Callable] = {}

def register_audio_func_loader(name: str):
    def decorator(func: Callable):
        PRIM_AUD_FUNC_REGISTRY[name] = func
        return func
    return decorator

def _filter_kwargs(
    func : Callable, 
    kwargs: dict) -> dict:
    sig = inspect.signature(func)
    valid_params = sig.parameters.keys()
    return {k: v for k, v in kwargs.items() if k in valid_params}

@register_audio_func_loader("spectrogram")
def to_spectrogram(
    waveform: torch.Tensor,
    sample_rate: int,
    n_fft: int = 512,
    hop_length: int = 256,
    img_size: int | None = None
) -> torch.Tensor:

    transform = T.Spectrogram(n_fft=n_fft, hop_length=hop_length)
    spec = transform(waveform)

    if img_size:
        spec = torch.nn.functional.interpolate(
            spec.unsqueeze(0),
            size=(img_size, img_size),
            mode="bilinear",
            align_corners=False
        ).squeeze(0)

    return spec

@register_audio_func_loader("mfcc")
def to_mfcc(
    waveform: torch.Tensor,
    sample_rate: int,
    n_fft: int = 2048,
    hop_length: int = 512,
    n_mels: int = 128,
    n_mfcc: int = 40,
    img_size: int | None = None
)-> torch.Tensor:
    transform = T.MFCC(
        sample_rate=sample_rate,
        n_mfcc=n_mfcc,
        melkwargs={
            "n_fft": n_fft,
            "n_mels": n_mels,
            "hop_length": hop_length,
        }
    )
    return transform(waveform)
    
def get_label(
    cls: str, 
    prob_dist : bool = False) -> torch.Tensor:
    idx = {"real":0., "fake":1.}
    if prob_dist:
        lbl = [0.,0.]
        lbl[int(idx[cls])] = 1.
        return torch.tensor(lbl)
    return torch.tensor([idx[cls]])

def load_audio(
    audio_path: str,
    max_secs: float | None = None,
):
    waveform, sample_rate = torchaudio.load(audio_path)

    if max_secs:
        max_samples = int(sample_rate * max_secs)
        waveform = waveform[:, :max_samples]

    return waveform, sample_rate