# Copyright (c) IBM Corp. 2024. All rights reserved.
# Copyright 2024 Prithvi-EO-2.0 Authors and The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
"""Self-contained Prithvi-EO-2.0 model and config for trust_remote_code loading."""

from __future__ import annotations

import warnings
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
from einops import rearrange
from timm.layers import to_2tuple
from timm.models.vision_transformer import Block

from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
from transformers.modeling_outputs import BaseModelOutputWithPooling
from transformers.modeling_utils import PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, logging


logger = logging.get_logger(__name__)

NO_DATA_FLOAT = 0.0001


def get_3d_sincos_pos_embed(embed_dim: int, grid_size: tuple[int, int, int] | list[int], add_cls_token: bool = False):
    assert embed_dim % 16 == 0

    t_size, h_size, w_size = grid_size
    w_embed_dim = embed_dim // 16 * 6
    h_embed_dim = embed_dim // 16 * 6
    t_embed_dim = embed_dim // 16 * 4

    w_pos_embed = get_1d_sincos_pos_embed_from_grid(w_embed_dim, np.arange(w_size))
    h_pos_embed = get_1d_sincos_pos_embed_from_grid(h_embed_dim, np.arange(h_size))
    t_pos_embed = get_1d_sincos_pos_embed_from_grid(t_embed_dim, np.arange(t_size))

    w_pos_embed = np.tile(w_pos_embed, (t_size * h_size, 1))
    h_pos_embed = np.tile(np.repeat(h_pos_embed, w_size, axis=0), (t_size, 1))
    t_pos_embed = np.repeat(t_pos_embed, h_size * w_size, axis=0)

    pos_embed = np.concatenate((w_pos_embed, h_pos_embed, t_pos_embed), axis=1)
    if add_cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed


def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
    if embed_dim % 2 != 0:
        raise ValueError("embed_dim must be even")

    omega = np.arange(embed_dim // 2, dtype=float)
    omega /= embed_dim / 2.0
    omega = 1.0 / 10000**omega

    pos = pos.reshape(-1)
    out = np.einsum("m,d->md", pos, omega)
    emb_sin = np.sin(out)
    emb_cos = np.cos(out)
    return np.concatenate([emb_sin, emb_cos], axis=1)


def _get_1d_sincos_embed_from_grid_torch(embed_dim: int, pos: torch.Tensor) -> torch.Tensor:
    assert embed_dim % 2 == 0
    assert pos.dtype in [torch.float32, torch.float16, torch.bfloat16]

    omega = torch.arange(embed_dim // 2, dtype=pos.dtype, device=pos.device)
    omega /= embed_dim / 2.0
    omega = 1.0 / 10000**omega

    pos = pos.reshape(-1)
    out = torch.einsum("m,d->md", pos, omega)
    emb_sin = torch.sin(out)
    emb_cos = torch.cos(out)
    return torch.cat([emb_sin, emb_cos], dim=1)


def _init_weights(module: nn.Module) -> None:
    if isinstance(module, nn.Linear):
        nn.init.xavier_uniform_(module.weight)
        if module.bias is not None:
            module.bias.data.zero_()
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)


def _interpolate_pos_encoding(
    pos_embed: torch.Tensor,
    grid_size: tuple[int, int, int] | list[int],
    patch_size: tuple[int, int, int] | list[int],
    shape: tuple[int, int, int],
    embed_dim: int,
) -> torch.Tensor:
    t, h, w = shape
    t_patches = t // patch_size[0]
    h_patches = h // patch_size[1]
    w_patches = w // patch_size[2]

    if [t_patches, h_patches, w_patches] == list(grid_size):
        return pos_embed
    if t_patches != grid_size[0]:
        new_grid_size = (t_patches, *grid_size[1:])
        new_pos_embed = get_3d_sincos_pos_embed(pos_embed.shape[-1], new_grid_size, add_cls_token=True)
        new_pos_embed = torch.from_numpy(new_pos_embed).float().unsqueeze(0).to(pos_embed.device)
    else:
        new_grid_size = grid_size
        new_pos_embed = pos_embed

    class_pos_embed, patch_pos_embed = new_pos_embed[:, :1], new_pos_embed[:, 1:]
    patch_pos_embed = patch_pos_embed.reshape(*new_grid_size, embed_dim).permute(0, 3, 1, 2)
    patch_pos_embed = nn.functional.interpolate(
        patch_pos_embed,
        size=(h_patches, w_patches),
        mode="bicubic",
        align_corners=True,
    )
    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, embed_dim)
    return torch.cat((class_pos_embed, patch_pos_embed), dim=1)


class PrithviConfig(PreTrainedConfig):
    model_type = "prithvi_eo"

    def __init__(
        self,
        image_size: int = 224,
        num_frames: int = 4,
        patch_size: tuple[int, int, int] | list[int] | int = (1, 16, 16),
        num_channels: int = 6,
        hidden_size: int = 1024,
        num_hidden_layers: int = 24,
        num_attention_heads: int = 16,
        decoder_hidden_size: int = 512,
        decoder_num_hidden_layers: int = 8,
        decoder_num_attention_heads: int = 16,
        mlp_ratio: float = 4.0,
        coords_encoding: list[str] | None = None,
        coords_scale_learn: bool = False,
        mask_ratio: float = 0.75,
        norm_pix_loss: bool = False,
        bands: list[str] | None = None,
        image_mean: list[float] | None = None,
        image_std: list[float] | None = None,
        origin_url: str | None = None,
        architecture_name: str | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        if isinstance(patch_size, int):
            patch_size = (1, patch_size, patch_size)

        self.image_size = image_size
        self.num_frames = num_frames
        self.patch_size = list(patch_size)
        self.num_channels = num_channels
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.decoder_hidden_size = decoder_hidden_size
        self.decoder_num_hidden_layers = decoder_num_hidden_layers
        self.decoder_num_attention_heads = decoder_num_attention_heads
        self.mlp_ratio = mlp_ratio
        self.coords_encoding = coords_encoding or []
        self.coords_scale_learn = coords_scale_learn
        self.mask_ratio = mask_ratio
        self.norm_pix_loss = norm_pix_loss
        self.bands = bands or ["B02", "B03", "B04", "B05", "B06", "B07"]
        self.image_mean = image_mean or [1087.0, 1342.0, 1433.0, 2734.0, 1958.0, 1363.0]
        self.image_std = image_std or [2248.0, 2179.0, 2178.0, 1850.0, 1242.0, 1049.0]
        self.origin_url = origin_url
        self.architecture_name = architecture_name


class PatchEmbed(nn.Module):
    def __init__(
        self,
        input_size: tuple[int, int, int] = (1, 224, 224),
        patch_size: tuple[int, int, int] = (1, 16, 16),
        in_chans: int = 3,
        embed_dim: int = 768,
        norm_layer: nn.Module | None = None,
        flatten: bool = True,
        bias: bool = True,
    ):
        super().__init__()
        self.input_size = input_size
        self.patch_size = patch_size
        self.grid_size = [s // p for s, p in zip(self.input_size, self.patch_size)]
        assert self.grid_size >= [1, 1, 1], "Patch size is bigger than input size."
        self.num_patches = self.grid_size[0] * self.grid_size[1] * self.grid_size[2]
        self.flatten = flatten
        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _, _, t, h, w = x.shape
        if t / self.patch_size[0] % 1 or h / self.patch_size[1] % 1 or w / self.patch_size[2] % 1:
            warnings.warn(
                f"Input {x.shape[-3:]} is not divisible by patch size {self.patch_size}. "
                "The border will be ignored; add backbone_padding for pixel-wise tasks."
            )
        x = self.proj(x)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)
        return self.norm(x)


class TemporalEncoder(nn.Module):
    def __init__(self, embed_dim: int, trainable_scale: bool = False):
        super().__init__()
        self.embed_dim = embed_dim
        self.year_embed_dim = embed_dim // 2
        self.julian_day_embed_dim = embed_dim - self.year_embed_dim
        if trainable_scale:
            self.scale = nn.Parameter(torch.full((1,), 0.1))
        else:
            self.register_buffer("scale", torch.ones(1))

    def forward(self, temporal_coords: torch.Tensor, tokens_per_frame: int | None = None) -> torch.Tensor:
        shape = temporal_coords.shape[:2] + (-1,)
        year = _get_1d_sincos_embed_from_grid_torch(
            self.year_embed_dim, temporal_coords[:, :, 0].flatten()
        ).reshape(shape)
        julian_day = _get_1d_sincos_embed_from_grid_torch(
            self.julian_day_embed_dim, temporal_coords[:, :, 1].flatten()
        ).reshape(shape)
        embedding = self.scale * torch.cat([year, julian_day], dim=-1)
        if tokens_per_frame is not None:
            embedding = torch.repeat_interleave(embedding, tokens_per_frame, dim=1)
        return embedding


class LocationEncoder(nn.Module):
    def __init__(self, embed_dim: int, trainable_scale: bool = False):
        super().__init__()
        self.embed_dim = embed_dim
        self.lat_embed_dim = embed_dim // 2
        self.lon_embed_dim = embed_dim - self.lat_embed_dim
        if trainable_scale:
            self.scale = nn.Parameter(torch.full((1,), 0.1))
        else:
            self.register_buffer("scale", torch.ones(1))

    def forward(self, location_coords: torch.Tensor) -> torch.Tensor:
        shape = location_coords.shape[:1] + (1, -1)
        lat = _get_1d_sincos_embed_from_grid_torch(
            self.lat_embed_dim, location_coords[:, 0].flatten()
        ).reshape(shape)
        lon = _get_1d_sincos_embed_from_grid_torch(
            self.lon_embed_dim, location_coords[:, 1].flatten()
        ).reshape(shape)
        return self.scale * torch.cat([lat, lon], dim=-1)


class PrithviViT(nn.Module):
    def __init__(self, config: PrithviConfig):
        super().__init__()
        patch_size = tuple(config.patch_size)
        img_size = to_2tuple(config.image_size)

        self.in_chans = config.num_channels
        self.num_frames = config.num_frames
        self.embed_dim = config.hidden_size
        self.img_size = img_size
        self.patch_embed = PatchEmbed(
            input_size=(config.num_frames,) + self.img_size,
            patch_size=patch_size,
            in_chans=config.num_channels,
            embed_dim=config.hidden_size,
        )
        self.out_channels = [config.hidden_size * self.patch_embed.grid_size[0]] * config.num_hidden_layers

        coords_encoding = config.coords_encoding or []
        self.temporal_encoding = "time" in coords_encoding
        self.location_encoding = "location" in coords_encoding
        if self.temporal_encoding:
            assert patch_size[0] == 1, f"With temporal encoding, patch_size[0] must be 1, received {patch_size[0]}"
            self.temporal_embed_enc = TemporalEncoder(config.hidden_size, config.coords_scale_learn)
        if self.location_encoding:
            self.location_embed_enc = LocationEncoder(config.hidden_size, config.coords_scale_learn)

        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
        self.register_buffer("pos_embed", torch.zeros(1, self.patch_embed.num_patches + 1, config.hidden_size))
        self.blocks = nn.ModuleList(
            [
                Block(
                    config.hidden_size,
                    config.num_attention_heads,
                    config.mlp_ratio,
                    qkv_bias=True,
                    norm_layer=nn.LayerNorm,
                )
                for _ in range(config.num_hidden_layers)
            ]
        )
        self.norm = nn.LayerNorm(config.hidden_size)
        self.initialize_weights()

    def initialize_weights(self) -> None:
        pos_embed = get_3d_sincos_pos_embed(
            self.pos_embed.shape[-1], self.patch_embed.grid_size, add_cls_token=True
        )
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
        w = self.patch_embed.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
        torch.nn.init.normal_(self.cls_token, std=0.02)
        self.apply(_init_weights)

    def random_masking(self, sequence: torch.Tensor, mask_ratio: float, noise: torch.Tensor | None = None):
        batch_size, seq_length, dim = sequence.shape
        len_keep = int(seq_length * (1 - mask_ratio))
        if noise is None:
            noise = torch.rand(batch_size, seq_length, device=sequence.device)
        ids_shuffle = torch.argsort(noise, dim=1).to(sequence.device)
        ids_restore = torch.argsort(ids_shuffle, dim=1).to(sequence.device)
        ids_keep = ids_shuffle[:, :len_keep]
        sequence_unmasked = torch.gather(sequence, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, dim))
        mask = torch.ones([batch_size, seq_length], device=sequence.device)
        mask[:, :len_keep] = 0
        mask = torch.gather(mask, dim=1, index=ids_restore)
        return sequence_unmasked, mask, ids_restore

    def interpolate_pos_encoding(self, sample_shape: tuple[int, int, int]) -> torch.Tensor:
        return _interpolate_pos_encoding(
            pos_embed=self.pos_embed,
            grid_size=self.patch_embed.grid_size,
            patch_size=self.patch_embed.patch_size,
            shape=sample_shape,
            embed_dim=self.embed_dim,
        )

    def forward(
        self,
        x: torch.Tensor,
        temporal_coords: torch.Tensor | None = None,
        location_coords: torch.Tensor | None = None,
        mask_ratio: float = 0.75,
    ):
        if len(x.shape) == 4 and self.patch_embed.input_size[0] == 1:
            x = x.unsqueeze(2)
        sample_shape = x.shape[-3:]
        x = self.patch_embed(x)
        pos_embed = self.interpolate_pos_encoding(sample_shape)
        x = x + pos_embed[:, 1:, :]

        if self.temporal_encoding and temporal_coords is not None:
            num_tokens_per_frame = x.shape[1] // self.num_frames
            temporal_encoding = self.temporal_embed_enc(temporal_coords, num_tokens_per_frame)
            x = x + temporal_encoding
        if self.location_encoding and location_coords is not None:
            location_encoding = self.location_embed_enc(location_coords)
            x = x + location_encoding

        x, mask, ids_restore = self.random_masking(x, mask_ratio)
        cls_token = self.cls_token + pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        for block in self.blocks:
            x = block(x)
        x = self.norm(x)
        return x, mask, ids_restore

    def forward_features(
        self,
        x: torch.Tensor,
        temporal_coords: torch.Tensor | None = None,
        location_coords: torch.Tensor | None = None,
        output_hidden_states: bool = False,
    ) -> tuple[torch.Tensor, list[torch.Tensor] | None]:
        if len(x.shape) == 4 and self.patch_embed.input_size[0] == 1:
            x = x.unsqueeze(2)
        sample_shape = x.shape[-3:]
        x = self.patch_embed(x)
        pos_embed = self.interpolate_pos_encoding(sample_shape)
        x = x + pos_embed[:, 1:, :]

        if self.temporal_encoding and temporal_coords is not None:
            num_tokens_per_frame = x.shape[1] // self.num_frames
            temporal_encoding = self.temporal_embed_enc(temporal_coords, num_tokens_per_frame)
            x = x + temporal_encoding
        if self.location_encoding and location_coords is not None:
            location_encoding = self.location_embed_enc(location_coords)
            x = x + location_encoding

        cls_token = self.cls_token + pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        hidden_states = [] if output_hidden_states else None
        for block in self.blocks:
            x = block(x)
            if output_hidden_states:
                hidden_states.append(x)
        x = self.norm(x)
        if output_hidden_states:
            hidden_states[-1] = x
        return x, hidden_states

    def prepare_features_for_image_model(self, features: list[torch.Tensor]) -> list[torch.Tensor]:
        out = []
        effective_time_dim = self.patch_embed.input_size[0] // self.patch_embed.patch_size[0]
        for x in features:
            x_no_token = x[:, 1:, :]
            number_of_tokens = x_no_token.shape[1]
            tokens_per_timestep = number_of_tokens // effective_time_dim
            h = int(np.sqrt(tokens_per_timestep))
            encoded = rearrange(
                x_no_token,
                "batch (t h w) e -> batch (t e) h w",
                e=self.embed_dim,
                t=effective_time_dim,
                h=h,
            )
            out.append(encoded)
        return out


class MAEDecoder(nn.Module):
    def __init__(self, config: PrithviConfig, grid_size: list[int]):
        super().__init__()
        patch_size = tuple(config.patch_size)
        self.decoder_embed = nn.Linear(config.hidden_size, config.decoder_hidden_size, bias=True)
        self.decoder_embed_dim = config.decoder_hidden_size
        self.grid_size = grid_size
        self.patch_size = patch_size
        self.num_frames = self.grid_size[0] * patch_size[0]
        num_patches = self.grid_size[0] * self.grid_size[1] * self.grid_size[2]

        coords_encoding = config.coords_encoding or []
        self.temporal_encoding = "time" in coords_encoding
        self.location_encoding = "location" in coords_encoding
        if self.temporal_encoding:
            self.temporal_embed_dec = TemporalEncoder(config.decoder_hidden_size, config.coords_scale_learn)
        if self.location_encoding:
            self.location_embed_dec = LocationEncoder(config.decoder_hidden_size, config.coords_scale_learn)

        self.mask_token = nn.Parameter(torch.zeros(1, 1, config.decoder_hidden_size))
        self.register_buffer("decoder_pos_embed", torch.zeros(1, num_patches + 1, config.decoder_hidden_size))
        self.decoder_blocks = nn.ModuleList(
            [
                Block(
                    config.decoder_hidden_size,
                    config.decoder_num_attention_heads,
                    config.mlp_ratio,
                    qkv_bias=True,
                    norm_layer=nn.LayerNorm,
                )
                for _ in range(config.decoder_num_hidden_layers)
            ]
        )
        self.decoder_norm = nn.LayerNorm(config.decoder_hidden_size)
        self.decoder_pred = nn.Linear(
            config.decoder_hidden_size,
            patch_size[0] * patch_size[1] * patch_size[2] * config.num_channels,
            bias=True,
        )
        self.initialize_weights()

    def initialize_weights(self) -> None:
        decoder_pos_embed = get_3d_sincos_pos_embed(
            self.decoder_pos_embed.shape[-1], self.grid_size, add_cls_token=True
        )
        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))
        torch.nn.init.normal_(self.mask_token, std=0.02)
        self.apply(_init_weights)

    def interpolate_pos_encoding(self, sample_shape: tuple[int, int, int]) -> torch.Tensor:
        return _interpolate_pos_encoding(
            pos_embed=self.decoder_pos_embed,
            grid_size=self.grid_size,
            patch_size=self.patch_size,
            shape=sample_shape,
            embed_dim=self.decoder_embed_dim,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        ids_restore: torch.Tensor,
        temporal_coords: torch.Tensor | None = None,
        location_coords: torch.Tensor | None = None,
        input_size: list[int] | None = None,
    ) -> torch.Tensor:
        x = self.decoder_embed(hidden_states)
        cls_token = x[:, :1, :]
        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
        x = torch.cat([x[:, 1:, :], mask_tokens], dim=1)
        x = torch.gather(x, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]).to(x.device))

        decoder_pos_embed = self.interpolate_pos_encoding(input_size[-3:])
        cls_token = cls_token + decoder_pos_embed[:, :1, :]
        x = x + decoder_pos_embed[:, 1:, :]

        if self.temporal_encoding and temporal_coords is not None:
            num_tokens_per_frame = x.shape[1] // self.num_frames
            temporal_encoding = self.temporal_embed_dec(temporal_coords, num_tokens_per_frame)
            x = x + temporal_encoding
        if self.location_encoding and location_coords is not None:
            location_encoding = self.location_embed_dec(location_coords)
            x = x + location_encoding

        x = torch.cat([cls_token, x], dim=1)
        for block in self.decoder_blocks:
            x = block(x)
        x = self.decoder_norm(x)
        pred = self.decoder_pred(x)
        return pred[:, 1:, :]


class PrithviPreTrainedModel(PreTrainedModel):
    config_class = PrithviConfig
    config: PrithviConfig
    base_model_prefix = "prithvi"
    main_input_name = "pixel_values"
    input_modalities = ("image",)
    supports_gradient_checkpointing = False
    _no_split_modules = ["Block"]


class PrithviMAEModel(PrithviPreTrainedModel):
    def __init__(self, config: PrithviConfig):
        super().__init__(config)
        self.encoder = PrithviViT(config)
        self.decoder = MAEDecoder(config, grid_size=self.encoder.patch_embed.grid_size)
        self.post_init()

    def patchify(self, pixel_values: torch.Tensor) -> torch.Tensor:
        patch_size_t, patch_size_h, patch_size_w = self.encoder.patch_embed.patch_size
        num_channels = self.encoder.in_chans
        return rearrange(
            pixel_values,
            "b c (t s) (h p) (w q) -> b (t h w) (s p q c)",
            c=num_channels,
            s=patch_size_t,
            p=patch_size_h,
            q=patch_size_w,
        )

    def unpatchify(self, patchified_pixel_values: torch.Tensor, image_size: tuple[int, int] | None = None) -> torch.Tensor:
        patch_size_t, patch_size_h, patch_size_w = self.encoder.patch_embed.patch_size
        image_size = to_2tuple(image_size) if image_size is not None else self.encoder.img_size
        original_height, original_width = image_size
        num_patches_h = original_height // patch_size_h
        num_patches_w = original_width // patch_size_w
        num_channels = self.encoder.in_chans
        return rearrange(
            patchified_pixel_values,
            "b (t h w) (s p q c) -> b c (t s) (h p) (w q)",
            c=num_channels,
            h=num_patches_h,
            w=num_patches_w,
            s=patch_size_t,
            p=patch_size_h,
            q=patch_size_w,
        )

    def forward_loss(self, pixel_values: torch.Tensor, pred: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        target = self.patchify(pixel_values)
        if self.config.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            var = target.var(dim=-1, keepdim=True)
            target = (target - mean) / (var + 1.0e-6) ** 0.5
        loss = (pred - target) ** 2
        loss = loss.mean(dim=-1)
        return (loss * mask).sum() / mask.sum()

    def forward(
        self,
        pixel_values: torch.Tensor,
        temporal_coords: torch.Tensor | None = None,
        location_coords: torch.Tensor | None = None,
        mask_ratio: float | None = None,
        return_dict: bool | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ):
        if len(pixel_values.shape) == 4 and self.encoder.patch_embed.input_size[0] == 1:
            pixel_values = pixel_values.unsqueeze(2)

        mask_ratio = self.config.mask_ratio if mask_ratio is None else mask_ratio
        latent, mask, ids_restore = self.encoder(pixel_values, temporal_coords, location_coords, mask_ratio)
        pred = self.decoder(latent, ids_restore, temporal_coords, location_coords, input_size=list(pixel_values.shape))
        loss = self.forward_loss(pixel_values, pred, mask)

        if not return_dict:
            return loss, pred, mask
        return {"loss": loss, "pred": pred, "mask": mask}


class PrithviModel(PrithviPreTrainedModel):
    """Encoder-only Prithvi ViT for spatiotemporal feature extraction."""

    def __init__(self, config: PrithviConfig):
        super().__init__(config)
        self.encoder = PrithviViT(config)
        self.post_init()

    def forward(
        self,
        pixel_values: Optional[torch.Tensor] = None,
        temporal_coords: Optional[torch.Tensor] = None,
        location_coords: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutputWithPooling:
        if pixel_values is None:
            raise ValueError("You must specify `pixel_values`")

        pixel_values = pixel_values.to(dtype=self.dtype)
        if return_dict is None:
            return_dict = self.config.use_return_dict
        if output_hidden_states is None:
            output_hidden_states = False

        last_hidden_state, hidden_states = self.encoder.forward_features(
            pixel_values,
            temporal_coords=temporal_coords,
            location_coords=location_coords,
            output_hidden_states=output_hidden_states,
        )
        pooler_output = last_hidden_state[:, 0]

        if not return_dict:
            outputs = (last_hidden_state, pooler_output)
            if output_hidden_states:
                outputs = outputs + (hidden_states,)
            return outputs

        return BaseModelOutputWithPooling(
            last_hidden_state=last_hidden_state,
            pooler_output=pooler_output,
            hidden_states=tuple(hidden_states) if hidden_states is not None else None,
        )


__all__ = [
    "PrithviConfig",
    "PrithviMAEModel",
    "PrithviModel",
    "PrithviPreTrainedModel",
]