import pandas as pd
import tensorflow as tf
import numpy as np
from matplotlib import pyplot as plt
import cv2
import wget
import os
import glob

print("opencv v is:" + cv2.__version__)

interpreter = tf.lite.Interpreter(model_path='lite-model_movenet_singlepose_lightning_3.tflite')
interpreter.allocate_tensors()


def draw_keypoints(frame, keypoints, confidence_threshold):
    y, x, c = frame.shape
    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))

    for kp in shaped:
        ky, kx, kp_conf = kp
        if kp_conf > confidence_threshold:
            cv2.circle(frame, (int(kx), int(ky)), 4, (0, 255, 0), -1)


EDGES = {
    (11, 13): 'm',
}


def draw_connections(frame, keypoints, edges, confidence_threshold):
    y, x, c = frame.shape
    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))

    for edge, color in edges.items():
        p1, p2 = edge
        y1, x1, c1 = shaped[p1]
        y2, x2, c2 = shaped[p2]

        if (c1 > confidence_threshold) & (c2 > confidence_threshold):
            cv2.line(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)


data = np.array([[0, 0, 0, 0, 0]])
path = "./squat_down/*.*"
for file in glob.glob(path):
    cap = cv2.imread(file)
    height, width, c = cap.shape
    #result = cv2.imwrite('trialData1.jpg', cap, (width, height))

    # Reshape image
    img = cap.copy()
    img = tf.image.resize_with_pad(np.expand_dims(cap, axis=0), 192, 192)
    input_image = tf.cast(img, dtype=tf.float32)

    # Setup input and output
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    # Make predictions
    interpreter.set_tensor(input_details[0]['index'], np.array(input_image))
    interpreter.invoke()
    keypoints_with_scores = interpreter.get_tensor(output_details[0]['index'])
    #print(keypoints_with_scores)


    # Rendering
    draw_connections(cap, keypoints_with_scores, EDGES, 0.4)
    draw_keypoints(cap, keypoints_with_scores, 0.4)
    #result.write(cap)
    cv2.imshow('MoveNet Lightning', cap)

    #result = cv2.imwrite('trialData1.jpg', cap, (width, height))


    cv2.destroyAllWindows()

    left_hip = keypoints_with_scores[0][0][11]
    left_knee = keypoints_with_scores[0][0][13]

    shaped = np.squeeze(np.multiply(interpreter.get_tensor(interpreter.get_output_details()[0]['index']), [480, 640, 1]))

    for kp in shaped:
        ky, kx, kp_conf = kp
        #print(int(ky), int(kx), kp_conf)

    shaped[0], shaped[1]

    for edge, color in EDGES.items():
        p1, p2 = edge
        y1, x1, c1 = shaped[p1]
        y2, x2, c2 = shaped[p2]
        #print((int(x2), int(y2))
        input = np.array([[x1, y1, x2, y2, 0]])
        data = np.concatenate([data, input], axis=0)


data = np.delete(data, 0, 0)
DF = pd.DataFrame(data)
DF.to_csv("squat_down2.csv")
print(data)