# coding=utf-8
#
# SPDX-License-Identifier: Apache-2.0
#
# Minimal reproduction of FGN (Functional Generative Networks, Google DeepMind
# "WeatherNext2", arXiv 2506.14285, 2025) following the paper's architecture:
#
#   * Grid encoder: a GNN that maps the two-prior-state gridded input onto a
#     latent mesh (a coarse regular lat/lon grid).
#   * Processor: a graph-transformer that operates on the latent mesh nodes
#     with conditional layer-norm layers.
#   * Grid decoder: a GNN that maps the latent mesh back onto the target grid.
#
# The probabilistic core of FGN is preserved:
#   * a global noise vector n ~ N(0, I)^32 is sampled per ensemble member and
#     per autoregressive step, embedded by a single matrix multiplication and
#     passed into *all* conditional layer-norm layers (learned functional
#     perturbations). This models aleatoric uncertainty.
#   * epistemic uncertainty is modelled by an ensemble of independently
#     trained models (deep ensembles); the mini constant model here uses one
#     seed by default (see conf/config.yaml).
#   * training objective is the fair CRPS estimator (Eq. 4) with N=2 samples.
#
# Differences from the paper (documented in README.md): the paper uses a
# spherical 6-times-refined icosahedral mesh and a full 768-latent / 24-layer
# / 6-head processor (about 180M params per seed); here the latent mesh is a
# fixed regular grid with a wrap-around neighbor graph, and the hyper
# parameters are reduced to CPU-friendly sizes for connectivity validation.
# sampled-per-step noise inside a single model plus the AR rollout are kept.
import math

import torch
import torch.nn as nn
import torch.nn.functional as F


def _latlon_grid(shape):
    """Regular lat/lon coordinates for a (H, W) grid, North-to-South rows."""
    H, W = shape
    lat = torch.linspace(90.0, -90.0, H)
    lon = torch.linspace(0.0, 360.0 - 360.0 / W, W)
    return lat, lon


def _haversine(lat1, lon1, lat2, lon2):
    """Haversine distance in metres given points in degrees."""
    R = 6371000.0
    p1 = torch.deg2rad(lat1)
    p2 = torch.deg2rad(lat2)
    dp = torch.deg2rad(lat2 - lat1)
    dl = torch.deg2rad(lon2 - lon1)
    a = torch.sin(dp / 2) ** 2 + torch.cos(p1) * torch.cos(p2) * torch.sin(dl / 2) ** 2
    return 2 * R * torch.asin(torch.sqrt(a.clamp(0, 1)))


def _bearing(lat1, lon1, lat2, lon2):
    """Initial forward bearing in radians from point 1 to point 2."""
    p1 = torch.deg2rad(lat1)
    p2 = torch.deg2rad(lat2)
    dl = torch.deg2rad(lon2 - lon1)
    y = torch.sin(dl) * torch.cos(p2)
    x = torch.cos(p1) * torch.sin(p2) - torch.sin(p1) * torch.cos(p2) * torch.cos(dl)
    return torch.atan2(y, x)


def build_mesh_graph(mesh_shape):
    """
    Build a fixed 8-neighbourhood graph over a regular latent mesh.
    Longitude wraps around; edge features are (forward bearing [rad],
    haversine distance [km]).
    """
    H, W = mesh_shape
    lat, lon = _latlon_grid(mesh_shape)
    lat = lat.view(-1, 1).expand(H, W)
    lon = lon.view(1, -1).expand(H, W)

    src_list, dst_list, feat_list = [], [], []
    for i in range(H):
        for j in range(W):
            for di, dj in ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)):
                ni, nj = i + di, (j + dj) % W
                if not (0 <= ni < H):
                    continue
                s = i * W + j
                d = ni * W + nj
                dist_km = _haversine(lat[i, j], lon[i, j], lat[ni, nj], lon[ni, nj]) / 1000.0
                bear = _bearing(lat[i, j], lon[i, j], lat[ni, nj], lon[ni, nj])
                src_list.append(s)
                dst_list.append(d)
                feat_list.append(torch.stack([bear / math.pi, dist_km / 1000.0]))
    edge_index = torch.stack([torch.as_tensor(src_list), torch.as_tensor(dst_list)], dim=0)
    edge_attr = torch.stack(feat_list)
    return edge_index, edge_attr


def _mlp(in_dim, out_dim, hidden_dim, n_layers=2):
    dims = [in_dim] + [hidden_dim] * (n_layers - 1) + [out_dim]
    layers = []
    for i in range(len(dims) - 1):
        layers.append(nn.Linear(dims[i], dims[i + 1]))
        if i < len(dims) - 2:
            layers.append(nn.GELU())
    return nn.Sequential(*layers)


class ConditionalLayerNorm(nn.Module):
    """
    Conditional layer-norm as used by FGN: a global noise vector is embedded
    (single matrix multiplication) and injected, via learned scale/shift, into
    every normalised module of the network. Sampling different noise vectors
    n for each ensemble member / timestep is what generates the variance
    across the ensemble (learned functional perturbations in weight space).
    """

    def __init__(self, dim, noise_dim=32):
        super().__init__()
        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
        self.gamma = nn.Linear(noise_dim, dim)
        self.beta = nn.Linear(noise_dim, dim)
        nn.init.zeros_(self.gamma.weight)
        nn.init.zeros_(self.beta.weight)
        nn.init.zeros_(self.gamma.bias)
        nn.init.zeros_(self.beta.bias)

    def forward(self, x, noise_emb):
        # x: [B, N, D]; noise_emb: [B, noise_dim]
        scale = self.gamma(noise_emb).unsqueeze(1)          # [B, 1, D]
        shift = self.beta(noise_emb).unsqueeze(1)           # [B, 1, D]
        return self.norm(x) * (1.0 + scale) + shift


class GNNLayer(nn.Module):
    """
    Message-passing layer with edge features (mean-aggregate, residual).
    For the grid->mesh encoder the message function knly conditions on the
    sender node features and the edge features, mirroring FGN's removal of
    the receiver-mesh-node conditioning in the encoder message function.
    """

    def __init__(self, dim, edge_dim=2, hidden_dim=64, noise_dim=32, receiver_cond=True):
        super().__init__()
        msg_in = 2 * dim + edge_dim if receiver_cond else dim + edge_dim
        self.edge_mlp = _mlp(msg_in, dim, hidden_dim)
        self.node_mlp = _mlp(dim, dim, hidden_dim)
        self.norm = ConditionalLayerNorm(dim, noise_dim)
        self.receiver_cond = receiver_cond

    def forward(self, x, edge_index, edge_attr, noise_emb):
        B, N, D = x.shape
        src, dst = edge_index
        offsets = torch.arange(B, device=x.device) * N
        src_b = (src.unsqueeze(0) + offsets.view(B, 1)).reshape(-1)
        dst_b = (dst.unsqueeze(0) + offsets.view(B, 1)).reshape(-1)
        edge_attr_b = edge_attr.unsqueeze(0).expand(B, -1, -1).reshape(-1, edge_attr.size(1))
        xb = x.reshape(B * N, D)
        if self.receiver_cond:
            msg = self.edge_mlp(torch.cat([xb[src_b], xb[dst_b], edge_attr_b], dim=1))
        else:
            msg = self.edge_mlp(torch.cat([xb[src_b], edge_attr_b], dim=1))
        agg = torch.zeros_like(xb)
        agg.index_add_(0, dst_b, msg)
        cnt = torch.bincount(dst_b, minlength=B * N).clamp(min=1).unsqueeze(1)
        agg = agg / cnt
        agg = agg.reshape(B, N, D)
        return self.norm(x + self.node_mlp(agg), noise_emb)


class GridEncoder(nn.Module):
    """
    Maps the gridded input states onto the latent mesh with a per-cell input
    MLP, an adaptive pooling to the mesh resolution, and graph message passing
    on the mesh graph (encoder message function without receiver conditioning).
    """

    def __init__(self, in_channels, latent_dim, mesh_shape, num_layers=2, hidden_dim=64,
                 edge_dim=2, noise_dim=32):
        super().__init__()
        self.input_mlp = _mlp(in_channels, latent_dim, hidden_dim)
        self.gnn = nn.ModuleList([
            GNNLayer(latent_dim, edge_dim=edge_dim, hidden_dim=hidden_dim,
                     noise_dim=noise_dim, receiver_cond=False)
            for _ in range(num_layers)
        ])
        self.mesh_shape = mesh_shape
        self.edge_index, self.edge_attr = build_mesh_graph(mesh_shape)

    def forward(self, x, noise_emb):
        B, C, H, W = x.shape
        feat = x.permute(0, 2, 3, 1).reshape(-1, C)
        feat = self.input_mlp(feat).reshape(B, H, W, -1).permute(0, 3, 1, 2)
        mesh = F.adaptive_avg_pool2d(feat, self.mesh_shape)     # [B, D, Hm, Wm]
        mesh = mesh.permute(0, 2, 3, 1).reshape(B, self.mesh_shape[0] * self.mesh_shape[1], -1)
        edge_index, edge_attr = self.edge_index.to(x.device), self.edge_attr.to(x.device)
        for layer in self.gnn:
            mesh = layer(mesh, edge_index, edge_attr, noise_emb)
        return mesh


class GraphTransformerBlock(nn.Module):
    """
    One graph-transformer block of the processor: multi-head self-attention
    over mesh tokens plus edge message passing, each with residual connection
    and conditioned layer-norm.
    """

    def __init__(self, latent_dim, n_heads, hidden_dim, edge_dim=2, noise_dim=32):
        super().__init__()
        self.attn = nn.MultiheadAttention(
            latent_dim, n_heads, batch_first=True, dropout=0.0
        )
        self.attn_norm = ConditionalLayerNorm(latent_dim, noise_dim)
        self.gnn = GNNLayer(latent_dim, edge_dim=edge_dim, hidden_dim=hidden_dim,
                            noise_dim=noise_dim, receiver_cond=True)
        self.gnn_norm = ConditionalLayerNorm(latent_dim, noise_dim)
        self.ff = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, latent_dim),
        )
        self.ff_norm = ConditionalLayerNorm(latent_dim, noise_dim)

    def forward(self, x, edge_index, edge_attr, noise_emb):
        # self-attention over mesh tokens
        h = self.attn_norm(x, noise_emb)
        h, _ = self.attn(h, h, h)
        x = x + self.gnn_norm(h, noise_emb)
        # edge message passing on the mesh graph
        x = self.gnn(x, edge_index, edge_attr, noise_emb)
        # feed-forward
        h = self.ff_norm(x, noise_emb)
        x = x + self.ff(h)
        return x


class GridDecoder(nn.Module):
    """
    Maps the latent mesh back onto the target grid (bilinear upsample) and
    predicts per-channel fields with an output MLP.
    """

    def __init__(self, latent_dim, out_channels, grid_shape, mesh_shape, num_layers=2,
                 hidden_dim=64, edge_dim=2, noise_dim=32):
        super().__init__()
        self.gnn = nn.ModuleList([
            GNNLayer(latent_dim, edge_dim=edge_dim, hidden_dim=hidden_dim,
                     noise_dim=noise_dim, receiver_cond=True)
            for _ in range(num_layers)
        ])
        self.grid_shape = grid_shape
        self.mesh_shape = mesh_shape
        self.edge_index, self.edge_attr = build_mesh_graph(mesh_shape)
        self.output_mlp = _mlp(latent_dim, out_channels, hidden_dim)

    def forward(self, mesh, noise_emb):
        B, N, D = mesh.shape
        edge_index = self.edge_index.to(mesh.device)
        edge_attr = self.edge_attr.to(mesh.device)
        for layer in self.gnn:
            mesh = layer(mesh, edge_index, edge_attr, noise_emb)
        H, W = self.grid_shape
        Hm, Wm = self.mesh_shape
        mesh = mesh.transpose(1, 2).reshape(B, D, Hm, Wm)
        grid = F.interpolate(mesh, size=self.grid_shape, mode="bilinear", align_corners=False)
        grid = grid.permute(0, 2, 3, 1).reshape(B, H * W, D)
        return self.output_mlp(grid).reshape(B, H, W, -1).permute(0, 3, 1, 2)


class FGN(nn.Module):
    """
    Config-driven FGN (Functional Generative Networks) wrapper.

    Args:
        in_channels: Number of state channels per frame (concatenated prior
            states fed to the grid encoder).
        out_channels: Number of forecast channels per frame.
        input_steps: Number of input (prior weather state) frames. FGN uses a
            second-order Markov assumption, input_steps=2.
        output_steps: Number of autoregressive forecast frames.
        grid_shape: Spatial shape of the (gridded) input state.
        mesh_shape: Latent mesh resolution (each dimension).
        latent_dim: Feature dimension of latent mesh tokens.
        num_encoder_layers / num_decoder_layers: GNN message-passing layers.
        num_processor_blocks: Graph-transformer blocks in the processor.
        n_heads: Attention heads of the processor.
        hidden_dim: Feed-forward / MLP hidden size.
        noise_dim: Dimension of the global noise vector injected through the
            conditional layer-norm layers (paper: 32).
        channel_weights: Per-channel weights for the fair-CRPS objective
            (taken from the GenCast/GraphCast loss weighting by default).
    """

    def __init__(
        self,
        in_channels=6,
        out_channels=6,
        input_steps=2,
        output_steps=2,
        grid_shape=(32, 32),
        mesh_shape=(8, 8),
        latent_dim=64,
        num_encoder_layers=2,
        num_decoder_layers=2,
        num_processor_blocks=1,
        n_heads=4,
        hidden_dim=64,
        noise_dim=32,
        channel_weights=None,
    ):
        super().__init__()
        self.in_channels = int(in_channels)
        self.out_channels = int(out_channels)
        self.input_steps = int(input_steps)
        self.output_steps = int(output_steps)
        self.grid_shape = (int(grid_shape[0]), int(grid_shape[1]))
        self.mesh_shape = (int(mesh_shape[0]), int(mesh_shape[1]))
        self.noise_dim = int(noise_dim)

        self.noise_embed = nn.Linear(self.noise_dim, self.noise_dim)

        self.encoder = GridEncoder(
            self.in_channels * self.input_steps, int(latent_dim), self.mesh_shape,
            num_layers=int(num_encoder_layers), hidden_dim=int(hidden_dim),
            noise_dim=self.noise_dim,
        )
        self.processor = nn.ModuleList([
            GraphTransformerBlock(
                int(latent_dim), int(n_heads), int(hidden_dim), noise_dim=self.noise_dim
            )
            for _ in range(int(num_processor_blocks))
        ])
        self.decoder = GridDecoder(
            int(latent_dim), self.out_channels, self.grid_shape, self.mesh_shape,
            num_layers=int(num_decoder_layers), hidden_dim=int(hidden_dim),
            noise_dim=self.noise_dim,
        )

        if channel_weights is None:
            channel_weights = torch.ones(self.out_channels)
        self.register_buffer("channel_weights", torch.as_tensor(channel_weights, dtype=torch.float32))

    def _rollout(self, x, noise):
        """
        Autoregressive rollout: at each output step sample the next state
        conditional on the last `input_steps` prior states x_{t-2}, x_{t-1}
        (second-order Markov), sampling a fresh global noise vector per step.
        """
        B, S, C, H, W = x.shape
        state = list(torch.unbind(x, dim=1))
        outs = []
        for t in range(self.output_steps):
            embrace = self.noise_embed(noise[t])              # [B, noise_dim]
            inp = torch.cat(state, dim=1)                     # [B, S*C, H, W]
            latent = self.encoder(inp, embrace)
            edge_index, edge_attr = self.encoder.edge_index, self.encoder.edge_attr
            edge_index = edge_index.to(x.device)
            edge_attr = edge_attr.to(x.device)
            for block in self.processor:
                latent = block(latent, edge_index, edge_attr, embrace)
            frame = self.decoder(latent, embrace)             # [B, C, H, W]
            outs.append(frame)
            state.append(frame)
            state = state[-self.input_steps:]
        return torch.stack(outs, dim=1)                       # [B, output_steps, C, H, W]

    def forward(self, x, num_members=1):
        """
        Args:
            x: Input state frames, shape [batch, input_steps, C, H, W].
            num_members: Number of independent ensemble members to generate
                (each member samples independent global noise per step).
        Returns:
            Forecast frames, shape [batch, num_members, output_steps, C, H, W].
        """
        device = x.device
        members = []
        for _ in range(int(num_members)):
            noise = torch.randn(self.output_steps, x.size(0), self.noise_dim, device=device)
            members.append(self._rollout(x, noise))
        return torch.stack(members, dim=1)

    def crps_loss(self, pred, target):
        """
        Fair CRPS objective (Eq. 4 of the paper) with an N-member ensemble,
        averaged over all locations, variables, levels and output steps:

            fCRPS(F_1:N, y) = 1/N sum_i |F_i - y|
                              - 1/(2 N (N-1)) sum_{i != i'} |F_i - F_i'|

        With N=2 this reduces to 0.5(|F1-y|+|F2-y|) - 0.5|F1-F2|. The loss is
        weighted per channel to match the GenCast/GraphCast loss weighting.
        """
        N = pred.size(1)
        mae = torch.abs(pred - target.unsqueeze(1)).mean(dim=1)      # (1/N) sum_i |F_i - y|
        per = torch.abs(pred.unsqueeze(2) - pred.unsqueeze(1)).sum(dim=(1, 2)) / (N * (N - 1))
        crps = mae - 0.5 * per                                        # [B, T, C, H, W]
        w = self.channel_weights.view(1, 1, self.out_channels, 1, 1)
        return (crps * w).mean()