#!/usr/bin/env python3
"""
Debug script to compare spectral convolution alone.
"""

import torch
import numpy as np
from safetensors.torch import save_file

def debug_spectral_conv():
    """Create a minimal test case for spectral convolution comparison."""
    torch.manual_seed(42)

    # Simple configuration
    in_ch, out_ch = 2, 2
    modes_h, modes_w = 4, 4
    height, width = 16, 16

    # Create a simple input
    x = torch.randn(1, in_ch, height, width)

    # Create weights [in_ch, out_ch, modes_h, modes_w]
    scale = 0.1
    w1_real = torch.randn(in_ch, out_ch, modes_h, modes_w) * scale
    w1_imag = torch.randn(in_ch, out_ch, modes_h, modes_w) * scale
    w2_real = torch.randn(in_ch, out_ch, modes_h, modes_w) * scale
    w2_imag = torch.randn(in_ch, out_ch, modes_h, modes_w) * scale

    # Step 1: FFT
    x_ft = torch.fft.fft2(x)
    print(f"Input shape: {x.shape}")
    print(f"FFT shape: {x_ft.shape}")
    print(f"FFT[0,0,0,0]: {x_ft[0,0,0,0]}")

    # Step 2: Create complex weights
    w1 = torch.complex(w1_real, w1_imag)
    w2 = torch.complex(w2_real, w2_imag)

    # Step 3: Apply weights to frequency domain
    out_ft = torch.zeros_like(x_ft)

    # Upper portion
    x_upper = x_ft[:, :, :modes_h, :modes_w]
    out_upper = torch.einsum('bihw,iohw->bohw', x_upper, w1)
    out_ft[:, :, :modes_h, :modes_w] = out_upper

    # Lower portion
    x_lower = x_ft[:, :, -modes_h:, :modes_w]
    out_lower = torch.einsum('bihw,iohw->bohw', x_lower, w2)
    out_ft[:, :, -modes_h:, :modes_w] = out_lower

    print(f"out_ft[0,0,0,0]: {out_ft[0,0,0,0]}")
    print(f"out_ft[0,0,-1,0]: {out_ft[0,0,-1,0]}")

    # Step 4: Inverse FFT
    output = torch.fft.ifft2(out_ft).real

    print(f"\nOutput shape: {output.shape}")
    print(f"Output range: [{output.min():.6f}, {output.max():.6f}]")
    print(f"Output mean: {output.mean():.6f}")
    print(f"Output[0,0,0,:4]: {output[0,0,0,:4]}")

    # Save for Rust comparison
    save_file({
        "input": x.contiguous(),
        "w1_real": w1_real.contiguous(),
        "w1_imag": w1_imag.contiguous(),
        "w2_real": w2_real.contiguous(),
        "w2_imag": w2_imag.contiguous(),
        "output": output.contiguous(),
    }, "/tmp/spectral_conv_debug.safetensors")

    print("\nSaved debug data to /tmp/spectral_conv_debug.safetensors")

    # Also print the raw weight values for first element
    print(f"\nw1[0,0,0,0] = {w1_real[0,0,0,0]:.6f} + {w1_imag[0,0,0,0]:.6f}i")
    print(f"x_ft[0,0,0,0] = {x_ft[0,0,0,0].real:.6f} + {x_ft[0,0,0,0].imag:.6f}i")
    print(f"After einsum out_upper[0,0,0,0] = {out_upper[0,0,0,0].real:.6f} + {out_upper[0,0,0,0].imag:.6f}i")


if __name__ == "__main__":
    debug_spectral_conv()