import cv2
import numpy as np
import pandas as pd
from include.PoseExtractor import PoseExtractor
from include.FeaturesExtractor import FeaturesExtractor

DF = pd.DataFrame()

videos_dict = {"hands_down1.mp4": "hands_small_down",
               "hands_down2.mp4": "hands_small_down",
               "hands_downest2.mp4": "hands_down",
               "hands_stop1.mp4": "hands_stop",
               "hands_stop2.mp4": "hands_stop",
               "hands_to_sides1.mp4": "hands_to_sides",
               "hands_to_sides2.mp4": "hands_to_sides",
               "hands_up1.mp4": "hands_up",
               "hands_up2.mp4": "hands_up",
               "hands_upper1.mp4": "hands_upper",
               "hands_upper2.mp4": "hands_upper",
               "hands_downest1.mp4": "hands_down"}

pose_extractor = PoseExtractor()
features_extractor = FeaturesExtractor()

for video_name, gesture_class in videos_dict.items():
    print(video_name)
    print(gesture_class)
    videos_dir = "videos"
    video_input = videos_dir + "/" + video_name

    stream = cv2.VideoCapture(video_input)
    stream.set(cv2.CAP_PROP_FPS, 60)

    if stream.isOpened():
        RUN = True
    else:
        RUN = False
        print("Can't open file '{}'".format(video_name))


    while RUN:
        ret, img = stream.read()
        if ret:
            pose_extractor.update_op_datum(img)
            drawed_pose = pose_extractor.draw_pose()
            poses = pose_extractor.get_poses()
            if len(poses) == 1:
                actual_pose = poses[0]
                k_points_list = pose_extractor.pose_to_k_points(actual_pose)

                k_points_features, scale_factor_points = features_extractor.extract_pose_k_points(k_points_list)
                scale_factor = features_extractor.get_scale_factor(scale_factor_points)

                x_f = features_extractor.get_x_features(k_points_features, scale_factor)
                y_f = features_extractor.get_y_features(k_points_features, scale_factor)

                if x_f is not None and y_f is not None:
                    data = {}
                    for i in range(10):
                        x_column = "x_{}".format(i)
                        y_column = "y_{}".format(i)
                        data[x_column] = []
                        data[y_column] = []
                    data["gestures_class"] = []

                    for i, x_val in enumerate(x_f):
                        data["x_{}".format(i)] = x_val
                    for i, y_val in enumerate(y_f):
                        data["y_{}".format(i)] = y_val
                    data["gestures_class"] = gesture_class

                    DS = pd.Series(data)
                    DF = DF.append(DS, ignore_index=True)

            cv2.imshow("i", drawed_pose)
            if cv2.waitKey(1) & 0xFF == 27:
                RUN = False
        else:
            RUN = False

DF.to_csv("scaled_data.csv")

cv2.destroyAllWindows()
stream.release()
