import cv2
import mediapipe as mp
import matplotlib
matplotlib.use("Agg")  # <-- ważne: wyłącza GUI
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

# ---------------------
# Wideo wejściowe
# ---------------------
cap = cv2.VideoCapture("input.mp4")
fps = cap.get(cv2.CAP_PROP_FPS)
width = 640
height = 640

# ---------------------
# Wideo wyjściowe
# ---------------------
fourcc = cv2.VideoWriter_fourcc(*"MJPG")


out = cv2.VideoWriter("output.mp4", fourcc, fps, (width, height))

# ---------------------
# MediaPipe Pose
# ---------------------
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=False, model_complexity=1)

frame_id = 0

while True:
    ok, frame = cap.read()
    if not ok:
        break

    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = pose.process(rgb)

    # -----------------------------------------
    # 3D landmarki: pose_world_landmarks
    # -----------------------------------------
    if results.pose_world_landmarks:
        lm = results.pose_world_landmarks.landmark

        xs = np.array([p.x for p in lm])
        ys = np.array([p.y for p in lm])
        zs = np.array([p.z for p in lm])

        # -----------------------------
        # RYSOWANIE 3D w Matplotlib
        # -----------------------------
        fig = plt.figure(figsize=(6.4, 6.4), dpi=100)
        ax = fig.add_subplot(111, projection="3d")

        ax.scatter(xs, zs, ys, s=20)

        ax.set_xlim([-1, 1])
        ax.set_ylim([-1, 1])
        ax.set_zlim([-1, 1])

        ax.set_xlabel("X")
        ax.set_ylabel("Y")
        ax.set_zlabel("Z")
        ax.invert_zaxis()

        # -----------------------------------------
        # Konwersja wykresu Matplotlib → klatka do MP4
        # -----------------------------------------
        fig.canvas.draw()
        renderer = fig.canvas.get_renderer()

        w, h = fig.canvas.get_width_height()

        buf = renderer.buffer_rgba()
        plot_img = np.frombuffer(buf, dtype=np.uint8).reshape((h, w, 4))[:, :, :3]


        plt.close(fig)

        # Dopasowanie rozmiaru do wideo
        plot_img = cv2.resize(plot_img, (width, height))
        plot_img = cv2.cvtColor(plot_img, cv2.COLOR_RGB2BGR)

        out.write(plot_img)

    frame_id += 1

cap.release()
out.release()
print("Zapisano: output.mp4")