import pandas as pd
from sklearn.model_selection import StratifiedKFold
import os
import glob 
import numpy as np 
import random
from PIL import Image
import matplotlib.pyplot as plt
import seaborn as sns
import sys
import cv2
import tensorflow as tf 

sys.path.append('./')
from utils.train_config import config
from utils.datagenerator import _denorm
from utils.datagenerator import get_tf_test_data
from utils.datagenerator import plot_batch_samples,view_image, debug_batch, get_eyeball_set_from_tf_data
from utils.tfexplain_utils import get_tfexplain_callbacks, get_post_gradcam_heatmap, get_peak_location
import utils.vis as vis_utl
from utils.tfexplain_utils import get_thresholded_img_contours

# preprocessing function
pre_func = tf.keras.applications.efficientnet.preprocess_input

IMG_WIDTH = config.MODEL.win
IMG_HEIGHT = config.MODEL.hin
img_shape = (IMG_HEIGHT, IMG_WIDTH, 3)

def load_model(model_path):
    new_model = tf.keras.models.load_model(model_path)
    return new_model

def  get_feature_extractor_model(keras_model):
    feature_extractor =  tf.keras.Model(keras_model.inputs,keras_model.layers[-3].output)
    return feature_extractor

def open_gray(fn):
    img = cv2.cvtColor(cv2.imread(fn), cv2.COLOR_BGR2GRAY)
    img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    return img

def read_image_3_channel_v2(path):
    img = open_gray(path)
    img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT), cv2.INTER_LINEAR)
    # img = img/255.
     # Apply model-specific preprocessing function
    img = pre_func(img).astype(float)
    return img

def get_predictions_and_roi(tmp_path, new_model):
    LABELS = ['CNV', 'DRUSEN', 'DME', 'NORMAL']
    img3 = read_image_3_channel_v2(tmp_path)
    # pre_procc_img_2 = read_image_3_channel_v2(row['path'])
    pre_procc_img_2 = np.expand_dims(img3,axis=0)
    class_prob = new_model.predict(pre_procc_img_2)
    pred_lbl = LABELS[np.argmax(class_prob)]
    cls_index = np.argmax(class_prob)
    heat_map = get_post_gradcam_heatmap(new_model,img3,class_index=cls_index)
    # get thresholded and bounding boox
    clos_img,cv2_bbox = get_thresholded_img_contours(heat_map)
    # tmp_img = _denorm(tmp_img, np.min(tmp_img), np.max(tmp_img))
    tmp_img = open_gray(tmp_path)
    tmp_img = cv2.resize(tmp_img, (IMG_WIDTH, IMG_HEIGHT), cv2.INTER_LINEAR)
    #merge map and frame
    tmp_img = cv2.addWeighted(heat_map.astype(float), 0.4, tmp_img.astype(float), 1, 0)
    tmp_img = vis_utl.draw_bbox(tmp_img,[cv2_bbox[0],cv2_bbox[1],cv2_bbox[0]+cv2_bbox[2],cv2_bbox[1]+cv2_bbox[3]])
    return _denorm(tmp_img, np.min(tmp_img), np.max(tmp_img)), class_prob

def get_feature_vector(tmp_path, feature_extractor):
    img3 = read_image_3_channel_v2(tmp_path)
    pre_procc_img_2 = np.expand_dims(img3,axis=0)
    # new_model = tf.keras.Model(new_model.inputs, new_model.layers[-3].output)
    feature = feature_extractor.predict(pre_procc_img_2)
    emb = feature.astype(np.float16)
    return emb



# LABELS = ['CNV', 'DRUSEN', 'DME', 'NORMAL']
# img3 = read_image_3_channel_v2(tmp_path.replace("F:/","E:/"))
# heat_map = get_post_gradcam_heatmap(new_model,img3,class_index=0)
# # pre_procc_img_2 = read_image_3_channel_v2(row['path'])
# pre_procc_img_2 = np.expand_dims(img3,axis=0)
# class_prob = new_model.predict(pre_procc_img_2)
# print("done")
# # new_model = tf.keras.Model(new_model.inputs, new_model.layers[-3].output)
# feature = feature_extractor.predict(pre_procc_img_2)
# print(feature.astype(np.float16).shape)
# pred_prob_max = np.max(class_prob)
# pred_lbl = LABELS[np.argmax(class_prob)]
# print(np.argmax(class_prob))
# print(pred_lbl)
# # get thresholded and bounding boox
# clos_img,cv2_bbox = get_thresholded_img_contours(heat_map)
# # tmp_img = _denorm(tmp_img, np.min(tmp_img), np.max(tmp_img))
# tmp_img = open_gray(tmp_path.replace("F:/","E:/"))
# tmp_img = cv2.resize(tmp_img, (IMG_WIDTH, IMG_HEIGHT), cv2.INTER_LINEAR)
# #merge map and frame
# tmp_img = cv2.addWeighted(heat_map.astype(float), 0.4, tmp_img.astype(float), 1, 0)
# tmp_img = vis_utl.draw_bbox(tmp_img,[cv2_bbox[0],cv2_bbox[1],cv2_bbox[0]+cv2_bbox[2],cv2_bbox[1]+cv2_bbox[3]])