"""
Verify that the OpenCL FFT convolution in solve_opencl / solve_opencl_batch2
produces the same result as scipy.signal.fftconvolve(f, G, mode='same').

This script simulates the OpenCL logic with numpy FFTs so we can compare
directly without needing a GPU.
"""
import numpy as np
from scipy.signal import fftconvolve

def make_green_2d(H, W, dx=1.0):
    """Reproduce PoissonSolver._create_radialGrid for 2D."""
    rdims = (2*H - 1, 2*W - 1)
    R = np.indices(rdims, dtype=np.float64)
    R[0] += 1 - H
    R[1] += 1 - W
    R[0] *= dx
    R[1] *= dx
    np.hypot(R[1], R[0], R[0])
    R = R[0]
    zval = 0.5 * (dx + dx) / 2.0
    R[H-1, W-1] = zval
    # Apply 2D Green's function: log(r)/(2*pi)
    G = np.log(R) * (0.5 / np.pi)
    return G

def test_equivalence(H, W, dx=1.0, seed=42):
    np.random.seed(seed)
    f = np.random.randn(H, W).astype(np.float64)
    G = make_green_2d(H, W, dx)
    gdx_prod = dx ** 2

    # --- CPU reference: scipy fftconvolve ---
    p_cpu = fftconvolve(f, G, mode='same') * gdx_prod

    # --- Simulate what OpenCL solve_opencl does ---
    f32 = f.astype(np.float32)
    G32 = G.astype(np.float32)
    full_shape = (2*H - 1, 2*W - 1)

    # G already has shape full_shape, no padding needed
    # f is zero-padded to full_shape at [0:H, 0:W]
    f_pad = np.zeros(full_shape, dtype=np.float32)
    f_pad[:H, :W] = f32

    # FFT at full_shape, multiply, iFFT
    F_f = np.fft.rfftn(f_pad)
    F_g = np.fft.rfftn(G32)
    conv_circ = np.fft.irfftn(F_f * F_g, s=full_shape)

    # OLD OpenCL extraction (BUGGY): same_starts = [(full_shape[i] - fshape[i]) // 2]
    same_starts_old = [(full_shape[0] - H) // 2, (full_shape[1] - W) // 2]
    p_old = conv_circ[
        same_starts_old[0]:same_starts_old[0]+H,
        same_starts_old[1]:same_starts_old[1]+W
    ].astype(np.float32) * np.float32(gdx_prod)

    # FIXED OpenCL extraction: same_starts = [s - 1 for s in fshape]
    same_starts_fixed = [H - 1, W - 1]
    p_fixed = conv_circ[
        same_starts_fixed[0]:same_starts_fixed[0]+H,
        same_starts_fixed[1]:same_starts_fixed[1]+W
    ].astype(np.float32) * np.float32(gdx_prod)

    # --- Also test with correct FFT size (3H-2) ---
    correct_fft_shape = (H + (2*H-1) - 1, W + (2*W-1) - 1)  # (3H-2, 3W-2)
    F_f_full = np.fft.rfftn(f32, s=correct_fft_shape)
    F_g_full = np.fft.rfftn(G32, s=correct_fft_shape)
    conv_full = np.fft.irfftn(F_f_full * F_g_full, s=correct_fft_shape)
    # scipy-style 'same' extraction from full linear conv
    start_full = [(correct_fft_shape[0] - H) // 2, (correct_fft_shape[1] - W) // 2]
    p_full_linear = conv_full[
        start_full[0]:start_full[0]+H,
        start_full[1]:start_full[1]+W
    ].astype(np.float32) * np.float32(gdx_prod)

    # --- Compare ---
    p_cpu32 = p_cpu.astype(np.float32)
    scale = max(np.abs(p_cpu32).max(), 1e-10)

    diff_old   = np.abs(p_cpu32 - p_old)
    diff_fixed = np.abs(p_cpu32 - p_fixed)
    diff_full  = np.abs(p_cpu32 - p_full_linear)

    print(f"=== Test: f({H},{W}), G({2*H-1},{2*W-1}), full_shape={full_shape} ===")
    print(f"  scipy fftconvolve full linear conv shape: {correct_fft_shape}")
    print(f"  p_cpu range: [{p_cpu32.min():.4f}, {p_cpu32.max():.4f}]")
    print()
    print(f"  [OLD BUG] same_starts={same_starts_old}:")
    print(f"      max_abs_diff = {diff_old.max():.6e}  rel_max = {diff_old.max()/scale:.6e}")
    print(f"  [FIXED]   same_starts={same_starts_fixed}:")
    print(f"      max_abs_diff = {diff_fixed.max():.6e}  rel_max = {diff_fixed.max()/scale:.6e}")
    print(f"  [REF]     Full linear conv (FFT size {correct_fft_shape}):")
    print(f"      max_abs_diff = {diff_full.max():.6e}  rel_max = {diff_full.max()/scale:.6e}")
    print()

    return diff_old.max(), diff_fixed.max(), diff_full.max()

if __name__ == '__main__':
    for size in [16, 32, 64, 128, 256, 384]:
        d_ocl, d_cor, d_full = test_equivalence(size, size)
        print()