{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Trajectory** on the video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Detection**\n",
    "\n",
    "From the Detection phase I want to have:\n",
    "|What I want|Where it should be|\n",
    "|-----------|------------------|\n",
    "|The video with the detection|*.../data/recording...*|\n",
    "|Data about the posizion of the ball for each frame|*.../data/auxiliary_data/Circle_posizions_clean.csv*|\n",
    "|(optional) The video only of the ball detected|*.../data/recording...*|\n",
    "\n",
    "If not present the file */data/auxiliary_data/Circle_positions.csv* run the Detection code at */notebook/ball_detection/Ball_detection_analysis.ipynb* in order to produce it"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Trajectory computation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "from pathlib import Path\n",
    "import csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading files paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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",
    "    / \"circle_positions\"\n",
    "    / \"Circle_positions_2.0_clean_radius.csv\"\n",
    ")\n",
    "OUTPUT_VIDEO_PATH = str(\n",
    "    PROJECT_ROOT / \"data\" / \"recording_2\" / \"Tracked_output_2.0_radius.mp4\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating video *Tracked_output.mp4*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "# Get video properties\n",
    "frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\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, (frame_width, frame_height))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read circle positions from CSV\n",
    "positions = {}\n",
    "radius_values = {}\n",
    "with open(CSV_FILE_PATH, mode=\"r\") as file:\n",
    "    reader = csv.reader(file)\n",
    "    next(reader)  # Skip header\n",
    "    for radius in reader:\n",
    "        frame_num, x, y, r = radius\n",
    "        positions[int(frame_num)] = (int(x), int(y)) if x and y else None\n",
    "        radius_values[int(frame_num)] = int(r) if r else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tracking video saved to /home/davic/projects/IACV_project/bowling-analysis/data/recording_2/Tracked_output_2.0_radius.mp4\n"
     ]
    }
   ],
   "source": [
    "frame_count = 0\n",
    "trajectory = []  # Store trajectory points\n",
    "\n",
    "while True:\n",
    "    ret, frame = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "\n",
    "    # Draw tracking path\n",
    "    if frame_count in positions and positions[frame_count] is not None:\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 and positions[frame_count] is not None:\n",
    "        x, y = positions[frame_count]\n",
    "        r = radius_values[frame_count]\n",
    "        cv2.circle(frame, (x, y), 10, (0, 0, 255), -1)  # Red dot for current position\n",
    "        cv2.circle(\n",
    "            frame, (x, y), r, (0, 255, 0), 3\n",
    "        )  # Green circle for the ball profile\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}\")"
   ]
  }
 ],
 "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
}