"""Independent PyTorch implementation of Appendix C.1, equations C.1--C.7.

The jet path is forward automatic differentiation carrying value, dx, dy and
Laplacian. Parameter derivatives are reverse AD through these exact operations;
no finite differences, finite-volume targets or reference values enter the PDE.
"""
import math
from typing import Tuple

import torch
from torch import Tensor, nn
import torch.nn.functional as F


def jet_swish(z: Tensor) -> Tensor:
    a, ax, ay, al = z.unbind(0)
    s = torch.sigmoid(a)
    sp = s * (1.0 - s)
    first = s + a * sp
    second = sp * (2.0 + a * (1.0 - 2.0 * s))
    return torch.stack((a * s, first * ax, first * ay,
                        first * al + second * (ax.square() + ay.square())))


def jet_product(a: Tensor, b: Tensor) -> Tensor:
    return torch.stack((a[0] * b[0], a[1] * b[0] + a[0] * b[1],
                        a[2] * b[0] + a[0] * b[2],
                        a[3] * b[0] + a[0] * b[3]
                        + 2.0 * (a[1] * b[1] + a[2] * b[2])))


class Dense(nn.Module):
    """LeCun truncated-normal initialization; layout is [in, out]."""
    def __init__(self, nin: int, nout: int, bias: bool = True):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(nin, nout))
        sigma = (1.0 / math.sqrt(nin)) / 0.87962566103423978
        nn.init.trunc_normal_(self.weight, std=sigma, a=-2 * sigma, b=2 * sigma)
        self.bias = nn.Parameter(torch.zeros(nout)) if bias else None

    def forward(self, x: Tensor) -> Tensor:
        y = x @ self.weight
        if self.bias is not None:
            y = y + self.bias
        return y

    @torch.jit.export
    def jet(self, x: Tensor) -> Tensor:
        z = x @ self.weight
        if self.bias is not None:
            z = torch.stack((z[0] + self.bias, z[1], z[2], z[3]))
        return z


class ResidualBlock(nn.Module):
    def __init__(self, width: int):
        super().__init__()
        self.first = Dense(width, width)
        self.second = Dense(width, width)
        self.third = Dense(width, width)
        self.alpha = nn.Parameter(torch.zeros(1))

    def forward(self, h: Tensor, u: Tensor, v: Tensor) -> Tensor:
        f = F.silu(self.first(h))
        z = v + f * (u - v)
        g = F.silu(self.second(z))
        z = v + g * (u - v)
        q = F.silu(self.third(z))
        return self.alpha * q + (1.0 - self.alpha) * h

    @torch.jit.export
    def jet(self, h: Tensor, u: Tensor, v: Tensor) -> Tensor:
        f = jet_swish(self.first.jet(h))
        z = v + jet_product(f, u - v)
        g = jet_swish(self.second.jet(z))
        z = v + jet_product(g, u - v)
        q = jet_swish(self.third.jet(z))
        return self.alpha * q + (1.0 - self.alpha) * h


class PseudoTimePINN(nn.Module):
    """BFS PirateNet mapping physical (x,y) to (u,v,p)."""
    def __init__(self, width: int = 256, depth: int = 3,
                 rff_scale: float = 2.0, rff_trainable: bool = True,
                 activation: str = "swish", output_bias: bool = False,
                 length: float = 15.0, height: float = 1.0):
        super().__init__()
        if width % 2 or activation != "swish":
            raise ValueError("Table 4 BFS requires an even RFF width and Swish")
        self.register_buffer("coordinate_scale", torch.tensor([length, height]))
        self.frequencies = nn.Parameter(torch.randn(2, width // 2) * rff_scale,
                                        requires_grad=rff_trainable)
        self.gate_u = Dense(width, width)
        self.gate_v = Dense(width, width)
        self.blocks = nn.ModuleList([ResidualBlock(width) for _ in range(depth)])
        self.readout = Dense(width, 3, output_bias)

    def forward(self, xy: Tensor) -> Tensor:
        phase = (xy / self.coordinate_scale) @ self.frequencies
        h = torch.cat((torch.cos(phase), torch.sin(phase)), dim=-1)
        u, v = F.silu(self.gate_u(h)), F.silu(self.gate_v(h))
        for block in self.blocks:
            h = block(h, u, v)
        return self.readout(h)

    @torch.jit.export
    def derivatives(self, xy: Tensor) -> Tensor:
        b = self.frequencies / self.coordinate_scale[:, None]
        phase = xy @ b
        cs, sn = torch.cos(phase), torch.sin(phase)
        h0 = torch.cat((cs, sn), -1)
        hx = torch.cat((-sn * b[0], cs * b[0]), -1)
        hy = torch.cat((-sn * b[1], cs * b[1]), -1)
        bnorm = b.square().sum(0)
        hl = torch.cat((-cs * bnorm, -sn * bnorm), -1)
        h = torch.stack((h0, hx, hy, hl))
        u, v = jet_swish(self.gate_u.jet(h)), jet_swish(self.gate_v.jet(h))
        for block in self.blocks:
            h = block.jet(h, u, v)
        return self.readout.jet(h)

    @torch.jit.export
    def physics(self, xy: Tensor, viscosity: float) -> Tuple[Tensor, Tensor]:
        d = self.derivatives(xy)
        value, dx, dy, lap = d.unbind(0)
        ru = value[:, 0] * dx[:, 0] + value[:, 1] * dy[:, 0] + dx[:, 2] - viscosity * lap[:, 0]
        rv = value[:, 0] * dx[:, 1] + value[:, 1] * dy[:, 1] + dy[:, 2] - viscosity * lap[:, 1]
        rc = dx[:, 0] + dy[:, 1]
        return value, torch.stack((ru, rv, rc), dim=-1)