# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

from pandas import DataFrame as PdDataFrame
from pyspark.sql.dataframe import DataFrame as SparkDataFrame
from ..dataaccess._open_image_blob_info import OpenImageBlobInfo
from azure.storage.blob import BlockBlobService
from ..environ import SparkEnv, PandasEnv
from multimethods import multimethod
import os
from pyspark.sql import functions as F
from pyspark.sql import SparkSession
import json


class ImageDownloader:
    """Open Image Dataset image downloader class."""

    @staticmethod
    def DownloadImages(elements, blobInfo, env, path):
        """
        Download images, used in GetImages function

        :param elements: each element is a tuple, it consists of (urlPrefix, subset/imageId, savePath)
         for example ('https://openimagestorage.blob.core.windows.net/openimagedataset/',
         'validation/0001eeaf4aed83f9.jpg', 'tmp/')
        :type elements: iterator
        :param runType: local or spark
        :type runType: str
        """
        containerName = blobInfo.blob_container_name
        blobService = BlockBlobService(
            account_name=blobInfo.blob_account_name,
            sas_token=blobInfo.blob_sas_token.lstrip('?'))
        for relativePath in elements:
            if type(env) == SparkEnv:
                # save image in dbfs, path format : "/dbfs/tmp..." -> tmp folder at dbfs
                imagePath = os.path.join("/dbfs", path, relativePath.split("/")[-1])
                blobService.get_blob_to_path(containerName, relativePath, imagePath)
            elif type(env) == PandasEnv:
                imagePath = os.path.join(path, relativePath.split("/")[-1])
                if not os.path.exists(imagePath):
                    blobService.get_blob_to_path(containerName, relativePath, imagePath)
            else:
                raise ValueError("Input is invalid")

    @multimethod(OpenImageBlobInfo, PdDataFrame, PdDataFrame, str, dict, str, str)
    def GetImages(self, blobInfo, bbox, urlTable, path, imageDict, storeType, tokens):
        """
        Download images (named by image ids) and a image attribute file (imageid, bb-info-list) into local path
        or customer azure blob.

        :param urlPrefix: prefix of url
        :type urlPrefix: str
        :param bbox: pandans dataframe contains bounding box information
        :type bbox: PdDataFrame
        :param urlTable: pandas dataframe contains imageID and relative path on azure blob.
        :type urlTable: PdDataFrame
        :param path: download path
        :type path: str
        :param imageDict: dict of imageID and bounding boxID list
        :type imageDict: dict
        :param storeType: local or azureblob.
        :type storeType: str
        :param tokens: for authentication on azure blob.
        :type tokens: str

        """
        if not imageDict:
            raise ValueError("Input can not be empty.")

        imageSavePath = os.path.join(path, "image")
        if not os.path.exists(imageSavePath):
            os.makedirs(imageSavePath)
        imageIds = list(imageDict.keys())

        # annotation file
        bboxLists = [v for imageID in imageIds for v in imageDict.get(imageID)]
        bbox = bbox[bbox.index.isin(bboxLists)].copy()
        bbox.reset_index(inplace=True)
        bbox.loc[:, 'BoxIDList'] = bbox.apply(lambda x: {'BoxID': x['BoxID'], 'XMin': x['XMin'], 'XMax': x['XMax'],
                                                         'YMin': x['YMin'], 'YMax': x['YMax'],
                                                         'IsOccluded': x['IsOccluded'],
                                                         'IsTruncated': x['IsTruncated'], 'IsGroupOf': x['IsGroupOf'],
                                                         'IsInside': x['IsInside'],
                                                         'Category': x['Category']},axis=1)
        df = bbox.groupby("ImageID", observed=True)['BoxIDList'].apply(list)
        with open(os.path.join(path, "annotation.json"), "w") as f:
            json.dump(df.reset_index().to_dict(orient='records'), f)

        print('download start...')
        # download images
        elements = []
        for imageId in imageIds:
            relativePath = urlTable.loc[imageId, "Subset"] + "/" + imageId + ".jpg"
            elements.append(relativePath)
        ImageDownloader.DownloadImages(elements, blobInfo, PandasEnv(), imageSavePath)
        print('finish.')

    @multimethod(OpenImageBlobInfo, SparkDataFrame, dict, str, dict, str, str)
    def GetImages(self, blobInfo, bbox, urlTable, path, imageDict, storeType, tokens):
        """
        Download images (named by image ids) and a image attribute file (imageid, categorylist, bb-info-list)
        into local path or customer azure blob.

        :param urlPrefix: prefix of url
        :type urlPrefix: str
        :param bbox: spark dataframe contains bounding box information
        :type bbox: SparkDataFrame
        :param urlTable: dict contains imageID and relative path on azure blob.
        :type urlTable: dict
        :param path: download path
        :type path: str
        :param imageDict: dict of imageID and bounding boxID list
        :type imageDict: dict
        :param storeType: local or azureblob.
        :type storeType: str
        :param tokens: for authentication on azure blob.
        :type tokens: str

        """

        if not imageDict:
            raise ValueError("Input can not be null.")

        # validate input path, starting with / or dbfs: is both valid
        if path.startswith("dbfs:"):
            path = path.split(":")[1][1:]
        elif path.startswith("/"):
            path = path[1:]
        else:
            raise ValueError("The path is invalid")

        spark = SparkSession.builder.getOrCreate()
        # get image list and box list
        imageIds = list(imageDict.keys())
        bboxLists = [v for imageID in imageIds for v in imageDict.get(imageID)]

        # annotaiton file
        bboxListDF = spark.sparkContext.parallelize(bboxLists).map(lambda x: (x, )).toDF(["BoxID"])
        fileterDF = bboxListDF.join(bbox, ["BoxID"], "inner")
        fileterDF = fileterDF.withColumn("BoxList", F.struct(F.col("BoxID"), F.col("XMin"), F.col("XMax"),
                                         F.col("YMin"), F.col("YMax"), F.col("IsOccluded"), F.col("IsTruncated"),
                                         F.col("IsGroupOf"), F.col("IsInside"), F.col("Category")))
        fileterDF.groupBy("ImageID").agg(F.collect_list("BoxList").alias("BoxList")) \
            .write.json(os.path.join(path, "annoatation"))

        # download image
        urlList = []
        for imageId in imageIds:
            url = os.path.join(urlTable.get(imageId), imageId + ".jpg")
            urlList.append(url)
        func = ImageDownloader.DownloadImages
        spark.sparkContext.parallelize(urlList).foreachPartition(lambda x: func(x, blobInfo, SparkEnv(), path))

    @multimethod(PdDataFrame, str, int, list, int, str, str)
    def GetImagesByBatch(self, urlPrefix, imageInfoTable, df, path, partitionSize,
                         imageIds, batchsize, storeType, tokens):
        """
        Return a list of images in image classes for training/testing
        and at the same time a file described imageInstance list will be created in local or blob
        image class will contain:
        (1) image bits
        (2) image tensors in RGB format and will consider other type of format
        (3) image category list (may use number only, like 0 means car, 1 means human, ...)
        (4) image bb list (each element in the list will be like (xMin, yMin, bbcategory, isTruncated, isOcclued, ...))
        (5) ...
        :param storeType: local or azureblob.
        :type: str
        :param path: download path.
        :type path: str
        :param tokens: for authentication on azure blob.
        :type tokens: str
        :param partitionSize: ...
        :type partitionSize: int
        :param imageIds: if None, means all images.
        :type imageIds: list
        :param batchsize: download images at the same time when training.
        :type batchsize: int
        :return: a list of images in image classes for training/testing.
        :rtype: list
        """

        pass

    @multimethod(SparkDataFrame, str, int, list, int, str, str)
    def GetImagesByBatch(self, urlPrefix, imageInfoTable, df, path, partitionSize,
                         imageIds, batchsize, storeType, tokens):
        """
        Return a list of images in image classes for training/testing
        and at the same time a file described imageInstance list will be created in local or blob
        image class will contain:
        (1) image bits
        (2) image tensors in RGB format and will consider other type of format
        (3) image category list (may use number only, like 0 means car, 1 means human, ...)
        (4) image bb list (each element in the list will be like (xMin, yMin, bbcategory, isTruncated, isOcclued, ...))
        (5) ...
        :param storeType: local or azureblob.
        :type: str
        :param path: download path.
        :type path: str
        :param tokens: for authentication on azure blob.
        :type tokens: str
        :param partitionSize: ...
        :type partitionSize: int
        :param imageIds: if None, means all images.
        :type imageIds: list
        :param batchsize: download images at the same time when training.
        :type batchsize: int
        :return: a list of images in image classes for training/testing.
        :rtype: list
        """

        pass
