# 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.
import os
from abc import ABC, abstractmethod
from typing import Union

import cupy as cp
import torch

from nemo_curator.datasets import ImageTextPairDataset
from nemo_curator.utils.cudf_utils import create_list_series_from_1d_or_2d_ar
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
    """

    def __init__(
        self,
        model_name: str,
        embedding_column: str,
        pred_column: str,
        pred_type: Union[str, type],
        batch_size: int,
        embedding_size: int,
    ) -> None:
        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:
        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):
        pass

    def postprocess(self, series):
        return series
