#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
滤波器工具类
包含双二阶滤波器计算器和各种滤波器类型的系数计算函数
"""

import math
import numpy as np
from typing import Dict, List, Union

class BiquadFilterCalculator:
    """双二阶滤波器计算器"""
    def __init__(self):
        self.coefficients = {
            'a0': 0.0,
            'a1': 0.0,
            'a2': 0.0,
            'b1': 0.0,
            'b2': 0.0
        }

    def calculate(self, filter_type: str, Fc: float, Fs: float, Q: float, peak_gain: float, bandwidth: float = None) -> Dict[str, float]:
        """
        计算滤波器系数
        :param filter_type: 滤波器类型
        :param Fc: 截止频率
        :param Fs: 采样率
        :param Q: 品质因数
        :param peak_gain: 峰值增益(dB)
        :return: 滤波器系数字典
        """
        V = math.pow(10, abs(peak_gain) / 20)
        K = math.tan(math.pi * Fc / Fs)
        norm = 0.0

        filter_type = filter_type.lower()
        
        if filter_type == "lowpass":
            norm = 1 / (1 + K / Q + K * K)
            self.coefficients['a0'] = K * K * norm
            self.coefficients['a1'] = 2 * self.coefficients['a0']
            self.coefficients['a2'] = self.coefficients['a0']
            self.coefficients['b1'] = 2 * (K * K - 1) * norm
            self.coefficients['b2'] = (1 - K / Q + K * K) * norm

        elif filter_type == "highpass":
            norm = 1 / (1 + K / Q + K * K)
            self.coefficients['a0'] = 1 * norm
            self.coefficients['a1'] = -2 * self.coefficients['a0']
            self.coefficients['a2'] = self.coefficients['a0']
            self.coefficients['b1'] = 2 * (K * K - 1) * norm
            self.coefficients['b2'] = (1 - K / Q + K * K) * norm

        elif filter_type == "bandpass":
            # 使用lib_audio_dsp中biquad.c的算法
            # 使用bandwidth参数，如果没有提供则使用Q作为近似值
            bw = bandwidth if bandwidth is not None else Q
            w0 = 2.0 * math.pi * Fc / Fs
            sin_w0 = math.sin(w0)
            alpha = sin_w0 * math.sinh(math.log(2) / 2.0 * bw * w0 / sin_w0)
            
            self.coefficients['a0'] = alpha
            self.coefficients['a1'] = 0
            self.coefficients['a2'] = -alpha
            self.coefficients['b1'] = -2.0 * math.cos(w0)
            self.coefficients['b2'] = 1.0 - alpha
            
            # 归一化
            a0 = 1.0 + alpha
            inv_a0 = 1.0 / a0
            self.coefficients['a0'] *= inv_a0
            self.coefficients['a1'] *= inv_a0
            self.coefficients['a2'] *= inv_a0
            self.coefficients['b1'] *= inv_a0
            self.coefficients['b2'] *= inv_a0

        elif filter_type == "notch":
            norm = 1 / (1 + K / Q + K * K)
            self.coefficients['a0'] = (1 + K * K) * norm
            self.coefficients['a1'] = 2 * (K * K - 1) * norm
            self.coefficients['a2'] = self.coefficients['a0']
            self.coefficients['b1'] = self.coefficients['a1']
            self.coefficients['b2'] = (1 - K / Q + K * K) * norm

        elif filter_type == "peak":
            if peak_gain >= 0:
                norm = 1 / (1 + 1/Q * K + K * K)
                self.coefficients['a0'] = (1 + V/Q * K + K * K) * norm
                self.coefficients['a1'] = 2 * (K * K - 1) * norm
                self.coefficients['a2'] = (1 - V/Q * K + K * K) * norm
                self.coefficients['b1'] = self.coefficients['a1']
                self.coefficients['b2'] = (1 - 1/Q * K + K * K) * norm
            else:
                norm = 1 / (1 + V/Q * K + K * K)
                self.coefficients['a0'] = (1 + 1/Q * K + K * K) * norm
                self.coefficients['a1'] = 2 * (K * K - 1) * norm
                self.coefficients['a2'] = (1 - 1/Q * K + K * K) * norm
                self.coefficients['b1'] = self.coefficients['a1']
                self.coefficients['b2'] = (1 - V/Q * K + K * K) * norm

        elif filter_type == "lowshelf":
            if peak_gain >= 0:
                norm = 1 / (1 + math.sqrt(2) * K + K * K)
                self.coefficients['a0'] = (1 + math.sqrt(2*V) * K + V * K * K) * norm
                self.coefficients['a1'] = 2 * (V * K * K - 1) * norm
                self.coefficients['a2'] = (1 - math.sqrt(2*V) * K + V * K * K) * norm
                self.coefficients['b1'] = 2 * (K * K - 1) * norm
                self.coefficients['b2'] = (1 - math.sqrt(2) * K + K * K) * norm
            else:
                norm = 1 / (1 + math.sqrt(2*V) * K + V * K * K)
                self.coefficients['a0'] = (1 + math.sqrt(2) * K + K * K) * norm
                self.coefficients['a1'] = 2 * (K * K - 1) * norm
                self.coefficients['a2'] = (1 - math.sqrt(2) * K + K * K) * norm
                self.coefficients['b1'] = 2 * (V * K * K - 1) * norm
                self.coefficients['b2'] = (1 - math.sqrt(2*V) * K + V * K * K) * norm

        elif filter_type == "highshelf":
            if peak_gain >= 0:
                norm = 1 / (1 + math.sqrt(2) * K + K * K)
                self.coefficients['a0'] = (V + math.sqrt(2*V) * K + K * K) * norm
                self.coefficients['a1'] = 2 * (K * K - V) * norm
                self.coefficients['a2'] = (V - math.sqrt(2*V) * K + K * K) * norm
                self.coefficients['b1'] = 2 * (K * K - 1) * norm
                self.coefficients['b2'] = (1 - math.sqrt(2) * K + K * K) * norm
            else:
                norm = 1 / (V + math.sqrt(2*V) * K + K * K)
                self.coefficients['a0'] = (1 + math.sqrt(2) * K + K * K) * norm
                self.coefficients['a1'] = 2 * (K * K - 1) * norm
                self.coefficients['a2'] = (1 - math.sqrt(2) * K + K * K) * norm
                self.coefficients['b1'] = 2 * (K * K - V) * norm
                self.coefficients['b2'] = (V - math.sqrt(2*V) * K + K * K) * norm

        else:
            raise ValueError("不支持的滤波器类型")

        return self.coefficients

def _check_filter_freq(filter_freq: float, fs: int) -> float:
    """检查滤波器频率是否有效"""
    if filter_freq >= fs / 2:
        raise ValueError(f"滤波器频率 ({filter_freq} Hz) 必须小于采样率的一半 ({fs/2} Hz)")
    return filter_freq

def make_biquad_bypass(fs: int) -> List[float]:
    """创建旁路滤波器系数"""
    return [1.0, 0.0, 0.0, 0.0, 0.0]

def make_biquad_lowpass(fs: int, filter_freq: float, q_factor: float) -> List[float]:
    """创建低通滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    K = math.tan(math.pi * filter_freq / fs)
    norm = 1 / (1 + K / q_factor + K * K)
    
    b0 = K * K * norm
    b1 = 2 * b0
    b2 = b0
    a1 = 2 * (K * K - 1) * norm
    a2 = (1 - K / q_factor + K * K) * norm
    
    return [b0, b1, b2, -a1, -a2]

def make_biquad_highpass(fs: int, filter_freq: float, q_factor: float) -> List[float]:
    """创建高通滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    K = math.tan(math.pi * filter_freq / fs)
    norm = 1 / (1 + K / q_factor + K * K)
    
    b0 = norm
    b1 = -2 * b0
    b2 = b0
    a1 = 2 * (K * K - 1) * norm
    a2 = (1 - K / q_factor + K * K) * norm
    
    return [b0, b1, b2, -a1, -a2]

def make_biquad_bandpass(fs: int, filter_freq: float, bandwidth: float) -> List[float]:
    """创建带通滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    w0 = 2.0 * math.pi * filter_freq / fs
    alpha = math.sin(w0) * math.sinh(math.log(2) / 2 * bandwidth * w0 / math.sin(w0))
    
    b0 = alpha
    b1 = 0.0
    b2 = -alpha
    a0 = 1.0 + alpha
    a1 = -2.0 * math.cos(w0)
    a2 = 1.0 - alpha
    
    return [b0/a0, b1/a0, b2/a0, -a1/a0, -a2/a0]

def make_biquad_bandstop(fs: int, filter_freq: float, bandwidth: float) -> List[float]:
    """创建带阻滤波器系数 - 使用lib_audio_dsp中biquad.c的算法"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    w0 = 2.0 * math.pi * filter_freq / fs
    sin_w0 = math.sin(w0)
    alpha = sin_w0 * math.sinh(math.log(2) / 2.0 * bandwidth * w0 / sin_w0)
    
    b0 = 1.0
    b1 = -2.0 * math.cos(w0)
    b2 = 1.0
    a0 = 1.0 + alpha
    a1 = b1  # 与lib_audio_dsp一致
    a2 = 1.0 - alpha
    
    return [b0/a0, b1/a0, b2/a0, -a1/a0, -a2/a0]

def make_biquad_notch(fs: int, filter_freq: float, q_factor: float) -> List[float]:
    """创建陷波滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    K = math.tan(math.pi * filter_freq / fs)
    norm = 1 / (1 + K / q_factor + K * K)
    
    b0 = (1 + K * K) * norm
    b1 = 2 * (K * K - 1) * norm
    b2 = b0
    a1 = b1
    a2 = (1 - K / q_factor + K * K) * norm
    
    return [b0, b1, b2, -a1, -a2]

def make_biquad_allpass(fs: int, filter_freq: float, q_factor: float) -> List[float]:
    """创建全通滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    K = math.tan(math.pi * filter_freq / fs)
    norm = 1 / (1 + K / q_factor + K * K)
    
    b0 = (1 - K / q_factor + K * K) * norm
    b1 = 2 * (K * K - 1) * norm
    b2 = 1.0
    a1 = b1
    a2 = b0
    
    return [b0, b1, b2, -a1, -a2]

def make_biquad_peaking(fs: int, filter_freq: float, q_factor: float, boost_db: float) -> List[float]:
    """创建峰值滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    A = math.sqrt(10 ** (boost_db / 20))
    w0 = 2.0 * math.pi * filter_freq / fs
    alpha = math.sin(w0) / (2.0 * q_factor)
    
    b0 = 1.0 + alpha * A
    b1 = -2.0 * math.cos(w0)
    b2 = 1.0 - alpha * A
    a0 = 1.0 + alpha / A
    a1 = -2.0 * math.cos(w0)
    a2 = 1.0 - alpha / A
    
    return [b0/a0, b1/a0, b2/a0, -a1/a0, -a2/a0]

def make_biquad_constant_q(fs: int, filter_freq: float, q_factor: float, boost_db: float) -> List[float]:
    """创建恒定Q值滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    V = 10 ** (boost_db / 20)
    w0 = 2.0 * math.pi * filter_freq / fs
    K = math.tan(w0 / 2)
    
    if boost_db > 0:
        b0 = 1 + V * K / q_factor + K**2
        b1 = 2 * (K**2 - 1)
        b2 = 1 - V * K / q_factor + K**2
        a0 = 1 + K / q_factor + K**2
        a1 = 2 * (K**2 - 1)
        a2 = 1 - K / q_factor + K**2
    else:
        V = 1 / V
        b0 = 1 + K / q_factor + K**2
        b1 = 2 * (K**2 - 1)
        b2 = 1 - K / q_factor + K**2
        a0 = 1 + V * K / q_factor + K**2
        a1 = 2 * (K**2 - 1)
        a2 = 1 - V * K / q_factor + K**2
    
    return [b0/a0, b1/a0, b2/a0, -a1/a0, -a2/a0]

def make_biquad_lowshelf(fs: int, filter_freq: float, q_factor: float, boost_db: float) -> List[float]:
    """创建低架滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    A = math.pow(10, boost_db / 40)
    w0 = 2.0 * math.pi * filter_freq / fs
    alpha = math.sin(w0) / 2 * math.sqrt((A + 1/A) * (1/q_factor - 1) + 2)
    
    if boost_db >= 0:
        b0 = A * ((A+1) + (A-1)*math.cos(w0) + 2*math.sqrt(A)*alpha)
        b1 = -2*A * ((A-1) + (A+1)*math.cos(w0))
        b2 = A * ((A+1) + (A-1)*math.cos(w0) - 2*math.sqrt(A)*alpha)
        a0 = (A+1) - (A-1)*math.cos(w0) + 2*math.sqrt(A)*alpha
        a1 = 2*((A-1) - (A+1)*math.cos(w0))
        a2 = (A+1) - (A-1)*math.cos(w0) - 2*math.sqrt(A)*alpha
    else:
        b0 = (A+1) - (A-1)*math.cos(w0) + 2*math.sqrt(A)*alpha
        b1 = 2*((A-1) - (A+1)*math.cos(w0))
        b2 = (A+1) - (A-1)*math.cos(w0) - 2*math.sqrt(A)*alpha
        a0 = A*((A+1) + (A-1)*math.cos(w0) + 2*math.sqrt(A)*alpha)
        a1 = -2*A*((A-1) + (A+1)*math.cos(w0))
        a2 = A*((A+1) + (A-1)*math.cos(w0) - 2*math.sqrt(A)*alpha)
    
    return [b0/a0, b1/a0, b2/a0, -a1/a0, -a2/a0]

def make_biquad_highshelf(fs: int, filter_freq: float, q_factor: float, boost_db: float) -> List[float]:
    """创建高架滤波器系数"""
    filter_freq = _check_filter_freq(filter_freq, fs)
    A = math.pow(10, boost_db / 40)  # 将dB转换为线性增益
    w0 = 2.0 * math.pi * filter_freq / fs  # 归一化角频率
    alpha = math.sin(w0) / 2 * math.sqrt((A + 1/A) * (1/q_factor - 1) + 2)
    
    if boost_db >= 0:  # 增益提升
        b0 = A*((A+1) + (A-1)*math.cos(w0) + 2*math.sqrt(A)*alpha)
        b1 = -2*A*((A-1) + (A+1)*math.cos(w0))
        b2 = A*((A+1) + (A-1)*math.cos(w0) - 2*math.sqrt(A)*alpha)
        a0 = (A+1) - (A-1)*math.cos(w0) + 2*math.sqrt(A)*alpha
        a1 = 2*((A-1) - (A+1)*math.cos(w0))
        a2 = (A+1) - (A-1)*math.cos(w0) - 2*math.sqrt(A)*alpha
    else:  # 增益衰减
        b0 = (A+1) - (A-1)*math.cos(w0) + 2*math.sqrt(A)*alpha
        b1 = 2*((A-1) - (A+1)*math.cos(w0))
        b2 = (A+1) - (A-1)*math.cos(w0) - 2*math.sqrt(A)*alpha
        a0 = A*((A+1) + (A-1)*math.cos(w0) + 2*math.sqrt(A)*alpha)
        a1 = -2*A*((A-1) + (A+1)*math.cos(w0))
        a2 = A*((A+1) + (A-1)*math.cos(w0) - 2*math.sqrt(A)*alpha)
    
    return coeffs 