{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Trajectory** on the reconstructed lane"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Trajectory computation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "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": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "VIDEO_NUMBER = \"7\"\n",
    "PROJECT_ROOT = Path().resolve().parent.parent\n",
    "INPUT_VIDEO_PATH = str(\n",
    "    PROJECT_ROOT\n",
    "    / \"data\"\n",
    "    / f\"recording_{VIDEO_NUMBER}\"\n",
    "    / f\"Recording_{VIDEO_NUMBER}.mp4\"\n",
    ")\n",
    "CSV_FILE_PATH = str(\n",
    "    PROJECT_ROOT\n",
    "    / \"data\"\n",
    "    / \"auxiliary_data\"\n",
    "    / \"reconstructed_positions\"\n",
    "    / f\"Transformed_positions_processed_{VIDEO_NUMBER}.csv\"\n",
    ")\n",
    "OUTPUT_VIDEO_PATH = str(\n",
    "    PROJECT_ROOT\n",
    "    / \"data\"\n",
    "    / f\"recording_{VIDEO_NUMBER}\"\n",
    "    / f\"Reconstructed_trajectory_processed_{VIDEO_NUMBER}.mp4\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating video *Tracked_output.mp4*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create image\n",
    "width = 106  # Approximate pixels for 1.0668m\n",
    "height = 1829  # Approximate pixels for 18.29m\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": 29,
   "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:\n",
    "            frame_str, x_str, y_str = row[:3]\n",
    "            if frame_str.strip() and x_str.strip() and y_str.strip():\n",
    "                try:\n",
    "                    frame_num = int(frame_str)\n",
    "                    x = int(float(x_str))  # scale down to fit image\n",
    "                    y = int(float(y_str))\n",
    "                    positions[frame_num] = (x, y)\n",
    "                except ValueError:\n",
    "                    continue  # skip rows with non-numeric values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tracking video saved to /home/davic/projects/IACV_project/bowling-analysis/data/recording_7/Reconstructed_trajectory_processed_7.mp4\n"
     ]
    }
   ],
   "source": [
    "frame_count = 0\n",
    "trajectory = []\n",
    "\n",
    "while cap.isOpened():\n",
    "    ret, _ = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "\n",
    "    frame = alley.copy()\n",
    "\n",
    "    # Check if this frame has a position\n",
    "    if frame_count in positions:\n",
    "        x, y = positions[frame_count]\n",
    "        trajectory.append((x, y))\n",
    "        # print(f\"Frame {frame_count}: Ball at ({x}, {y})\")\n",
    "\n",
    "    # Draw trajectory path\n",
    "    for i in range(1, len(trajectory)):\n",
    "        cv2.line(frame, trajectory[i - 1], trajectory[i], (0, 0, 255), 2)\n",
    "\n",
    "    # Draw current ball position\n",
    "    if frame_count in positions:\n",
    "        cv2.circle(frame, (x, y), 10, (0, 0, 255), -1)  # red dot\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
}