import cv2

from PIL import Image
import tensorflow as tf
import numpy as np
import time 
from datetime import datetime
import glob
import os
from datetime import datetime
import shutil
import csv
import re

today = time.strftime("%Y-%m-%d")

###
# Load the TFLite model and allocate tensors.
models= ['models/efficientdet/efficientdet_lite3x_640_ptq.tflite',
         'models/tf2_ssd_mobilenet_v2_coco17_ptq.tflite',
         'models/ssdlite_mobiledet_coco_qat_postprocess.tflite',
         'models/efficientdet/lite-model_efficientdet_lite4_detection_default_2.tflite',
         'models/modelv1-4.tflite' # 4
         ]

model_path  = models[4]
labels_path = 'models/efficientdet/coco_labels.txt'
print(model_path)

# Load the labels into a list
#classes = open(labels_path, 'r').read().splitlines()
classes = ['Phone', 'Not Phone']
##########
print('Files are exists : ',os.path.exists(model_path), os.path.exists(labels_path), classes)

# Define a list of colors for visualization
COLORS = np.random.randint(0, 255, size=(len(classes), 3), dtype=np.uint8)

def preprocess_image(image_path, input_size):
  """Preprocess the input image to feed to the TFLite model"""
    
  img = tf.image.convert_image_dtype(image_path, tf.uint8)  # commented as image generated from cv2
  original_image = img
  resized_img = tf.image.resize(img, input_size)
  resized_img = resized_img[tf.newaxis, :]
  return resized_img, original_image


def set_input_tensor(interpreter, image):
  """Set the input tensor."""
  tensor_index = interpreter.get_input_details()[0]['index']
  input_tensor = interpreter.tensor(tensor_index)()[0]
  input_tensor[:, :] = image


def get_output_tensor(interpreter, index):
  """Retur the output tensor at the given index."""
  output_details = interpreter.get_output_details()[index]
  tensor = np.squeeze(interpreter.get_tensor(output_details['index']))
  return tensor


def detect_objects(interpreter, image, threshold):
  """Returns a list of detection results, each a dictionary of object info."""
  set_input_tensor(interpreter, image)
  interpreter.invoke()
  scores = get_output_tensor(interpreter, 0)
  boxes = get_output_tensor(interpreter, 1)
  count = get_output_tensor(interpreter, 2)
  class_ = get_output_tensor(interpreter, 3) 
  
  #print('\nboxes : ', boxes, '\nclasses', class_,'\nscores : ', scores, '\ncount : ', count)
  #input('Press any key to continue...')
  results = []
  for i in range(len(class_)):
    # check the detection score and class_id equals to 76 (cell phone)
    if scores[i] >= threshold and class_[i] == 0.0: # new model class 0.0 76.0:
      result = {
        'bounding_box': boxes[i],
        'class_id': class_[i],
        'score': scores[i]
      }
      results.append(result)
  #print('\nboxes : ', results['bounding_box'], '\nclasses', results['class_'],'\nscores : ', results['score'])
  
  return results


def run_odt_and_draw_results(image_path, interpreter, threshold=0.4, counter=0):
  """Run object detection on an image and visualize the detection results""" 
  # Load the input image and preprocess it
  preprocessed_image, original_image = preprocess_image(
      image_path, (input_height, input_width)
    )
  # Plot the detection results on the input image
  original_image_np = original_image.numpy().astype(np.uint8)

  # Run object detection on the input image
  detections = []  # np.array([])
  results = detect_objects(interpreter, preprocessed_image, threshold=threshold)

  class_ids = []
  detection_scors = []
  for obj in results:
    element = []
    ymin, xmin, ymax, xmax = obj['bounding_box']
    class_ids.append(obj['class_id'])
    detection_scors.append(obj['score'])
    xmin = int(xmin * original_image_np.shape[1])
    xmax = int(xmax * original_image_np.shape[1])
    ymin = int(ymin * original_image_np.shape[0])
    ymax = int(ymax * original_image_np.shape[0])

    element = np.array([xmin, ymin, xmax, ymax, obj['score'], obj['class_id'] ])
    detections.append(element) 

  return detections, original_image_np

def normalize(x, y, x1, y1, w,h):
  # 960x440
  return (round(x/w, 8), round(y/h,8), round(x1/w,8), round(y1/h,8))
  
  
  
#########################################
DETECTION_THRESHOLD = 0.85
ROI=0.6


# Load the TFLite model
interpreter = tf.lite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()

_, input_height, input_width, _ = interpreter.get_input_details()[0]['shape']
 

t = datetime.now()
counter = 0
# project name to be used for the input and output file
project = 'negative'

image_path = 'data/%s' % project

files=glob.glob(image_path + '/*.*')
i=0
for count, image in enumerate(files):

    image_np = cv2.imread(image)
    imagex = image # keep the original image name
    # replace space and brackets with underscore in the image name with regex
    image = re.sub(r'\(|\)|\s+', '_', image)
    image = re.sub('.JPG', '.jpg', image)
            
    if image_np is None:    
        print('Image not found: ', image)
        continue
    
    hh,ww,alf = image_np.shape
    ww_hh=ww/hh
    image_np = cv2.resize(image_np, (960, int(960/ww_hh)))
    hh,ww,alf = image_np.shape
    # write the new shape image to the disk replacing extantion with jpg
    new_shapes_path = image.replace(project, 'new_shapes/%s' % project)
    
    cv2.imwrite(new_shapes_path, image_np)
  
    # Run inference and draw detection result on the local copy of the original file
    start = time.monotonic()
    detections, original_image_np = run_odt_and_draw_results(
        # Image.fromarray(image_np), #tf.convert_to_tensor(image_np, dtype=tf.string ), 
        image_np,
        interpreter,
        threshold=DETECTION_THRESHOLD, counter=i)
    inference_time = time.monotonic()-start
        
    # if no any detection, move image to no_phone folder:
    if len(detections) == 0:
      # make a copy of the image if no phone found in the image
      shutil.copy(imagex, image.replace(project, 'no_phone/%s' % project) )
      print('No phone found in the image: ', image)
      
    else:
      i += 1
      print(len(detections), 'Phone found in the image: ', image)
      # if any detection, save the image with bounding boxes
      for tr in  detections[:1]: 
          xmin, ymin, xmax, ymax, score, label_id = int(tr[0].item()), int(tr[1].item()), int(tr[2].item()), int(tr[3].item()), int(tr[4].item()*100), int(tr[5].item())
          color = [0,0,255]
          cv2.rectangle(original_image_np, (xmin, ymin), (xmax, ymax), color, 1)

          # Draw the label background
          label = "{0:}: {1:2d}%".format(classes[int(label_id)], score ) 
          cv2.putText(original_image_np, label, (xmin, ymin),   cv2.FONT_HERSHEY_SIMPLEX, 0.3, color, 1)
          
          # Seperate the data into 80 train, 10 test and 10 validation 
          if i % 8 == 0: grup = 'TEST'  
          elif i % 8 == 1:
              grup = 'VALIDATION'
          else:
              grup = 'TRAIN'
              
          # if cell phone detected
          if classes[int(label_id)] == 'Phone':
              # crop the image from detected object size
              crop_img = original_image_np[ymin:ymax, xmin:xmax]
              if crop_img.size > 0:
                  # save the cropped image
                  cv2.imwrite(image.replace(project, 'cropped/%s' % project), crop_img)
                  xminn, yminn, xmaxx, ymaxx = normalize(xmin, ymin, xmax, ymax, ww,hh )                
                  # line = [grup, new_shapes_path, classes[int(label_id)], xminn, yminn, xmaxx, yminn, xmaxx, ymaxx, xminn, ymaxx]
                  
                  # new line to csv : filename,width,height,class,xmin,ymin,xmax,ymax
                  line = [new_shapes_path.split('/')[-1], ww, hh, 'Phone', xmin, ymin, xmax, ymax]
                  # write line to csv file
                  with open('/home/cappittall/Documents/onarbeni/data/annotations.csv', 'a') as f:
                    csv.writer(f).writerow(line)
                    
    cv2.imwrite(image.replace(project, '/results/%s' % project), original_image_np)
    cv2.imshow('Images-1', original_image_np )
    print('%s Resmi %.2f Saniyede detect edildi ' % (image,  inference_time) )
    #input('Press any key to continue...')
    # Show the detection result
    

    key = cv2.waitKey(1) & 0xFF
        # if the `q` key was pressed, break from the loop

    if key == ord("q"):
        break