From 7d57d40991ba41eb6848dc73f323a6aa9037525b Mon Sep 17 00:00:00 2001
From: Michael Kennedy <34204085+mlists@users.noreply.github.com>
Date: Thu, 14 Mar 2024 14:36:54 +1100
Subject: [PATCH] Add shot logging for tuning at comp

---
 controllers/shooter.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/controllers/shooter.py b/controllers/shooter.py
index 5e06304e..256a2667 100644
--- a/controllers/shooter.py
+++ b/controllers/shooter.py
@@ -1,6 +1,8 @@
 import math
 
 from wpimath.geometry import Translation2d
+from wpiutil.log import DataLog, FloatArrayLogEntry, FloatLogEntry
+from wpilib import DriverStation
 
 from magicbot import StateMachine, state, timed_state, feedback, tunable
 
@@ -18,6 +20,8 @@ class Shooter(StateMachine):
     intake_component: IntakeComponent
     status_lights: LightStrip
 
+    data_log: DataLog
+
     SPEED_LIMIT = tunable(1)
     SPINNING_SPEED_LIMIT = tunable(1)
 
@@ -26,6 +30,13 @@ def __init__(self):
         self.bearing_tolerance = 0.0
         self.bearing_to_speaker = 0.0
 
+    def setup(self):
+        self.shot_time_entry = FloatLogEntry(self.data_log, "Shooter: Shot match times")
+        self.shot_range_entry = FloatLogEntry(self.data_log, "Shooter: Shot ranges")
+        self.shot_pos_entry = FloatArrayLogEntry(
+            self.data_log, "Shooter: Field translation from target"
+        )
+
     def translation_to_goal(self) -> Translation2d:
         return (
             get_goal_speaker_position().toTranslation2d()
@@ -48,11 +59,17 @@ def update_range(self) -> None:
         self.range = self.translation_to_goal().norm()
         self.shooter_component.set_range(self.range)
 
-    def try_jettison(self):
+    def try_jettison(self) -> None:
         self.engage(self.preparing_to_jettison)
 
+    def log_shot(self) -> None:
+        translation = self.translation_to_goal()
+        self.shot_time_entry.append(DriverStation.getMatchTime())
+        self.shot_range_entry.append(self.range)
+        self.shot_pos_entry.append([translation.x, translation.y])
+
     @feedback
-    def in_range(self):
+    def in_range(self) -> bool:
         return self.shooter_component.is_range_in_bounds(self.range)
 
     @state(first=True)
@@ -67,6 +84,7 @@ def aiming(self, initial_call) -> None:
                 and self.is_below_speed_limit()
                 and self.is_below_spinning_limit()
             ):
+                self.log_shot()
                 self.next_state(self.firing)
             else:
                 self.aim()