from typing import Annotated, Union, Literal, Tuple, Dict
WIDTH = DICT_CONSTANTS['YUNET']['WIDTH']  # Width of the image to diuplay
HEIGHT = DICT_CONSTANTS['YUNET']['HEIGHT']  # Height of the image to display
# --------------------------------------------------------------------------------
# FUNCTION TO PLOT THE PREPROCESSING TRANSFORMATION
# --------------------------------------------------------------------------------
def plot_preprocessing_transformation(
    df: pd.DataFrame = None,
    cols_to_apply: list = [],
    input_size: Tuple[int, int] = (WIDTH, HEIGHT), 
    output_size: Tuple[int, int] = (128, 128),
    input_mode: Literal['BGR', 'RGB'] = 'BGR',
    output_mode: Literal['BGR', 'RGB'] = 'RGB',
    expand_ratio: float = 1.0,
    force_square: bool = True,
    suptitle: str = "PREPROCESSING PIPELINE"
):
    
    """
    Takes a sample list of files from the df, shows original images (both correct and faulty), amended images with the bounding boxes and 
    resulting imagees to be channeled to the neuronet model.

    Paramters:
    - df: pd.DataFrame, a Pandas DataFrame instance with obligatory default integer index and 'file_path' column and df_file_register structure 
        (see FP_1_1_Data_Preparation.ipynb). Default None.
    - cols_to_apply: list, a  list of columns named as faults to decide what is to be displayed. Default []
        These must be a list of any of DICT_CONSTANTS['LIST_FAULTS'], relating to images (not file consistency). 
        If left [] only 1 sample of correct image will be shown.
   - input_size: Tuple[int, int], image size to conver an image before porcessing. Deafult (WIDTH, HEIGHT) constants. 
        No need to change in normal circumstatnces.
    - output_size: Tuple[int, int], image size for resulting output image.  Deafult (128, 128).
    - input_mode: Literal['BGR', 'RGB'], whether to expect an input of and image np.ndarray in RGB or BGR colour mode. 
        Deafult 'BGR' (openCV stnadard).
    - output_mode: Literal['BGR', 'RGB'], OpenCV colour mode - the colour scheem for resulting output np.ndarrays representing the image. 
        Defalut 'RGB'. NOTE! All internal methods are expecting and returning BGR format. 
    - expand_ratio: float, ratio to shift each of the bounding box boundaries outwards 
        respective to both width and height multiplied by expand_ratio. Default 1.0.
    - force_square: bool, whether to force the bounding box to be square, taking its largest side for the side of the square. 
         Default False. 
    - suptitle: str, a suptitle string to the plot. Default "PREPROCESSING PIPELINE"
    Returns:
    - a matplotlib figure with pictures of images orginized in rows by each category

    Usage:
    >>> plot_preprocessing_transformation(
    >>> df=df_file_register,
    >>> cols_to_apply = ['no_face', 'many_faces', 'rotated', 'low_quality', 'too_small', 'fragmented', 'truncated']
    >>> )    
    """
    # Initialize a preprocessor instance
    ipc = ipp.ImageProcessor(
        input_size = input_size, 
        output_size = output_size,
        input_mode= input_mode,
        output_mode = output_mode,
        expand_ratio = expand_ratio,
        force_square = force_square,
        initial_score_threshold = 1.0
    )

  
    # Append a category for a correct file
    cols_to_apply.append('no_fault')  

    # Define the figure    
    n_stages = 4
    n_faults = len(cols_to_apply)
    fig = plt.figure(
        figsize=(n_stages * 4, n_faults * 4),
        constrained_layout=True
    ) 

    
    # Take a sample of files indices of 1 for each fault category + 1 correct
    for fig_row, fault in enumerate(cols_to_apply):
        if fault == 'no_fault':
            file_idx = np.random.choice(df[(df[cols_to_apply[:-1]] == False).all(axis=1)].index)
        else:
            file_idx = np.random.choice(df[df[fault]].index)
        
        # Find the file path in case if it was moved out the DS
        if df.at[file_idx, 'moved_out']:
            image_file = df.path_to_original[file_idx]
        else:
            image_file = df.file_path[file_idx]
        
        # Read all image faults related to the file
        # dict_image_faults = dict(df.loc[file_idx, cols_to_apply[:-1]])
        dict_image_faults = None
        n_rotations = df.n_rotations[file_idx]
        best_face = df.best_face[file_idx]

        # Get the image array as is
        arr_image_input = ipc.load_image_from_path(image_file)  # BGR format

        # Get preprocessed data and image
        result = ipc.process_image(
            image=image_file,
            best_face=best_face,
            n_rotations=n_rotations,
            image_faults = dict_image_faults
        )

        arr_image_corrected, arr_coords_xy, score, arr_image_final = result
        arr_bbox = arr_coords_xy[:2]
        
        arr_image_w_bbox = cv.rectangle(
            arr_image_corrected, 
            arr_bbox[0],
            arr_bbox[1],
            (0, 255, 0), 
            2, 
            cv.LINE_AA
        )
        # Draw landmarks
        landmarks = arr_coords_xy[2:]
        for i, landmark in enumerate(landmarks):
            if i in (0, 1):
                color = (255, 0, 0)
            elif i == 2:
                color = (0, 255, 0)
            else:
                color = (0, 0, 255)
            radius = 4
            thickness = 2
            arr_image_w_bbox = cv.circle(
                arr_image_w_bbox, 
                tuple(landmark), 
                radius, 
                color, 
                thickness, 
                cv.LINE_AA
            )
        
        # Add confidence value
        confidence = f"{score:.2f}"
        position = (50, 50)
        font = cv.FONT_HERSHEY_SIMPLEX
        scale = 0.45
        thickness = 1
        arr_image_w_bbox = cv.putText(
            arr_image_w_bbox, 
            confidence, 
            position, 
            font, 
            scale, 
            color, 
            thickness, 
            cv.LINE_AA
        )
        arr_image_cropped = improc.crop_image(
            arr_image = arr_image_corrected,
            arr_coords = arr_bbox
        )
         
        list_arr_to_show = [arr_image_input, arr_image_w_bbox, arr_image_cropped, arr_image_final]
        list_titles = [
            'Original image',
            'Corrected image with a bounding box',
            'Cropped image',
            'Preprocessed output image'
        ]
        for stage, (arr_im, title) in enumerate(zip(list_arr_to_show, list_titles)):
            ax = fig.add_subplot(n_faults, n_stages, fig_row * n_stages + (stage + 1))
            if stage == 0:
                ax.imshow(cv.cvtColor(arr_im, cv.COLOR_BGR2RGB))
            else:
                ax.imshow(arr_im)  # we defind RGB mode for output!
            dummy = ax.set_title(title, size=14)
    
    dummy = fig.suptitle(suptitle, size=16)
    plt.show()


cols_to_show = [
    'no_face',
    'many_faces',
    'rotated',
    'low_quality',
    'too_small',
    'fragmented',
    'truncated'
]
plot_preprocessing_transformation(df_file_register, cols_to_show)