import glob
import numpy as np
import cv2
import tqdm
from ultralytics import YOLO
import matplotlib.pyplot as plt

# --- Wczytanie kalibracji ---
data = np.load("calib_relative_to_A_3cams.npz")

K_A, D_A = data["K_A"], data["D_A"]
K_B, D_B = data["K_B"], data["D_B"]
K_C, D_C = data["K_C"], data["D_C"]

R_AA, T_AA = data["R_AA"], data["T_AA"]
R_BA, T_BA = data["R_BA"], data["T_BA"]
R_CA, T_CA = data["R_CA"], data["T_CA"]

# reshape translacji
T_AA = T_AA.reshape(3,1)
T_BA = T_BA.reshape(3,1)
T_CA = T_CA.reshape(3,1)

# Kamera A = układ odniesienia
P_A = K_A @ np.hstack((np.eye(3), np.zeros((3,1))))
P_B = K_B @ np.hstack((R_BA, T_BA))
P_C = K_C @ np.hstack((R_CA, T_CA))


def triangulate_three_views(pA, pB, pC):
    pA = pA.reshape(2,1)
    pB = pB.reshape(2,1)
    pC = pC.reshape(2,1)

    XAB_h = cv2.triangulatePoints(P_A, P_B, pA, pB)
    XAB = (XAB_h / XAB_h[3])[:3].reshape(3)

    XAC_h = cv2.triangulatePoints(P_A, P_C, pA, pC)
    XAC = (XAC_h / XAC_h[3])[:3].reshape(3)

    return (XAB + XAC)/2


# --- YOLO Pose ---
model = YOLO("yolo11x-pose.pt")

skeleton = [
    [0,1],[0,2],[1,3],[2,4],[0,5],[0,6],
    [5,7],[7,9],[6,8],[8,10],[5,6],
    [11,12],[12,14],[14,16],[11,13],[13,15]
]

points3DList = {}

frames = sorted(glob.glob("video/camA/*.jpg"), key=lambda f: int(__import__("re").search(r"\d+", f).group()))

for frame in tqdm.tqdm(frames):
    name = frame.replace('video/camA\\',"")
    imgA = cv2.imread(f"video/camA/{name}")
    imgB = cv2.imread(f"video/camB/{name}")
    imgC = cv2.imread(f"video/camC/{name}")

    rA = model(imgA, verbose=False)[0]
    rB = model(imgB, verbose=False)[0]
    rC = model(imgC, verbose=False)[0]

    if len(rA.keypoints.xy)==0: continue
    if len(rB.keypoints.xy)==0: continue
    if len(rC.keypoints.xy)==0: continue

    kpA = rA.keypoints.xy[0].cpu().numpy()
    kpB = rB.keypoints.xy[0].cpu().numpy()
    kpC = rC.keypoints.xy[0].cpu().numpy()

    pts = []
    for i in range(kpA.shape[0]):
        X = triangulate_three_views(kpA[i], kpB[i], kpC[i])
        pts.append(X)

    pts = np.array(pts)
    points3DList[name] = pts

import pickle
with open("replay_tpose.pkl", "wb") as f:
    pickle.dump(points3DList, f)