package com.njcn.advance.event.cause.algorithm;

import org.apache.commons.math3.complex.Complex;
import org.apache.commons.math3.transform.DftNormalization;
import org.apache.commons.math3.transform.FastFourierTransformer;
import org.apache.commons.math3.transform.TransformType;

/**
 * FFT工具类
 * 使用Apache Commons Math实现FFT变换
 */
public class FFTUtils {
    
    private static final FastFourierTransformer transformer = 
        new FastFourierTransformer(DftNormalization.STANDARD);

    /**
     * 执行FFT变换
     * @param input 输入实数数组
     * @return 复数结果数组
     */
    public static Complex[] fft(float[] input) {
        // 确保输入长度是2的幂
        int n = nextPowerOfTwo(input.length);
        double[] paddedInput = new double[n];
        
        // 复制输入数据并用零填充
        for (int i = 0; i < input.length; i++) {
            paddedInput[i] = input[i];
        }
        for (int i = input.length; i < n; i++) {
            paddedInput[i] = 0.0;
        }
        
        return transformer.transform(paddedInput, TransformType.FORWARD);
    }

    /**
     * 执行FFT变换（复数输入）
     * @param input 输入复数数组
     * @return 复数结果数组
     */
    public static Complex[] fft(Complex[] input) {
        // 确保输入长度是2的幂
        int n = nextPowerOfTwo(input.length);
        Complex[] paddedInput = new Complex[n];
        
        // 复制输入数据并用零填充
        System.arraycopy(input, 0, paddedInput, 0, input.length);
        for (int i = input.length; i < n; i++) {
            paddedInput[i] = Complex.ZERO;
        }
        
        return transformer.transform(paddedInput, TransformType.FORWARD);
    }

    /**
     * 执行IFFT逆变换
     * @param input 输入复数数组
     * @return 复数结果数组
     */
    public static Complex[] ifft(Complex[] input) {
        return transformer.transform(input, TransformType.INVERSE);
    }

    /**
     * 计算复数数组的模
     * @param complexArray 复数数组
     * @param output 输出模值数组
     * @param harmonicCount 需要计算的谐波个数
     * @param N FFT点数
     */
    public static void calculateMagnitude(Complex[] complexArray, float[] output, 
                                        int harmonicCount, int N) {
        int count = Math.min(harmonicCount, output.length);
        count = Math.min(count, complexArray.length);
        
        for (int i = 0; i < count; i++) {
            double magnitude = complexArray[i].abs();
            // 归一化处理，与C代码保持一致
            output[i] = (float) (magnitude / (N / 2.0 * Math.sqrt(2)));
        }
    }

    /**
     * 找到下一个2的幂
     * @param n 输入数字
     * @return 大于等于n的最小2的幂
     */
    private static int nextPowerOfTwo(int n) {
        if (n <= 0) return 1;
        if ((n & (n - 1)) == 0) return n; // 已经是2的幂
        
        int power = 1;
        while (power < n) {
            power <<= 1;
        }
        return power;
    }

    /**
     * 创建复数数组（从实数数组）
     * @param realArray 实数数组
     * @return 复数数组
     */
    public static Complex[] createComplexArray(float[] realArray) {
        Complex[] complexArray = new Complex[realArray.length];
        for (int i = 0; i < realArray.length; i++) {
            complexArray[i] = new Complex(realArray[i], 0.0);
        }
        return complexArray;
    }

    /**
     * 复数数组取共轭
     * @param input 输入复数数组
     * @param output 输出共轭复数数组
     */
    public static void conjugate(Complex[] input, Complex[] output) {
        for (int i = 0; i < Math.min(input.length, output.length); i++) {
            output[i] = input[i].conjugate();
        }
    }
}