import json
from pathlib import Path
from typing import Optional, Union

import torch
from torch import nn
import torch.nn.functional as F


class BatAudioProcessor(nn.Module):
    def __init__(
        self,
        sample_rate: int = 16000,
        n_fft: int = 1024,
        hop_length: int = 160,
        n_mels: int = 128,
        f_min: int = 0,
        top_db: int = 80,
        time_frame_size: int = 1024,
        crop: bool = True,
    ):
        super().__init__()
        try:
            import torchaudio
        except ImportError as exc:
            raise ImportError("BatAudioProcessor requires torchaudio for waveform preprocessing.") from exc

        self.sample_rate = sample_rate
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.n_mels = n_mels
        self.f_min = f_min
        self.top_db = top_db
        self.time_frame_size = time_frame_size
        self.crop = crop
        self.mel = torchaudio.transforms.MelSpectrogram(
            sample_rate=sample_rate,
            n_fft=n_fft,
            hop_length=hop_length,
            n_mels=n_mels,
            f_min=f_min,
        )
        self.db = torchaudio.transforms.AmplitudeToDB("power", top_db=top_db)

    @classmethod
    def from_pretrained(cls, model_id_or_path: Optional[Union[str, Path]] = None, **kwargs):
        model_id_or_path = Path(model_id_or_path or Path(__file__).resolve().parent)
        if model_id_or_path.exists():
            config_path = model_id_or_path / "config.json"
        else:
            from huggingface_hub import hf_hub_download

            config_path = Path(hf_hub_download(str(model_id_or_path), filename="config.json"))

        with config_path.open("r", encoding="utf-8") as f:
            config = json.load(f)
        processor_kwargs = {
            "sample_rate": config.get("sample_rate", 16000),
            "n_fft": config.get("n_fft", 1024),
            "hop_length": config.get("hop_length", 160),
            "n_mels": config.get("n_mels", 128),
            "f_min": config.get("f_min", 0),
            "top_db": config.get("top_db", 80),
            "time_frame_size": config.get("input_shape", [1024, 128])[0],
        }
        processor_kwargs.update(kwargs)
        return cls(**processor_kwargs)

    def minmax_normalize(self, x: torch.Tensor) -> torch.Tensor:
        shape = x.shape
        x = x.flatten(1)
        min_, max_ = x.aminmax(dim=1, keepdim=True)
        x = (x - min_) / (max_ - min_ + 1e-8)
        return x.reshape(shape)

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        if waveform.ndim == 1:
            waveform = waveform.unsqueeze(0)
        if waveform.ndim != 2:
            raise ValueError(f"Expected waveform shape [batch, samples] or [samples], got {tuple(waveform.shape)}.")

        x = self.mel(waveform).unsqueeze(1)
        x = self.db(x)
        x = self.minmax_normalize(x)

        frames = x.shape[-1]
        if frames < self.time_frame_size:
            x = F.pad(x, (0, self.time_frame_size - frames), mode="constant", value=0)
        elif frames > self.time_frame_size:
            if not self.crop:
                raise ValueError(
                    f"Waveform produced {frames} frames, longer than configured {self.time_frame_size} frames."
                )
            x = x[..., : self.time_frame_size]

        return x.transpose(2, 3)