# Copyright 2025 The Galileo Authors and The HuggingFace Inc. team.
"""Custom Galileo feature extraction pipeline for trust_remote_code loading."""

from __future__ import annotations

from typing import Any, Union

import torch

from transformers.pipelines.base import GenericTensor, Pipeline, build_pipeline_init_args
from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
from transformers.utils import add_end_docstrings


@add_end_docstrings(
    build_pipeline_init_args(has_processor=True),
    """
        s1 (`torch.Tensor` or `numpy.ndarray`, *optional*):
            Sentinel-1 tensor with shape `(H, W, T, C)` or `(H, W, C)`.
        s2 (`torch.Tensor` or `numpy.ndarray`, *optional*):
            Sentinel-2 tensor with shape `(H, W, T, C)` or `(H, W, C)`.
        normalize (`bool`, *optional*, defaults to `True`):
            Whether to apply pretraining normalization statistics.
        patch_size (`int`, *optional*):
            Spatial patch size for the flexible patch embedder (1-8).
        pool (`bool`, *optional*, defaults to `False`):
            Whether to return pooled features instead of sequence hidden states.
    """,
)
class GalileoImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
    """
    Galileo multimodal remote sensing feature extraction pipeline.

    Load with:

    ```python
    >>> from transformers import pipeline
    >>> pipe = pipeline(
    ...     task="galileo-feature-extraction",
    ...     model="/path/to/galileo-nano-patch8",
    ...     trust_remote_code=True,
    ... )
    >>> features = pipe(s2=s2_array, pool=True, return_tensors=True)
    ```
    """

    _load_processor = True
    _load_image_processor = False
    _load_feature_extractor = False
    _load_tokenizer = False

    def _sanitize_parameters(
        self,
        processor_kwargs=None,
        image_processor_kwargs=None,
        return_tensors=None,
        pool=None,
        normalize=None,
        patch_size=None,
        **kwargs,
    ):
        preprocess_params = {}
        if processor_kwargs is not None:
            preprocess_params.update(processor_kwargs)
        if image_processor_kwargs is not None:
            preprocess_params.update(image_processor_kwargs)

        for key in (
            "s1",
            "s2",
            "era5",
            "tc",
            "viirs",
            "srtm",
            "dw",
            "wc",
            "landscan",
            "latlon",
            "months",
        ):
            if key in kwargs:
                preprocess_params[key] = kwargs.pop(key)

        if normalize is not None:
            preprocess_params["normalize"] = normalize
        if patch_size is not None:
            preprocess_params["patch_size"] = patch_size

        postprocess_params = {}
        if pool is not None:
            postprocess_params["pool"] = pool
        if return_tensors is not None:
            postprocess_params["return_tensors"] = return_tensors

        return preprocess_params, {}, postprocess_params

    def preprocess(self, image=None, timeout=None, **processor_kwargs) -> dict[str, GenericTensor]:
        del image, timeout
        model_inputs = self.processor(return_tensors="pt", **processor_kwargs)
        model_inputs = model_inputs.to(self.dtype)
        return model_inputs

    def _forward(self, model_inputs):
        patch_size = model_inputs.pop("patch_size", None)
        if torch.is_tensor(patch_size):
            patch_size = int(patch_size.item())
        return self.model(patch_size=patch_size, **model_inputs)

    def __call__(self, *args: Union[Any, list[Any]], **kwargs: Any) -> list[Any]:
        if not args:
            args = (None,)
        return super().__call__(*args, **kwargs)


__all__ = ["GalileoImageFeatureExtractionPipeline"]