import tensorflow as tf
import numpy as np
from PIL import Image

def get_odometer_xy(model_path, image_path):
    #model_path = 'odo_detector.tflite'
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()

    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    # Obtain the height and width of the corresponding image from the input tensor
    image_height = input_details[0]['shape'][1] # 640
    image_width = input_details[0]['shape'][2] # 640

    # Image Preparation
    # image_name = 'car.jpg'
    image = Image.open(image_path)
    image_resized = image.resize((image_width, image_height)) # Resize the image to the corresponding size of the input tensor and store it in a new variable

    image_np = np.array(image_resized) #
    image_np = np.true_divide(image_np, 255, dtype=np.float32) 
    image_np = image_np[np.newaxis, :]

    # inference
    interpreter.set_tensor(input_details[0]['index'], image_np)
    interpreter.invoke()

    # Obtaining output results
    output = interpreter.get_tensor(output_details[0]['index'])
    output = output[0]
    output = output.T

    boxes_xywh = output[:, :4] #Get coordinates of bounding box, first 4 columns of output tensor
    scores = output[:, 4]#np.max(output[..., 5:], axis=1) #Get score value, 5th column of output tensor
    classes = np.zeros(len(scores))#np.argmax(output[..., 5:], axis=1) # Get the class value, get the 6th and subsequent columns of the output tensor, and store the largest value in the output tensor.

    # Threshold Setting
    # threshold = 0.7
    final_score = 0
    x_center, y_center, width, height = 0, 0, 0, 0
    class_name = 'odometer'

    # Bounding boxes, scores, and classes are drawn on the image
    # draw = ImageDraw.Draw(image_resized)

    for box, score, cls in zip(boxes_xywh, scores, classes):
        if score >= final_score:
            x_center, y_center, width, height = box
            final_score = score
            class_name = cls
        else:
            pass
            
    x1 = int((x_center - width / 2) * image_width)
    y1 = int((y_center - height / 2) * image_height)
    x2 = int((x_center + width / 2) * image_width)
    y2 = int((y_center + height / 2) * image_height)

    # draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
    # text = f"Class: {class_name}, Score: {final_score:.2f}"
    # draw.text((x1, y1), text, fill="red")

     # Saving Images
    # image_resized.save('test_img.jpg')

    return x1, y1, x2, y2, final_score