"""
This is the utils module related to the main
editor.
"""
from yta_validation import PythonValidator
from yta_programming.decorators.requires_dependency import requires_dependency
from typing import Union


# TODO: I think I have this methods in other library
# maybe yta-image-utils (?)
@requires_dependency('cv2', 'yta_editor_utils', 'opencv-python')
def scale_numpy(
    input: 'np.ndarray',
    size: tuple[int, int]
) -> 'np.ndarray':
    """
    *Optional dependency `opencv-python` required*

    Resize the 'input' numpy array to the provided 'size'
    if needed, using a rescaling method with 'opencv-python'
    (cv2).

    The 'size' provided must be (width, height).
    """
    import cv2

    return cv2.resize(input, size, interpolation = cv2.INTER_LINEAR)

@requires_dependency('PIL', 'yta_editor_utils', 'pillow')
@requires_dependency('numpy', 'yta_editor_utils', 'numpy')
def scale_numpy_pillow(
    input: 'np.ndarray',
    size: tuple[int, int]
) -> 'np.ndarray':
    """
    *Optional dependency `pillow` (imported as `PIL`) required*

    *Optional dependency `numpy` (imported as `numpy`) required*

    Resize the 'input' numpy array to the provided 'size'
    if needed, using a rescaling method with 'pillow'
    (PIL).

    The 'size' provided must be (width, height).
    """
    from PIL import Image
    import numpy as np

    return np.array(Image.fromarray(input).resize(size, Image.BILINEAR))

@requires_dependency('PIL', 'yta_editor_utils', 'pillow')
@requires_dependency('numpy', 'yta_editor_utils', 'numpy')
def read_image_as_numpy(
    path: str,
    do_read_as_rgba: bool = True,
    size: Union[tuple[int, int], None] = None
) -> 'np.ndarray':
    """
    *Optional dependency `pillow` (imported as `PIL`) required*

    *Optional dependency `numpy` (imported as `numpy`) required*

    Read an image from a file and transform it into a
    numpy array. It will force the 'size' if provided,
    or leave the original one if it is None.
    """
    from PIL import Image
    import numpy as np

    mode = (
        'RGBA'
        if do_read_as_rgba else
        'RGB'
    )

    with Image.open(path) as img:
        img = img.convert(mode)
        np_img = np.array(img, dtype = np.uint8)

    return (
        scale_numpy_pillow(
            input = np_img,
            size = size
        ) if size is not None else
        np_img
    )

@requires_dependency('PIL', 'yta_editor_utils', 'pillow')
@requires_dependency('numpy', 'yta_editor_utils', 'numpy')
def numpy_to_file(
    input: 'np.ndarray',
    output_filename: str
) -> str:
    """
    *Optional dependency `pillow` (imported as `PIL`) required*

    Export the provided 'array' numpy array as a file
    with the given 'output_filename' name.
    """
    from PIL import Image

    Image.fromarray(input).save(output_filename)

    return output_filename

class TextureUtils:
    """
    Class to be used as a static class, including
    some utils related to handling and converting
    OpenGL textures.

    Check the optional libraries that are needed
    to use the different methods included in this
    class.
    """

    # TODO: I should make different methods to
    # obtain a VideoFrame or a numpy array frame
    @staticmethod
    @requires_dependency('numpy', 'yta_editor_utils', 'numpy')
    @requires_dependency('moderngl', 'yta_editor_utils', 'moderngl')
    def texture_to_frame(
        texture: 'moderngl.Texture',
        do_include_alpha: bool = True
    ) -> 'np.ndarray':
        """
        *Optional dependency `numpy` (imported as `numpy`) required*

        *Optional dependency `moderngl` (imported as `moderngl`) required*

        Transform an opengl texture into numpy array.

        The `do_include_alpha` will include the
        alpha channel if True.
        """
        import numpy as np
        import moderngl

        data = texture.read(alignment = 1)
        # Read 4 channels always (RGBA8)
        frame = np.frombuffer(data, dtype = np.uint8).reshape((texture.size[1], texture.size[0], 4))
        # TODO: Do this with a utils
        frame = (
            frame
            if do_include_alpha else
            # Discard alpha channel if not needed
            frame[:, :, :3]
        )
        # Opengl gives it with the 'y' inverted
        frame = np.flipud(frame)
        # TODO: This can be returned as a numpy frame

        # This is if we need an 'av' VideoFrame (to
        # export through the demuxer, for example)
        # TODO: I avoid this by now because we don't
        # want to import the pyav library, so this
        # below has to be done with the numpy array
        # received...
        # frame = av.VideoFrame.from_ndarray(frame, format = 'rgba')
        # # TODO: Make this customizable
        # frame = frame.reformat(format = 'yuv420p')

        return frame
    
    # TODO: The code of this one is similar to the
    # 'texture_to_frame', so we should keep only one
    # (maybe this one)
    @staticmethod
    @requires_dependency('numpy', 'yta_editor_utils', 'numpy')
    @requires_dependency('moderngl', 'yta_editor_utils', 'moderngl')
    def texture_to_numpy(
        texture: 'moderngl.Texture',
        do_include_alpha: bool = True
    ) -> 'np.ndarray':
        """
        *Optional dependency `numpy` (imported as `numpy`) required*

        *Optional dependency `moderngl` (imported as `moderngl`) required*

        Transform the provided OpenGL `texture` into a numpy
        array to be able to validate the values that are
        contained in the texture.

        The alpha layer will be removed if the parameter
        `do_include_alpha` is set as False.
        """
        import numpy as np
        import moderngl

        data = texture.read()

        # OpenGL uses only f4 and u1 for the textures so
        # we don't need a mapping
        dtype = (
            np.float32
            if texture.dtype == 'f4' else
            np.uint8
        )

        frame = np.frombuffer(
            buffer = data,
            dtype = dtype
        ).reshape(
            texture.height,
            texture.width,
            texture.components
        )

        # Discard alpha channel if not needed
        frame = (
            frame
            if do_include_alpha else
            frame[:, :, :3]
        )

        # Opengl gives it with the 'y' inverted
        frame = np.flipud(frame)

        return frame

    @staticmethod
    @requires_dependency('PIL', 'yta_editor_utils', 'pillow')
    def texture_to_file(
        texture: 'moderngl.Texture',
        output_filename: str
    ) -> str:
        """
        *Optional dependency `pillow` (imported as `PIL`) required*

        Export the provided OpenGL texture 'texture' as a
        file with the given 'output_filename' name.
        """
        return numpy_to_file(
            input = TextureUtils.texture_to_frame(texture),
            output_filename = output_filename
        )
    
    @requires_dependency('moderngl', 'yta_editor_utils', 'moderngl')
    @requires_dependency('numpy', 'yta_editor_utils', 'numpy')
    def from_frame(
        frame: Union['VideoFrame', 'np.ndarray'],
        context: 'moderngl.Context',
        numpy_format: str = 'rgb24'
    ):
        """
        *Optional dependency `numpy` (imported as `numpy`) required*

        *Optional dependency `moderngl` (imported as `moderngl`) required*
        
        Transform the given 'frame' to an opengl
        texture. The frame can be a VideoFrame
        instance (from pyav library) or a numpy
        array.

        (!) This method is useful to transform a
        frame into a texture quick and for a single
        use, but we have the GPUTextureHandler class
        to handle it in an specific contexto to 
        optimize the performance and avoid creating
        textures but rewriting on them.
        """
        import numpy as np
        import moderngl

        # To numpy RGB inverted for opengl
        frame: np.ndarray = (
            np.flipud(frame.to_ndarray(format = numpy_format))
            if PythonValidator.is_instance_of(frame, 'VideoFrame') else
            np.flipud(frame)
        )

        # TODO: This is not ok as it is not considering the
        # different types of dtypes we can have and handle
        # and we are doing that in other part of the code...

        # Sometimes we have 'float32' values but we need to
        # force 'uint8' to be able to work with
        if frame.dtype != np.uint8:
            frame = np.clip(frame, 0, 255).astype(np.uint8)

        return context.texture(
            size = (frame.shape[1], frame.shape[0]),
            components = frame.shape[2],
            data = frame.tobytes()
        )
    
class NumpyTextureUtils:
    """
    A class to wrap utils related to the numpy and
    texture conversion (with the dtypes and values).
    """

    @staticmethod
    @requires_dependency('numpy', 'yta_editor_utils', 'numpy')
    def numpy_to_float32(
        input: 'np.ndarray'
    ) -> 'np.ndarray':
        """
        Transform the provided 'input' numpy array to the 
        np.float32 dtype if needed. The np.float32 has values
        in the range [0.0, 1.0] and is necesary for some
        operations.

        The formula:
        - `(input.astype(np.float32) / 255.0)`
        (if the dtype is `np.uint8`)
        """
        import numpy as np

        NumpyTextureUtils._validate_numpy_input(input)
        
        return (
            # np.uint8 has values in [0, 255]
            # np.floate2 has values in [0.0, 1.0]
            (input.astype(np.float32) / 255.0)
            if input.dtype == np.uint8 else
            input
        )

    @staticmethod
    @requires_dependency('numpy', 'yta_editor_utils', 'numpy')
    def numpy_to_uint8(
        input: 'np.ndarray'
    ) -> 'np.ndarray':
        """
        Transform the provided 'input' numpy array to the 
        np.uint8 dtype if needed. The np.uint8 has values
        in the range [0, 255] and is necesary for some
        operations.

        The formula:
        - `np.clip(input * 255.0, 0, 255).astype(np.uint8)`
        (if the dtype is `np.float32`)
        """
        import numpy as np

        NumpyTextureUtils._validate_numpy_input(input)
        
        return (
            # np.uint8 has values in [0, 255]
            # np.floate2 has values in [0.0, 1.0]
            np.clip(input * 255.0, 0, 255).astype(np.uint8)
            if input.dtype == np.float32 else
            input
        )

    @staticmethod
    @requires_dependency('numpy', 'yta_editor_utils', 'numpy')
    def _validate_numpy_input(
        input: 'np.ndarray'
    ) -> None:
        """
        Validate that the `input` numpy array has one of
        the accepted dtypes or raise an exception if not.
        """
        import numpy as np

        if input.dtype not in [np.uint8, np.float32]:
            raise Exception(f'The type of the "input" array provided is not "np.uint8" nor "np.float32" but "{input.dtype}".')