#!/usr/bin/env python3
"""
Debug script to compare lifting MLP alone.
"""

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


def debug_lifting_mlp():
    """Create a minimal test case for lifting MLP comparison."""
    torch.manual_seed(42)

    # Configuration matching the FNO
    in_channels = 3  # After positional encoding: data_channels(1) + 2
    width = 32
    hidden_dim = 2 * width  # 64
    height, width_spatial = 8, 8

    # Create input: [batch, in_channels, height, width]
    x = torch.randn(1, in_channels, height, width_spatial)

    # Create weights
    scale_fc1 = np.sqrt(2.0 / (in_channels + hidden_dim))
    scale_fc2 = np.sqrt(2.0 / (hidden_dim + width))

    fc1_weight = torch.randn(hidden_dim, in_channels) * scale_fc1
    fc1_bias = torch.zeros(hidden_dim)
    fc2_weight = torch.randn(width, hidden_dim) * scale_fc2
    fc2_bias = torch.zeros(width)

    print(f"Input shape: {x.shape}")
    print(f"FC1 weight shape: {fc1_weight.shape}")
    print(f"FC2 weight shape: {fc2_weight.shape}")

    # Apply lifting MLP (same as in run_fno_forward_manual)
    # Reshape for linear: [B, C, H, W] -> [B, H, W, C]
    x_perm = x.permute(0, 2, 3, 1)
    print(f"After permute: {x_perm.shape}")

    # FC1 with GELU
    x1 = torch.nn.functional.linear(x_perm, fc1_weight, fc1_bias)
    print(f"After FC1: {x1.shape}, range=[{x1.min():.4f}, {x1.max():.4f}]")

    x1_gelu = torch.nn.functional.gelu(x1)
    print(f"After GELU: {x1_gelu.shape}, range=[{x1_gelu.min():.4f}, {x1_gelu.max():.4f}]")

    # FC2
    x2 = torch.nn.functional.linear(x1_gelu, fc2_weight, fc2_bias)
    print(f"After FC2: {x2.shape}, range=[{x2.min():.4f}, {x2.max():.4f}]")

    # Reshape back: [B, H, W, C] -> [B, C, H, W]
    output = x2.permute(0, 3, 1, 2)
    print(f"Final output: {output.shape}, range=[{output.min():.4f}, {output.max():.4f}]")

    # Save for Rust comparison
    save_file({
        "input": x.contiguous(),
        "fc1_weight": fc1_weight.contiguous(),
        "fc1_bias": fc1_bias.contiguous(),
        "fc2_weight": fc2_weight.contiguous(),
        "fc2_bias": fc2_bias.contiguous(),
        "output": output.contiguous(),
    }, "/tmp/lifting_mlp_debug.safetensors")

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


if __name__ == "__main__":
    debug_lifting_mlp()