import cv2
import yaml
import numpy as np
import mediapipe as mp
import pyrealsense2 as rs
from collections import deque

from biohead.algo import *
from biohead.utils import calc_feature
from biohead.utils import update_text
from biohead.utils import norm


def main(calib_file):
    with open(r"./config/config.yaml", "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    with open(calib_file, "r") as f:
        calib = yaml.load(f, Loader=yaml.FullLoader)

    w, h, fps = config['Camera']['image_width'], config['Camera']['image_height'], config['Camera']['fps']
    pipe = rs.pipeline()
    cfg = rs.config()
    cfg.enable_stream(rs.stream.color, w, h, rs.format.bgr8, fps)
    pipe.start(cfg)
    align = rs.align(rs.stream.color)

    mp_mesh = mp.solutions.face_mesh
    mesh = mp_mesh.FaceMesh(
        max_num_faces=1,
        refine_landmarks=True,
        min_detection_confidence=config['MediaPipe']['min_detection_confidence'],
        min_tracking_confidence=config['MediaPipe']['min_tracking_confidence']
    )

    ray_origins = deque(maxlen=config['Smooth'])
    ray_directions = deque(maxlen=config['Smooth'])
    result = HeadJoints()

    while True:
        frames = align.process(pipe.wait_for_frames())
        color_f = frames.get_color_frame()
        if not color_f: continue

        color = np.asanyarray(color_f.get_data())
        h, w, _ = color.shape
        rgb = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)

        res = mesh.process(rgb)
        if not res.multi_face_landmarks: continue

        uv = calc_feature(color, res, ray_origins, ray_directions)
        result = calc_eyebrow(uv, result)
        result = calc_eyelid(uv, result)
        result = calc_eyeball(uv, result)
        result = calc_mouth(uv, result)
        result = calc_jaw(uv, result)
        result = norm(result, calib)

        update_text(color, result)
        cv2.imshow("RGB", color)

        # canvas = np.zeros((400, 400, 3), dtype=np.uint8)
        # arr = np.array(list(uv.values()))
        # if len(arr) > 0:
        #     umin, vmin = arr.min(0)
        #     umax, vmax = arr.max(0)
        #     for (u, v) in arr:
        #         x2 = int((u - umin) / (umax - umin + 1e-6) * 360 + 20)
        #         y2 = int((vmax - v) / (vmax - vmin + 1e-6) * 360 + 20)
        #         cv2.circle(canvas, (x2, y2), 3, (0, 255, 0), -1)
        # cv2.imshow("Projected Plane", canvas)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    pipe.stop()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main("./config/calibrated_interval.yaml")