"""Independent implementation of RONOM, arXiv:2507.12814v2, Sec. 2 and SM3.

The projection solves weighted regularized least squares in a Gaussian basis;
it is not a learned Fourier operator. All geometric matrices are path independent.
"""
import math
from typing import List

import numpy as np
from scipy.special import erf
import torch
from torch import nn


def gaussian(x, centers, sigma):
    return np.exp(-0.5 * ((x[:, None] - centers[None, :]) / sigma) ** 2)


def laplacian_gram(centers, sigma):
    """Closed-form integral of phi_i'' phi_j'' on [0,1], in float64.

    Complete the square at the midpoint and integrate Gaussian moments 0,2,4.
    The recurrence follows integration by parts, without numerical quadrature.
    """
    mid = (centers[:, None] + centers[None, :]) / 2
    dist = centers[:, None] - centers[None, :]
    lo, hi = -mid, 1 - mid
    ea, eb = np.exp(-(lo / sigma)**2), np.exp(-(hi / sigma)**2)
    j0 = sigma * math.sqrt(math.pi) / 2 * (erf(hi / sigma) - erf(lo / sigma))
    j2 = sigma**2 / 2 * (j0 - hi * eb + lo * ea)
    j4 = sigma**2 / 2 * (3 * j2 - hi**3 * eb + lo**3 * ea)
    value = np.exp(-dist**2 / (4 * sigma**2)) / sigma**8 * (
        j4 - (dist**2 / 2 + 2 * sigma**2) * j2
        + (dist**2 / 4 - sigma**2)**2 * j0)
    return (value + value.T) / 2


def projection_matrix(x_in, x_measure, centers, sigma, alpha, sobolev_reg, bc_reg):
    basis = gaussian(x_in, centers, sigma)
    # Trapezoidal integral weights; all input-generalization points are uniform.
    w = np.empty(len(x_in), dtype=np.float64)
    w[0], w[-1] = (x_in[1] - x_in[0]) / 2, (x_in[-1] - x_in[-2]) / 2
    w[1:-1] = (x_in[2:] - x_in[:-2]) / 2
    edge = gaussian(np.array([0., 1.]), centers, sigma)
    # Official residual convention scales the derivative residual by sigma.
    deriv = -(np.array([0., 1.])[:, None] - centers) / sigma * edge
    residual = np.stack([edge[0] - edge[1], deriv[0] - deriv[1]])
    rhs = basis.T * w
    normal = (rhs @ basis + sobolev_reg * laplacian_gram(centers, sigma)
              + bc_reg * (residual.T @ residual) + alpha * np.eye(len(centers)))
    return gaussian(x_measure, centers, sigma) @ np.linalg.solve(normal, rhs)


class LatentFlow(nn.Module):
    """Autonomous neural ODE with the explicit fourth order 3/8 tableau."""
    def __init__(self, dimension, widths, step_size):
        super().__init__()
        layers = []
        prev = dimension
        for width in widths:
            layers += [nn.Linear(prev, width), nn.GELU()]
            prev = width
        layers.append(nn.Linear(prev, dimension))
        self.velocity = nn.Sequential(*layers)
        self.step_size = float(step_size)

    def forward(self, initial: torch.Tensor, times: torch.Tensor) -> torch.Tensor:
        state = initial
        states = torch.jit.annotate(List[torch.Tensor], [])
        states.append(state)
        for i in range(times.numel() - 1):
            # Sampling times are 0,.1,...,1. Substeps also allow a coarser query.
            interval = float(times[i + 1] - times[i])
            count = max(1, int(math.ceil(interval / self.step_size - 1e-6)))
            h = interval / count
            for _ in range(count):
                a = self.velocity(state)
                b = self.velocity(state + (h / 3.) * a)
                c = self.velocity(state + h * (b - a / 3.))
                d = self.velocity(state + h * (a - b + c))
                state = state + (h / 8.) * (a + 3. * b + 3. * c + d)
            states.append(state)
        return torch.stack(states, dim=1)


class RONOM(nn.Module):
    def __init__(self, grid_size=33, features=(64, 128, 128, 256),
                 latent_channels=4, fourier_features=32, velocity_widths=(64, 64, 64),
                 padding_size=2, alpha=1e-5, sobolev_reg=1e-8, bc_reg=10.,
                 step_size=0.1):
        super().__init__()
        self.grid_size = int(grid_size)
        self.latent_channels = int(latent_channels)
        self.alpha, self.sobolev_reg, self.bc_reg = alpha, sobolev_reg, bc_reg
        self.sigma = 1 / ((grid_size - 1) * math.sqrt(-2 * math.log(0.5)))
        self.register_buffer('frequencies', torch.randn(fourier_features, 1))
        f = list(features)
        self.lift = nn.Sequential(nn.Conv1d(1 + 2 * fourier_features, 64, 3, padding=1),
                                  nn.ReLU(), nn.Conv1d(64, f[0] // 2, 3, padding=1))
        self.output_project = nn.Sequential(
            nn.ConvTranspose1d(f[0] // 2, 64, 3, padding=1), nn.ReLU(),
            nn.ConvTranspose1d(64, 1, 3, padding=1))
        enc = []
        prev = f[0] // 2
        for width in f:
            enc += [nn.Conv1d(prev, width, 3, padding=1), nn.ReLU()]
            prev = width
        self.encoder = nn.Sequential(*enc)
        self.to_latent = nn.Conv1d(f[-1], latent_channels, 3, padding=1)
        dec = []
        back = list(reversed(f)) + [f[0] // 2]
        for left, right in zip(back[:-1], back[1:]):
            dec += [nn.ConvTranspose1d(left, right, 3, padding=1), nn.ReLU()]
        self.decoder = nn.Sequential(*dec)
        self.from_latent = nn.ConvTranspose1d(latent_channels, f[-1], 3, padding=1)
        x = np.linspace(0., 1., grid_size)
        centers = np.arange(-padding_size, grid_size + padding_size) / (grid_size - 1)
        proj = projection_matrix(x, x, centers, self.sigma, alpha, sobolev_reg, bc_reg)
        kernel = gaussian(x, x, self.sigma)
        recovery = np.linalg.solve(kernel + alpha * np.eye(grid_size), kernel.T).T
        self.register_buffer('grid', torch.tensor(x, dtype=torch.float32))
        self.register_buffer('centers', torch.tensor(centers, dtype=torch.float64))
        self.register_buffer('projection', torch.tensor(proj, dtype=torch.float32))
        self.register_buffer('recovery', torch.tensor(recovery, dtype=torch.float32))
        self.latent_flow = LatentFlow(grid_size * latent_channels, velocity_widths, step_size)

    def encode(self, initial, x_input=None):
        if initial.ndim == 3:
            initial = initial.squeeze(-1)
        if x_input is None:
            projected = initial @ self.projection.T
        else:
            matrix = projection_matrix(x_input.detach().cpu().double().numpy(),
                self.grid.detach().cpu().double().numpy(), self.centers.cpu().numpy(),
                self.sigma, self.alpha, self.sobolev_reg, self.bc_reg)
            projected = initial @ torch.as_tensor(matrix.T, device=initial.device, dtype=initial.dtype)
        phase = 2 * math.pi * self.grid[:, None] @ self.frequencies.T
        positional = torch.cat([phase.cos(), phase.sin()], dim=-1).T
        lifted_input = torch.cat([projected[:, None, :],
                                  positional[None].expand(initial.shape[0], -1, -1)], dim=1)
        return self.to_latent(self.encoder(self.lift(lifted_input)))

    def decode(self, latent, x_output=None):
        b, nt = latent.shape[:2]
        features = latent.reshape(b * nt, self.latent_channels, self.grid_size)
        values = self.output_project(self.decoder(self.from_latent(features)))
        values = values.reshape(b, nt, self.grid_size)
        if x_output is None:
            return values @ self.recovery.T
        train_x = self.grid.detach().cpu().double().numpy()
        target_x = x_output.detach().cpu().double().numpy()
        kernel = gaussian(train_x, train_x, self.sigma)
        matrix = np.linalg.solve(kernel + self.alpha * np.eye(self.grid_size),
                                  gaussian(target_x, train_x, self.sigma).T).T
        return values @ torch.as_tensor(matrix.T, device=values.device, dtype=values.dtype)

    def forward(self, initial, times, return_latent=False):
        # CPU time coordinates avoid a host/device synchronization at every RK stage.
        z0 = self.encode(initial)
        zt = self.latent_flow(z0.flatten(1), times)
        zt = zt.reshape(initial.shape[0], times.numel(), self.latent_channels, self.grid_size)
        prediction = self.decode(zt)
        if return_latent:
            return prediction, zt
        return prediction