import numpy as np
import torch
from typing import cast
from torch import fft
from torch.nn import functional as F

from zyplib.utils.ensure import ensure_3dims


def _periodogram(X: torch.Tensor, fs, detrend, scaling):
    """内部函数：使用 pytorch 计算周期图

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入信号
    - `fs` : `int`
        - 采样率
    - `detrend` : `bool`
        - 是否去趋势
    - `scaling` : `str`
        - 缩放方式

    Returns
    ----------
    - `tuple[torch.Tensor, torch.Tensor]`
        - 频率和功率谱
    """
    if X.dim() > 2:
        X = torch.squeeze(X)
    elif X.dim() == 1:
        X = X.unsqueeze(0)

    if detrend:
        X -= X.mean(-1, keepdim=True)

    N = X.size(-1)
    assert N % 2 == 0

    df = fs / N
    dt = df
    f = torch.arange(0, N / 2 + 1) * df  # [0:df:f/2]

    dual_side = fft.fft(X)  # 双边谱
    half_idx = int(N / 2 + 1)
    single_side = dual_side[:, 0:half_idx]
    win = torch.abs(single_side)

    ps = win**2
    if scaling == 'density':  # 计算功率谱密度
        scale = N * fs
    elif scaling == 'spectrum':  # 计算功率谱
        scale = N**2
    elif scaling is None:  # 不做缩放
        scale = 1
    else:
        raise ValueError('Unknown scaling: %r' % scaling)
    Pxy = ps / scale

    Pxy[:, 1:-1] *= 2  # 能量2倍;直流分量不用二倍, 中心频率点不用二倍

    return f, Pxy.squeeze()


def periodogram(X: torch.Tensor, fs=256, detrend=True, scaling='density', no_grad=True):
    """使用 pytorch 计算信号单边 PSD

    基本等价于 scipy.signal.periodogram，计算信号的功率谱密度或功率谱。

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - EEG 信号, 形状为 [T] 或 [N, T]
    - `fs` : `int`, optional
        - 采样率 (Hz), by default 256
    - `detrend` : `bool`, optional
        - 是否去趋势 (去除直流分量), by default True
    - `scaling` : `str`, optional
        - 缩放方式, by default 'density'
        - 'density': 计算功率谱密度 (V ** 2 / Hz)
        - 'spectrum': 计算功率谱 (V ** 2)
    - `no_grad` : `bool`, optional
        - 是否启用 no_grad() 模式, by default True

    Returns
    ----------
    - `tuple[torch.Tensor, torch.Tensor]`
        - 频率和单边功率谱
    """
    if no_grad:
        with torch.no_grad():
            return _periodogram(X, fs, detrend, scaling)
    else:
        return _periodogram(X, fs, detrend, scaling)


def _get_window(window, nwlen, device):
    """内部函数：获取窗函数

    Parameters
    ----------
    - `window` : `str`
        - 窗函数类型
    - `nwlen` : `int`
        - 窗长度
    - `device` : `torch.device`
        - 设备

    Returns
    ----------
    - `torch.Tensor`
        - 窗函数

    Raises
    ----------
    - `ValueError`
        - 当窗函数类型无效时抛出
    """
    if window == 'hann':
        window = torch.hann_window(
            nwlen, dtype=torch.float32, device=device, periodic=False
        )
    elif window == 'hamming':
        window = torch.hamming_window(
            nwlen, dtype=torch.float32, device=device, periodic=False
        )
    elif window == 'blackman':
        window = torch.blackman_window(
            nwlen, dtype=torch.float32, device=device, periodic=False
        )
    elif window == 'boxcar':
        window = torch.ones(nwlen, dtype=torch.float32, device=device)
    else:
        raise ValueError('Invalid Window {}' % window)
    return window


def _pwelch(X: torch.Tensor, fs, detrend, scaling, window, nwlen, nhop):
    """内部函数：Welch 方法计算功率谱

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入信号
    - `fs` : `int`
        - 采样率
    - `detrend` : `bool`
        - 是否去趋势
    - `scaling` : `str`
        - 缩放方式
    - `window` : `torch.Tensor`
        - 窗函数
    - `nwlen` : `int`
        - 窗长度
    - `nhop` : `int`
        - 窗移动步长

    Returns
    ----------
    - `tuple[torch.Tensor, torch.Tensor]`
        - 频率和功率谱
    """
    X = cast(torch.Tensor, ensure_3dims(X, newaxis='channel'))
    if scaling == 'density':
        scale = fs * (window * window).sum().item()
    elif scaling == 'spectrum':
        scale = window.sum().item() ** 2
    else:
        raise ValueError('Unknown scaling: %r' % scaling)
    # --------------- Fold and windowing --------------- #
    N, T = X.size(0), X.size(-1)
    X = X.view(N, 1, 1, T)
    X_fold = F.unfold(X, (1, nwlen), stride=nhop)  # [N, 1, 1, T] -> [N, nwlen, win_cnt]
    if detrend:
        X_fold -= X_fold.mean(1, keepdim=True)  # 各个窗口各自detrend
    window = window.view(1, -1, 1)  # [1, nwlen, 1]
    X_windowed = X_fold * window  # [N, nwlen, win_cnt]
    win_cnt = X_windowed.size(-1)

    # --------------- Pwelch --------------- #
    X_windowed = X_windowed.transpose(1, 2).contiguous()  # [N, win_cnt, nwlen]
    X_windowed = X_windowed.view(N * win_cnt, nwlen)  # [N * win_cnt, nwlen]
    f, pxx = _periodogram(
        X_windowed, fs, detrend=False, scaling=None
    )  # [N * win_cnt, nwlen // 2 + 1]
    pxx /= scale
    pxx = pxx.view(N, win_cnt, -1)  # [N, win_cnt, nwlen // 2 + 1]
    pxx = torch.mean(pxx, dim=1)  # [N, nwlen // 2 + 1]
    return f, pxx


def pwelch(
    X: torch.Tensor,
    fs=256,
    detrend=True,
    scaling='density',
    window='hann',
    nwlen=128,
    nhop=None,
    noverlap=None,
    no_grad=True,
):
    """使用 pytorch 计算功率谱

    大致相当于 scipy.signal.welch，使用 Welch 方法计算信号的功率谱密度或功率谱。

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - EEG 信号, 形状为 [T] 或 [N, T]
    - `fs` : `int`, optional
        - 采样率 (Hz), by default 256
    - `detrend` : `bool`, optional
        - 是否去趋势 (去除直流分量), by default True
    - `scaling` : `str`, optional
        - 缩放方式, by default 'density'
        - 'density': 计算功率谱密度 (V ** 2 / Hz)
        - 'spectrum': 计算功率谱 (V ** 2)
    - `window` : `str`, optional
        - 窗函数名称, by default 'hann'
    - `nwlen` : `int`, optional
        - 窗函数长度 (点的个数), by default 128
    - `nhop` : `int`, optional
        - 窗函数移动步长, 即 nwlen - noverlap (点的个数)
        - 如果为 None 且 noverlap 为 None，则默认为 nwlen // 2
    - `noverlap` : `int`, optional
        - 仅当 nhop 为 None 时有效
    - `no_grad` : `bool`, optional
        - 是否启用 no_grad() 模式, by default True

    Returns
    ----------
    - `tuple[torch.Tensor, torch.Tensor]`
        - (频率, 单边功率谱)
    """
    if nhop is None:
        if noverlap is not None:
            nhop = nwlen - noverlap
        else:
            nhop = nwlen // 2
    window = _get_window(window, nwlen, X.device)
    if no_grad:
        with torch.no_grad():
            return _pwelch(X, fs, detrend, scaling, window, nwlen, nhop)
    else:
        return _pwelch(X, fs, detrend, scaling, window, nwlen, nhop)


def _stft(X: torch.Tensor, fs, window, nhop, nfft):
    """内部函数：短时傅里叶变换

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入信号
    - `fs` : `int`
        - 采样率
    - `window` : `torch.Tensor`
        - 窗函数
    - `nhop` : `int`
        - 窗移动步长
    - `nfft` : `int`
        - FFT 长度

    Returns
    ----------
    - `tuple[torch.Tensor, torch.Tensor, torch.Tensor]`
        - (时间点, 频率点, 功率谱)
    """
    nfft = window.size(-1)
    # [batch, freq_points, time_points]
    stft_ans = torch.stft(
        X, n_fft=nfft, window=window, hop_length=nhop, return_complex=True
    )
    pxx = torch.abs(stft_ans) ** 2
    f = torch.arange(0, stft_ans.size(-2))
    ts = nhop / fs
    t = torch.arange(0, stft_ans.size(-1)) * ts
    return t, f, pxx


def stft(
    X: torch.Tensor, fs, window='hann', nwlen=256, nhop=None, nfft=512, no_grad=True
):
    """使用 pytorch 计算短时傅里叶变换 (Short Time Fourier Transform)

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - EEG 信号, 形状为 [T] 或 [N, T]
    - `fs` : `int`
        - 采样率 (Hz)
    - `window` : `str`, optional
        - 窗函数名称, by default 'hann'
    - `nwlen` : `int`, optional
        - 窗函数长度 (点的个数), by default 256
    - `nhop` : `int`, optional
        - 窗函数移动步长 (点的个数), by default None
        - 如果为 None，则默认为 nwlen // 4
    - `nfft` : `int`, optional
        - FFT 长度, by default 512
    - `no_grad` : `bool`, optional
        - 是否启用 no_grad() 模式, by default True

    Returns
    ----------
    - `tuple[torch.Tensor, torch.Tensor, torch.Tensor]`
        - (时间点, 频率点, 功率谱)
    """
    if nhop is None:
        nhop = nwlen // 4
    if nfft < nwlen:
        nfft = nwlen
    window = _get_window(window, nwlen, X.device)
    if no_grad:
        with torch.no_grad():
            return _stft(X, fs, window, nhop, nfft)
    else:
        return _stft(X, fs, window, nhop, nfft)


def sinc(x: torch.Tensor):
    """Sinc 函数

    Parameters
    ----------
    - `x` : `torch.Tensor`
        - 输入张量

    Returns
    ----------
    - `torch.Tensor`
        - Sinc 函数值
    """
    return torch.where(
        x == 0, torch.tensor(1.0, device=x.device, dtype=x.dtype), torch.sin(x) / (x)
    )


def gate_signal(cutoff, half_size=256):
    """生成门控信号

    Parameters
    ----------
    - `cutoff` : `float`
        - 截止频率
    - `half_size` : `int`, optional
        - 半窗长度, by default 256

    Returns
    ----------
    - `torch.Tensor`
        - 门控信号
    """
    window = torch.hann_window(2 * half_size, periodic=False)
    time = torch.arange(-half_size, half_size)
    gate = 2 * cutoff * window * sinc(2 * cutoff * np.pi * time)
    gate /= gate.sum()
    return gate
