# 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.
"""Prithvi-EO-2.0 spatiotemporal feature extraction pipeline."""

from typing import Any, Union

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


if is_vision_available():
    from transformers.image_utils import load_image


@add_end_docstrings(
    build_pipeline_init_args(has_image_processor=True),
    """
        pool (`bool`, *optional*, defaults to `False`):
            Whether to return the CLS token (`True`) or the full token sequence (`False`).
        temporal_coords (`list[list[float]]`, *optional*):
            Year and Julian day for each frame, shape `(T, 2)`.
        location_coords (`list[float]`, *optional*):
            Center latitude and longitude, shape `(2,)`.
    """,
)
class PrithviFeatureExtractionPipeline(Pipeline):
    """
    Prithvi-EO-2.0 spatiotemporal feature extraction pipeline.

    Example:

    ```python
    >>> from transformers import pipeline
    >>> import numpy as np

    >>> pipe = pipeline(
    ...     task="prithvi-eo-feature-extraction",
    ...     model="prithvi-eo-v2-300m-tl",
    ...     trust_remote_code=True,
    ... )
    >>> frames = [np.random.rand(224, 224, 6).astype("float32") for _ in range(4)]
    >>> features = pipe(frames, pool=True, return_tensors=True)
    ```
    """

    def _sanitize_parameters(
        self,
        image_processor_kwargs=None,
        return_tensors=None,
        pool=None,
        temporal_coords=None,
        location_coords=None,
        **kwargs,
    ):
        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
        if "timeout" in kwargs:
            preprocess_params["timeout"] = kwargs["timeout"]
        if temporal_coords is not None:
            preprocess_params["temporal_coords"] = temporal_coords
        if location_coords is not None:
            preprocess_params["location_coords"] = location_coords

        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, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
            image = load_image(image, timeout=timeout)
        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
        model_inputs = model_inputs.to(self.dtype)
        return model_inputs

    def _forward(self, model_inputs):
        temporal_coords = model_inputs.pop("temporal_coords", None)
        location_coords = model_inputs.pop("location_coords", None)
        return self.model(
            temporal_coords=temporal_coords,
            location_coords=location_coords,
            **model_inputs,
        )

    def postprocess(self, model_outputs, pool=False, return_tensors=False):
        if pool:
            output = model_outputs.pooler_output
        else:
            output = model_outputs.last_hidden_state

        if return_tensors:
            return output
        return output.tolist()

    def _is_temporal_frame_list(self, value) -> bool:
        if not isinstance(value, list) or len(value) == 0:
            return False
        if not all(hasattr(item, "shape") for item in value):
            return False
        shapes = [tuple(item.shape) for item in value]
        return all(len(shape) == 3 for shape in shapes) and len(set(shapes)) == 1

    def __call__(
        self,
        *args: Union[str, Any, list[Any]],
        **kwargs: Any,
    ) -> list[Any]:
        if len(args) == 1 and self._is_temporal_frame_list(args[0]):
            results = super().__call__([args[0]], **kwargs)
            return results[0] if len(results) == 1 else results
        results = super().__call__(*args, **kwargs)
        if len(args) == 1 and not isinstance(args[0], list) and len(results) == 1:
            return results[0]
        return results


__all__ = ["PrithviFeatureExtractionPipeline"]