"""
Pipeline modules for basic image operations.
"""

import time

from typing import Tuple

from typeguard import typechecked
from scipy.ndimage import rotate

from pynpoint.core.processing import ProcessingModule
from pynpoint.util.module import progress, memory_frames


class SubtractImagesModule(ProcessingModule):
    """
    Pipeline module for subtracting two sets of images.
    """

    __author__ = "Tomas Stolker"

    @typechecked
    def __init__(
        self,
        name_in: str,
        image_in_tags: Tuple[str, str],
        image_out_tag: str,
        scaling: float = 1.0,
    ) -> None:
        """
        Parameters
        ----------
        name_in : str
            Unique name of the module instance.
        image_in_tags : tuple(str, str)
            Tuple with two tags of the database entry that are read as input.
        image_out_tag : str
            Tag of the database entry with the subtracted images that are written as output.
        scaling : float
            Additional scaling factor.

        Returns
        -------
        NoneType
            None
        """

        super().__init__(name_in)

        self.m_image_in1_port = self.add_input_port(image_in_tags[0])
        self.m_image_in2_port = self.add_input_port(image_in_tags[1])
        self.m_image_out_port = self.add_output_port(image_out_tag)

        self.m_scaling = scaling

    @typechecked
    def run(self) -> None:
        """
        Run method of the module. Subtracts the images from the second database tag from the images
        of the first database tag, on a frame-by-frame basis.

        Returns
        -------
        NoneType
            None
        """

        if self.m_image_in1_port.get_shape() != self.m_image_in2_port.get_shape():
            raise ValueError("The shape of the two input tags has to be the same.")

        memory = self._m_config_port.get_attribute("MEMORY")
        nimages = self.m_image_in1_port.get_shape()[0]
        frames = memory_frames(memory, nimages)

        start_time = time.time()

        for i, _ in enumerate(frames[:-1]):
            progress(i, len(frames[:-1]), "Subtracting images...", start_time)

            images1 = self.m_image_in1_port[frames[i] : frames[i + 1],]
            images2 = self.m_image_in2_port[frames[i] : frames[i + 1],]

            self.m_image_out_port.append(
                self.m_scaling * (images1 - images2), data_dim=3
            )

        history = f"scaling = {self.m_scaling}"
        self.m_image_out_port.add_history("SubtractImagesModule", history)
        self.m_image_out_port.copy_attributes(self.m_image_in1_port)
        self.m_image_out_port.close_port()


class AddImagesModule(ProcessingModule):
    """
    Pipeline module for adding two sets of images.
    """

    __author__ = "Tomas Stolker"

    @typechecked
    def __init__(
        self,
        name_in: str,
        image_in_tags: Tuple[str, str],
        image_out_tag: str,
        scaling: float = 1.0,
    ) -> None:
        """
        Parameters
        ----------
        name_in : str
            Unique name of the module instance.
        image_in_tags : tuple(str, str)
            Tuple with two tags of the database entry that are read as input.
        image_out_tag : str
            Tag of the database entry with the added images that are written as output.
        scaling: float
            Additional scaling factor.

        Returns
        -------
        NoneType
            None
        """

        super().__init__(name_in)

        self.m_image_in1_port = self.add_input_port(image_in_tags[0])
        self.m_image_in2_port = self.add_input_port(image_in_tags[1])
        self.m_image_out_port = self.add_output_port(image_out_tag)

        self.m_scaling = scaling

    @typechecked
    def run(self) -> None:
        """
        Run method of the module. Add the images from the two database tags on a frame-by-frame
        basis.

        Returns
        -------
        NoneType
            None
        """

        if self.m_image_in1_port.get_shape() != self.m_image_in2_port.get_shape():
            raise ValueError("The shape of the two input tags has to be the same.")

        memory = self._m_config_port.get_attribute("MEMORY")
        nimages = self.m_image_in1_port.get_shape()[0]
        frames = memory_frames(memory, nimages)

        start_time = time.time()

        for i, _ in enumerate(frames[:-1]):
            progress(i, len(frames[:-1]), "Adding images...", start_time)

            images1 = self.m_image_in1_port[frames[i] : frames[i + 1],]
            images2 = self.m_image_in2_port[frames[i] : frames[i + 1],]

            self.m_image_out_port.append(
                self.m_scaling * (images1 + images2), data_dim=3
            )

        history = f"scaling = {self.m_scaling}"
        self.m_image_out_port.add_history("AddImagesModule", history)
        self.m_image_out_port.copy_attributes(self.m_image_in1_port)
        self.m_image_out_port.close_port()


class RotateImagesModule(ProcessingModule):
    """
    Pipeline module for rotating images.
    """

    __author__ = "Tomas Stolker"

    @typechecked
    def __init__(
        self, name_in: str, image_in_tag: str, image_out_tag: str, angle: float
    ) -> None:
        """
        Parameters
        ----------
        name_in : str
            Unique name of the module instance.
        image_in_tag : str
            Tag of the database entry that is read as input.
        image_out_tag : str
            Tag of the database entry that is written as output.
        angle : float
            Rotation angle (deg). Rotation is clockwise for positive values.
        Returns
        -------
        NoneType
            None
        """

        super().__init__(name_in)

        self.m_image_in_port = self.add_input_port(image_in_tag)
        self.m_image_out_port = self.add_output_port(image_out_tag)

        self.m_angle = angle

    @typechecked
    def run(self) -> None:
        """
        Run method of the module. Rotates all images by a constant angle.

        Returns
        -------
        NoneType
            None
        """

        memory = self._m_config_port.get_attribute("MEMORY")
        nimages = self.m_image_in_port.get_shape()[0]
        frames = memory_frames(memory, nimages)

        start_time = time.time()

        for i, _ in enumerate(frames[:-1]):
            progress(i, len(frames[:-1]), "Rotating images...", start_time)

            images = self.m_image_in_port[frames[i] : frames[i + 1],]

            for j in range(frames[i + 1] - frames[i]):
                im_tmp = images[j,]

                # ndimage.rotate rotates in clockwise direction for positive angles
                im_tmp = rotate(im_tmp, self.m_angle, reshape=False)

                self.m_image_out_port.append(im_tmp, data_dim=3)

        history = f"angle [deg] = {self.m_angle}"
        self.m_image_out_port.add_history("RotateImagesModule", history)
        self.m_image_out_port.copy_attributes(self.m_image_in_port)
        self.m_image_out_port.close_port()


class RepeatImagesModule(ProcessingModule):
    """
    Pipeline module for repeating the images from a dataset.
    """

    __author__ = "Tomas Stolker"

    @typechecked
    def __init__(
        self, name_in: str, image_in_tag: str, image_out_tag: str, repeat: int
    ) -> None:
        """
        Parameters
        ----------
        name_in : str
            Unique name of the module instance.
        image_in_tag : str
            Tag of the database entry that is read as input.
        image_out_tag : str
            Tag of the database entry with the added images that are written as output.
        repeat: int
            The number of times the input images get repeated.

        Returns
        -------
        NoneType
            None
        """

        super().__init__(name_in)

        self.m_image_in_port = self.add_input_port(image_in_tag)
        self.m_image_out_port = self.add_output_port(image_out_tag)

        self.m_repeat = repeat

    @typechecked
    def run(self) -> None:
        """
        Run method of the module. Repeats the stack of input images a specified number of times.

        Returns
        -------
        NoneType
            None
        """

        nimages = self.m_image_in_port.get_shape()[0]
        memory = self._m_config_port.get_attribute("MEMORY")

        frames = memory_frames(memory, nimages)

        start_time = time.time()

        for i in range(self.m_repeat):
            progress(i, self.m_repeat, "Repeating images...", start_time)

            for j, _ in enumerate(frames[:-1]):
                images = self.m_image_in_port[frames[j] : frames[j + 1],]
                self.m_image_out_port.append(images, data_dim=3)

        history = f"repeat = {self.m_repeat}"
        self.m_image_out_port.add_history("RepeatImagesModule", history)
        self.m_image_out_port.copy_attributes(self.m_image_in_port)
        self.m_image_out_port.close_port()
