"""
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

from magicgui import magic_factory
from magicgui.widgets import CheckBox, Container, create_widget
from qtpy.QtWidgets import QHBoxLayout, QPushButton, QWidget
from skimage.util import img_as_float
from napari.types import ImageData, LabelsData
if TYPE_CHECKING:
    import napari

from segtree.sam2.build_sam import build_sam2
from segtree.sam2.sam2_image_predictor import SAM2ImagePredictor
import torch
from ultralytics import YOLO
import torchvision.transforms as transforms
from PIL import Image 
import numpy as np
from skimage.io import imread
import matplotlib.pyplot as plt
from segtree.utils import *
import scipy
import os
from transformers import pipeline
import sys
sys.path.append(os.path.join(get_base_dir(),'segtree'))
print(os.getcwd())

def trunk_segmentation(bgr_data):
    pth_trunk = os.path.join(get_base_dir(),'segtree','best (1).pt') # trunk best model Yolov8 CV1 Aug_200_epochs-9
    my_new_model = YOLO(pth_trunk) 
    new_results = my_new_model.predict([bgr_data], conf=0.2)
    if new_results[0].masks is None:
        h,w = new_results[0].orig_shape
        labeled_mask = np.zeros((h,w),dtype='uint8')
        return labeled_mask

    data_msk = new_results[0].masks.data
    h,w = new_results[0].orig_shape
    resize_transform = transforms.Resize((h,w))
    resized_image = resize_transform(data_msk)

    data_msk = resized_image.cpu().numpy()
    data_msk = np.sum(data_msk,axis=0)
    data_msk = (data_msk!=0)
    labeled_mask, _ = scipy.ndimage.label(data_msk)
    return labeled_mask


def prompt_engineering(rgb_data,msks):
    image = Image.fromarray(rgb_data)
    depth = pipe(image)["depth"]
    msks_depth = np.array(depth)

    msks_front = frontground_part(msks_depth)
    msks = keep_largest_object(msks)
    n,m = msks.shape
    msks = np.where(msks!=0,1,0)
    dico_c = centroid_msk(msks)
    msks,dico_v = reattribute_data(dico_c,msks)    
    dict_coords = {}
    if len(np.unique(msks))!=1:
        dico = area_pts(msks,dico_v) # cle label / valeur espace droite et gauche et cx, cy
        L = list(dico.keys())
        dico_out = angle_assessment(msks)
        
        for curr_label in L:
            coords_diamond = {}
            I = dico[curr_label]
            curr_msk = np.where(msks==curr_label,255,0)
            angle_deg = dico_out[curr_label][0]
            curr_reg_angl = regression_with_correction_angle(curr_msk,angle_deg,n,m)
            curr_reg = curr_sum_one(curr_reg_angl,msks_front)
            if len(np.unique(curr_reg))>1:
                top,mid,mid_droite,mid_gauche,bottom = diamant_star(curr_reg,I) 
                coords_diamond['TOPX'] = top[0]
                coords_diamond['TOPY'] = top[1]
                coords_diamond['MIDX'] = mid[0]
                coords_diamond['MIDY'] = mid[1]
                coords_diamond['MIDXD'] = mid_droite[0]
                coords_diamond['MIDYD'] = mid_droite[1]
                coords_diamond['MIDXG'] = mid_gauche[0]
                coords_diamond['MIDYG'] = mid_gauche[1]
                coords_diamond['BOTTOMX'] = bottom[0]
                coords_diamond['BOTTOMY'] = bottom[1]
                dict_coords[curr_label] = [coords_diamond]
    return dict_coords

def foundation_model(image,dict_coords,checkpoint,model_cfg):
    L = list(dict_coords.keys())
    msk_output = np.zeros((image.shape[0],image.shape[1]),dtype='uint8')
    for ix in L:
        coords_diamond = dict_coords[ix][0]
        input_point = np.array([[coords_diamond['TOPX'],coords_diamond['TOPY']],
                                            [coords_diamond['MIDX'],coords_diamond['MIDY']],
                                            [coords_diamond['MIDXD'],coords_diamond['MIDYD']],
                                            [coords_diamond['MIDXG'],coords_diamond['MIDYG']],
                                            [coords_diamond['BOTTOMX'],coords_diamond['BOTTOMY']],
                                            ])

        input_label = np.array([1,1,1,1,1])

        predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint,device="cpu")) # Change this line
        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
            predictor.set_image(image)
            masks, scores, logits = predictor.predict(point_coords=input_point,
                                                    point_labels=input_label,
                                                    multimask_output=False, hq_token_only=False)

        sorted_ind = np.argsort(scores)[::-1]
        output_pred = masks[sorted_ind]
        scores = scores[sorted_ind]
        logits = logits[sorted_ind]
        n,_,_ = output_pred.shape
        resized_masks = []
        for i in range(n):
            output_pred_i = output_pred[i,...]
            output_pred_i = np.where(output_pred_i!=0,i+1,0)
            resized_masks.append(output_pred_i)

        msk = np.stack(resized_masks, axis=0)  # Stack resized masks back together
        msk = np.squeeze(msk)

        msk_output = np.where(msk!=0,ix,msk_output)
    return msk_output

checkpoint = os.path.join(get_base_dir(),"segtree","sam2.1_hq_hiera_large.pt")
# model_cfg = os.path.join(get_base_dir(),"segtree","sam2","configs","sam2.1","sam2.1_hq_hiera_l.yaml") 
model_cfg = os.path.join("configs","sam2.1","sam2.1_hq_hiera_l.yaml") 

def codes(rgb_data,data_msk):
    if len(np.unique(data_msk))==1:
        bgr_data = rgb_data[..., ::-1]  # Inverse l'ordre des canaux
        bgr_data = bgr_data.astype('uint8')
        labeled_mask = trunk_segmentation(bgr_data)
    else:
        labeled_mask = data_msk.copy()
    msks = labeled_mask.copy()
    if len(np.unique(msks))>1:
        dict_coords = prompt_engineering(rgb_data,msks)
        rgb_data = rgb_data[:,:,:3]
        msks = foundation_model(rgb_data,dict_coords,checkpoint,model_cfg)
    msks = msks.astype('uint8')
    return msks

@magic_factory(call_button="Run")
def individual_tree(rgb_data: ImageData, trunk_data: LabelsData) -> LabelsData:
    if rgb_data.ndim==3:
        rgb_data = rgb_data[:,:,:3]
        data_msk = trunk_data[:,:]
        msks = codes(rgb_data,data_msk)
        return msks
    elif rgb_data.ndim==4:
        n = rgb_data.shape[0]
        msks = np.zeros((n,rgb_data.shape[1],rgb_data.shape[2]),dtype='uint8')
        for i in range(n):
            data_img = rgb_data[i,:,:,:]
            data_msk = trunk_data[i,:,:]
            data_img = np.array(data_img)
            data_msk = np.array(data_msk)
            msks[i,:,:] = codes(data_img,data_msk)
        return msks
    else:
        msks = np.zeros((rgb_data.shape[1],rgb_data.shape[2]),dtype='uint8')
        return msks
    

# Uses the `autogenerate: true` flag in the plugin manifest
# to indicate it should be wrapped as a magicgui to autogenerate
# a widget.
def threshold_autogenerate_widget(
    img: "napari.types.ImageData",
    threshold: "float",
) -> "napari.types.LabelsData":
    return img_as_float(img) > threshold

# if we want even more control over our widget, we can use
# magicgui `Container`
class ImageThreshold(Container):
    def __init__(self, viewer: "napari.viewer.Viewer"):
        super().__init__()
        self._viewer = viewer
        # use create_widget to generate widgets from type annotations
        self._image_layer_combo = create_widget(
            label="Image", annotation="napari.layers.Image"
        )
        self._threshold_slider = create_widget(
            label="Threshold", annotation=float, widget_type="FloatSlider"
        )
        self._threshold_slider.min = 0
        self._threshold_slider.max = 1
        # use magicgui widgets directly
        self._invert_checkbox = CheckBox(text="Keep pixels below threshold")

        # connect your own callbacks
        self._threshold_slider.changed.connect(self._threshold_im)
        self._invert_checkbox.changed.connect(self._threshold_im)

        # append into/extend the container with your widgets
        self.extend(
            [
                self._image_layer_combo,
                self._threshold_slider,
                self._invert_checkbox,
            ]
        )

    def _threshold_im(self):
        image_layer = self._image_layer_combo.value
        if image_layer is None:
            return

        image = img_as_float(image_layer.data)
        name = image_layer.name + "_thresholded"
        threshold = self._threshold_slider.value
        if self._invert_checkbox.value:
            thresholded = image < threshold
        else:
            thresholded = image > threshold
        if name in self._viewer.layers:
            self._viewer.layers[name].data = thresholded
        else:
            self._viewer.add_labels(thresholded, name=name)


class ExampleQWidget(QWidget):
    # your QWidget.__init__ can optionally request the napari viewer instance
    # use a type annotation of 'napari.viewer.Viewer' for any parameter
    def __init__(self, viewer: "napari.viewer.Viewer"):
        super().__init__()
        self.viewer = viewer

        btn = QPushButton("Click me!")
        btn.clicked.connect(self._on_click)

        self.setLayout(QHBoxLayout())
        self.layout().addWidget(btn)

    def _on_click(self):
        print("napari has", len(self.viewer.layers), "layers")
