import cv2
from time import time
import pickle as pk
import mediapipe as mp
import pandas as pd
import multiprocessing as mtp

from recommendations import check_pose_angle
from landmarks import extract_landmarks
from calc_angles import rangles

def init_cam():
    cam = cv2.VideoCapture(0)  # Use 0 for default camera
    if not cam.isOpened():
        raise ValueError("Could not open camera.")
    cam.set(cv2.CAP_PROP_AUTOFOCUS, 0)
    cam.set(cv2.CAP_PROP_FOCUS, 360)
    cam.set(cv2.CAP_PROP_BRIGHTNESS, 130)
    cam.set(cv2.CAP_PROP_SHARPNESS, 125)
    cam.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
    return cam

def get_pose_name(index):
    names = {
        0: "Adho Mukha Svanasana",
        1: "Phalakasana",
        2: "Utkata Konasana",
        3: "Vrikshasana",
    }
    return str(names.get(index, "Unknown Pose"))

def init_dicts():
    landmarks_points = {
        "nose": 0,
        "left_shoulder": 11, "right_shoulder": 12,
        "left_elbow": 13, "right_elbow": 14,
        "left_wrist": 15, "right_wrist": 16,
        "left_hip": 23, "right_hip": 24,
        "left_knee": 25, "right_knee": 26,
        "left_ankle": 27, "right_ankle": 28,
        "left_heel": 29, "right_heel": 30,
        "left_foot_index": 31, "right_foot_index": 32,
    }
    landmarks_points_array = {key: [] for key in landmarks_points.keys()}
    col_names = [f"{name}_{axis}" for name in landmarks_points.keys() for axis in ["x", "y", "z", "v"]]
    return col_names, landmarks_points_array

def cv2_put_text(image, message):
    cv2.putText(
        image,
        message,
        (50, 50),
        cv2.FONT_HERSHEY_SIMPLEX,
        2,
        (255, 0, 0),
        5,
        cv2.LINE_AA
    )

def destroy(cam):
    cam.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    try:
        cam = init_cam()
        model = pk.load(open("./models/4_poses.model", "rb"))
        cols, landmarks_points_array = init_dicts()
        angles_df = pd.read_csv("./csv_files/4_angles_poses_angles.csv")
        mp_drawing = mp.solutions.drawing_utils
        mp_pose = mp.solutions.pose

        while True:
            result, image = cam.read()
            if not result:
                print("Failed to grab frame.")
                break

            flipped = cv2.flip(image, 1)
            resized_image = cv2.resize(
                flipped,
                (640, 360),
                interpolation=cv2.INTER_AREA
            )

            err, df, landmarks = extract_landmarks(
                resized_image,
                mp_pose,
                cols
            )

            if not err:
                prediction = model.predict(df)
                probabilities = model.predict_proba(df)

                mp_drawing.draw_landmarks(
                    flipped,
                    landmarks,
                    mp_pose.POSE_CONNECTIONS
                )

                if probabilities[0, prediction[0]] > 0.85:
                    cv2_put_text(
                        flipped,
                        get_pose_name(prediction[0])
                    )

                    angles = rangles(df, landmarks_points_array)
                    suggestions = check_pose_angle(
                        prediction[0], angles, angles_df)

                    cv2_put_text(
                        flipped,
                        suggestions[0]
                    )

                else:
                    cv2_put_text(
                        flipped,
                        "No Pose Detected"
                    )
            cv2.imshow("Frame", flipped)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    finally:
        destroy(cam)