import torch
import torch.nn as nn
import torch.fft
import numpy as np
import matplotlib.pyplot as plt

# ------------------- SPDE-based weight generator -------------------
class SPDEWeightLinear(nn.Module):
    def __init__(self, out_features, in_features, grid_size, base_kappa=1.0):
        """
        out_features, in_features: dimensions of the linear layer.
        grid_size: (H,W) of the 2D grid that will be flattened to (out_features * in_features).
        The weight matrix W is a sample from the SPDE on that grid.
        """
        super().__init__()
        self.out_f = out_features
        self.in_f  = in_features
        self.H, self.W = grid_size
        assert self.H * self.W == out_features * in_features, "grid_size must match product of features"

        # Learnable length scale (log scale for positivity)
        self.log_kappa = nn.Parameter(torch.tensor(np.log(base_kappa)))
        # Fixed noise seed for deterministic training (one random field per run)
        self.register_buffer('noise_seed', torch.randn(self.H, self.W, dtype=torch.float32))
        # Precompute wave numbers
        kx = torch.fft.fftfreq(self.H) * 2 * np.pi  # length H
        ky = torch.fft.fftfreq(self.W) * 2 * np.pi  # length W
        KX, KY = torch.meshgrid(kx, ky, indexing='ij')
        self.register_buffer('k_sq', KX**2 + KY**2)   # ω_x^2 + ω_y^2

    def forward(self, x):
        # Generate weight field from SPDE
        kappa = torch.exp(self.log_kappa)
        # Fourier transform of fixed noise
        eps_hat = torch.fft.fft2(self.noise_seed)
        # SPDE solution in Fourier domain
        f_hat = eps_hat / (kappa**2 + self.k_sq)
        # Inverse FFT -> real valued weight matrix
        W_field = torch.fft.ifft2(f_hat).real
        # Flatten to 2D weight matrix (out_features x in_features)
        W = W_field.reshape(self.out_f, self.in_f)
        return torch.nn.functional.linear(x, W)

# ------------------- Synthetic data -------------------
torch.manual_seed(42)
N = 200
x = torch.linspace(0, 1, N).unsqueeze(1)   # input 1D
y = torch.sin(5 * np.pi * x.squeeze()) + 0.1 * torch.randn(N)
x, y = x.float(), y.float().unsqueeze(1)

# ------------------- Model setup -------------------
in_dim = 1
out_dim = 1
# Choose grid size: 20x20 = 400 weights for a 1->1 linear layer? We need 1*1=1 weight,
# so let's make a tiny layer (1->1) and use a 1x1 grid (trivial). For demonstration,
# let's make a 10->10 layer (100 weights) and learn to reconstruct an identity mapping?
# Better: create a wider layer that maps 1->1 but with a larger grid to show the field.
# We'll map 1->1 but still use a 20x20 grid, flattening to (1,1) means out_f=1, in_f=1 -> 1*1=1,
# which doesn't fit grid 20x20=400. So let's map 1->1 with a weight matrix that is a scalar?
# The SPDE field would be a scalar function evaluated at one point, not interesting.
# Redesign: Let's map 1->1 but using an intermediate layer with structured weights:
# We'll create a linear layer that maps 1 -> 50, then 50 -> 1, and apply SPDE weights to the
# 50->50 hidden layer (2500 weights on a 50x50 grid). That shows the concept nicely.

hidden_size = 50
grid_size = (50, 50)   # 2500 weights

model = nn.Sequential(
    nn.Linear(1, hidden_size, bias=False),
    SPDEWeightLinear(hidden_size, hidden_size, grid_size, base_kappa=10.0),
    nn.Linear(hidden_size, 1, bias=False)
)

# Freeze the first and last layers, train only the SPDE layer's kappa
for p in model[0].parameters(): p.requires_grad = False
for p in model[2].parameters(): p.requires_grad = False

# ------------------- Training -------------------
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

for epoch in range(500):
    optimizer.zero_grad()
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        print(f"Epoch {epoch:3d}  Loss: {loss.item():.4f}  kappa: {torch.exp(model[1].log_kappa).item():.3f}")

# ------------------- Evaluate & show continuous weight field -------------------
with torch.no_grad():
    # Original grid weight field
    W_orig = model[1].forward(torch.eye(hidden_size)).detach()  # not needed; we can just get W_field
    # But easier: regenerate the field on a finer grid
    kappa = torch.exp(model[1].log_kappa).item()
    H_fine, W_fine = 200, 200
    # Build fine grid wave numbers
    kx_f = torch.fft.fftfreq(H_fine) * 2 * np.pi
    ky_f = torch.fft.fftfreq(W_fine) * 2 * np.pi
    KX_f, KY_f = torch.meshgrid(kx_f, ky_f, indexing='ij')
    k_sq_f = KX_f**2 + KY_f**2
    # Sample new noise? To keep field consistent we'd need to interpolate noise.
    # For demonstration, we can just sample a fresh field with the learned kappa.
    eps_f = torch.randn(H_fine, W_fine)
    f_hat_f = torch.fft.fft2(eps_f) / (kappa**2 + k_sq_f)
    W_field_fine = torch.fft.ifft2(f_hat_f).real

    plt.figure(figsize=(10,4))
    plt.subplot(1,2,1)
    plt.imshow(model[1].noise_seed, cmap='RdBu')
    plt.title("Fixed noise seed (50×50)")
    plt.colorbar()
    plt.subplot(1,2,2)
    plt.imshow(W_field_fine, cmap='RdBu')
    plt.title(f"Upsampled weight field (200×200)\nLearned κ = {kappa:.2f}")
    plt.colorbar()
    plt.tight_layout()
    plt.show()