{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "57439bfb",
   "metadata": {},
   "source": [
    "# New approaches to axis detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "491b8d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.spatial.transform import Rotation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84cc5269",
   "metadata": {},
   "source": [
    "Import video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "200992e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opened: True, FPS: 30.0, Total Frames: 258.0\n",
      "Width: 1920, Height: 1080\n"
     ]
    }
   ],
   "source": [
    "video_number = \"5\"\n",
    "# Define the relative path to the video file\n",
    "notebook_dir = Path().resolve()\n",
    "project_root = notebook_dir.parent.parent\n",
    "video_path = (\n",
    "    project_root\n",
    "    / \"data\"\n",
    "    / f\"recording_{video_number}\"\n",
    "    / f\"Recording_{video_number}.mp4\"\n",
    ")\n",
    "video_path = str(video_path)\n",
    "\n",
    "# Load the video\n",
    "cap = cv2.VideoCapture(video_path)\n",
    "\n",
    "# Check\n",
    "print(\n",
    "    f\"Opened: {cap.isOpened()}, FPS: {cap.get(cv2.CAP_PROP_FPS)}, Total Frames: {cap.get(cv2.CAP_PROP_FRAME_COUNT)}\"\n",
    ")\n",
    "fps = round(cap.get(cv2.CAP_PROP_FPS))\n",
    "# size of the video\n",
    "width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "print(f\"Width: {width}, Height: {height}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41d9befd",
   "metadata": {},
   "source": [
    "Import data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "da051d27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the path to the CSV file\n",
    "input_data_path = (\n",
    "    project_root\n",
    "    / \"data\"\n",
    "    / \"auxiliary_data\"\n",
    "    / \"circle_positions\"\n",
    "    / f\"Adjusted_positions_new_{video_number}.csv\"\n",
    ")\n",
    "\n",
    "# Load the CSV file into a DataFrame\n",
    "df = pd.read_csv(input_data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b3c6d44",
   "metadata": {},
   "source": [
    "Cut the frames roi (gray)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "3e83fe6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_two_roi(x, y, x_next, y_next, r, curr_frame, next_frame, width, height, off=2):\n",
    "    x0, x1 = max(x - r - off, 0), min(x + r + off, width)\n",
    "    y0, y1 = max(y - r - off, 0), min(y + r + off, height)\n",
    "\n",
    "    x0_next, x1_next = max(x_next - r - off, 0), min(x_next + r + off, width)\n",
    "    y0_next, y1_next = max(y_next - r - off, 0), min(y_next + r + off, height)\n",
    "\n",
    "    roi1 = curr_frame[y0:y1, x0:x1]\n",
    "    roi2 = next_frame[y0_next:y1_next, x0_next:x1_next]\n",
    "    g1 = cv2.cvtColor(roi1, cv2.COLOR_BGR2GRAY)\n",
    "    g2 = cv2.cvtColor(roi2, cv2.COLOR_BGR2GRAY)\n",
    "    c_roi = np.array([x - x0, y - y0])\n",
    "\n",
    "    return g1, g2, c_roi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11380edc",
   "metadata": {},
   "source": [
    "Sample points random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "338e4b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_points_random(frame, center_roi, r):\n",
    "    # Randomly sample points within the circle\n",
    "    pts = []\n",
    "    while len(pts) < 1000:\n",
    "        xi = np.random.randint(0, frame.shape[1])\n",
    "        yi = np.random.randint(0, frame.shape[0])\n",
    "        if (xi - center_roi[0]) ** 2 + (yi - center_roi[1]) ** 2 <= (r * 0.7) ** 2:\n",
    "            pts.append([xi, yi])\n",
    "    p0 = np.array(pts, dtype=np.float32).reshape(-1, 1, 2)\n",
    "\n",
    "    if p0.size == 0:\n",
    "        return None\n",
    "    return p0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ebab63e",
   "metadata": {},
   "source": [
    "Sample points manually"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "9783376d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_points_manually(center_roi, r):\n",
    "    # selezione 100 punti manuali distribuiti su 10 cerchi concentrici\n",
    "    manual_points = []\n",
    "    num_rings = 10  # numero di cerchi concentrici\n",
    "    points_per_ring = 10  # punti per ogni cerchio (10 × 10 = 100)\n",
    "    max_radius = r * 0.99  # lascia un margine vicino al bordo\n",
    "\n",
    "    for i in range(1, num_rings + 1):\n",
    "        r = (i / num_rings) * max_radius\n",
    "        for j in range(points_per_ring):\n",
    "            if j < 20 or j > 9:\n",
    "                theta = 2 * np.pi * j / points_per_ring\n",
    "                x = int(center_roi[0] + r * np.cos(theta))\n",
    "                y = int(center_roi[1] + r * np.sin(theta))\n",
    "                manual_points.append((x, y))\n",
    "\n",
    "    # Converti in formato richiesto da calcOpticalFlowPyrLK\n",
    "    p0 = np.array(manual_points, dtype=np.float32).reshape(-1, 1, 2)\n",
    "\n",
    "    if p0.size == 0:\n",
    "        return None\n",
    "    return p0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac8b6bd6",
   "metadata": {},
   "source": [
    "Compute the new points with optical flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "5cfd3cb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def new_points_optical_flow(g1, g2, p0):\n",
    "    p1, st, _ = cv2.calcOpticalFlowPyrLK(\n",
    "        g1,\n",
    "        g2,\n",
    "        p0,\n",
    "        None,\n",
    "        winSize=(15, 15),\n",
    "        maxLevel=2,\n",
    "        criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03),\n",
    "    )\n",
    "\n",
    "    return p1, st"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79c25e69",
   "metadata": {},
   "source": [
    "Get the 3d points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "c6e554cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_3d_points(p0, p1, st, c_roi, r):\n",
    "    old3d, new3d = [], []\n",
    "    for o, n, s in zip(p0.reshape(-1, 2), p1.reshape(-1, 2), st.reshape(-1)):\n",
    "        if s:\n",
    "            ox, oy = o - c_roi\n",
    "            nx, ny = n - c_roi\n",
    "            oz = np.sqrt(max(r * r - ox * ox - oy * oy, 0))\n",
    "            nz = np.sqrt(max(r * r - nx * nx - ny * ny, 0))\n",
    "            old3d.append([ox, oy, oz])\n",
    "            new3d.append([nx, ny, nz])\n",
    "\n",
    "    if len(old3d) < 3:\n",
    "        return None, None\n",
    "\n",
    "    old3d = np.array(old3d)\n",
    "    new3d = np.array(new3d)\n",
    "\n",
    "    return old3d, new3d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2912c47a",
   "metadata": {},
   "source": [
    "Compute the axis with SVD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "0e869cd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_axis_svd(old3d, new3d):\n",
    "    # Compute Rotation matrix using SVD\n",
    "    Hmat = old3d.T @ new3d\n",
    "    U, S, Vt = np.linalg.svd(Hmat)\n",
    "    R = Vt.T @ U.T\n",
    "    if np.linalg.det(R) < 0:\n",
    "        Vt[-1, :] *= -1\n",
    "        R = Vt.T @ U.T\n",
    "\n",
    "    # Compute the rotation axis and angle\n",
    "    theta = np.arccos(np.clip((np.trace(R) - 1) / 2, -1, 1))\n",
    "    if np.sin(theta) != 0:\n",
    "        ax = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])\n",
    "        ax /= 2 * np.sin(theta)\n",
    "        ax /= np.linalg.norm(ax)\n",
    "        return ax, theta\n",
    "    return None, None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddfa9cd1",
   "metadata": {},
   "source": [
    "Compute axis with Kabsch SVD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "aeee310c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def axis_kabsch(P1, P2):\n",
    "    H = P1.T @ P2\n",
    "    U, S, Vt = np.linalg.svd(H)\n",
    "    R = Vt.T @ U.T\n",
    "    if np.linalg.det(R) < 0:\n",
    "        Vt[-1, :] *= -1\n",
    "        R = Vt.T @ U.T\n",
    "    # Eigenvector of R corresponding to eigenvalue 1\n",
    "    axis = np.real_if_close(np.linalg.svd(R - np.eye(3))[2][-1])\n",
    "    return axis / np.linalg.norm(axis)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d78cd042",
   "metadata": {},
   "source": [
    "Get axis with Rodriguez formula"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "f2528896",
   "metadata": {},
   "outputs": [],
   "source": [
    "def axis_rodrigues(p1, p2):\n",
    "    cross = np.cross(p1, p2)\n",
    "    if np.linalg.norm(cross) < 1e-6:\n",
    "        raise ValueError(\"Vectors are nearly colinear or identical.\")\n",
    "    axis = cross / np.linalg.norm(cross)\n",
    "    return axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "c2bd3ed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def average_rodrigues_axis(P1, P2):\n",
    "    axes = []\n",
    "    for i in range(len(P1)):\n",
    "        try:\n",
    "            a = axis_rodrigues(P1[i], P2[i])\n",
    "            axes.append(a)\n",
    "        except:\n",
    "            pass\n",
    "    return np.mean(axes, axis=0) / np.linalg.norm(np.mean(axes, axis=0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c35ac87",
   "metadata": {},
   "source": [
    "Get axis with quaternions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "08dfbe59",
   "metadata": {},
   "outputs": [],
   "source": [
    "def axis_quaternion(P1, P2):\n",
    "    rot, _ = Rotation.align_vectors(P2, P1)\n",
    "    rotvec = rot.as_rotvec()\n",
    "    axis = rotvec / np.linalg.norm(rotvec)\n",
    "    return axis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1ca990c",
   "metadata": {},
   "source": [
    "Get axis from velocities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "03162a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "def axis_from_velocity(P1, P2, delta_t):\n",
    "    V = (P2 - P1) / delta_t  # Nx3 velocity vectors\n",
    "    A = np.zeros((3 * len(P1), 3))\n",
    "    b = np.zeros((3 * len(P1), 1))\n",
    "\n",
    "    for i in range(len(P1)):\n",
    "        r = P1[i]\n",
    "        vx, vy, vz = V[i]\n",
    "        # ω x r = v => [r]_x ω = v\n",
    "        skew = np.array([[0, -r[2], r[1]], [r[2], 0, -r[0]], [-r[1], r[0], 0]])\n",
    "        A[3 * i : 3 * i + 3, :] = skew\n",
    "        b[3 * i : 3 * i + 3, 0] = V[i]\n",
    "\n",
    "    # Least squares solution\n",
    "    omega, _, _, _ = np.linalg.lstsq(A, b, rcond=None)\n",
    "    axis = omega.flatten() / np.linalg.norm(omega)\n",
    "    return axis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55172fa5",
   "metadata": {},
   "source": [
    "Main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "6b082aa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_axis(\n",
    "    x, y, x_next, y_next, r, curr_frame, next_frame, width, height, max_iterations=3\n",
    "):\n",
    "    grey_curr, grey_next, center_roi = get_two_roi(\n",
    "        x, y, x_next, y_next, r, curr_frame, next_frame, width, height\n",
    "    )\n",
    "\n",
    "    axes = []\n",
    "    angles = []\n",
    "\n",
    "    for _ in range(max_iterations):\n",
    "        p0 = sample_points_random(grey_curr, center_roi, r)\n",
    "        # p0 = sample_points_manually(center_roi, r)\n",
    "\n",
    "        if p0 is None:\n",
    "            break\n",
    "        p1, st = new_points_optical_flow(grey_curr, grey_next, p0)\n",
    "        curr_3d_points, next_3d_points = get_3d_points(p0, p1, st, center_roi, r)\n",
    "        if curr_3d_points is None or next_3d_points is None:\n",
    "            break\n",
    "\n",
    "        axis, angle = get_axis_svd(curr_3d_points, next_3d_points)\n",
    "        # axis = axis_kabsch(curr_3d_points, next_3d_points) # similar to get_axis_svd\n",
    "        # axis = average_rodrigues_axis(curr_3d_points, next_3d_points) # bad results\n",
    "        # axis = axis_quaternion(curr_3d_points, next_3d_points) # similar to get_axis_svd\n",
    "        # axis = axis_from_velocity(curr_3d_points, next_3d_points, 1/fps) # good results\n",
    "\n",
    "        if axis is None:\n",
    "            break\n",
    "        axes.append(axis)\n",
    "        angles.append(angle)\n",
    "\n",
    "    if axes:\n",
    "        avg_axis = np.mean(np.vstack(axes), axis=0)\n",
    "        avg_axis /= np.linalg.norm(avg_axis)\n",
    "        avg_angle = np.mean(angles)\n",
    "        return avg_axis, avg_angle\n",
    "    return None, None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dea8b0eb",
   "metadata": {},
   "source": [
    "Generate video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "cb582356",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- prepare writer ---\n",
    "fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "out_path = str(\n",
    "    project_root\n",
    "    / \"data\"\n",
    "    / f\"recording_{video_number}\"\n",
    "    / \"test_video\"\n",
    "    / f\"Spin_video_{video_number}.mp4\"\n",
    ")\n",
    "fourcc = cv2.VideoWriter_fourcc(*\"mp4v\")\n",
    "out = cv2.VideoWriter(str(out_path), fourcc, fps, (W, H))\n",
    "\n",
    "# circle video\n",
    "# Define the output path for the new video\n",
    "output_path = str(\n",
    "    project_root\n",
    "    / \"data\"\n",
    "    / f\"recording_{video_number}\"\n",
    "    / \"test_video\"\n",
    "    / f\"Red_Circle_Video_{video_number}.mp4\"\n",
    ")\n",
    "\n",
    "# Create a VideoWriter object\n",
    "circle_video = cv2.VideoWriter(output_path, fourcc, fps, (300, 300))\n",
    "\n",
    "# Calculate the total number of frames in the original video\n",
    "total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "\n",
    "# Prepare list to log axis endpoints\n",
    "log = []\n",
    "\n",
    "# Read frame0 and frame1\n",
    "ret_curr, curr = cap.read()\n",
    "ret_next, next_ = cap.read()\n",
    "frame_idx = 0\n",
    "\n",
    "while ret_curr and ret_next:\n",
    "    vis = curr.copy()\n",
    "    pA, pB = (np.nan, np.nan), (np.nan, np.nan)\n",
    "    axis = None\n",
    "    # Only compute if we have valid circle data\n",
    "    if (\n",
    "        frame_idx < len(df) - 1\n",
    "        and not df.iloc[frame_idx][[\"x\", \"y\", \"radius\"]].isna().any()\n",
    "        and not df.iloc[frame_idx + 1][[\"x\", \"y\", \"radius\"]].isna().any()\n",
    "    ):\n",
    "        x, y, r = df.iloc[frame_idx][[\"x\", \"y\", \"radius\"]].astype(int)\n",
    "        x_next, y_next, r_next = df.iloc[frame_idx + 1][[\"x\", \"y\", \"radius\"]].astype(\n",
    "            int\n",
    "        )\n",
    "        # Get the axis of rotation\n",
    "        axis, angle = compute_axis(x, y, x_next, y_next, r, curr, next_, W, H)\n",
    "\n",
    "        # Get the Intersection points with the sphere\n",
    "        if axis is not None:\n",
    "            # add axis to df\n",
    "            df.at[frame_idx, \"axis_x\"] = axis[0]\n",
    "            df.at[frame_idx, \"axis_y\"] = axis[1]\n",
    "            df.at[frame_idx, \"axis_z\"] = axis[2]\n",
    "            df.at[frame_idx, \"angle\"] = angle\n",
    "            # Calculate the intersection points with the sphere\n",
    "            pA = axis * r\n",
    "            pB = -axis * r\n",
    "            pA = (x + pA[0], y + pA[1])\n",
    "            pB = (x + pB[0], y + pB[1])\n",
    "\n",
    "            # Draw the axis on the current frame\n",
    "            cv2.line(\n",
    "                vis, (int(pA[0]), int(pA[1])), (int(pB[0]), int(pB[1])), (0, 255, 0), 2\n",
    "            )\n",
    "            cv2.circle(vis, (int(pA[0]), int(pA[1])), 5, (0, 255, 255), -1)\n",
    "            cv2.circle(vis, (int(pB[0]), int(pB[1])), 5, (0, 0, 255), -1)\n",
    "            cv2.circle(vis, (int(x), int(y)), 5, (255, 0, 0), 2)\n",
    "\n",
    "    # log the endpoints for this frame\n",
    "    log.append(\n",
    "        {\"frame\": frame_idx, \"pA_x\": pA[0], \"pA_y\": pA[1], \"pB_x\": pB[0], \"pB_y\": pB[1]}\n",
    "    )\n",
    "\n",
    "    out.write(vis)\n",
    "\n",
    "    # Create a black square\n",
    "    frame = np.zeros((300, 300, 3), dtype=np.uint8)\n",
    "\n",
    "    # Draw a red circle in the center\n",
    "    center = (150, 150)  # Center of the square\n",
    "    radius = 100  # Radius of the circle\n",
    "    color = (0, 0, 255)  # Red color in BGR\n",
    "    thickness = -1  # Fill the circle\n",
    "    cv2.circle(frame, center, radius, color, thickness)\n",
    "    if axis is not None:\n",
    "        # convert points to circle coordinates\n",
    "        pA = axis * 100\n",
    "        pB = -axis * 100\n",
    "        pA = (150 + pA[0], 150 + pA[1])\n",
    "        pB = (150 + pB[0], 150 + pB[1])\n",
    "        cv2.circle(frame, (int(pA[0]), int(pA[1])), 5, (255, 255, 0), -1)\n",
    "        cv2.circle(frame, (int(pB[0]), int(pB[1])), 5, (0, 255, 255), -1)\n",
    "        cv2.circle(frame, center, 5, (255, 0, 0), 2)\n",
    "        cv2.line(\n",
    "            frame, (int(pA[0]), int(pA[1])), (int(pB[0]), int(pB[1])), (0, 255, 0), 2\n",
    "        )\n",
    "\n",
    "    # Write the frame to the video\n",
    "    circle_video.write(frame)\n",
    "\n",
    "    # advance frames\n",
    "    frame_idx += 1\n",
    "    curr = next_\n",
    "    ret_next, next_ = cap.read()\n",
    "    ret_curr = curr is not None\n",
    "\n",
    "# write last frame if exists\n",
    "if ret_curr:\n",
    "    vis = curr.copy()\n",
    "    log.append(\n",
    "        {\n",
    "            \"frame\": frame_idx,\n",
    "            \"pA_x\": np.nan,\n",
    "            \"pA_y\": np.nan,\n",
    "            \"pB_x\": np.nan,\n",
    "            \"pB_y\": np.nan,\n",
    "        }\n",
    "    )\n",
    "    out.write(vis)\n",
    "\n",
    "cap.release()\n",
    "out.release()\n",
    "circle_video.release()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32a83764",
   "metadata": {},
   "source": [
    "Save in csv file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "98623d61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DataFrame saved to C:\\Users\\miche\\OneDrive\\Documenti\\GitHub\\bowling-analysis\\notebook\\spin\\intermediate_data\\Axis_5.csv\n"
     ]
    }
   ],
   "source": [
    "output_data_path = (\n",
    "    project_root\n",
    "    / \"notebook\"\n",
    "    / \"spin\"\n",
    "    / \"intermediate_data\"\n",
    "    / f\"Axis_{video_number}.csv\"\n",
    ")\n",
    "df.to_csv(output_data_path, index=False)\n",
    "print(f\"DataFrame saved to {output_data_path}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}