package com.yfd.platform.utils;

import org.apache.poi.ss.usermodel.Cell;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.jtransforms.fft.DoubleFFT_1D;

import javax.sound.sampled.AudioFormat;
import javax.sound.sampled.AudioInputStream;
import javax.sound.sampled.AudioSystem;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.List;

public class AudioSpectrumAnalyzer {

    public static void main(String[] args) throws Exception {
        File audioFile = new File("D:\\riis\\video\\010203040506-02-20250313094706.wav");
        AudioData audioData = readAudioData(audioFile);
        SpectrumData spectrumData = calculateSpectrum(audioData);
        String[] headers = {"Hz", "dB"};
        // 写入excel
        writerExcel(headers, spectrumData);

    }

    // 读取音频数据
    static AudioData readAudioData(File file) throws Exception {
        AudioInputStream audioStream = AudioSystem.getAudioInputStream(file);
        AudioFormat format = audioStream.getFormat();

        // 转换为PCM格式
        if (format.getEncoding() != AudioFormat.Encoding.PCM_SIGNED) {
            AudioFormat newFormat = new AudioFormat(
                    AudioFormat.Encoding.PCM_SIGNED,
                    format.getSampleRate(),
                    16,
                    format.getChannels(),
                    format.getChannels() * 2,
                    format.getSampleRate(),
                    false);
            audioStream = AudioSystem.getAudioInputStream(newFormat, audioStream);
            format = newFormat;
        }

        // 2. 读取PCM数据
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        byte[] buffer = new byte[4096];
        int bytesRead;
        try {
            while ((bytesRead = audioStream.read(buffer)) != -1) {
                baos.write(buffer, 0, bytesRead);
            }
        } finally {
            audioStream.close();
        }
        byte[] bytes = baos.toByteArray();
        audioStream.close();

        // 转换为double数组（处理16-bit样本）
        int sampleSize = format.getSampleSizeInBits() / 8;
        double[] samples = new double[bytes.length / sampleSize];
        ByteBuffer bb = ByteBuffer.wrap(bytes)
                .order(format.isBigEndian() ? ByteOrder.BIG_ENDIAN : ByteOrder.LITTLE_ENDIAN);

        for (int i = 0; i < samples.length; i++) {
            samples[i] = bb.getShort() / 32768.0; // 16-bit PCM归一化
        }

        return new AudioData(samples, format.getSampleRate());
    }

    // 计算频谱
    static SpectrumData calculateSpectrum(AudioData audioData) {
        double[] samples = audioData.samples;
        double sampleRate = audioData.sampleRate;
        int n = samples.length;

        // 应用汉明窗
        applyHammingWindow(samples);

        // 执行FFT
        DoubleFFT_1D fft = new DoubleFFT_1D(n);
        double[] fftData = new double[n * 2];
        System.arraycopy(samples, 0, fftData, 0, n);
        fft.realForwardFull(fftData);

        // 计算幅度和分贝
        List<Double> frequencies = new ArrayList<>();
        List<Double> amplitudesDb = new ArrayList<>();
        double maxAmplitude = 0;

        for (int k = 0; k < n / 2; k++) {
            double real = fftData[2 * k];
            double imag = fftData[2 * k + 1];
            double magnitude = Math.sqrt(real * real + imag * imag);
            // 归一化分贝值
            double db = 20 * Math.log10(magnitude / n);

            double frequency = k * sampleRate / n;
            frequencies.add(frequency);
            amplitudesDb.add(db);
        }

        // 转换为数组
        SpectrumData result = new SpectrumData();
        result.frequencies = frequencies.stream().mapToDouble(d -> d).toArray();
        result.amplitudesDb = amplitudesDb.stream().mapToDouble(d -> d).toArray();

        return result;
    }

    public static void writerExcel(String[] headers, SpectrumData spectrumData) {
        // 创建Excel工作簿
        try (Workbook workbook = new XSSFWorkbook()) {
            Sheet sheet = workbook.createSheet("数据表");

            // 写入标题行
            Row headerRow = sheet.createRow(0);
            for (int i = 0; i < headers.length; i++) {
                Cell cell = headerRow.createCell(i);
                cell.setCellValue(headers[i]);
            }
            // 写入数据行
            for (int i = 0; i < spectrumData.frequencies.length; i++) {
                Row row = sheet.createRow(i + 1);
                Cell cell = row.createCell(0);
                cell.setCellValue(spectrumData.frequencies[i]);
                Cell cell1 = row.createCell(1);
                cell1.setCellValue(spectrumData.amplitudesDb[i]);
            }

            // 保存到文件
            try (FileOutputStream outputStream = new FileOutputStream("D:\\output.xlsx")) {
                workbook.write(outputStream);
                System.out.println("Excel文件已保存！");
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // 应用汉明窗
    static void applyHammingWindow(double[] data) {
        int n = data.length;
        for (int i = 0; i < n; i++) {
            data[i] *= 0.54 - 0.46 * Math.cos(2 * Math.PI * i / (n - 1));
        }
    }

    // 数据载体类
    static class AudioData {

        double[] samples;
        double sampleRate;

        AudioData(double[] samples, double sampleRate) {
            this.samples = samples;
            this.sampleRate = sampleRate;
        }
    }

    static class SpectrumData {

        double[] frequencies;
        double[] amplitudesDb;
    }
}
