from .base import BaseModule
import numpy as np
import cv2
from .topology.graph import NerveImage, NerveGraph


def _graphify(image: np.ndarray, binary_image: np.ndarray, skeleton_image: np.ndarray):
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    nerve_image = NerveImage(image, binary_image, skeleton_image)
    nerve_graph = NerveGraph(nerve_image)
    return nerve_graph


class GraphifyModule(BaseModule):
    """
    A module for constructing CCM images into multiple undirected graphs
    for morphological or topological structure analysis.
    This module is expected to accept three images: the original image,
    the binary image, and the skeleton image.
    All of their shapes should be (384, 384), and they will all be converted to grayscale images.
    It will output a NerveGraph object, which contains all the results generated by the analysis.
    """

    def __init__(self):
        super().__init__()
        self.name = 'graphify'
        self.output_name = 'nerve_graph'

    def __call__(self, *args, **kwargs) -> NerveGraph:
        if not args:
            raise ValueError("An input is required.")
        return _graphify(*args, **kwargs)
