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



def init_interval_dict():
    interval = {}
    for key in HeadJoints().to_dict().keys():
        interval[key] = [np.inf, -np.inf]
    return interval


def update_interval(interval, result):
    for key, val in result.to_dict().items():
        if key in interval:
            pre_min, pre_max = interval[key]
            cur_min = min(getattr(result, key), pre_min)
            cur_max = max(getattr(result, key), pre_max)
            interval[key] = [cur_min, cur_max]


def save_interval_yaml(interval, filename="./config/calibrated_interval.yaml"):
    # Wrap in AngleInterval
    out = {}
    for k, v in interval.items():
        out[k] = [float(v[0]), float(v[1])]
    with open(filename, "w") as f:
        yaml.dump(out, f)


def calibrate():
    with open(r"./config/config.yaml", "r") as f:
        config = 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'])

    interval = init_interval_dict()
    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 = HeadJoints()
        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)

        update_interval(interval, result)

        cv2.imshow("Calibration", color)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    pipe.stop()
    cv2.destroyAllWindows()
    save_interval_yaml(interval)


if __name__ == '__main__':
    calibrate()