# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/09_e2e_pipeline.ipynb (unless otherwise specified).

__all__ = ['E2EInference']

# Cell

import math
import json
import os
import cv2
from tqdm import tqdm
from skimage import io
import pandas as pd
import numpy as np
from glob import glob
import matplotlib.pyplot as plt
from fastcore.transform import store_attr
from PIL import Image

from .inference_pipeline import KeypointInference
from .tv_detection import InferLogo
from .utils import replacenth
from .perspective_correction import four_point_transform,order_points,transform_keypoints,get_perspective_transform

# Cell

class E2EInference():
    def __init__(self, infer_config: str):
        store_attr(self, "infer_config")

        self.kp_outer_edge = KeypointInference(infer_config, 'kp_outer_edge')
        self.kp_inner_edge = KeypointInference(infer_config, 'kp_inner_edge')
        self.kp_sticky_note = KeypointInference(infer_config, 'kp_sticky_note')
        self.logo_identifier = InferLogo(config)

    def infer(self, img_loc: str, sticky_w: float, sticky_h: float, visualize: bool, save_output="output.jpg"):
        self.warped_outer = False
        self.warped_inner = False
        self.note_detected = False

        self.size_inner = False #size using inner edge
        self.size_outer = False #size using outer edge

        self.size_inner_updated = False #size using inner edge and single perspective correction
        self.size_outer_updated = False #size using outer edge and single perspective correction

        self.screen_width = False
        self.screen_height = False
        self.tv_width = False
        self.tv_height = False

        self.screen_width_updated = False
        self.screen_height_updated = False
        self.tv_width_updated = False
        self.tv_height_updated = False

        self.img = io.imread(img_loc)
        self.outer_kps, outer_vis = self.kp_outer_edge.predict_keypoints(self.img, visualize)
        self.inner_kps, inner_vis = self.kp_inner_edge.predict_keypoints(self.img, visualize)
        self.sticky_note_kps, note_vis = self.kp_sticky_note.predict_keypoints(self.img, visualize)
        self.apply_perspective_correction()
        self.brand, self.score = self.logo_detection(visualize)

        #if sticky note found, run size estimation
        if self.note_detected:
            ratio_h = self.warped_note.shape[0]/sticky_h
            ratio_w = self.warped_note.shape[1]/sticky_w
            self.size_inner, self.screen_width, self.screen_height = self.estimate_tv_size(self.warped_inner, ratio_w, ratio_h)
            self.size_outer, self.tv_width, self.tv_height = self.estimate_tv_size(self.warped_outer, ratio_w, ratio_h)

        #if outer kps and note found project sticky note onto perspectively corrected tv
        if len(self.outer_kps)>0 and len(self.sticky_note_kps)>0:
            sticky_note_x_projection = abs(self.transformed_stickyKPs[0][0] - self.transformed_stickyKPs[1][0])
            sticky_note_y_projection = abs(self.transformed_stickyKPs[0][1] - self.transformed_stickyKPs[-1][1])
            ratio_h = sticky_note_y_projection/sticky_h
            ratio_w = sticky_note_x_projection/sticky_w
            self.size_inner_updated, self.screen_width_updated, self.screen_height_updated = self.estimate_tv_size(self.warped_inner, ratio_w, ratio_h)
            self.size_outer_updated, self.tv_width_updated, self.tv_height_updated = self.estimate_tv_size(self.warped_outer, ratio_w, ratio_h)

        pil_img = self.visualize_keypoints(self.img)
        return pil_img, self.export()

    def apply_perspective_correction(self):
        #TODO : consider keypoints with highest score instead of keypoints from 1st instance
        #if keypoints are detected then appply perspective correction
        if len(self.outer_kps)>0:
            self.warped_outer = four_point_transform(self.img, self.outer_kps[0])
        if len(self.inner_kps)>0:
            self.warped_inner = four_point_transform(self.img, self.inner_kps[0])
        if len(self.sticky_note_kps)>0:
            self.warped_note = four_point_transform(self.img, self.sticky_note_kps[0])
            self.note_detected = True

        if len(self.outer_kps)>0 and len(self.sticky_note_kps)>0:
            sticky_note_kps = order_points(self.sticky_note_kps[0])
            M, transformed_size , warped_image = get_perspective_transform(self.img, np.array(self.outer_kps[0]))
            self.transformed_stickyKPs = transform_keypoints(warped_image, M, sticky_note_kps)

    def logo_detection(self, visualize):
            return self.logo_identifier.predict(img_loc, visualize)

    def estimate_tv_size(self, warped_tv, ratio_w, ratio_h):
        #if keypoints found return estimated sizes, else return False
        if warped_tv.any():
            if ratio_w==0:
                tv_width = 0
            else:
                tv_width = (warped_tv.shape[1]/ratio_w)*10

            if ratio_h==0:
                tv_height = 0
            else:
                tv_height = (warped_tv.shape[0]/ratio_h)*10

            l_tv = tv_width*0.0393701
            b_tv = tv_height*0.0393701
            size = math.sqrt(l_tv**2 + b_tv**2)

            return size, tv_width, tv_height
        else:
            return False, False

    def save_output(self, output, outer_vis, inner_vis, note_vis):
        fig = plt.figure(figsize=(10,10))
        ax = []
        ax.append(fig.add_subplot(2,2,1))
        plt.imshow(self.img)
        v = 2
        if len(self.outer_kps)>0:
            ax.append(fig.add_subplot(2,2,v))
            v = v + 1
            plt.imshow(outer_vis)
        if len(self.inner_kps)>0:
            ax.append(fig.add_subplot(2,2,v))
            v = v + 1
            plt.imshow(inner_vis)
        if self.note_detected:
            ax.append(fig.add_subplot(2,2,v))
            v = v + 1
            plt.imshow(note_vis)
        plt.savefig(f"../data/{output}")
        plt.close()

    def visualize_keypoints(self, img, thickness = 7, color = (255, 0, 0)):
        if len(self.outer_kps)>0:
            pred_kp = self.outer_kps[0]
            for j in range(-1,3):
                x1,y1,x2,y2 = pred_kp[j][0],pred_kp[j][1],pred_kp[j+1][0], pred_kp[j+1][1]
                cv2.circle(img,(x1, y1), 9, (0,255,0), -1)
                cv2.line(img,(x1,y1),(x2,y2),color,thickness)
        if len(self.inner_kps)>0:
            pred_kp = self.inner_kps[0]
            for j in range(-1,3):
                x1,y1,x2,y2 = pred_kp[j][0],pred_kp[j][1],pred_kp[j+1][0], pred_kp[j+1][1]
                cv2.circle(img,(x1, y1), 9, (0,255,0), -1)
                cv2.line(img,(x1,y1),(x2,y2),color,thickness)
        if len(self.sticky_note_kps)>0:
            pred_kp = self.sticky_note_kps[0]
            for j in range(-1,3):
                x1,y1,x2,y2 = pred_kp[j][0],pred_kp[j][1],pred_kp[j+1][0], pred_kp[j+1][1]
                cv2.circle(img,(x1, y1), 9, (0,255,0), -1)
                cv2.line(img,(x1,y1),(x2,y2),color,thickness)

        pil_img = Image.fromarray(img)
        return pil_img

    def export(self):
        result = {}
        result["sticky_note_detected"] = self.note_detected
        result["is_outer_detected"] = True if len(self.outer_kps)>0 else False
        result["is_inner_detected"] = True if len(self.inner_kps)>0 else False
        result["brand_predited"] = False if len(self.brand)<1 else self.brand
        result["size_inner"] = self.size_inner_updated
        result["size_outer"] = self.size_outer_updated
        result["tv_height"] = self.tv_height_updated
        result["tv_width"] = self.tv_width_updated
        result["screen_height"] = self.screen_height_updated
        result["screen_width"] = self.screen_width_updated
#         result["outer_kps"] = [self.outer_kps[0].tolist()] if len(self.outer_kps)>0 else False
#         result["inner_kps"] = [self.inner_kps[0].tolist()] if len(self.inner_kps)>0 else False
#         result["note_kps"] = [self.sticky_note_kps[0].tolist()] if self.note_detected else False
        return result