from typing import Optional
import os.path
try:
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
except ImportError:
    plt = mpatches = None
import numpy as np
import cv2
from scipy.ndimage import gaussian_filter1d

from .utils import crop_image_inside_box
from .utils.rotate import rotate_image_different
from .utils.resize import resize_image

class EynollahPlotter:
    """
    Class collecting all the plotting and image writing methods
    """

    def __init__(
        self,
        *,
        dir_out,
        dir_of_all: Optional[str] = None,
        dir_save_page: Optional[str] = None,
        dir_of_deskewed: Optional[str] = None,
        dir_of_layout: Optional[str] = None,
        dir_of_cropped_images: Optional[str] = None,
    ):
        self.dir_out = dir_out
        self.dir_of_all = dir_of_all
        self.dir_save_page = dir_save_page
        self.dir_of_layout = dir_of_layout
        self.dir_of_cropped_images = dir_of_cropped_images
        self.dir_of_deskewed = dir_of_deskewed

    def save_plot_of_layout_main(self, text_regions_p, image_page, name=None):
        if self.dir_of_layout is not None:
            values = np.unique(text_regions_p[:, :])
            # pixels=['Background' , 'Main text' , 'Heading' , 'Marginalia' ,'Drop capitals' , 'Images' , 'Seperators' , 'Tables', 'Graphics']
            pixels=['Background' , 'Main text'  , 'Image' , 'Separator','Marginalia']
            values_indexes = [0, 1, 2, 3, 4]
            plt.figure(figsize=(40, 40))
            plt.rcParams["font.size"] = "40"
            im = plt.imshow(text_regions_p[:, :])
            colors = [im.cmap(im.norm(value)) for value in values]
            patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values]
            plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=40)
            plt.savefig(os.path.join(self.dir_of_layout,
                                     (name or "page") + "_layout_main.png"))

    def save_plot_of_layout_main_all(self, text_regions_p, image_page, name=None):
        if self.dir_of_all is not None:
            values = np.unique(text_regions_p[:, :])
            # pixels=['Background' , 'Main text' , 'Heading' , 'Marginalia' ,'Drop capitals' , 'Images' , 'Seperators' , 'Tables', 'Graphics']
            pixels=['Background' , 'Main text'  , 'Image' , 'Separator','Marginalia']
            values_indexes = [0, 1, 2, 3, 4]
            plt.figure(figsize=(80, 40))
            plt.rcParams["font.size"] = "40"
            plt.subplot(1, 2, 1)
            plt.imshow(image_page)
            plt.subplot(1, 2, 2)
            im = plt.imshow(text_regions_p[:, :])
            colors = [im.cmap(im.norm(value)) for value in values]
            patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values]
            plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=60)
            plt.savefig(os.path.join(self.dir_of_all,
                                     (name or "page") + "_layout_main_and_page.png"))

    def save_plot_of_layout(self, text_regions_p, image_page, name=None):
        if self.dir_of_layout is not None:
            values = np.unique(text_regions_p[:, :])
            # pixels=['Background' , 'Main text' , 'Heading' , 'Marginalia' ,'Drop capitals' , 'Images' , 'Seperators' , 'Tables', 'Graphics']
            pixels = ["Background", "Main text", "Header", "Marginalia", "Drop capital", "Image", "Separator", "Tables"]
            values_indexes = [0, 1, 2, 8, 4, 5, 6, 10]
            plt.figure(figsize=(40, 40))
            plt.rcParams["font.size"] = "40"
            im = plt.imshow(text_regions_p[:, :])
            colors = [im.cmap(im.norm(value)) for value in values]
            patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values]
            plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=40)
            plt.savefig(os.path.join(self.dir_of_layout,
                                     (name or "page") + "_layout.png"))

    def save_plot_of_layout_all(self, text_regions_p, image_page, name=None):
        if self.dir_of_all is not None:
            values = np.unique(text_regions_p[:, :])
            # pixels=['Background' , 'Main text' , 'Heading' , 'Marginalia' ,'Drop capitals' , 'Images' , 'Seperators' , 'Tables', 'Graphics']
            pixels = ["Background", "Main text", "Header", "Marginalia", "Drop capital", "Image", "Separator", "Tables"]
            values_indexes = [0, 1, 2, 8, 4, 5, 6, 10]
            plt.figure(figsize=(80, 40))
            plt.rcParams["font.size"] = "40"
            plt.subplot(1, 2, 1)
            plt.imshow(image_page)
            plt.subplot(1, 2, 2)
            im = plt.imshow(text_regions_p[:, :])
            colors = [im.cmap(im.norm(value)) for value in values]
            patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values]
            plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=60)
            plt.savefig(os.path.join(self.dir_of_all,
                                     (name or "page") + "_layout_and_page.png"))

    def save_plot_of_textlines(self, textline_mask_tot_ea, image_page, name=None):
        if self.dir_of_all is not None:
            values = np.unique(textline_mask_tot_ea[:, :])
            pixels = ["Background", "Textlines"]
            values_indexes = [0, 1]
            plt.figure(figsize=(80, 40))
            plt.rcParams["font.size"] = "40"
            plt.subplot(1, 2, 1)
            plt.imshow(image_page)
            plt.subplot(1, 2, 2)
            im = plt.imshow(textline_mask_tot_ea[:, :])
            colors = [im.cmap(im.norm(value)) for value in values]
            patches = [mpatches.Patch(color=colors[np.where(values == i)[0][0]], label="{l}".format(l=pixels[int(np.where(values_indexes == i)[0][0])])) for i in values]
            plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=60)
            plt.savefig(os.path.join(self.dir_of_all,
                                     (name or "page") + "_textline_and_page.png"))

    def save_deskewed_image(self, slope_deskew, image_org, name=None):
        if self.dir_of_all is not None:
            cv2.imwrite(os.path.join(self.dir_of_all,
                                     (name or "page") + "_org.png"), image_org)
        if self.dir_of_deskewed is not None:
            img_rotated = rotate_image_different(image_org, slope_deskew)
            cv2.imwrite(os.path.join(self.dir_of_deskewed,
                                     (name or "page") + "_deskewed.png"), img_rotated)

    def save_page_image(self, image_page, name=None):
        if self.dir_of_all is not None:
            cv2.imwrite(os.path.join(self.dir_of_all,
                                     (name or "page") + "_page.png"), image_page)
        if self.dir_save_page is not None:
            cv2.imwrite(os.path.join(self.dir_save_page,
                                     (name or "page") + "_page.png"), image_page)

    def save_enhanced_image(self, img_res, name=None):
        cv2.imwrite(os.path.join(self.dir_out,
                                 (name or "page") + "_enhanced.png"), img_res)
        
    def save_plot_of_textline_density(self, img_patch_org, name=None):
        if self.dir_of_all is not None:
            plt.figure(figsize=(80,40))
            plt.rcParams['font.size']='50'
            plt.subplot(1,2,1)
            plt.imshow(img_patch_org)
            plt.subplot(1,2,2)
            plt.plot(gaussian_filter1d(img_patch_org.sum(axis=1), 3),np.array(range(len(gaussian_filter1d(img_patch_org.sum(axis=1), 3)))),linewidth=8)
            plt.xlabel('Density of textline prediction in direction of X axis',fontsize=60)
            plt.ylabel('Height',fontsize=60)
            plt.yticks([0,len(gaussian_filter1d(img_patch_org.sum(axis=1), 3))])
            plt.gca().invert_yaxis()
            plt.savefig(os.path.join(self.dir_of_all,
                                     (name or "page") + '_density_of_textline.png'))

    def save_plot_of_rotation_angle(self, angels, var_res, name=None):
        if self.dir_of_all is not None:
            plt.figure(figsize=(60,30))
            plt.rcParams['font.size']='50'
            plt.plot(angels,np.array(var_res),'-o',markersize=25,linewidth=4)
            plt.xlabel('angle',fontsize=50)
            plt.ylabel('variance of sum of rotated textline in direction of x axis',fontsize=50)
            plt.plot(angels[np.argmax(var_res)],var_res[np.argmax(np.array(var_res))]  ,'*',markersize=50,label='Angle of deskewing=' +str("{:.2f}".format(angels[np.argmax(var_res)]))+r'$\degree$')
            plt.legend(loc='best')
            plt.savefig(os.path.join(self.dir_of_all,
                                     (name or "page") + '_rotation_angle.png'))

    def write_images_into_directory(self, img_contours, image_page, scale_x=1.0, scale_y=1.0, name=None):
        if self.dir_of_cropped_images is not None:
            index = 0
            for cont_ind in img_contours:
                x, y, w, h = cv2.boundingRect(cont_ind)
                box = [x, y, w, h]
                image, _ = crop_image_inside_box(box, image_page)
                if not image.size:
                    continue
                image = resize_image(image,
                                     int(image.shape[0] / scale_y),
                                     int(image.shape[1] / scale_x))
                cv2.imwrite(os.path.join(self.dir_of_cropped_images,
                                         (name or "page") + f"_{index:03d}.jpg"), image)
                index += 1

