from typing import List, Union
from PIL import Image
import torch
import numpy as np

from diffusers.modular_pipelines import (
    PipelineState,
    ModularPipelineBlocks,
    InputParam,
    ComponentSpec,
    OutputParam,
)
from controlnet_aux import CannyDetector
import numpy as np


class CannyBlock(ModularPipelineBlocks):
    @property
    def expected_components(self):
        return []

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "image",
                type_hint=Union[Image.Image, np.ndarray],
                required=True,
                description="Image to compute canny filter on",
            ),
            InputParam(
                "low_threshold",
                type_hint=int,
                default=50,
            ),
            InputParam("high_threshold", type_hint=int, default=200),
            InputParam(
                "detect_resolution",
                type_hint=int,
                default=1024,
                description="Resolution to resize to when running the Canny filtering process.",
            ),
            InputParam(
                "image_resolution",
                type_hint=int,
                default=1024,
                description="Resolution to resize the detected Canny edge map to.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        return [
            OutputParam(
                "control_image",
                type_hint=Image,
                description="Canny map for input image",
            )
        ]

    def compute_canny(self, image, low_threshold, high_threshold, detect_resolution, image_resolution):
        canny_detector = CannyDetector()
        canny_map = canny_detector(
            input_image=image,
            low_threshold=low_threshold,
            high_threshold=high_threshold,
            detect_resolution=detect_resolution,
            image_resolution=image_resolution,
        )
        return canny_map

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        block_state.control_image = self.compute_canny(
            block_state.image,
            block_state.low_threshold,
            block_state.high_threshold,
            block_state.detect_resolution,
            block_state.image_resolution,
        )
        self.set_block_state(state, block_state)

        return components, state