import cv2, datetime, os
import numpy as np
from copy import deepcopy
from matplotlib import pyplot as plt

from . import Algorithm

class InvisCloak (Algorithm):

    """ init function """
    def __init__(self):
        # Number of buffered images
        self.n = 5

        # Picture buffer
        self.picture_buffer = []

        # Middle value image built of buffer images
        # Includes noice reduction and histogram spread
        self.middle_value_picture = None

        # Buffer for background image
        self.background = None

        # Clean up results folder
        folder_path = os.path.join("results")

        for filename in os.listdir(folder_path):
            file_path = os.path.join(folder_path, filename)

            if os.path.isfile(file_path):
                os.unlink(file_path)


    """ Processes the input image"""
    def process(self, img):

        """ 2.1 Vorverarbeitung """
        """ 2.1.1 Rauschreduktion """
        plotNoise = False   # Schaltet die Rauschvisualisierung ein
        if plotNoise:
            self._plotNoise(img, "Rauschen vor Korrektur")
        img = self._211_Rauschreduktion(img)
        if plotNoise:
            self._plotNoise(img, "Rauschen nach Korrektur")
        """ 2.1.2 HistogrammSpreizung """
        img = self._212_HistogrammSpreizung(img)


        """ 2.2 Farbanalyse """
        """ 2.2.1 RGB """
        #self._221_RGB(img)
        """ 2.2.2 HSV """
        #self._222_HSV(img)


        """ 2.3 Segmentierung und Bildmodifikation """
        img = self._23_SegmentUndBildmodifizierung(img)

        return img

    """ Reacts on mouse callbacks """
    def mouse_callback(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONUP:
            print("A Mouse click happend! at position", x, y)

            # Stores the current image to data folder
            cv2.imwrite(f"results/{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}_original_image.png",
                        self.picture_buffer[self.n - 1])

            # Create RGB histogram
            self._221_RGB(self.middle_value_picture)

            # Create HSV histogram
            self._222_HSV(self.middle_value_picture)

            # Get binary mask and write it to file
            binary_mask = self._23_SegmentUndBildmodifizierung(self.middle_value_picture, True)
            cv2.imwrite(f"results/{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}_binary_mask.png",
                        binary_mask)
        elif event == cv2.EVENT_MBUTTONUP:
            # Save current image as background
            cv2.imwrite(f"results/background.png", self.picture_buffer[self.n - 1])
            self.background = cv2.imread("results/background.png")


    def _plotNoise(self, img, name:str):
        height, width = np.array(img.shape[:2])
        centY = (height / 2).astype(int)
        centX = (width / 2).astype(int)

        cutOut = 5
        tmpImg = deepcopy(img)
        tmpImg = tmpImg[centY - cutOut:centY + cutOut, centX - cutOut:centX + cutOut, :]

        outSize = 500
        tmpImg = cv2.resize(tmpImg, (outSize, outSize), interpolation=cv2.INTER_NEAREST)

        cv2.imshow(name, tmpImg)
        cv2.waitKey(1)

    def _211_Rauschreduktion(self, img):
        """
            Hier steht Ihr Code zu Aufgabe 2.1.1 (Rauschunterdrückung)
            - Implementierung Mittelwertbildung über N Frames
        """
        self.picture_buffer.append(img)

        if len(self.picture_buffer) < self.n:
            # If number of buffered images < defined buffer size n, return current image
            return img
        elif len(self.picture_buffer) > self.n:
            # If number of buffered images > defined buffer size n, remove oldest image
            self.picture_buffer.pop(0)

        # Reduce noise, return result image
        return np.mean(self.picture_buffer, axis=0).astype(np.uint8)


    def _212_HistogrammSpreizung(self, img):
        """
            Hier steht Ihr Code zu Aufgabe 2.1.2 (Histogrammspreizung)
            - Transformation HSV
            - Histogrammspreizung berechnen
            - Transformation BGR
        """
        # Convert brg image to hsv image
        hsv_image = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

        # Get hsv parts
        h, s, v = cv2.split(hsv_image)

        # Calc histogram spread
        v = cv2.equalizeHist(v)

        # Merge histogram spread to image
        hsv_stretched = cv2.merge([h, s, v])

        # Convert hsv image to brg and store result to member variable
        self.middle_value_picture = cv2.cvtColor(hsv_stretched, cv2.COLOR_HSV2BGR)

        # Return brg result image
        return self.middle_value_picture


    def _221_RGB(self, img, colorspectrum = "bgr"):
        """
            Hier steht Ihr Code zu Aufgabe 2.2.1 (RGB)
            - Histogrammberechnung und Analyse
        """
        # Names of the colors in histogram
        channels = ["b", "g", "r"]

        # Calc histogram
        for index, channel_name in enumerate(channels):
            hist = cv2.calcHist([img], [index], None, [256], [0, 256])
            plt.plot(hist, color=channel_name)
            plt.xlim([0, 256])

        # Save histogram, clear cache
        plt.savefig(f"results/{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}_histogram_{colorspectrum}.png")
        plt.clf()


    def _222_HSV(self, img):
        """
            Hier steht Ihr Code zu Aufgabe 2.2.2 (HSV)
            - Histogrammberechnung und Analyse im HSV-Raum
        """
        # Convert image to hsv, call _221_RGB to create histogram
        self._221_RGB(cv2.cvtColor(img, cv2.COLOR_BGR2HSV), "hsv")


    def _23_SegmentUndBildmodifizierung (self, img, save_binary_mask = False):
        """
            Hier steht Ihr Code zu Aufgabe 2.3.1 (StatischesSchwellwertverfahren)
            - Binärmaske erstellen
        """
        # Convert BGR -> HSV
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

        # Area 1: H = 0-10 (strong Rot)
        lower_red1 = np.array([0, 100, 50])
        upper_red1 = np.array([0, 255, 255])

        # Area 2: H = 169-179 (red-violet)
        lower_red2 = np.array([171, 100, 50])
        upper_red2 = np.array([179, 255, 255])

        # Create binary mask for both red areas
        mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
        mask2 = cv2.inRange(hsv, lower_red2, upper_red2)

        # Combine both masks
        mask = cv2.bitwise_or(mask1, mask2)

        """
            Hier steht Ihr Code zu Aufgabe 2.3.2 (Binärmaske)
            - Binärmaske optimieren mit Opening/Closing
            - Wahl größte zusammenhängende Region
        """

        # Optimizing mask with opening and closing
        kernel = np.ones((5, 5), np.uint8)

        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        # Select biggest connected area
        cnts, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

        if cnts:
            c = max(cnts, key=cv2.contourArea)
            new_mask = np.zeros_like(mask)
            cv2.drawContours(new_mask, [c], -1, color = 255, thickness = -1)
            mask = new_mask
        else:
            mask = np.zeros_like(mask)

        """
            Hier steht Ihr Code zu Aufgabe 2.3.3 (Bildmodifizerung)
            - Hintergrund mit Mausklick definieren
            - Ersetzen des Hintergrundes
        """

        # Return image if no background is set
        if self.background is None:
            return img

        # 3-channel mask for binary operations
        mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)

        # Background are: Get from save background image
        background_part = cv2.bitwise_and(self.background, mask_3ch)

        # Foreground area: Extract from current image
        foreground_part = cv2.bitwise_and(img, cv2.bitwise_not(mask_3ch))

        # Merge both areas
        output = cv2.add(background_part, foreground_part)
        return output