"""
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/building_a_plugin/guides.html#widgets
- magicgui docs: https://pyapp-kit.github.io/magicgui/

Replace code below according to your needs.
"""

from typing import TYPE_CHECKING

import numpy as np
import torch
from magicgui import magic_factory
from PIL import Image
from scipy import ndimage
from transformers import pipeline

if TYPE_CHECKING:
    import napari
from frontveg.utils import frontground_part, ground_dino, sam2

pipe = pipeline(
    task="depth-estimation", model="depth-anything/Depth-Anything-V2-Large-hf"
)


@magic_factory(call_button="Run")
def vegetation(
    input_data: "napari.types.ImageData",
) -> "napari.types.LabelsData":
    device = "cuda"

    if input_data.ndim == 4:
        output_data = np.zeros(
            (input_data.shape[0], input_data.shape[1], input_data.shape[2]),
            dtype="uint8",
        )
        INPUT = []
        for i in range(len(input_data)):
            rgb_data = input_data[i, :, :, :].compute()
            image = Image.fromarray(rgb_data)
            INPUT.append(image)
    else:
        output_data = np.zeros(
            (1, input_data.shape[0], input_data.shape[1]), dtype="uint8"
        )
        rgb_data = input_data
        image = Image.fromarray(rgb_data)
        INPUT = [image]
    depth = pipe(INPUT)
    n = len(depth)

    model, processor = ground_dino()
    predictor, text_labels = sam2()

    for i in range(n):
        depth_pred = depth[i]["depth"]
        msks_depth = np.array(depth_pred)
        msks_front = frontground_part(msks_depth)
        msks_front = msks_front.astype(np.uint8) * 255

        image = INPUT[i]
        inputs = processor(
            images=image, text=text_labels, return_tensors="pt"
        ).to(device)
        with torch.no_grad():
            outputs = model(**inputs)

        results = processor.post_process_grounded_object_detection(
            outputs,
            inputs.input_ids,
            box_threshold=0.4,
            text_threshold=0.3,
            target_sizes=[image.size[::-1]],
        )

        # Retrieve the first image result
        result = results[0]
        for box, score, labels in zip(
            result["boxes"], result["scores"], result["labels"], strict=False
        ):
            box = [round(x, 2) for x in box.tolist()]
            print(
                f"Detected {labels} with confidence {round(score.item(), 3)} at location {box}"
            )
        if len(result["boxes"]) == 0:
            masks = np.zeros(image.size[::-1], dtype="uint8")
        else:
            with (
                torch.inference_mode(),
                torch.autocast("cuda", dtype=torch.bfloat16),
            ):
                predictor.set_image(image)
                masks_sam, _, _ = predictor.predict(
                    box=result["boxes"],
                    point_labels=result["labels"],
                    multimask_output=False,
                )
            if masks_sam.ndim == 4:
                masks = np.sum(masks_sam, axis=0)
                masks = masks[0, :, :]
            else:
                masks = masks_sam[0, :, :]

        msks_veg = masks.astype(np.uint8) * 255

        mask1 = msks_front.copy()  # Masque 1
        mask2 = msks_veg.copy()  # Masque 2
        mask2 = ndimage.binary_fill_holes(mask2)  # Fill holes
        mask1 = (mask1 > 0).astype(np.uint8)  # Convertir en binaire
        mask2 = (mask2 > 0).astype(np.uint8)  # Convertir en binaire
        if len(np.unique(mask2)) == 2:
            intersection = (
                mask1 & mask2
            )  # Intersection : les pixels qui sont 1 dans les deux masques
            intersection = intersection > 0
        else:
            intersection = mask1.copy()
        intersection = (intersection * 255).astype(
            np.uint8
        )  # Si tu veux un masque avec des 0 et 255 (ex. pour OpenCV)
        output_data[i, :, :] = intersection
    return output_data
