#include <cstdio>
#include <cufft.h>
#include <cmath>
#include <cuda_runtime.h>
#include <iostream>
#include <memory>
#include <cmath>
#include <complex>
#include <device_launch_parameters.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cuComplex.h>

#include "BaseConstVariable.h"
#include "GPUTool.cuh"
#include "GPUBPTool.cuh"
#include "BPBasic0_CUDA.cuh"



__global__ void phaseCompensationKernel(cufftComplex* phdata, const float* Freq, float r, int K, int Na) {
    int freqIdx = blockIdx.x * blockDim.x + threadIdx.x;
    int pulseIdx = blockIdx.y * blockDim.y + threadIdx.y;

    if (freqIdx >= K || pulseIdx >= Na) return;

    int idx = pulseIdx * K + freqIdx;
    float phase = 4 * PI * Freq[freqIdx] * r / c;
    float cos_phase = cosf(phase);
    float sin_phase = sinf(phase);

    cufftComplex ph = phdata[idx];
    float new_real = ph.x * cos_phase - ph.y * sin_phase;
    float new_imag = ph.x * sin_phase + ph.y * cos_phase;
    phdata[idx] = make_cuComplex(new_real, new_imag);
}

__global__ void fftshiftKernel(cufftComplex* data, int Nfft, int Np) {
    int pulse = blockIdx.x * blockDim.x + threadIdx.x;
    if (pulse >= Np) return;

    int half = Nfft / 2;
    for (int i = 0; i < half; ++i) {
        cufftComplex temp = data[pulse * Nfft + i];
        data[pulse * Nfft + i] = data[pulse * Nfft + i + half];
        data[pulse * Nfft + i + half] = temp;
    }
}

__global__ void processPulseKernel(int nx, int ny, const float* x_mat, const float* y_mat, const float* z_mat,
    float AntX, float AntY, float AntZ, float R0, float minF,
    const cufftComplex* rc_pulse, float r_start, float dr, int nR,
    cufftComplex* im_final) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= nx || y >= ny) return;

    int idx = x * ny + y;
    float dx = AntX - x_mat[idx];
    float dy = AntY - y_mat[idx];
    float dz = AntZ - z_mat[idx];
    float dR = sqrtf(dx * dx + dy * dy + dz * dz) - R0;

    // Range check
    if (dR < r_start || dR >= (r_start + dr * (nR - 1))) return;

    // Linear interpolation
    float pos = (dR - r_start) / dr;
    int index = (int)floorf(pos);
    float weight = pos - index;

    if (index < 0 || index >= nR - 1) return;

    cufftComplex rc_low = rc_pulse[index];
    cufftComplex rc_high = rc_pulse[index + 1];
    cufftComplex rc_interp;
    rc_interp.x = rc_low.x * (1 - weight) + rc_high.x * weight;
    rc_interp.y = rc_low.y * (1 - weight) + rc_high.y * weight;

    // Phase correction
    float phase = 4 * PI * minF * dR / c;
    float cos_phase = cosf(phase);
    float sin_phase = sinf(phase);

    cufftComplex phCorr;
    phCorr.x = rc_interp.x * cos_phase - rc_interp.y * sin_phase;
    phCorr.y = rc_interp.x * sin_phase + rc_interp.y * cos_phase;

    // Accumulate
    im_final[idx].x += phCorr.x;
    im_final[idx].y += phCorr.y;
}

void bpBasic0CUDA(GPUDATA& data, int flag) {
    // Phase compensation
    if (flag == 1) {
        dim3 block(16, 16);
        dim3 grid((data.K + 15) / 16, (data.Np + 15) / 16);
        phaseCompensationKernel << <grid, block >> > (data.phdata, data.Freq, data.R0, data.K, data.Np);
        cudaCheckError(cudaDeviceSynchronize());
        data.R0 = data.r;  // 假设data.r已正确设置
    }

    // FFT处理
    cufftHandle plan;
    cufftPlan1d(&plan, data.Nfft, CUFFT_C2C, data.Np);
    cufftExecC2C(plan, data.phdata, data.phdata, CUFFT_INVERSE);
    cufftDestroy(plan);

    // FFT移位
    dim3 blockShift(256);
    dim3 gridShift((data.Np + 255) / 256);
    fftshiftKernel << <gridShift, blockShift >> > (data.phdata, data.Nfft, data.Np);
    cudaCheckError(cudaDeviceSynchronize());

    // 图像重建
    float r_start = data.r_vec[0];
    float dr = (data.r_vec[data.Nfft - 1] - r_start) / (data.Nfft - 1);

    dim3 block(16, 16);
    dim3 grid((data.nx + 15) / 16, (data.ny + 15) / 16);

    for (int ii = 0; ii < data.Np; ++ii) {
        processPulseKernel << <grid, block >> > (
            data.nx, data.ny,
            data.x_mat, data.y_mat, data.z_mat,
            data.AntX[ii], data.AntY[ii], data.AntZ[ii],
            data.R0, data.minF[ii],
            data.phdata + ii * data.Nfft,
            r_start, dr, data.Nfft,
            data.im_final
            );
        cudaCheckError(cudaPeekAtLastError());
    }
    cudaCheckError(cudaDeviceSynchronize());
}


void initGPUData(GPUDATA& h_data, GPUDATA& d_data) {
    
}

void freeGPUData(GPUDATA& d_data) {
 
}



//int main() {
//    GPUDATA h_data, d_data;
//
//    // 初始化主机数据
//    h_data.Nfft = 1024;
//    h_data.K = 512;
//    // ... 其他参数初始化
//
//    // 初始化设备内存
//    initGPUData(h_data, d_data);
//
//    // 执行算法
//    bpBasic0CUDA(d_data, 0);
//
//    // 拷贝结果回主机
//    cudaCheckError(cudaMemcpy(h_data.im_final, d_data.im_final,
//        sizeof(cufftComplex) * h_data.nx * h_data.ny, cudaMemcpyDeviceToHost));
//
//    // 释放资源
//    freeGPUData(d_data);
//
//    return 0;
//}