"""
Collection of various utils 
"""

import numpy as np

import imageio.v3 as iio
from PIL import Image
# we may have very large images (e.g. panoramic SEM images), allow to read them w/o warnings
Image.MAX_IMAGE_PIXELS = 933120000

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.lines import Line2D


import math


###
### load SEM images
### 
def load_image(filename : str) -> np.ndarray :
    """Load an SEM image 

    Args:
        filename (str): full path and name of the image file to be loaded

    Returns:
        np.ndarray: file as numpy ndarray
    """
    image =  iio.imread(filename,mode='F')

    return image



###
### show SEM image with boxes in various colours around each damage site
###
def show_boxes(image : np.ndarray, damage_sites : dict, box_size = [250,250],
               save_image = False, image_path : str = None) :
    """_summary_

    Args:
        image (np.ndarray): SEM image to be shown
        damage_sites (dict): python dictionary using the coordinates as key (x,y), and the label as value
        box_size (list, optional): size of the rectangle drawn around each centroid. Defaults to [250,250].
        save_image (bool, optional): save the image with the boxes or not. Defaults to False.
        image_path (str, optional) : Full path and name of the output file to be saved
    """

    _, ax = plt.subplots(1)
    ax.imshow(image, cmap='gray')  # show image on correct axis
    ax.set_xticks([])
    ax.set_yticks([])

    for key, label in damage_sites.items():
        position = [key[0], key[1]]
        edgecolor = {
            'Inclusion': 'b',
            'Interface': 'g',
            'Martensite': 'r',
            'Notch': 'y',
            'Shadowing': 'm'
        }.get(label, 'k')  # default: black

        rect = patches.Rectangle((position[1] - box_size[1] / 2., position[0] - box_size[0] / 2),
                                 box_size[1], box_size[0],
                                 linewidth=1, edgecolor=edgecolor, facecolor='none')
        ax.add_patch(rect)

    legend_elements = [
        Line2D([0], [0], color='b', lw=4, label='Inclusion'),
        Line2D([0], [0], color='g', lw=4, label='Interface'),
        Line2D([0], [0], color='r', lw=4, label='Martensite'),
        Line2D([0], [0], color='y', lw=4, label='Notch'),
        Line2D([0], [0], color='m', lw=4, label='Shadow'),
        Line2D([0], [0], color='k', lw=4, label='Not Classified')
    ]
    ax.legend(handles=legend_elements, bbox_to_anchor=(1.04, 1), loc="upper left")

    fig = ax.figure
    fig.tight_layout(pad=0)

    if save_image and image_path:
        fig.savefig(image_path, dpi=1200, bbox_inches='tight')

    canvas = fig.canvas
    canvas.draw()

    data = np.frombuffer(canvas.buffer_rgba(), dtype=np.uint8).reshape(
        canvas.get_width_height()[::-1] + (4,))
    data = data[:, :, :3]  # RGB only

    plt.close(fig)

    return data 


###
### cut out small images from panorama, append colour information
###
def prepare_classifier_input(panorama: np.ndarray, centroids: list, window_size=[250, 250]) -> list:
    """
    Extracts square image patches centered at each given centroid from a grayscale panoramic SEM image.
    
    Each extracted patch is resized to the specified window size and converted into a 3-channel (RGB-like)
    normalized image suitable for use with classification neural networks that expect color input.

    Parameters
    ----------
    panorama : np.ndarray
        Input SEM image. Should be a 2D array (H, W) or a 3D array (H, W, 1) representing grayscale data.
    
    centroids : list of [int, int]
        List of (y, x) coordinates marking the centers of regions of interest. These are typically damage sites
        identified in preprocessing (e.g., clustering).

    window_size : list of int, optional
        Size [height, width] of each extracted image patch. Defaults to [250, 250].

    Returns
    -------
    list of np.ndarray
        List of extracted and normalized 3-channel image patches, each with shape (height, width, 3). Only
        centroids that allow full window extraction within image bounds are used.
    """
    if panorama.ndim == 2:
        panorama = np.expand_dims(panorama, axis=-1)  # (H, W, 1)

    H, W, _ = panorama.shape
    win_h, win_w = window_size
    images = []

    for (cy, cx) in centroids:
        x1 = int(cx - win_w / 2)
        y1 = int(cy - win_h / 2)
        x2 = x1 + win_w
        y2 = y1 + win_h

        # Skip if patch would go out of bounds
        if x1 < 0 or y1 < 0 or x2 > W or y2 > H:
            continue

        # Extract and normalize patch
        patch = panorama[y1:y2, x1:x2, 0].astype(np.float32)
        patch = patch * 2. / 255. - 1.

        # Replicate grayscale channel to simulate RGB
        patch_color = np.repeat(patch[:, :, np.newaxis], 3, axis=2)
        images.append(patch_color)

    return images