# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from abc import ABC, abstractmethod
from typing import Callable, Union

import cudf
import cupy as cp
import torch
from crossfit.backend.cudf.series import create_list_series_from_1d_or_2d_ar

from nemo_curator.datasets import ImageTextPairDataset
from nemo_curator.utils.distributed_utils import load_object_on_worker


class ImageClassifier(ABC):
    """
    An abstract base class that represents a classifier on top
    of embeddings generated by a CLIP vision encoder.

    Subclasses only need to define how a model is loaded.
    They may also override the postprocess method if they would like
    to modify output series of predictions before it gets combined into
    the dataset. The classifier must be able to fit on a single GPU.
    """

    def __init__(
        self,
        model_name: str,
        embedding_column: str,
        pred_column: str,
        pred_type: Union[str, type],
        batch_size: int,
        embedding_size: int,
    ) -> None:
        """
        Constructs an image classifier.

        Args:
            model_name (str): A unqiue name to identify the model on each worker
                and in the logs.
            embedding_column (str): The column name that stores the image
                embeddings.
            pred_column (str): The column name to be added where the classifier's
                predictions will be stored.
            pred_type (Union[str, type]): The datatype of the pred_column.
            batch_size (int): If greater than 0, the image embeddings
                will be processed in batches of at most this size. If less than 0,
                all embeddings will be processed at once.
        """
        self.model_name = model_name
        self.embedding_column = embedding_column
        self.pred_column = pred_column
        self.pred_type = pred_type
        self.batch_size = batch_size
        self.embedding_size = embedding_size

    def __call__(self, dataset: ImageTextPairDataset) -> ImageTextPairDataset:
        """
        Classifies all embeddings in the dataset.

        Args:
            dataset (ImageTextPairDataset): The dataset to classify.

        Returns:
            ImageTextPairDataset: A dataset with classifier scores.
        """
        meta = dataset.metadata.dtypes.to_dict()
        meta[self.pred_column] = self.pred_type
        embedding_df = dataset.metadata.map_partitions(self._run_inference, meta=meta)

        return ImageTextPairDataset(
            dataset.path,
            metadata=embedding_df,
            tar_files=dataset.tar_files,
            id_col=dataset.id_col,
        )

    def _run_inference(self, partition, partition_info=None):
        device = "cuda"

        model = load_object_on_worker(
            self.model_name,
            self.load_model,
            {"device": device},
        )

        embeddings = torch.as_tensor(
            partition[self.embedding_column].list.leaves.values.reshape(
                len(partition), -1
            ),
            device=device,
        )

        if self.embedding_size != embeddings.shape[-1]:
            raise RuntimeError(
                f"{self.model_name} expects embedding size {self.embedding_size} but column "
                f"'{self.embedding_column}' has embedding size {embeddings.shape[-1]}. Ensure your "
                "classifier is compatible with the CLIP model you used to generate the embeddings."
            )

        with torch.no_grad():
            if self.batch_size > 0:
                batches = torch.split(embeddings, self.batch_size)
                model_results = []
                for batch in batches:
                    batch_results = model(batch)
                    model_results.append(batch_results)
                scores = torch.cat(model_results, dim=0)
            else:
                scores = model(embeddings)

        scores = cp.asarray(scores)

        series = create_list_series_from_1d_or_2d_ar(scores, index=partition.index)
        postprocessed_series = self.postprocess(series)
        partition[self.pred_column] = postprocessed_series

        print(
            f"Partition {partition_info['number']} - Classification with {self.model_name} completed for {len(scores)} samples."
        )

        return partition

    @abstractmethod
    def load_model(self, device: str) -> Callable:
        """
        Loads the classifier model.

        Args:
            device (str): A PyTorch device identifier that specifies what GPU
                to load the model on.

        Returns:
            Callable: A callable model, usually a torch.nn.Module.
                The input to this model will be the batches of images output
                by the ImageEmbedder.load_dataset_shard.
        """
        pass

    def postprocess(self, series: cudf.Series) -> cudf.Series:
        """
        Postprocesses the predictions of the classifier before saving
        them to the metadata.

        Args:
            series (cudf.Series): The cuDF series of raw model predictions.

        Returns:
            cudf.Series: The same series unmodified. Override in your classifier
                if needed.
        """
        return series
