#!/usr/bin/env python3
"""
Debug script to trace through FNO forward pass step by step.
Saves intermediate outputs at each stage for comparison with Rust.
"""

import json
import numpy as np
import torch
from pathlib import Path
from safetensors.torch import save_file


def create_grid_positional_encoding(height: int, width: int) -> torch.Tensor:
    """Create 2D grid positional encoding."""
    y = torch.linspace(-1, 1, height)
    x = torch.linspace(-1, 1, width)
    grid_y, grid_x = torch.meshgrid(y, x, indexing='ij')
    grid = torch.stack([grid_x, grid_y], dim=0)
    return grid


def debug_fno_step_by_step():
    """Run FNO forward pass with intermediate outputs."""
    torch.manual_seed(42)

    # Simple config
    batch_size = 1
    data_channels = 1
    height = 32
    width = 32
    model_width = 16  # Smaller for faster debugging
    n_modes = (8, 8)
    n_layers = 2  # Fewer layers

    # Create deterministic input
    y = torch.linspace(0, 2 * np.pi, height)
    x = torch.linspace(0, 2 * np.pi, width)
    yy, xx = torch.meshgrid(y, x, indexing='ij')
    input_tensor = torch.sin(xx) * torch.cos(yy)
    input_tensor = input_tensor.unsqueeze(0).unsqueeze(0)  # [1, 1, H, W]

    print(f"Input: shape={input_tensor.shape}, range=[{input_tensor.min():.4f}, {input_tensor.max():.4f}]")

    # Step 1: Positional encoding
    grid = create_grid_positional_encoding(height, width)
    grid = grid.unsqueeze(0).expand(batch_size, -1, -1, -1)
    with_pos = torch.cat([input_tensor, grid], dim=1)
    print(f"After pos encoding: shape={with_pos.shape}, range=[{with_pos.min():.4f}, {with_pos.max():.4f}]")

    # Create weights
    in_with_pos = data_channels + 2
    hidden_dim = 2 * model_width

    # Lifting weights
    scale_fc1 = np.sqrt(2.0 / (in_with_pos + hidden_dim))
    scale_fc2 = np.sqrt(2.0 / (hidden_dim + model_width))
    lifting_fc1_weight = torch.randn(hidden_dim, in_with_pos) * scale_fc1
    lifting_fc1_bias = torch.zeros(hidden_dim)
    lifting_fc2_weight = torch.randn(model_width, hidden_dim) * scale_fc2
    lifting_fc2_bias = torch.zeros(model_width)

    # Step 2: Lifting MLP
    x = with_pos.permute(0, 2, 3, 1)  # [B, H, W, C]
    x = torch.nn.functional.linear(x, lifting_fc1_weight, lifting_fc1_bias)
    x = torch.nn.functional.gelu(x)
    x = torch.nn.functional.linear(x, lifting_fc2_weight, lifting_fc2_bias)
    x = x.permute(0, 3, 1, 2)  # [B, C, H, W]
    print(f"After lifting: shape={x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")
    after_lifting = x.clone()

    # Spectral and skip weights
    spectral_weights = []
    skip_weights = []
    scale_spectral = 1.0 / (model_width * n_modes[0] * n_modes[1])
    scale_conv = np.sqrt(2.0 / (model_width + model_width))

    for i in range(n_layers):
        sw = {
            'w1_real': torch.randn(model_width, model_width, n_modes[0], n_modes[1]) * scale_spectral,
            'w1_imag': torch.randn(model_width, model_width, n_modes[0], n_modes[1]) * scale_spectral,
            'w2_real': torch.randn(model_width, model_width, n_modes[0], n_modes[1]) * scale_spectral,
            'w2_imag': torch.randn(model_width, model_width, n_modes[0], n_modes[1]) * scale_spectral,
        }
        spectral_weights.append(sw)
        skip_weights.append({
            'weight': torch.randn(model_width, model_width) * scale_conv,
            'bias': torch.zeros(model_width),
        })

    # Step 3: Fourier layers
    for i in range(n_layers):
        # Spectral conv path
        x_ft = torch.fft.fft2(x)

        sw = spectral_weights[i]
        w1 = torch.complex(sw['w1_real'], sw['w1_imag'])
        w2 = torch.complex(sw['w2_real'], sw['w2_imag'])

        modes_h, modes_w = n_modes
        out_ft = torch.zeros_like(x_ft)

        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

        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

        x_spectral = torch.fft.ifft2(out_ft).real
        print(f"  Layer {i} spectral: range=[{x_spectral.min():.4f}, {x_spectral.max():.4f}]")

        # Skip connection path
        skip_w = skip_weights[i]['weight']
        skip_b = skip_weights[i]['bias']
        x_skip = x.permute(0, 2, 3, 1)
        x_skip = torch.nn.functional.linear(x_skip, skip_w, skip_b)
        x_skip = x_skip.permute(0, 3, 1, 2)
        print(f"  Layer {i} skip: range=[{x_skip.min():.4f}, {x_skip.max():.4f}]")

        # Combine
        x = x_spectral + x_skip
        print(f"  Layer {i} combined: range=[{x.min():.4f}, {x.max():.4f}]")

        if i < n_layers - 1:
            x = torch.nn.functional.gelu(x)
            print(f"  Layer {i} after GELU: range=[{x.min():.4f}, {x.max():.4f}]")

    after_fourier = x.clone()
    print(f"After all Fourier layers: shape={x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")

    # Step 4: Projection
    proj_hidden = 128
    scale_proj1 = np.sqrt(2.0 / (model_width + proj_hidden))
    scale_proj2 = np.sqrt(2.0 / (proj_hidden + 1))

    proj1_weight = torch.randn(proj_hidden, model_width) * scale_proj1
    proj1_bias = torch.zeros(proj_hidden)
    proj2_weight = torch.randn(1, proj_hidden) * scale_proj2
    proj2_bias = torch.zeros(1)

    x = x.permute(0, 2, 3, 1)
    x = torch.nn.functional.linear(x, proj1_weight, proj1_bias)
    x = torch.nn.functional.gelu(x)
    x = torch.nn.functional.linear(x, proj2_weight, proj2_bias)
    x = x.permute(0, 3, 1, 2)

    print(f"Final output: shape={x.shape}, range=[{x.min():.4f}, {x.max():.4f}]")
    print(f"Final output mean: {x.mean():.6f}")

    # Save all data
    output_dir = Path("/tmp/fno_debug")
    output_dir.mkdir(exist_ok=True)

    tensors = {
        "input": input_tensor.contiguous(),
        "after_pos_encoding": with_pos.contiguous(),
        "lifting_fc1_weight": lifting_fc1_weight.contiguous(),
        "lifting_fc1_bias": lifting_fc1_bias.contiguous(),
        "lifting_fc2_weight": lifting_fc2_weight.contiguous(),
        "lifting_fc2_bias": lifting_fc2_bias.contiguous(),
        "after_lifting": after_lifting.contiguous(),
        "after_fourier": after_fourier.contiguous(),
        "proj1_weight": proj1_weight.contiguous(),
        "proj1_bias": proj1_bias.contiguous(),
        "proj2_weight": proj2_weight.contiguous(),
        "proj2_bias": proj2_bias.contiguous(),
        "output": x.contiguous(),
    }

    # Add spectral and skip weights
    for i in range(n_layers):
        sw = spectral_weights[i]
        for key, tensor in sw.items():
            tensors[f"spectral_{i}_{key}"] = tensor.contiguous()
        tensors[f"skip_{i}_weight"] = skip_weights[i]['weight'].contiguous()
        tensors[f"skip_{i}_bias"] = skip_weights[i]['bias'].contiguous()

    save_file(tensors, output_dir / "debug.safetensors")

    # Save config
    config = {
        "batch_size": batch_size,
        "data_channels": data_channels,
        "height": height,
        "width": width,
        "model_width": model_width,
        "n_modes": n_modes,
        "n_layers": n_layers,
    }
    with open(output_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    print(f"\nSaved debug data to {output_dir}")


if __name__ == "__main__":
    debug_fno_step_by_step()