# coding=utf-8
#
# SPDX-License-Identifier: Apache-2.0
#
# Minimal reproduction of GraphDOP (ECMWF, "Towards skilful medium-range
# forecasts learnt directly from observations", 2025 preprint) following the
# encoder -- processor -- decoder design:
#
#   * Encoder: a GNN that projects gridded "observations" inside the input
#     window onto a latent mesh (a coarse regular lat/lon grid), using graph
#     edges with (forward bearing, haversine distance) features.
#   * Processor: a transformer that advances the latent atmospheric state
#     forward in time, once per output frame (latent-space rollout).
#   * Decoder: a GNN that maps the latent mesh back onto the target grid and
#     predicts per-channel observations with instrument-like output MLPs.
#
# Differences from the paper (documented in README.md): the paper consumes
# irregular, instrument-specific Level-1 observations with dynamic graphs built
# per batch (PyTorch Geometric); here the OneScience ERA5-h5 gridded pipeline is
# used as the observation placeholder, and the graphs are fixed regular-grid
# meshes. The weighted MSE objective is kept (per-channel weights).
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 between points given 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 GNNLayer(nn.Module):
    """Message-passing layer with edge features (mean-aggregate, residual)."""

    def __init__(self, dim, edge_dim=2, hidden_dim=64):
        super().__init__()
        self.edge_mlp = _mlp(2 * dim + edge_dim, dim, hidden_dim)
        self.node_mlp = _mlp(dim, dim, hidden_dim)
        self.norm = nn.LayerNorm(dim)

    def forward(self, x, edge_index, edge_attr):
        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)
        msg = self.edge_mlp(torch.cat([xb[src_b], xb[dst_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))


class ObsEncoder(nn.Module):
    """
    Maps the observation grid onto the latent mesh with a per-cell input MLP,
    an adaptive pooling to the mesh resolution, and graph message passing.
    """

    def __init__(self, in_channels, latent_dim, mesh_shape, num_layers=2, hidden_dim=64):
        super().__init__()
        self.in_channels = in_channels
        self.input_mlp = _mlp(in_channels, latent_dim, hidden_dim)
        self.gnn = nn.ModuleList([GNNLayer(latent_dim, hidden_dim=hidden_dim) 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):
        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)
        mesh = mesh.permute(0, 2, 3, 1).reshape(B, -1, mesh.size(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)
        return mesh


class LatentProcessor(nn.Module):
    """Transformer over latent mesh tokens that advances the state in time."""

    def __init__(self, latent_dim, mesh_shape, num_blocks=1, n_heads=4, hidden_dim=128):
        super().__init__()
        n_nodes = mesh_shape[0] * mesh_shape[1]
        self.pos_emb = nn.Parameter(torch.zeros(1, n_nodes, latent_dim))
        nn.init.trunc_normal_(self.pos_emb, std=0.02)
        block = nn.TransformerEncoderLayer(
            d_model=latent_dim, nhead=n_heads, dim_feedforward=hidden_dim,
            dropout=0.0, activation="gelu", batch_first=True, norm_first=True,
        )
        self.blocks = nn.ModuleList([block for _ in range(num_blocks)])

    def forward(self, mesh):
        tokens = mesh + self.pos_emb
        for block in self.blocks:
            tokens = block(tokens)
        return tokens


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

    def __init__(self, latent_dim, out_channels, grid_shape, mesh_shape, num_layers=2, hidden_dim=64):
        super().__init__()
        self.gnn = nn.ModuleList([GNNLayer(latent_dim, hidden_dim=hidden_dim) 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):
        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)
        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 GraphDOP(nn.Module):
    """
    Config-driven GraphDOP wrapper.

    Args:
        in_channels: Number of observation channels per frame.
        out_channels: Number of forecast channels per frame.
        input_steps: Number of input (observation window) frames.
        output_steps: Number of forecast frames.
        grid_shape: Spatial shape of the (gridded) observation field.
        mesh_shape: Latent mesh resolution (each dimension, powers of two fine).
        latent_dim: Feature dimension of latent mesh tokens.
        num_encoder_layers / num_decoder_layers: GNN message-passing layers.
        num_processor_blocks: Transformer blocks in the processor.
        n_heads: Attention heads of the processor.
        channel_weights: Per-channel weights for the weighted MSE objective.
    """

    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,
        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.encoder = ObsEncoder(
            self.in_channels, int(latent_dim), self.mesh_shape, num_layers=int(num_encoder_layers), hidden_dim=int(hidden_dim)
        )
        self.processor = LatentProcessor(
            int(latent_dim), self.mesh_shape, num_blocks=int(num_processor_blocks), n_heads=int(n_heads), hidden_dim=int(hidden_dim)
        )
        self.decoder = ObsDecoder(
            int(latent_dim), self.out_channels, self.grid_shape, self.mesh_shape,
            num_layers=int(num_decoder_layers), hidden_dim=int(hidden_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 forward(self, x):
        """
        Args:
            x: Observation frames, shape [batch, input_steps, C, H, W].
        Returns:
            Forecast frames, shape [batch, output_steps, C, H, W].
        """
        latents = torch.stack([self.encoder(x[:, t]) for t in range(self.input_steps)], dim=0)
        latent = latents.mean(dim=0)
        outs = []
        for _ in range(self.output_steps):
            latent = self.processor(latent)
            outs.append(self.decoder(latent))
        return torch.stack(outs, dim=1)

    def wmse_loss(self, pred, target):
        """Weighted mean squared error objective (Eq. 1 of the paper)."""
        diff = (pred - target) ** 2
        w = self.channel_weights.view(1, 1, self.out_channels, 1, 1)
        return (diff * w).mean()