#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""
.. attribute:: ImageSchema

    An attribute of this module that contains the instance of :class:`_ImageSchema`.

.. autoclass:: _ImageSchema
   :members:
"""

import sys
from typing import Any, Dict, List, NoReturn, cast
from functools import cached_property

import numpy as np

from pyspark.sql.types import Row, StructType, _create_row, _parse_datatype_json_string
from pyspark.sql import SparkSession

__all__ = ["ImageSchema"]


class _ImageSchema:
    """
    Internal class for `pyspark.ml.image.ImageSchema` attribute. Meant to be private and
    not to be instantized. Use `pyspark.ml.image.ImageSchema` attribute to access the
    APIs of this class.
    """

    @cached_property
    def imageSchema(self) -> StructType:
        """
        Returns the image schema.

        Returns
        -------
        :class:`StructType`
            with a single column of images named "image" (nullable)
            and having the same type returned by :meth:`columnSchema`.

        .. versionadded:: 2.3.0
        """
        from pyspark.core.context import SparkContext

        ctx = SparkContext._active_spark_context
        assert ctx is not None and ctx._jvm is not None
        jschema = getattr(ctx._jvm, "org.apache.spark.ml.image.ImageSchema").imageSchema()
        return cast(StructType, _parse_datatype_json_string(jschema.json()))

    @cached_property
    def ocvTypes(self) -> Dict[str, int]:
        """
        Returns the OpenCV type mapping supported.

        Returns
        -------
        dict
            a dictionary containing the OpenCV type mapping supported.

        .. versionadded:: 2.3.0
        """
        from pyspark.core.context import SparkContext

        ctx = SparkContext._active_spark_context
        assert ctx is not None and ctx._jvm is not None
        return dict(getattr(ctx._jvm, "org.apache.spark.ml.image.ImageSchema").javaOcvTypes())

    @cached_property
    def columnSchema(self) -> StructType:
        """
        Returns the schema for the image column.

        Returns
        -------
        :class:`StructType`
            a schema for image column,
            ``struct<origin:string, height:int, width:int, nChannels:int, mode:int, data:binary>``.

        .. versionadded:: 2.4.0
        """
        from pyspark.core.context import SparkContext

        ctx = SparkContext._active_spark_context
        assert ctx is not None and ctx._jvm is not None
        jschema = getattr(ctx._jvm, "org.apache.spark.ml.image.ImageSchema").columnSchema()
        return cast(StructType, _parse_datatype_json_string(jschema.json()))

    @cached_property
    def imageFields(self) -> List[str]:
        """
        Returns field names of image columns.

        Returns
        -------
        list
            a list of field names.

        .. versionadded:: 2.3.0
        """
        from pyspark.core.context import SparkContext

        ctx = SparkContext._active_spark_context
        assert ctx is not None and ctx._jvm is not None
        return list(getattr(ctx._jvm, "org.apache.spark.ml.image.ImageSchema").imageFields())

    @cached_property
    def undefinedImageType(self) -> str:
        """
        Returns the name of undefined image type for the invalid image.

        .. versionadded:: 2.3.0
        """
        from pyspark.core.context import SparkContext

        ctx = SparkContext._active_spark_context
        assert ctx is not None and ctx._jvm is not None
        return getattr(ctx._jvm, "org.apache.spark.ml.image.ImageSchema").undefinedImageType()

    def toNDArray(self, image: Row) -> np.ndarray:
        """
        Converts an image to an array with metadata.

        Parameters
        ----------
        image : :class:`Row`
            image: A row that contains the image to be converted. It should
            have the attributes specified in `ImageSchema.imageSchema`.

        Returns
        -------
        :class:`numpy.ndarray`
            that is an image.

        .. versionadded:: 2.3.0
        """

        if not isinstance(image, Row):
            raise TypeError(
                "image argument should be pyspark.sql.types.Row; however, "
                "it got [%s]." % type(image)
            )

        if any(not hasattr(image, f) for f in self.imageFields):
            raise ValueError(
                "image argument should have attributes specified in "
                "ImageSchema.imageSchema [%s]." % ", ".join(self.imageFields)
            )

        height = image.height
        width = image.width
        nChannels = image.nChannels
        return np.ndarray(
            shape=(height, width, nChannels),
            dtype=np.uint8,
            buffer=image.data,
            strides=(width * nChannels, nChannels, 1),
        )

    def toImage(self, array: np.ndarray, origin: str = "") -> Row:
        """
        Converts an array with metadata to a two-dimensional image.

        Parameters
        ----------
        array : :class:`numpy.ndarray`
            The array to convert to image.
        origin : str
            Path to the image, optional.

        Returns
        -------
        :class:`Row`
            that is a two dimensional image.

        .. versionadded:: 2.3.0
        """

        if not isinstance(array, np.ndarray):
            raise TypeError(
                "array argument should be numpy.ndarray; however, it got [%s]." % type(array)
            )

        if array.ndim != 3:
            raise ValueError("Invalid array shape")

        height, width, nChannels = array.shape
        ocvTypes = ImageSchema.ocvTypes
        if nChannels == 1:
            mode = ocvTypes["CV_8UC1"]
        elif nChannels == 3:
            mode = ocvTypes["CV_8UC3"]
        elif nChannels == 4:
            mode = ocvTypes["CV_8UC4"]
        else:
            raise ValueError("Invalid number of channels")

        data = bytearray(array.astype(dtype=np.uint8).ravel().tobytes())

        # Creating new Row with _create_row(), because Row(name = value, ... )
        # orders fields by name, which conflicts with expected schema order
        # when the new DataFrame is created by UDF
        return _create_row(self.imageFields, [origin, height, width, nChannels, mode, data])


ImageSchema = _ImageSchema()


# Monkey patch to disallow instantiation of this class.
def _disallow_instance(_: Any) -> NoReturn:
    raise RuntimeError("Creating instance of _ImageSchema class is disallowed.")


_ImageSchema.__init__ = _disallow_instance  # type: ignore[assignment]


def _test() -> None:
    import doctest
    import pyspark.ml.image

    globs = pyspark.ml.image.__dict__.copy()
    spark = SparkSession.builder.master("local[2]").appName("ml.image tests").getOrCreate()
    globs["spark"] = spark

    (failure_count, test_count) = doctest.testmod(
        pyspark.ml.image, globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE
    )
    spark.stop()
    if failure_count:
        sys.exit(-1)


if __name__ == "__main__":
    _test()