# AUTOGENERATED! DO NOT EDIT! File to edit: core.ipynb (unless otherwise specified).

__all__ = ['clahe_img', 'CorrespondenceAnnotator']

# Cell
import cv2
from .io import *
def clahe_img(rgb):
    lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB)
    lab_planes = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=10.0,tileGridSize=(16,16))
    lab_planes[0] = clahe.apply(lab_planes[0])
    lab = cv2.merge(lab_planes)
    rgb = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
    return rgb

# Cell
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import ipywidgets as wdg  # Using the ipython notebook widgets
import cv2
import os
import kornia as K
import kornia.feature as KF
import os
import torch
from copy import deepcopy
from typing import List
from kornia_moons.feature import *
from IPython.display import display
from .io import *

class CorrespondenceAnnotator():
    def __init__(self,
                 img_pair_list, save_on_next,
                 corrs_save_name = 'corrs.txt', max_pts = 100):
        self.save_on_next = save_on_next
        self.img_pairs = img_pair_list
        self.corrs_save_name = corrs_save_name
        self.pair_state = self.initialize_pairdict()
        self.state_dict = self.initialize_statedict()
        self.max_pts = max_pts

        return
    def initialize_pairdict(self):
        return {"img1": None,
                "img2": None,
                "pts1": [],
                "pts2": [],
                "F": None,
                "H": None}

    def initialize_ui(self, figsize=(15,8)):
        fig, ((self.ax1, self.ax2), (self.ax3, self.ax4)) = plt.subplots(2, 2, figsize=figsize)
        self.figure = fig
        gs = self.ax3.get_gridspec()
        self.ax3.remove()
        self.ax4.remove()
        self.ax3 = fig.add_subplot(gs[1, :])
        plt.tight_layout()
        self.ln1, = self.ax1.plot([], [], 'x', markersize=10, color='r')
        self.ln2, = self.ax2.plot([], [], 'x', markersize=10, color='r')
        self.texts1 = [self.ax1.text([], [], '', fontsize=12) for i in range(self.max_pts)]
        self.texts2 = [self.ax2.text([], [], '', fontsize=12) for i in range(self.max_pts)]
        #self.fig.tight_layout()

        self.next_button = wdg.Button(description="Next")
        self.prev_button = wdg.Button(description="Prev")
        self.save_button = wdg.Button(description="Save points")
        self.clahe_button = wdg.Button(description='CLAHE')
        self.showF_button = wdg.Button(description='ShowF')

        self.next_button.on_click(self.next_pair)
        self.prev_button.on_click(self.prev_pair)
        self.save_button.on_click(self.save_current_pts)
        self.clahe_button.on_click(self.show_clahed)
        self.showF_button.on_click(self.show_epipolar_lines)
        display(self.prev_button,
                self.next_button,
                self.save_button,
                self.showF_button,
                self.clahe_button)
        return

    def next_pair(self, b):
        if self.save_on_next:
            self.save_current_pts(b)
        if self.state_dict['idx'] >= len(self.img_pairs) -1:
            return
        self.state_dict['idx'] += 1
        self.load_pair_into_statedict(self.state_dict['idx'])
        self.show_current_pair()
        return

    def prev_pair(self, b):
        if self.save_on_next:
            self.save_current_pts(b)
        if self.state_dict['idx'] <= 0:
            return
        self.state_dict['idx'] -= 1
        self.load_pair_into_statedict(self.state_dict['idx'])
        self.show_current_pair()
        return

    def save_current_pts(self, b):
        idx = self.state_dict['idx']
        try:
            img1_fname, img2_fname, pts_fname = self.img_pairs[idx]
        except:
            img1_fname, img2_fname = self.img_pairs[idx]
            pts_fname = os.path.join(os.path.dirname(img1_fname), self.corrs_save_name)
        min_size = min(len(self.pair_state['pts1']), len(self.pair_state['pts2']))
        if min_size == 0:
            return
        pts_arr = np.concatenate([np.array(self.pair_state['pts1']),
                                  np.array(self.pair_state['pts2'])],axis=1)
        np.savetxt(pts_fname, pts_arr)
        return

    def initialize_statedict(self):
        return {"idx": -1}

    def load_pts_to_statedict(self, pts:np.ndarray):
        self.pair_state['pts1'].clear()
        self.pair_state['pts2'].clear()
        if len(pts)>0:
            for x,y,x2,y2 in pts:
                self.pair_state['pts1'].append((x,y))
                self.pair_state['pts2'].append((x2,y2))
        return

    def show_clahed(self, b):
        self.ax1.imshow(clahe_img(self.pair_state['img1']))
        self.ax2.imshow(clahe_img(self.pair_state['img2']))
        return

    def show_epipolar_lines(self, b):
        min_size = min(len(self.pair_state['pts1']), len(self.pair_state['pts2']))
        num = min_size
        if min_size < 8:
            return
        F, inliers = cv2.findFundamentalMat(np.array(self.pair_state['pts1'][:min_size]),
                                            np.array(self.pair_state['pts2'][:min_size]), cv2.FM_8POINT)

        lafs1 = KF.laf_from_center_scale_ori(torch.tensor(self.pair_state['pts1'][:num]).view(1, -1, 2).float(),
                                             torch.ones(1, num, 1, 1, dtype=torch.float),
                                             torch.zeros(1, num, 1,  dtype=torch.float))
        lafs2 = KF.laf_from_center_scale_ori(torch.tensor(self.pair_state['pts2'][:num]).view(1, -1, 2).float(),
                                             torch.ones(1, num, 1, 1, dtype=torch.float),
                                             torch.zeros(1, num, 1,  dtype=torch.float))

        idxs = torch.stack([torch.arange(min_size), torch.arange(min_size)], dim=-1)
        self.ax3.clear()
        draw_LAF_matches(lafs1, lafs2, idxs,
                         self.pair_state['img1'],
                         self.pair_state['img2'],
                         inliers, Fm=F.astype(float), ax=self.ax3)
        return
    def process_user_click(self, event):
        # We are interested only in LB Clicks in the images
        if event.inaxes not in [self.ax1, self.ax2]:
            return
        if str(event.button) not in  ['MouseButton.LEFT', 'MouseButton.RIGHT']:
            return
        # If left click: add point, if right click: delete
        picking_mode = str(event.button) == 'MouseButton.LEFT'

        # Not adding points when zooming
        zoom_pan = ['Cursors.SELECT_REGION', 'Cursors.MOVE']
        pick = 'Cursors.POINTER'
        ptr_type = str(self.figure.canvas.toolbar.cursor)
        if ptr_type != pick:
            txt.value = str('Non picking')
            return

        # Select proper subplot to work with
        ax = event.inaxes
        use_pts = self.pair_state['pts1'] if ax == self.ax1 else self.pair_state['pts2']
        other_pts = self.pair_state['pts2'] if ax == self.ax1 else self.pair_state['pts1']

        x,y = event.xdata, event.ydata
        if picking_mode: # Adding point
            use_pts.append((x,y))
        else:  # remove point, if rightclick near it
            if len(use_pts) > 0:
                already_pts = np.array(use_pts)
                diff = np.abs(already_pts - np.array([[x,y]])).mean(axis=1)
                closest_pts = np.argmin(diff)
                if diff[closest_pts] < 2.0: # We will update the point
                    use_pts.pop(closest_pts)
                    if len(other_pts) >= closest_pts:
                        other_pts.pop(closest_pts)
        # Redraw is set to false to avoid zoom reset
        self.show_current_pair(redraw=False)
        return


    def load_pair_into_statedict(self, idx):
        try:
            img1_fname, img2_fname, pts_fname = self.img_pairs[idx]
        except:
            img1_fname, img2_fname = self.img_pairs[idx]
            pts_fname = os.path.join(os.path.dirname(img1_fname), self.corrs_save_name)
        img1 = imread(img1_fname)
        img2 = imread(img2_fname)
        pts = load_pts(pts_fname)
        self.pair_state['img1'] = img1
        self.pair_state['img2'] = img2
        self.load_pts_to_statedict(pts)
        return

    def show_points(self, pts, line_artist, text_artists):
        if len(pts) == 0:
            return
        num_pts = len(pts)
        pts_np = np.array(pts)
        # We have to draw points and their idxs
        point_idxs = [str(i) for i in range(num_pts)]
        for i, coord in enumerate(pts_np):
            text_artists[i].set_position((coord[0], coord[1]))
            text_artists[i].set_text(point_idxs[i])
        for i in range(num_pts, self.max_pts):
            text_artists[i].set_text("")
        line_artist.set_data(pts_np[:,0], pts_np[:,1])
        return

    def show_current_pair(self, redraw=True):
        if redraw:
            self.ax1.imshow(np.zeros((2,2)))
            self.ax1.imshow(self.pair_state['img1'])
            self.ax2.imshow(np.zeros((2,2)))
            self.ax2.imshow(self.pair_state['img2'])
        self.show_points(self.pair_state['pts1'], self.ln1, self.texts1)
        self.show_points(self.pair_state['pts2'], self.ln2, self.texts2)
        return

    def start(self,  figsize=(15,8)):
        idx = 0
        self.state_dict['idx'] = 0
        self.load_pair_into_statedict(idx)
        self.initialize_ui(figsize=figsize)
        self.show_current_pair()
        self.callback_ref = self.figure.canvas.mpl_connect('button_press_event',
                                                        lambda event: self.process_user_click(event))
        return
