#! COCO image captioning dataset

import os

from clarifai.datasets.upload.base import ClarifaiDataLoader

from ..features import VisualClassificationFeatures

#pycocotools is a dependency for this loader
try:
  from pycocotools.coco import COCO
except ImportError:
  raise ImportError("Could not import pycocotools package. "
                    "Please do `pip install 'clarifai[all]'` to import pycocotools.")


class COCOCaptionsDataLoader(ClarifaiDataLoader):
  """COCO Image Captioning Dataset."""

  def __init__(self, images_dir, label_filepath):
    """
    Args:
      images_dir: Directory containing the images.
      label_filepath: Path to the COCO annotation file.
    """
    self.images_dir = images_dir
    self.label_filepath = label_filepath

    self.map_ids = {}
    self.load_data()

  @property
  def task(self):
    return "visual_captioning"

  def load_data(self) -> None:
    self.coco = COCO(self.label_filepath)
    self.map_ids = {i: img_id for i, img_id in enumerate(list(self.coco.imgs.keys()))}

  def __len__(self):
    return len(self.coco.imgs)

  def __getitem__(self, index):
    value = self.coco.imgs[self.map_ids[index]]
    image_path = os.path.join(self.images_dir, value['file_name'])
    annots = []

    input_ann_ids = self.coco.getAnnIds(imgIds=[value['id']])
    input_anns = self.coco.loadAnns(input_ann_ids)

    for ann in input_anns:
      annots.append(ann['caption'])

    return VisualClassificationFeatures(image_path, labels=annots[0], id=str(value['id']))
