"""
This module contains four napari widgets declared in
different ways:

- a pure Python function flagged with `autogenerate: true`
    in the plugin manifest. Type annotations are used by
    magicgui to generate widgets for each parameter. Best
    suited for simple processing tasks - usually taking
    in and/or returning a layer.
- a `magic_factory` decorated function. The `magic_factory`
    decorator allows us to customize aspects of the resulting
    GUI, including the widgets associated with each parameter.
    Best used when you have a very simple processing task,
    but want some control over the autogenerated widgets. If you
    find yourself needing to define lots of nested functions to achieve
    your functionality, maybe look at the `Container` widget!
- a `magicgui.widgets.Container` subclass. This provides lots
    of flexibility and customization options while still supporting
    `magicgui` widgets and convenience methods for creating widgets
    from type annotations. If you want to customize your widgets and
    connect callbacks, this is the best widget option for you.
- a `QWidget` subclass. This provides maximal flexibility but requires
    full specification of widget layouts, callbacks, events, etc.

References:
- Widget specification: https://napari.org/stable/plugins/guides.html?#widgets
- magicgui docs: https://pyapp-kit.github.io/magicgui/

Replace code below according to your needs.
"""
from typing import TYPE_CHECKING

import os
import cv2
import imutils
import tifffile
import numpy as np
from PIL import Image
from magicgui.widgets import create_widget, PushButton, Container, ToolBar, SpinBox, LineEdit, ComboBox, Label
from qtpy.QtWidgets import QFileDialog, QMessageBox

from .functions import color_analysis, bw_correction, rgb_correction, cmy_correction, set_cm_scale, lock_controls

import warnings
warnings.filterwarnings('ignore')


if TYPE_CHECKING:
    import napari


class FishAnalysis(Container):
    def __init__(self, viewer: "napari.viewer.Viewer"):
        super().__init__()
        self._viewer = viewer
        self.rect_layer = None
        self.line_layer = None
        self.fishmask_layer = None
        self.rotation = 0
        self.cmscale = None
        self.default = None
        self.previous_corrected_layer = None


        self.tb = ToolBar()

        ################################################################
        self.c1 = Container(name ="Color correction")
        # Select file
        self._image_layer_col = create_widget(
            label="Image to correct", annotation="napari.layers.Image"
        )
        # Button shape layer
        self.btn_rectangle = PushButton(value=True, text='Select the color area')
        self.btn_rectangle.clicked.connect(lambda: self._on_click_createlayer('color_rectangle'))
        # Color combobox
        self.colorbox = ComboBox(choices=['white', 'black', 'red', 'green', 'blue', 'cyan', 'magenta', 'yellow'])
        # Button correct color
        self.btn_correct = PushButton(value=True, text='Correct the color')
        self.btn_correct.clicked.connect(self._on_click_colorcorrection)
        # Button correct color
        self.btn_median = PushButton(value=True, text='Apply median filter')
        self.btn_median.clicked.connect(self._on_click_medianfilter)

        self.c1.extend(
            [
                self.tb,
                self._image_layer_col,
                self.btn_rectangle,
                self.colorbox,
                self.btn_correct,
                self.btn_median,
            ]
        )

        ################################################################
        self.c2 = Container(name="Set scale")
        # Button shape layer
        self.btn_line = PushButton(value=True, text='Create 1cm layer')
        self.btn_line.clicked.connect(lambda: self._on_click_createlayer('line'))
        # Button correct color
        self.btn_setscale = PushButton(value=True, text='Set scale')
        self.btn_setscale.clicked.connect(self._on_click_cmscale)

        self.c2.extend(
            [
                self.tb,
                self.btn_line,
                self.btn_setscale,
            ]
        )

        ################################################################
        self.c3 = Container(name='Segmentation')
        # Button label layer
        self.btn_rectangle1 = PushButton(value=True, text='Create segmentation layer')
        self.btn_rectangle1.clicked.connect(self._on_click_createmask)

        self.c3.extend(
            [
                self.tb,
                self.btn_rectangle1,
            ]
        )

        ################################################################
        self.c4 = Container(name='Analysis')
        # Select file
        self.label = Label(value='Use current image or upload a previous one')
        self.image_dir = LineEdit(label="Upload image")
        self.browseimage_btn = PushButton(value=True, text='Browse and Open')
        self.browseimage_btn.clicked.connect(self._add_image)
        # Select file
        self.image_toanalyse = create_widget(
            label="Image to analyse", annotation="napari.layers.Image"
        )
        self.btn_landmarks = PushButton(value=True, text='Create landmarks layer')
        self.btn_landmarks.clicked.connect(self._on_click_createlandmarks)
        # Define output directory
        self.saving_dir = LineEdit(label="Output folder")
        self.browse_btn = PushButton(value=True, text='Browse')
        self.browse_btn.clicked.connect(self._add_folder)
        # Button analysis
        self.btn_analyse = PushButton(value=True, text='Perform the analysis')
        self.btn_analyse.clicked.connect(self._on_click_analyse)

        self.c4.extend(
            [
                self.tb,
                self.label,
                self.image_dir,
                self.browseimage_btn,
                self.image_toanalyse,
                self.btn_landmarks,
                self.saving_dir,
                self.browse_btn,
                self.btn_analyse 
            ]
        )

        ################################################################
        self.extend(
            [
                self.c1,
                self.c2,
                self.c3,
                self.c4,
            ]
        )

    
    def _add_folder(self):
        # Add the selected mask as input
        self.saving_dir.value = QFileDialog.getExistingDirectory()


    def _add_image(self):
        self.image_dir.value, _ = QFileDialog.getOpenFileName(filter='*.tif')
        image = tifffile.imread(self.image_dir.value)
        self._viewer.add_image(image[0], name='Corrected_'+os.path.basename(self.image_dir.value).replace('.tif',''), rgb=True)
        if len(image[1].shape) > 2:
            mask = cv2.cvtColor(image[1], cv2.COLOR_BGR2GRAY).astype('uint8')
            self.fishmask_layer = self._viewer.add_labels(mask, name='Mask')
        else:
            self.fishmask_layer = self._viewer.add_labels(image[1].astype('uint8'), name='Mask')


    def _on_click_createlayer(self, t):
        if t == 'color_rectangle':
            self.rect_layer = self._viewer.add_shapes(name='Rectangle', shape_type='rectangle', edge_width=5, edge_color='red')
            widget_list_toblock = [
                'ellipse_button',
                'line_button',
                'path_button',
                'polygon_button',
                'polygon_lasso_button',
                'vertex_remove_button',
                'vertex_insert_button',
                'move_back_button',
                'move_front_button'
            ]
            lock_controls(self.rect_layer, self._viewer, widget_list_toblock)
        else:
            self.line_layer = self._viewer.add_shapes(name='Line', shape_type='line', edge_width=5, edge_color='red')
            widget_list_toblock = [
                'ellipse_button',
                'rectangle_button',
                'path_button',
                'polygon_button',
                'polygon_lasso_button',
                'vertex_remove_button',
                'vertex_insert_button',
                'move_back_button',
                'move_front_button'
            ]
            lock_controls(self.line_layer, self._viewer, widget_list_toblock)


    def _on_click_colorcorrection(self):
        # Read image and its name (cleaned)
        image_layer = self._image_layer_col.value
        if image_layer is None:
            print("Select image layer before correcting color")
            return
        image = np.asarray(image_layer.data)
        image_name = image_layer.name.replace('Current_corrected_', '').replace('Previous_corrected_', '')
        # Read the color to process
        color = self.colorbox.value
        # Process
        if color == 'white':
            new_image = bw_correction(image, self.rect_layer, 'white')
        elif color == 'black':
            new_image = bw_correction(image, self.rect_layer, 'black')
        elif color in ['red', 'green', 'blue']:
            new_image = rgb_correction(image, self.rect_layer, color)
        else:
            new_image = cmy_correction(image, self.rect_layer, color)

        # Add the new layer
        previous_name = 'Previous_corrected_'+image_name
        current_name = 'Current_corrected_'+image_name
        layers_names = self._viewer.layers
        
        ## If there are no images OR
        ## If there is only the previous (it means they deleted the current)
        ## just insert the current image
        if (previous_name not in layers_names and current_name not in layers_names) or (previous_name in layers_names and current_name not in layers_names):
            self._viewer.add_image(new_image, name=current_name, rgb=True, rotate=self.rotation)

        ## If there are is a currect images, insert as current and move the previous
        elif previous_name not in layers_names and current_name in layers_names:
            # Remove current and save it as previous
            self.previous_corrected_layer = self._viewer.layers[current_name]
            self._viewer.add_image(self.previous_corrected_layer.data, name=previous_name, rgb=True, rotate=self.rotation)
            self._viewer.layers.remove(current_name)
            # Save new current
            self._viewer.add_image(new_image, name=current_name, rgb=True, rotate=self.rotation)

        ## If there are both images, insert as current the new one, move the previous, delete the other
        elif previous_name in layers_names and current_name in layers_names:
            # Delete the previous
            self._viewer.layers.remove(previous_name)
            # Save the new previous
            self.previous_corrected_layer = self._viewer.layers[current_name]
            self._viewer.add_image(self.previous_corrected_layer.data, name=previous_name, rgb=True, rotate=self.rotation)
            # Detete the current (now previous)
            self._viewer.layers.remove(current_name)
            # Save new current
            self._viewer.add_image(new_image, name=current_name, rgb=True, rotate=self.rotation)
            
        # Delete the rectangle layer
        if self.rect_layer is not None:
            self._viewer.layers.remove('Rectangle')
            self.rect_layer = None
        

    def _on_click_medianfilter(self):
        image_layer = self._image_layer_col.value
        if image_layer is None:
            print("Select image layer before correcting color")
            msg = QMessageBox() 
            msg.setIcon(QMessageBox.Warning)
            msg.setText("Select image layer before correcting color.")
            msg.setStandardButtons(QMessageBox.Ok | QMessageBox.Cancel)
            msg.exec_()
            return
        image = np.asarray(image_layer.data)
        name = image_layer.name
        self._viewer.layers.remove(name)
        median = cv2.medianBlur(image, 5)
        self._viewer.add_image(median, name=name, rgb=True, rotate=self.rotation)


    def _on_click_cmscale(self):
        if self.line_layer.data is None:
            print("Create a line before setting scale")
            msg = QMessageBox() 
            msg.setIcon(QMessageBox.Warning)
            msg.setText("Create a line before setting scale.")
            msg.setStandardButtons(QMessageBox.Ok | QMessageBox.Cancel)
            msg.exec_()
            return
        self.cmscale = set_cm_scale(self.line_layer.data)
        self._viewer.layers.remove('Line')


    def _on_click_createmask(self):
        image = np.asarray(self._viewer.layers[0].data)
        self.fishmask_layer = self._viewer.add_labels(np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8), name='Fish mask')

    
    def _on_click_createlandmarks(self):
        image = np.asarray(self._viewer.layers[0].data)
        landmarks = np.asarray(Image.open(os.path.join(os.path.dirname(__file__), 'support_files', 'landmarks.png')))
        landmarks = imutils.resize(landmarks, width=int(image.shape[1]/3))
        self._viewer.add_image(landmarks, name='Reference', rgb=True)
        self.landmarks_layer = self._viewer.add_labels(np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8), name='Landmarks', opacity=1)
        self.landmarks_layer.brush_size = 25


    def _on_click_analyse(self):
        image_layer = self.image_toanalyse.value
        mask = self.fishmask_layer.data
        landmarks = self.landmarks_layer.data
        if image_layer is None:
            print("Select image layer before segmenting")
            msg = QMessageBox() 
            msg.setIcon(QMessageBox.Warning)
            msg.setText("Select image layer before segmenting.")
            msg.setStandardButtons(QMessageBox.Ok | QMessageBox.Cancel)
            msg.exec_()
            return
        if len(np.unique(mask)) < 2:
            print("Draw mask layer before segmenting")
            msg = QMessageBox() 
            msg.setIcon(QMessageBox.Warning)
            msg.setText("Draw mask layer before segmenting.")
            msg.setStandardButtons(QMessageBox.Ok | QMessageBox.Cancel)
            msg.exec_()
            return
        if len(np.unique(landmarks)) != 30:
            print("Select all the 29 landmarks with the correct id")
            msg = QMessageBox() 
            msg.setIcon(QMessageBox.Warning)
            msg.setText("Select all the 29 landmarks with the correct id.")
            msg.setStandardButtons(QMessageBox.Ok | QMessageBox.Cancel)
            msg.exec_()
            return
        
        image = np.asarray(image_layer.data)
        image_name = image_layer.name.replace('Current_corrected_', '').replace('Previous_corrected_', '').replace('Corrected_', '')
        if self.saving_dir.value != '':
            output_path = self.saving_dir.value
        else:
            output_path = os.path.dirname(image_layer.source.path)
        
        color_analysis(image, mask, landmarks, output_path, image_name ,self.cmscale)
        