{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Trajectory** on the reconstructed lane"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Trajectory computation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading files paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ROOT = Path().resolve().parent.parent\n",
    "INPUT_VIDEO_PATH = str(\n",
    "    PROJECT_ROOT / \"data\" / \"recording_2\" / \"Recording_2_normal_speed.mp4\"\n",
    ")\n",
    "CSV_FILE_PATH = str(\n",
    "    PROJECT_ROOT\n",
    "    / \"data\"\n",
    "    / \"auxiliary_data\"\n",
    "    / \"reconstructed_positions\"\n",
    "    / \"Transformed_positions_2.csv\"\n",
    ")\n",
    "OUTPUT_VIDEO_PATH = str(\n",
    "    PROJECT_ROOT / \"data\" / \"recording_2\" / \"Reconstructed_trajectory.mp4\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating video *Tracked_output.mp4*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create image\n",
    "width = 106  # Approximate pixels for 1.0668m\n",
    "height = 1829  # Approximate pixels for 18.29m\n",
    "\n",
    "brown_color = (135, 184, 222)  # RGB for burly wood\n",
    "alley = np.full((height, width, 3), brown_color, dtype=np.uint8)\n",
    "\n",
    "# Load video\n",
    "cap = cv2.VideoCapture(INPUT_VIDEO_PATH)\n",
    "if not cap.isOpened():\n",
    "    print(\"Error: Could not open video.\")\n",
    "    exit\n",
    "\n",
    "fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
    "\n",
    "# Setup video writer\n",
    "fourcc = cv2.VideoWriter_fourcc(*\"mp4v\")\n",
    "out = cv2.VideoWriter(OUTPUT_VIDEO_PATH, fourcc, fps, (width, height))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read circle positions from CSV\n",
    "positions = {}\n",
    "with open(CSV_FILE_PATH, mode=\"r\") as file:\n",
    "    reader = csv.reader(file)\n",
    "    next(reader)  # Skip header\n",
    "    for row in reader:\n",
    "        if len(row) >= 3 and row[1] and row[2]:  # Ensure valid data\n",
    "            frame_num, x, y = map(int, row[:3])\n",
    "            positions[frame_num] = (int(x / 10), int(y / 10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tracking video saved to /home/davic/projects/IACV_project/bowling-analysis/data/recording_2/Reconstructed_trajectory.mp4\n"
     ]
    }
   ],
   "source": [
    "frame_count = 0\n",
    "trajectory = []  # Store trajectory points\n",
    "\n",
    "while cap.isOpened():\n",
    "    ret, _ = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "\n",
    "    frame = alley.copy()\n",
    "\n",
    "    # Draw tracking path\n",
    "    if frame_count in positions:\n",
    "        x, y = positions[frame_count]\n",
    "        trajectory.append((x, y))\n",
    "\n",
    "    # Draw trajectory\n",
    "    for i in range(1, len(trajectory)):\n",
    "        cv2.line(frame, trajectory[i - 1], trajectory[i], (0, 255, 255), 2)\n",
    "\n",
    "    # Draw current ball position\n",
    "    if frame_count in positions:\n",
    "        x, y = positions[frame_count]\n",
    "        cv2.circle(frame, (x, y), 10, (0, 0, 255), -1)  # Red dot for current position\n",
    "\n",
    "    out.write(frame)\n",
    "    frame_count += 1\n",
    "\n",
    "cap.release()\n",
    "out.release()\n",
    "print(f\"Tracking video saved to {OUTPUT_VIDEO_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "IMPROVEMENTS: Make is smoothy, more visible and maybe use another recostructed lane to print che trajectory"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}