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

torch.manual_seed(42)

# ------------------- Learnable SPDE weight field -------------------
class SPDEWeightField(nn.Module):
    """
    Produces a 2D weight field f(x,y) by solving the SPDE:
        (kappa^2 - Δ) f = noise
    with periodic BC. Both the noise seed and kappa are trainable.
    """
    def __init__(self, H, W, initial_kappa=10.0):
        super().__init__()
        self.H, self.W = H, W
        # Trainable log‑kappa (ensures positivity)
        self.log_kappa = nn.Parameter(torch.tensor(initial_kappa).log())
        # Trainable noise source in Fourier space (complex representation)
        # Initialised with small random values so the field is smooth at start.
        self.noise_hat = nn.Parameter(
            torch.randn(H, W, dtype=torch.complex64) * 0.1
        )
        # Wave‑number grid (fixed, not trained)
        kx = torch.fft.fftfreq(H) * 2 * torch.pi
        ky = torch.fft.fftfreq(W) * 2 * torch.pi
        KX, KY = torch.meshgrid(kx, ky, indexing='ij')
        self.register_buffer('k_sq', KX**2 + KY**2)

    def forward(self):
        kappa = torch.exp(self.log_kappa)
        # SPDE solution in Fourier domain: f_hat = noise_hat / (kappa^2 + k_sq)
        f_hat = self.noise_hat / (kappa**2 + self.k_sq)
        # Inverse FFT -> real‑valued field
        f_real = torch.fft.ifft2(f_hat).real
        return f_real

# ------------------- Generate a target smooth weight matrix -------------------
# We'll use a Gaussian bump that the SPDE field should learn to approximate.
H, W = 64, 64
x = torch.linspace(-1, 1, H)
y = torch.linspace(-1, 1, W)
X, Y = torch.meshgrid(x, y, indexing='ij')
target = torch.exp(-(X**2 + Y**2) / 0.1)  # a sharp Gaussian

# ------------------- Model, loss, optimiser -------------------
model = SPDEWeightField(H, W, initial_kappa=5.0)
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# ------------------- Training loop -------------------
n_epochs = 2000
for epoch in range(n_epochs):
    optimizer.zero_grad()
    field = model()
    loss = loss_fn(field, target)
    loss.backward()
    optimizer.step()
    if epoch % 500 == 0:
        print(f"Epoch {epoch:4d}  Loss: {loss.item():.6f}  κ: {torch.exp(model.log_kappa).item():.3f}")

# ------------------- Show result & upsample -------------------
with torch.no_grad():
    learned_field = model().detach()
    print(f"Final loss: {loss_fn(learned_field, target).item():.6f}")

    # Upsample to a finer grid (e.g., 256×256) – the “more exact” version
    H_fine, W_fine = 256, 256
    kx_f = torch.fft.fftfreq(H_fine) * 2 * torch.pi
    ky_f = torch.fft.fftfreq(W_fine) * 2 * torch.pi
    KX_f, KY_f = torch.meshgrid(kx_f, ky_f, indexing='ij')
    k_sq_f = KX_f**2 + KY_f**2

    # Interpolate the learned noise Fourier coefficients to the fine grid
    # (nearest neighbour in frequency space – simpler but works for smooth fields)
    noise_hat_fine = nn.functional.interpolate(
        model.noise_hat[None, None, ...].real, size=(H_fine, W_fine), mode='bilinear'
    ).squeeze().to(torch.complex64)

    kappa = torch.exp(model.log_kappa)
    f_hat_fine = noise_hat_fine / (kappa**2 + k_sq_f)
    field_fine = torch.fft.ifft2(f_hat_fine).real

    # Plot everything
    fig, axes = plt.subplots(1, 4, figsize=(16, 4))
    axes[0].imshow(target, cmap='inferno')
    axes[0].set_title("Target (64×64)")
    axes[1].imshow(learned_field, cmap='inferno')
    axes[1].set_title(f"Learned (64×64)\nκ={kappa.item():.2f}")
    axes[2].imshow(field_fine, cmap='inferno')
    axes[2].set_title("Upsampled (256×256)")
    # Difference on the fine grid vs downsampled target
    target_fine = nn.functional.interpolate(
        target[None, None, ...], size=(H_fine, W_fine), mode='bilinear'
    ).squeeze()
    diff = field_fine - target_fine
    axes[3].imshow(diff, cmap='RdBu', vmin=-0.1, vmax=0.1)
    axes[3].set_title("Error (256×256)")
    for ax in axes:
        ax.axis('off')
    plt.tight_layout()
    plt.show()