import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import roc_auc_score, f1_score, jaccard_score, accuracy_score
import tensorflow as tf



# create prediction mask
def create_mask(pred_mask):
  if pred_mask.shape[-1] > 1:
      pred_mask = tf.argmax(pred_mask, axis=-1)
      pred_mask = pred_mask[..., tf.newaxis]
  
  return pred_mask[0]



def metric_copy(premask, groundtruth):
    seg_inv, gt_inv = np.logical_not(premask), np.logical_not(groundtruth)
    true_pos = float(np.logical_and(premask, groundtruth).sum())  # float for division
    true_neg = np.logical_and(seg_inv, gt_inv).sum()
    false_pos = np.logical_and(premask, gt_inv).sum()
    false_neg = np.logical_and(seg_inv, groundtruth).sum()
    f1 = 2 * true_pos / (2 * true_pos + false_pos + false_neg + 1e-6)
    cross = np.logical_and(premask, groundtruth)
    union = np.logical_or(premask, groundtruth)
    iou = np.sum(cross) / (np.sum(union) + 1e-6)
    if np.sum(cross) + np.sum(union) == 0:
        iou = 1
    return f1, iou



def show_prediction(img, pred):
    print("max_pred = ", np.max(pred), " min_pred = ", np.min(pred))
    plt.subplot(1,2,1)
    plt.imshow(img)
    plt.subplot(1,2,2)
    plt.imshow(pred, cmap='gray') #, vmin=0, vmax=1)
    plt.show()


def show_predictions(dataset=None, num=1):
  if dataset:
    for image, mask in dataset.take(num):
      pred_mask = model.predict(image)
      display([image[0], mask[0], create_mask(pred_mask)])
  else:
    print(sample_image.shape)
    print(sample_mask.shape)
    display([sample_image, sample_mask,
             create_mask(model.predict(sample_image[tf.newaxis, ...]))])



def display(display_list, reverseRGB = True):
  plt.figure(figsize=(4, 4))

  title = ['Input Image', 'True Mask', 'Predicted Mask']

  for i in range(len(display_list)):
    plt.subplot(1, len(display_list), i+1)
    plt.title(title[i])
    if reverseRGB:
        plt.imshow(tf.keras.utils.array_to_img(display_list[i][...,::-1]))
    else:
        plt.imshow(tf.keras.utils.array_to_img(display_list[i]))
    plt.axis('off')
  plt.show()


def get_gt_and_osn_folders(folder):
    folder_list = [folder]
    folder_list.append(folder+"_Facebook")
    folder_list.append(folder+"_Whatsapp")
    folder_list.append(folder+"_Weibo")
    folder_list.append(folder+"_Wechat")
    gt_folder = folder + "_GT"
    return gt_folder,folder_list

def get_gt_and_osn_folder(folder, osn):
    osn_folder = folder+osn
    gt_folder = folder + "_GT"
    return gt_folder,osn_folder


# plots the image + prediction + ground truth
def plot_img_pred_gt(img_path, pre_t, gt):
    print("INPUT plot_img_pred_gt:")
    print("  img_path: ", img_path)
    #get image
    img = cv2.imread(img_path)
    #plot image, prediction and mask
    plot_img_pred_gt_execute(img,pre_t, gt)

    
def plot_img_pred_gt_execute(img, pre_t, gt, DISCRETIZE_OUTPUT=True):
    #print("plot_img_pred_gt_execute(): pre_t.max: ", np.max(pre_t))
    #print("plot_img_pred_gt_execute(): pre_t.min: ", np.min(pre_t))
    if DISCRETIZE_OUTPUT:
        pre_t = pre_t.numpy()
        pre_t[pre_t > 0.5] = 1.0
        pre_t[pre_t <= 0.5] = 0.0
    plt.subplots(1,3,figsize=(10,10))
    plt.subplot(1,3,1)
    plt.imshow(img[...,::-1])
    plt.title("Original Image")
    plt.subplot(1,3,2)
    plt.imshow(pre_t, cmap='gray')
    #plt.imshow(pre_t>0.5, cmap='gray')
    plt.title("Prediction")
    plt.subplot(1,3,3)
    plt.imshow(gt, cmap='gray')
    plt.title("Ground Truth")
    plt.show()
    
    
def mask_bigger_fifty_perc(mask):
    mask_size = mask.size
    #print("mask.shape: ", mask.shape)
    #print("mask_size: ", mask_size)
    nr_points_in_mask = mask_size - (mask == 0.).sum()
    mask_cover_perc_of_img = nr_points_in_mask/mask_size
    #print("mask_cover_perc_of_img: ", mask_cover_perc_of_img)
    if mask_cover_perc_of_img>0.5:
        return True
    return False


#evaluation for one image (auc roc, f1, iou)
def eval_image(pre_t, gt, auc, f1, iou, acc):
    #df_out("pre_t_in evalimage",pre_t,True)

    pre = np.repeat(pre_t.numpy()[:,:,np.newaxis],3,2)
    H, W, _ = pre.shape
    Hg, Wg, C = gt.shape

    if mask_bigger_fifty_perc(gt):
        print("FLIP pre because mask > 50% of image")
        pre = 1 - pre
    
    if H != Hg or W != Wg:
        print("ERROR: values not matching:")
        print(f'H: {H}, W: {W}, C: {C}')
        print(f'Hg: {Hg}, Wg: {Wg}, C: {C}')
        gt = cv2.resize(gt, (W, H))
        gt[gt > 127] = 255
        gt[gt <= 127] = 0
            
    if np.max(gt) != np.min(gt):    
        auc.append(roc_auc_score((gt.reshape(H*W*C) / 255.).astype('int'), pre.reshape(H*W*C)))
    else:
        print("!!!!!!!!!!!!!! eval_image(): np.max(gt) != np.min(gt) !!!!!!!!!!!!")
    pre[pre>0.5] = 1.0
    pre[pre<=0.5] = 0.0
    
    #consider changing to: a, b = metric_copy(pre , gt > 127)
    #a, b = metric_copy(pre , gt / 255.) #old
    a, b = metric_copy(pre , gt)
        
    
    pre_ = tf.reshape(pre, [-1])
    gt_ = tf.reshape(gt / 255., [-1]).astype(tf.int32)
    acc_tmp = accuracy_score(pre_, gt_)
    acc.append(acc_tmp)
          
    f1.append(a)
    iou.append(b)
    #print('Evaluation: AUC: %5.4f, F1: %5.4f, IOU: %5.4f' % (np.mean(auc), np.mean(f1), np.mean(iou)))