"""
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 numpy as np
from magicgui.widgets import create_widget, PushButton, Container, ToolBar, SpinBox, LineEdit, ComboBox
from qtpy.QtWidgets import QFileDialog

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

from functools import partial


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="Rotation")
        # Select file
        self._image_layer_rot = create_widget(
            label="Image", annotation="napari.layers.Image"
        )
        # Select rotation angle
        self.rotationangle = SpinBox(label="Rotation angle", value=0, min=-180, max=180)
        # Button rotate image
        self.btn_rotate = PushButton(value=True, text='Rotate')
        self.btn_rotate.clicked.connect(self._on_click_rotate)

        self.c1.extend(
            [
                self.tb,
                self._image_layer_rot,
                self.rotationangle,
                self.btn_rotate,
            ]
        )

        ################################################################
        self.c2 = Container(name ="Color correction")
        # Select file
        self._image_layer_col = create_widget(
            label="Image", annotation="napari.layers.Image"
        )
        # Button shape layer
        self.btn_rectangle = PushButton(value=True, text='Select the 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 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.c2.extend(
            [
                self.tb,
                self._image_layer_col,
                self.btn_rectangle,
                self.colorbox,
                self.btn_correct,
                self.btn_median,
            ]
        )

        ################################################################
        self.c3 = Container(name="Set scale")
        # Button shape layer
        self.btn_line = PushButton(value=True, text='Select 1 cm')
        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.c3.extend(
            [
                self.tb,
                self.btn_line,
                self.btn_setscale,
            ]
        )

        ################################################################
        self.c4 = Container(name='Segmentation')
        # Select file
        self._image_layer_seg = create_widget(
            label="Image", annotation="napari.layers.Image"
        )
        # Button label layer
        self.btn_rectangle1 = PushButton(value=True, text='Select fish area')
        self.btn_rectangle1.clicked.connect(self._on_click_createmask)

        self.c4.extend(
            [
                self.tb,
                self._image_layer_seg,
                self.btn_rectangle1
            ]
        )

        ################################################################
        self.c5 = Container(name='Analysis')
        # 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='Analyse the fish')
        self.btn_analyse.clicked.connect(self._on_click_analyse)

        self.c5.extend(
            [
                self.saving_dir,
                self.browse_btn,
                self.btn_analyse,
            ]
        )
        ################################################################
        self.extend(
            [
                #self.c1,
                self.c2,
                self.c3,
                self.c4,
                self.c5,
            ]
        )

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


    def _on_click_rotate(self):
        """image_layer = self._image_layer_rot.value
        if image_layer is None:
            print("Select image layer before rotating")
            return"""

        names = [x.name for x in self._viewer.layers]

        for image_name in names:
            img = np.asarray(self._viewer.layers[image_name].data)
            self._viewer.layers.remove(image_name)
            self._viewer.add_image(img, name=image_name, rgb=True, rotate=self.rotationangle.value)
            self.rotation = self.rotationangle.value


    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")
            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")
            return
        self.cmscale = set_cm_scale(self.line_layer.data)
        self._viewer.layers.remove('Line')


    def _on_click_createmask(self):
        image_layer = self._image_layer_seg.value
        if image_layer is None:
            print("Select image layer before segmenting")
            return
        image = np.asarray(image_layer.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_analyse(self):
        image_layer = self._image_layer_seg.value
        if image_layer is None:
            print("Select image layer before segmenting")
            return
        mask = self.fishmask_layer.data
        image = np.asarray(image_layer.data)
        image_name = image_layer.name.replace('Current_corrected_', '').replace('Previous_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, output_path, image_name ,self.cmscale)
        