from typing import TYPE_CHECKING, Any, Optional

import librosa.display
import matplotlib.pyplot as plt
import numpy as np
from scipy import fft
from scipy import signal as ss

from wandas.utils.types import NDArrayComplex, NDArrayReal

from .base_channel import BaseChannel

if TYPE_CHECKING:
    from matplotlib.axes import Axes
    from matplotlib.collections import QuadMesh


class TimeFrequencyChannel(BaseChannel):
    def __init__(
        self,
        data: NDArrayReal,
        sampling_rate: int,
        n_fft: int,
        hop_length: int,
        win_length: int,
        window: str,
        # center: bool = None,
        # pad_mode: str,
        label: Optional[str] = None,
        unit: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
        previous: Optional["BaseChannel"] = None,
    ):
        """
        TimeFrequencyChannel オブジェクトを初期化します。

        Parameters:
            frequencies (numpy.ndarray): 周波数データ。
            times (numpy.ndarray): 時間データ。
            data (numpy.ndarray): スペクトルデータ（時間周波数成分）。
            label (str, optional): チャンネルのラベル。
            unit (str, optional):    データの単位。
            metadata (dict, optional): メタデータ。
        """
        super().__init__(
            data=data,
            sampling_rate=sampling_rate,
            label=label,
            unit=unit,
            metadata=metadata,
            previous=previous,
        )
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.window = window
        # self.center = center
        # self.pad_mode = pad_mode

    @classmethod
    def stft(
        cls,
        data: NDArrayReal,
        n_fft: Optional[int] = None,
        hop_length: Optional[int] = None,
        win_length: Optional[int] = None,
        window: str = "hann",
        # pad_mode: str = "constant",
    ) -> dict[str, Any]:
        """
        STFT（短時間フーリエ変換）を実行します。

        Parameters:
            data (numpy.ndarray): 入力データ。
            n_fft (int): FFT のサンプル数。
            hop_length (int): ホップサイズ（フレーム間の移動量）。
            win_length (int): ウィンドウの長さ。
            window (str): ウィンドウ関数の種類。
            center (bool): フレームを中央に配置するかどうか。
            pad_mode (str): パディングモード。

        Returns:
            numpy.ndarray: STFT の結果。
        """
        if win_length is None:
            win_length = 2048
        if n_fft is None:
            n_fft = win_length
        if hop_length is None:
            hop_length = win_length // 2

        spec_data: NDArrayComplex
        _, _, spec_data = ss.stft(
            data,
            nfft=n_fft,
            noverlap=win_length - hop_length,
            nperseg=win_length,
            window=window,
            detrend="constant",  # type: ignore[unused-ignore]
            # pad_mode=pad_mode,
        )
        spec_data[..., 1:-1, :] *= 2.0
        return dict(
            data=spec_data,
            n_fft=n_fft,
            hop_length=hop_length,
            win_length=win_length,
            window=window,
        )

    def melspectrogram(
        self,
        n_mels: int = 128,
        # pad_mode: str = "constant",
    ) -> "TimeMelFrequencyChannel":
        result = TimeMelFrequencyChannel.spec2melspec(
            sampling_rate=self.sampling_rate,
            data=np.abs(self.data),
            n_fft=self.n_fft,
            n_mels=n_mels,
        )
        result.update(
            dict(
                hop_length=self.hop_length,
                win_length=self.win_length,
                window=self.window,
            )
        )
        return TimeMelFrequencyChannel.from_channel(self, **result)

    def data_Aw(self, to_dB: bool = False) -> NDArrayReal:  # noqa: N802, N803
        """
        A特性を適用した振幅データを返します。
        """
        freqs = fft.rfftfreq(self.n_fft, 1 / self.sampling_rate)
        weighted: NDArrayReal = librosa.perceptual_weighting(
            np.abs(self._data) ** 2, freqs, kind="A", ref=self.ref**2
        )

        if to_dB:
            return weighted.astype(np.float64)

        return np.asarray(librosa.db_to_amplitude(weighted), dtype=np.float64)

    def hpss_harmonic(self, **kwargs: Any) -> "TimeFrequencyChannel":
        """
        ハーモニック成分を抽出します。
        """
        harmonic, _ = librosa.decompose.hpss(self.data, **kwargs)
        result = dict(
            data=harmonic,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window,
        )
        return TimeFrequencyChannel.from_channel(self, **result)

    def hpss_percussive(self, **kwargs: Any) -> "TimeFrequencyChannel":
        """
        パーカッシブ成分を抽出します。
        """
        _, percussive = librosa.decompose.hpss(self.data, **kwargs)
        result = dict(
            data=percussive,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window,
        )
        return TimeFrequencyChannel.from_channel(self, **result)

    def _plot(
        self,
        ax: "Axes",
        title: Optional[str] = None,
        db_scale: bool = True,
        fmin: Optional[float] = None,
        fmax: Optional[float] = None,
        Aw: bool = False,  # noqa: N803
        vmin: Optional[float] = None,
        vmax: Optional[float] = None,
    ) -> tuple["QuadMesh", NDArrayReal]:
        """
        時間周波数データをプロットします。

        Parameters:
            ax (matplotlib.axes.Axes, optional): 既存のプロット軸。
            title (str, optional): プロットのタイトル。
            db_scale (bool): dBスケールでプロットするかどうか。
        """

        if Aw:
            data_to_plot = self.data_Aw(to_dB=True)
        elif db_scale:
            data_to_plot = self._to_db()
        else:
            data_to_plot = np.abs(self.data)

        # 時間周波数データをプロット
        img = librosa.display.specshow(
            data=data_to_plot,
            sr=self.sampling_rate,
            hop_length=self.hop_length,
            n_fft=self.n_fft,
            win_length=self.win_length,
            x_axis="time",
            y_axis="linear",
            ax=ax,
            fmin=fmin,
            fmax=fmax,
            cmap="magma",
            vmin=vmin,
            vmax=vmax,
        )

        # ラベルとタイトルを設定
        ax.set_xlabel("Time [s]")
        ax.set_ylabel("Frequency [Hz]")
        ax.set_title(title or self.label or "Time-Frequency Representation")

        return img, data_to_plot

    def plot(
        self,
        ax: Optional["Axes"] = None,
        title: Optional[str] = None,
        db_scale: bool = True,
        fmin: Optional[float] = None,
        fmax: Optional[float] = None,
        Aw: bool = False,  # noqa: N803
        vmin: Optional[float] = None,
        vmax: Optional[float] = None,
    ) -> tuple["Axes", NDArrayReal]:
        """
        時間周波数データをプロットします。

        Parameters:
            ax (matplotlib.axes.Axes, optional): 既存のプロット軸。
            title (str, optional): プロットのタイトル。
            db_scale (bool): dBスケールでプロットするかどうか。
        """
        _ax = ax
        if _ax is None:
            _, _ax = plt.subplots(figsize=(10, 6))

        img, data_to_plot = self._plot(
            ax=_ax,
            title=title,
            db_scale=db_scale,
            fmin=fmin,
            fmax=fmax,
            Aw=Aw,
            vmin=vmin,
            vmax=vmax,
        )

        if _ax.figure is not None:
            if db_scale:
                cbar = _ax.figure.colorbar(img, ax=ax, format="%+2.0f")
                if Aw:
                    unit = "dBA"
                else:
                    unit = "dB"
                cbar.set_label(f"Spectrum level [{unit}]")
            else:
                cbar = _ax.figure.colorbar(img, ax=ax)
                cbar.set_label(f"Amplitude [{self.unit}]")

        if ax is None:
            plt.tight_layout()
            plt.show()

        return _ax, data_to_plot

    def _to_db(self) -> NDArrayReal:
        """
        スペクトルデータを dB スケールに変換した新しい TimeFrequencyChannel を返す。

        Returns:
            TimeFrequencyChannel: dBスケールに変換された新しい TimeFrequencyChannel。
        """
        return np.asarray(
            librosa.amplitude_to_db(np.abs(self.data), ref=self.ref, amin=1e-12),
            np.float64,
        )


class TimeMelFrequencyChannel(TimeFrequencyChannel):
    def __init__(
        self,
        data: NDArrayReal,
        sampling_rate: int,
        n_fft: int,
        hop_length: int,
        win_length: int,
        window: str,
        # center: bool = None,
        # pad_mode: str,
        n_mels: int = 128,
        label: Optional[str] = None,
        unit: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
        previous: Optional["BaseChannel"] = None,
    ):
        """
        TimeMelFrequencyChannel オブジェクトを初期化します。

        Parameters:
            frequencies (numpy.ndarray): 周波数データ。
            times (numpy.ndarray): 時間データ。
            data (numpy.ndarray): スペクトルデータ（時間周波数成分）。
            label (str, optional): チャンネルのラベル。
            unit (str, optional):    データの単位。
            metadata (dict, optional): メタデータ。

        """
        super().__init__(
            data=data,
            sampling_rate=sampling_rate,
            label=label,
            unit=unit,
            metadata=metadata,
            n_fft=n_fft,
            hop_length=hop_length,
            win_length=win_length,
            window=window,
            previous=previous,
        )
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.window = window
        # self.center = center
        # self.pad_mode = pad_mode

    @classmethod
    def spec2melspec(
        cls,
        sampling_rate: int,
        data: NDArrayReal,
        n_fft: int = 2048,
        fmin: float = 0.0,
        fmax: Optional[float] = None,
        n_mels: int = 128,
        # pad_mode: str = "constant",
    ) -> dict[str, Any]:
        melfb = librosa.filters.mel(
            sr=sampling_rate,
            n_fft=n_fft,
            n_mels=n_mels,
            fmin=fmin,
            fmax=fmax,
            norm=None,
        )

        melspec: NDArrayReal = np.einsum("...ft,mf->...mt", data, melfb, optimize=True)  # type: ignore[arg-type, unused-ignore]

        return dict(
            sampling_rate=sampling_rate,
            data=melspec,
            n_fft=n_fft,
            n_mels=n_mels,
        )

    def plot(
        self,
        ax: Optional["Axes"] = None,
        title: Optional[str] = None,
        db_scale: bool = True,
        fmin: Optional[float] = None,
        fmax: Optional[float] = None,
        Aw: bool = False,  # noqa: N803
        vmin: Optional[float] = None,
        vmax: Optional[float] = None,
    ) -> tuple["Axes", NDArrayReal]:
        """
        時間周波数データをプロットします。

        Parameters:
            ax (matplotlib.axes.Axes, optional): 既存のプロット軸。
            title (str, optional): プロットのタイトル。
            db_scale (bool): dBスケールでプロットするかどうか。
        """
        _ax = ax
        if _ax is None:
            _, _ax = plt.subplots(figsize=(10, 6))

        if Aw:
            data_to_plot = self.data_Aw(to_dB=True)
        elif db_scale:
            data_to_plot = self._to_db()
        else:
            data_to_plot = np.abs(self.data)

        # 時間周波数データをプロット
        img = librosa.display.specshow(
            data=data_to_plot,
            sr=self.sampling_rate,
            hop_length=self.hop_length,
            n_fft=self.n_fft,
            win_length=self.win_length,
            x_axis="time",
            y_axis="mel",
            ax=ax,
            fmin=fmin,
            fmax=fmax,
            cmap="magma",
            vmin=vmin,
            vmax=vmax,
        )

        # ラベルとタイトルを設定
        _ax.set_xlabel("Time [s]")
        _ax.set_ylabel("Frequency [Hz]")
        _ax.set_title(title or self.label or "Time-Frequency Representation")

        if _ax.figure is not None:
            if db_scale:
                cbar = _ax.figure.colorbar(img, ax=ax, format="%+2.0f")
                if Aw:
                    unit = "dBA"
                else:
                    unit = "dB"
                cbar.set_label(f"Spectrum level [{unit}]")
            else:
                cbar = _ax.figure.colorbar(img, ax=ax)
                cbar.set_label(f"Amplitude [{self.unit}]")

        if ax is None:
            plt.tight_layout()
            plt.show()

        return _ax, data_to_plot
