package cn.iocoder.yudao.framework.common.thealgorithms.audiofilters;


/**
 * N阶IIR滤波器，假设输入已经归一化至[-1, 1]
 *
 * 基于
 * <a href="https://en.wikipedia.org/wiki/Infinite_impulse_response">Wikipedia链接</a>
 * 中的差分方程
 */
public class IIRFilter {

    private final int order;  // 滤波器阶数
    private final double[] coeffsA;  // 分母系数
    private final double[] coeffsB;  // 分子系数
    private final double[] historyX;  // 输入历史记录
    private final double[] historyY;  // 输出历史记录

    /**
     * 构造一个IIR滤波器
     *
     * @param order 滤波器的阶数
     * @throws IllegalArgumentException 如果阶数为零或小于零
     */
    public IIRFilter(int order) throws IllegalArgumentException {
        if (order < 1) {
            throw new IllegalArgumentException("阶数必须大于零");
        }
        this.order = order;  // 设置阶数
        coeffsA = new double[order + 1];  // 初始化分母系数数组
        coeffsB = new double[order + 1];  // 初始化分子系数数组

        // 设置默认值
        coeffsA[0] = 1.0;
        coeffsB[0] = 1.0;

        historyX = new double[order];  // 初始化输入历史记录数组
        historyY = new double[order];  // 初始化输出历史记录数组
    }

    /**
     * 设置系数
     *
     * @param aCoeffs 分母系数
     * @param bCoeffs 分子系数
     * @throws IllegalArgumentException 如果aCoeffs或bCeffs的大小不为order，或者aCoeffs[0]为0.0
     */
    public void setCoeffs(double[] aCoeffs, double[] bCoeffs) throws IllegalArgumentException {
        if (aCoeffs.length != order) {
            throw new IllegalArgumentException("aCoeffs的大小必须为" + order + "，当前为" + aCoeffs.length);
        }

        if (aCoeffs[0] == 0.0) {
            throw new IllegalArgumentException("aCoeffs的第一个元素不能为零");
        }

        if (bCoeffs.length != order) {
            throw new IllegalArgumentException("bCoeffs的大小必须为" + order + "，当前为" + bCoeffs.length);
        }

        for (int i = 0; i <= order; i++) {  // 设置分母和分子系数
            coeffsA[i] = aCoeffs[i];
            coeffsB[i] = bCoeffs[i];
        }
    }

    /**
     * 处理单个样本
     *
     * @param sample 要处理的样本
     * @return 处理后的样本
     */
    public double process(double sample) {
        double result = 0.0;

        // 处理
        for (int i = 1; i <= order; i++) {
            result += (coeffsB[i] * historyX[i - 1] - coeffsA[i] * historyY[i - 1]);
        }
        result = (result + coeffsB[0] * sample) / coeffsA[0];

        // 反馈
        for (int i = order - 1; i > 0; i--) {
            historyX[i] = historyX[i - 1];
            historyY[i] = historyY[i - 1];
        }

        historyX[0] = sample;
        historyY[0] = result;

        return result;
    }
}
