# coding=utf-8
#
# SPDX-License-Identifier: MIT
#
# Copyright (c) 2024 Inter Agency Implementation and Advanced Concepts
#
# This module is a thin configuration wrapper around the official Prithvi WxC
# model (Schmude et al., arXiv:2409.13598), taken verbatim from the
# NASA-IMPACT/Prithvi-WxC repository (commit 79dabfcd17abe77e2d5c696707c0164a04f2ec01,
# MIT License). The official implementation is vendored as
# ``prithvi_wxc_official.py``; only scaler construction and tensor-level call
# signatures are added here for YAML-driven usage.
import torch
import torch.nn as nn

from model.prithvi_wxc_official import PrithviWxC as _PrithviWxC


def _identity_scalers(n, device="cpu"):
    return torch.zeros(n, device=device), torch.ones(n, device=device)


class PrithviWxC(nn.Module):
    """
    Config-driven Prithvi WxC wrapper.

    The official model is an encoder-decoder vision Transformer (Hiera + MaxViT,
    alternating local/global attention) that maps an input state history to a
    target state. This wrapper builds the official network with identity
    scalers so that a small YAML config can be exercised; loading the 2.3B
    checkpoints published on Hugging Face is possible by calling
    ``load_state_dict`` with the official checkpoint state (the buffer names,
    e.g. ``input_scalers_mu``, are identical).

    Args:
        in_channels: number of input/output parameters.
        input_size_time: number of input timestamps (paper uses 2).
        in_channels_static: number of static input channels (>= 1).
        n_lats_px / n_lons_px: grid size in pixels.
        patch_size_px: patch/token size in pixels.
        mask_unit_size_px: mask unit size in pixels.
        mask_ratio_inputs: input masking ratio (0..1, 0 disables).
        embed_dim / n_blocks_encoder / n_blocks_decoder / mlp_multiplier / n_heads:
            transformer architecture hyper-parameters (paper: 2560 / 25 / 5 / 4 / 16).
        dropout / drop_path / parameter_dropout: regularization ratios.
        residual: one of "none" / "temporal" / "climate".
        masking_mode: "global" / "local" / "both".
        positional_encoding: "absolute" or "fourier".
        encoder_shifting / decoder_shifting: Swin-style shifting.
    """

    def __init__(
        self,
        in_channels=4,
        input_size_time=2,
        in_channels_static=4,
        n_lats_px=32,
        n_lons_px=64,
        patch_size_px=(2, 2),
        mask_unit_size_px=(8, 8),
        mask_ratio_inputs=0.0,
        embed_dim=16,
        n_blocks_encoder=1,
        n_blocks_decoder=1,
        mlp_multiplier=4.0,
        n_heads=4,
        dropout=0.0,
        drop_path=0.0,
        parameter_dropout=0.0,
        residual="none",
        masking_mode="global",
        positional_encoding="absolute",
        encoder_shifting=False,
        decoder_shifting=False,
    ):
        super().__init__()
        self.in_channels = int(in_channels)
        self.input_size_time = int(input_size_time)
        self.in_channels_static = int(in_channels_static)
        self.n_lats_px = int(n_lats_px)
        self.n_lons_px = int(n_lons_px)
        self.residual = residual
        self.positional_encoding = positional_encoding

        mu_in, sigma_in = _identity_scalers(self.in_channels)
        mu_s, sigma_s = _identity_scalers(self.in_channels_static)

        self.model = _PrithviWxC(
            in_channels=self.in_channels,
            input_size_time=self.input_size_time,
            in_channels_static=self.in_channels_static,
            input_scalers_mu=mu_in,
            input_scalers_sigma=sigma_in,
            input_scalers_epsilon=1e-6,
            static_input_scalers_mu=mu_s,
            static_input_scalers_sigma=sigma_s,
            static_input_scalers_epsilon=1e-6,
            output_scalers=sigma_in,
            n_lats_px=self.n_lats_px,
            n_lons_px=self.n_lons_px,
            patch_size_px=tuple(patch_size_px),
            mask_unit_size_px=tuple(mask_unit_size_px),
            mask_ratio_inputs=float(mask_ratio_inputs),
            mask_ratio_targets=0.0,
            embed_dim=int(embed_dim),
            n_blocks_encoder=int(n_blocks_encoder),
            n_blocks_decoder=int(n_blocks_decoder),
            mlp_multiplier=int(mlp_multiplier),
            n_heads=int(n_heads),
            dropout=float(dropout),
            drop_path=float(drop_path),
            parameter_dropout=float(parameter_dropout),
            residual=residual,
            masking_mode=masking_mode,
            positional_encoding=positional_encoding,
            encoder_shifting=bool(encoder_shifting),
            decoder_shifting=bool(decoder_shifting),
        )

    def forward(self, x, static, input_time=None, lead_time=None, climate=None):
        """
        Args:
            x: Tensor of shape [batch, time, parameter, lat, lon].
            static: Tensor of shape [batch, static_channel, lat, lon].
            input_time: Tensor of shape [batch] (optional, default zeros).
            lead_time: Tensor of shape [batch] (optional, default zeros).
            climate: Optional Tensor of shape [batch, parameter, lat, lon].
        Returns:
            Tensor of shape [batch, parameter, lat, lon].
        """
        if input_time is None:
            input_time = torch.zeros(x.shape[0], device=x.device)
        if lead_time is None:
            lead_time = torch.zeros(x.shape[0], device=x.device)
        batch = {
            "x": x,
            "y": x[:, -1],
            "static": static,
            "input_time": input_time,
            "lead_time": lead_time,
        }
        if climate is not None:
            batch["climate"] = climate
        return self.model(batch)