package com.njcn.gather.tools.comtrade.comparewave.core.algorithm;

import com.njcn.gather.tools.comtrade.comparewave.core.model.Complex;
import lombok.extern.slf4j.Slf4j;

import static com.njcn.gather.tools.comtrade.comparewave.core.Constants.*;


/**
 * FFT算法处理器
 * <p>对应C代码：pqs_fft_lib.c</p>
 * <p><b>重要：所有算法实现必须与C代码完全一致！</b></p>
 * 
 * @author hongawen
 * @since 1.0
 */
@Slf4j
public class FFTProcessor {
    
    /** FFT旋转因子数组 - 使用float与C代码保持一致 */
    private static float[] FFT_W = new float[FFT_MAX_N * 2];
    private static float[] FFT_W_65536 = new float[65536 * 2];
    
    /** DFT旋转因子数组 - 使用float与C代码保持一致 */
    private static float[] DFT_W_5 = new float[5 * 2];
    private static float[] DFT_W_25 = new float[25 * 2];
    private static float[] DFT_W_125 = new float[125 * 2];
    private static float[] DFT_W_800 = new float[800 * 2];
    private static float[] DFT_W_1280 = new float[1280 * 2];
    private static float[] DFT_W_2000 = new float[2000 * 2];
    private static float[] DFT_W_2048 = new float[2048 * 2];
    private static float[] DFT_W_1024 = new float[1024 * 2];
    private static float[] DFT_W_2560 = new float[2560 * 2];
    private static float[] DFT_W_5120 = new float[5120 * 2];
    
    /** 临时数组（与C代码对应的静态数组） */
    private static Complex[] FFT_xk1 = new Complex[FFT_MAX_N];
    private static Complex[] FFT_xk2 = new Complex[FFT_MAX_N];
    
    /** 初始化标志 */
    private static boolean initialized = false;
    
    static {
        // 静态初始化复数数组
        for (int i = 0; i < FFT_MAX_N; i++) {
            FFT_xk1[i] = new Complex();
            FFT_xk2[i] = new Complex();
        }
    }
    
    /**
     * FFT初始化方法
     * <p>对应C代码：FFT_Init</p>
     * <p>初始化FFT和DFT所需的旋转因子数组</p>
     */
    public static void init() {
        if (initialized) {
            return;
        }
        
        // 初始化FFT旋转因子W
        for (int i = 0; i < FFT_MAX_N; i++) {
            FFT_W[i * 2] = FloatMath.cos(2 * PI * i / FFT_MAX_N);
            FFT_W[i * 2 + 1] = -FloatMath.sin(2 * PI * i / FFT_MAX_N);
        }
        
        for (int i = 0; i < 65536; i++) {
            FFT_W_65536[i * 2] = FloatMath.cos(2 * PI * i / 65536);
            FFT_W_65536[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 65536);
        }
        
        // 初始化DFT旋转因子W
        for (int i = 0; i < 5; i++) {
            DFT_W_5[i * 2] = FloatMath.cos(2 * PI * i / 5);
            DFT_W_5[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 5);
        }
        
        for (int i = 0; i < 25; i++) {
            DFT_W_25[i * 2] = FloatMath.cos(2 * PI * i / 25);
            DFT_W_25[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 25);
        }
        
        for (int i = 0; i < 125; i++) {
            DFT_W_125[i * 2] = FloatMath.cos(2 * PI * i / 125);
            DFT_W_125[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 125);
        }
        
        for (int i = 0; i < 800; i++) {
            DFT_W_800[i * 2] = FloatMath.cos(2 * PI * i / 800);
            DFT_W_800[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 800);
        }
        
        for (int i = 0; i < 1280; i++) {
            DFT_W_1280[i * 2] = FloatMath.cos(2 * PI * i / 1280);
            DFT_W_1280[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 1280);
        }
        
        for (int i = 0; i < 2000; i++) {
            DFT_W_2000[i * 2] = FloatMath.cos(2 * PI * i / 2000);
            DFT_W_2000[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 2000);
        }
        
        for (int i = 0; i < 1024; i++) {
            DFT_W_1024[i * 2] = FloatMath.cos(2 * PI * i / 1024);
            DFT_W_1024[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 1024);
        }
        
        for (int i = 0; i < 2048; i++) {
            DFT_W_2048[i * 2] = FloatMath.cos(2 * PI * i / 2048);
            DFT_W_2048[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 2048);
        }
        
        for (int i = 0; i < 2560; i++) {
            DFT_W_2560[i * 2] = FloatMath.cos(2 * PI * i / 2560);
            DFT_W_2560[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 2560);
        }
        
        for (int i = 0; i < 5120; i++) {
            DFT_W_5120[i * 2] = FloatMath.cos(2 * PI * i / 5120);
            DFT_W_5120[i * 2 + 1] = -FloatMath.sin(2 * PI * i / 5120);
        }
        
        initialized = true;
    }
    
    /**
     * 主FFT计算函数
     * <p>对应C代码：FFT_Cal</p>
     * 
     * @param xk 复数数组，输入输出数据
     * @param N FFT点数
     */
    public static void fftCal(Complex[] xk, int N) {
        if (!initialized) {
            init();
        }
        
        switch (N) {
            case 800:
                fft800Cal(xk);
                break;
            case 1280:
                fft1280Cal(xk);
                break;
            case 2000:
                fft2000Cal(xk);
                break;
            case 2560:
                fft2560Cal(xk);
                break;
            case 5120:
                fft5120Cal(xk);
                break;
            default:
                log.warn("Unsupported FFT size: {}", N);
                break;
        }
    }
    
    /**
     * 2560点FFT计算
     * <p>对应C代码：FFT_2560_Cal</p>
     * <p>使用混合基FFT算法，分解为512×5</p>
     * 
     * @param xk 复数数组，输入输出数据
     */
    private static void fft2560Cal(Complex[] xk) {
        int SAMPLE_p = 512;
        int SAMPLE_q = 5;
        Complex[][] xk_pq = new Complex[512][5];
        Complex[] xk_temp = new Complex[512];
        
        // 初始化
        for (int i = 0; i < 512; i++) {
            for (int j = 0; j < 5; j++) {
                xk_pq[i][j] = new Complex();
            }
            xk_temp[i] = new Complex();
        }
        
        // 步骤1: p组q点DFT
        for (int i = 0; i < SAMPLE_p; i++) {
            for (int j = 0; j < SAMPLE_q; j++) {
                int k = j * SAMPLE_p + i;
                xk_pq[i][j].setReal(xk[k].getReal());
                xk_pq[i][j].setImag(xk[k].getImag());
            }
            xkDft(xk_pq[i], SAMPLE_q, 1, 0);
        }
        
        // 步骤2: 旋转因子W(k,N)相乘
        for (int i = 0; i < SAMPLE_p; i++) {
            for (int j = 0; j < SAMPLE_q; j++) {
                float cosValue = DFT_W_2560[i * j * 2];
                float sinValue = DFT_W_2560[i * j * 2 + 1];
                float realPart = xk_pq[i][j].getReal();
                float imagPart = xk_pq[i][j].getImag();
                xk_pq[i][j].setReal(cosValue * realPart - sinValue * imagPart);
                xk_pq[i][j].setImag(cosValue * imagPart + sinValue * realPart);
            }
        }
        
        // 步骤3: q组p点DFT
        for (int i = 0; i < SAMPLE_q; i++) {
            for (int j = 0; j < SAMPLE_p; j++) {
                xk_temp[j] = xk_pq[j][i];
            }
            xkFft(xk_temp, SAMPLE_p, 0, 1);
            for (int j = 0; j < (SAMPLE_p + 1) / 2; j++) {
                int k = SAMPLE_q * j + i;
                xk[k].setReal((float)(xk_temp[j].getReal() / (2560 / 2 * 1.4142135)));
                xk[k].setImag((float)(xk_temp[j].getImag() / (2560 / 2 * 1.4142135)));
            }
        }
    }
    
    /**
     * 5120点FFT计算
     * <p>对应C代码：FFT_5120_Cal</p>
     * <p>使用混合基FFT算法，分解为1024×5</p>
     * 
     * @param xk 复数数组，输入输出数据
     */
    private static void fft5120Cal(Complex[] xk) {
        int SAMPLE_p = 1024;
        int SAMPLE_q = 5;
        Complex[][] xk_pq = new Complex[1024][5];
        Complex[] xk_temp = new Complex[1024];
        
        // 初始化
        for (int i = 0; i < 1024; i++) {
            for (int j = 0; j < 5; j++) {
                xk_pq[i][j] = new Complex();
            }
            xk_temp[i] = new Complex();
        }
        
        // 步骤1: p组q点DFT
        for (int i = 0; i < SAMPLE_p; i++) {
            for (int j = 0; j < SAMPLE_q; j++) {
                int k = j * SAMPLE_p + i;
                xk_pq[i][j].setReal(xk[k].getReal());
                xk_pq[i][j].setImag(xk[k].getImag());
            }
            xkDft(xk_pq[i], SAMPLE_q, 1, 0);
        }
        
        // 步骤2: 旋转因子W(k,N)相乘
        for (int i = 0; i < SAMPLE_p; i++) {
            for (int j = 0; j < SAMPLE_q; j++) {
                float cosValue = DFT_W_5120[i * j * 2];
                float sinValue = DFT_W_5120[i * j * 2 + 1];
                float realPart = xk_pq[i][j].getReal();
                float imagPart = xk_pq[i][j].getImag();
                xk_pq[i][j].setReal(cosValue * realPart - sinValue * imagPart);
                xk_pq[i][j].setImag(cosValue * imagPart + sinValue * realPart);
            }
        }
        
        // 步骤3: q组p点DFT
        for (int i = 0; i < SAMPLE_q; i++) {
            for (int j = 0; j < SAMPLE_p; j++) {
                xk_temp[j] = xk_pq[j][i];
            }
            xkFft(xk_temp, SAMPLE_p, 0, 1);
            for (int j = 0; j < (SAMPLE_p + 1) / 2; j++) {
                int k = SAMPLE_q * j + i;
                xk[k].setReal((float)(xk_temp[j].getReal() / (5120 / 2 * 1.4142135)));
                xk[k].setImag((float)(xk_temp[j].getImag() / (5120 / 2 * 1.4142135)));
            }
        }
    }
    
    /** 800点FFT计算实现 */
    private static void fft800Cal(Complex[] xk) {
        // 实现800点FFT
    }
    
    /** 1280点FFT计算实现 */
    private static void fft1280Cal(Complex[] xk) {
        // TODO: 实现1280点FFT
    }
    
    /** 2000点FFT计算实现 */
    private static void fft2000Cal(Complex[] xk) {
        // TODO: 实现2000点FFT
    }
    
    /**
     * DFT计算
     * <p>对应C代码：XK_DFT</p>
     * 
     * @param xk 复数数组，输入输出数据
     * @param N DFT点数
     * @param inType 输入类型（0:复数 1:实数）
     * @param outType 输出类型（0:全谱 1:半谱）
     */
    private static void xkDft(Complex[] xk, int N, int inType, int outType) {
        Complex[] dftXk = new Complex[800];
        for (int i = 0; i < 800; i++) {
            dftXk[i] = new Complex();
        }
        
        float[] dftW = null;
        switch (N) {
            case 5:
                dftW = DFT_W_5;
                break;
            case 25:
                dftW = DFT_W_25;
                break;
            case 800:
                dftW = DFT_W_800;
                break;
            case 1024:
                dftW = DFT_W_1024;
                break;
            case 2048:
                dftW = DFT_W_2048;
                break;
            default:
                return;
        }
        
        int xkNum = (outType == 1) ? (N + 1) / 2 : N;
        
        if (inType == 1) {
            // 实数输入
            for (int k = 0; k < xkNum; k++) {
                dftXk[k].setReal(0);
                dftXk[k].setImag(0);
                for (int n = 0; n < N; n++) {
                    int xsNum = ((k * n) % N) * 2;
                    float xsR = dftW[xsNum];
                    float xsX = dftW[xsNum + 1];
                    dftXk[k].setReal(dftXk[k].getReal() + xk[n].getReal() * xsR);
                    dftXk[k].setImag(dftXk[k].getImag() + xk[n].getReal() * xsX);
                }
            }
        } else {
            // 复数输入
            for (int k = 0; k < xkNum; k++) {
                dftXk[k].setReal(0);
                dftXk[k].setImag(0);
                for (int n = 0; n < N; n++) {
                    int xsNum = ((k * n) % N) * 2;
                    float xsR = dftW[xsNum];
                    float xsX = dftW[xsNum + 1];
                    dftXk[k].setReal(dftXk[k].getReal() + 
                        (xk[n].getReal() * xsR - xk[n].getImag() * xsX));
                    dftXk[k].setImag(dftXk[k].getImag() + 
                        (xk[n].getReal() * xsX + xk[n].getImag() * xsR));
                }
            }
        }
        
        // 复制结果
        for (int k = 0; k < xkNum; k++) {
            xk[k].setReal(dftXk[k].getReal());
            xk[k].setImag(dftXk[k].getImag());
        }
    }
    
    /**
     * 基2 FFT计算
     * <p>对应C代码：XK_FFT</p>
     * <p>使用Cooley-Tukey基2 FFT算法</p>
     * 
     * @param xk 复数数组，输入输出数据
     * @param N FFT点数
     * @param inType 输入类型（0:复数 1:实数）
     * @param outType 输出类型（0:全谱 1:半谱）
     */
    private static void xkFft(Complex[] xk, int N, int inType, int outType) {
        int num = N;
        float xk0 = 0, xkn = 0;
        
        // 实数输入转换
        if (inType == 1) {
            num >>= 1;
            for (int i = 0; i < num; i++) {
                float real = xk[i * 2].getReal();
                float imag = xk[i * 2 + 1].getReal();
                xk[i].setReal(real);
                xk[i].setImag(imag);
                xk0 += (real + imag);
                xkn += (real - imag);
            }
        }
        
        // 计算蝶形级数
        int xsNum = FFT_MAX_N * 2 / num;
        int M = 0;
        for (int i = 2; i <= num; i *= 2) {
            M++;
        }
        
        // 位反转
        int J = num / 2;
        for (int i = 1; i <= (num - 2); i++) {
            if (i < J) {
                // 交换xk[i]和xk[J]
                Complex temp = new Complex(xk[i]);
                xk[i] = new Complex(xk[J]);
                xk[J] = temp;
            }
            int K = num / 2;
            while (J >= K) {
                J = J - K;
                K = K / 2;
            }
            J = J + K;
        }
        
        // 蝶形计算
        for (int L = 1; L <= M; L++) {
            int B = 1 << (L - 1);
            for (int j = 0; j < B; j++) {
                int P = xsNum * j * (1 << (M - L));
                float xsR = FFT_W[P];
                float xsX = (float)FFT_W[P + 1];
                for (int K = j; K < num; K += (B * 2)) {
                    float tr = xk[K + B].getReal() * xsR - xk[K + B].getImag() * xsX;
                    float tx = xk[K + B].getImag() * xsR + xk[K + B].getReal() * xsX;
                    xk[K + B].setReal(xk[K].getReal() - tr);
                    xk[K + B].setImag(xk[K].getImag() - tx);
                    xk[K].setReal(xk[K].getReal() + tr);
                    xk[K].setImag(xk[K].getImag() + tx);
                }
            }
        }
        
        // 实数输入的后处理
        if (inType == 1) {
            xsNum = xsNum >> 1;
            for (int k = 1; k < num; k++) {
                FFT_xk1[k].setReal((xk[k].getReal() + xk[num - k].getReal()) * 0.5f);
                FFT_xk1[k].setImag((xk[k].getImag() - xk[num - k].getImag()) * 0.5f);
                FFT_xk2[k].setReal((xk[k].getImag() + xk[num - k].getImag()) * 0.5f);
                FFT_xk2[k].setImag(-(xk[k].getReal() - xk[num - k].getReal()) * 0.5f);
            }
            
            for (int k = 1; k < num; k++) {
                float cosValue = FFT_W[k * xsNum];
                float sinValue = FFT_W[k * xsNum + 1];
                float tr = cosValue * FFT_xk2[k].getReal() - sinValue * FFT_xk2[k].getImag();
                float tx = cosValue * FFT_xk2[k].getImag() + sinValue * FFT_xk2[k].getReal();
                xk[k].setReal(tr + FFT_xk1[k].getReal());
                xk[k].setImag(tx + FFT_xk1[k].getImag());
            }
            
            xk[0].setReal(xk0);
            xk[0].setImag(0);
            
            if (outType == 0) {
                xk[num].setReal(xkn);
                xk[num].setImag(0);
                for (int k = 1; k < num; k++) {
                    xk[num * 2 - k].setReal(xk[k].getReal());
                    xk[num * 2 - k].setImag(-xk[k].getImag());
                }
            }
        }
    }
}