Provide plotter for image correction debugging

Change-Id: Id2545c7520be06b45b59cf0214a7c9b61bc0cb10
Signed-off-by: James Kuszmaul <jabukuszmaul@gmail.com>
diff --git a/frc971/analysis/BUILD b/frc971/analysis/BUILD
index 36c450a..1ebf1ba 100644
--- a/frc971/analysis/BUILD
+++ b/frc971/analysis/BUILD
@@ -58,6 +58,7 @@
         "//y2022/control_loops/superstructure:turret_plotter",
         "//y2022/localizer:localizer_plotter",
         "//y2022/vision:vision_plotter",
+        "//y2023/localizer:corrections_plotter",
     ],
 )
 
diff --git a/frc971/analysis/plot_index.ts b/frc971/analysis/plot_index.ts
index ab00b99..1b62e42 100644
--- a/frc971/analysis/plot_index.ts
+++ b/frc971/analysis/plot_index.ts
@@ -54,6 +54,8 @@
     '../../y2022/localizer/localizer_plotter'
 import {plotVision as plot2022Vision} from
     '../../y2022/vision/vision_plotter'
+import {plotVision as plot2023Corrections} from
+    '../../y2023/localizer/corrections_plotter'
 import {plotDemo} from '../../aos/network/www/demo_plot';
 
 const rootDiv = document.createElement('div');
@@ -112,6 +114,7 @@
   ['Spline Debug', new PlotState(plotDiv, plotSpline)],
   ['Down Estimator', new PlotState(plotDiv, plotDownEstimator)],
   ['Robot State', new PlotState(plotDiv, plotRobotState)],
+  ['2023 Vision', new PlotState(plotDiv, plot2023Corrections)],
   ['2020 Finisher', new PlotState(plotDiv, plot2020Finisher)],
   ['2020 Accelerator', new PlotState(plotDiv, plot2020Accelerator)],
   ['2020 Hood', new PlotState(plotDiv, plot2020Hood)],
diff --git a/y2023/localizer/BUILD b/y2023/localizer/BUILD
index d97d4a3..70a4d57 100644
--- a/y2023/localizer/BUILD
+++ b/y2023/localizer/BUILD
@@ -1,5 +1,6 @@
 load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
 load("@com_github_google_flatbuffers//:typescript.bzl", "flatbuffer_ts_library")
+load("//tools/build_rules:js.bzl", "ts_project")
 
 flatbuffer_cc_library(
     name = "status_fbs",
@@ -176,3 +177,16 @@
         "//y2023/control_loops/drivetrain:drivetrain_base",
     ],
 )
+
+ts_project(
+    name = "corrections_plotter",
+    srcs = ["corrections_plotter.ts"],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":visualization_ts_fbs",
+        "//aos/network/www:aos_plotter",
+        "//aos/network/www:colors",
+        "//aos/network/www:proxy",
+    ],
+)
diff --git a/y2023/localizer/corrections_plotter.ts b/y2023/localizer/corrections_plotter.ts
new file mode 100644
index 0000000..a114f3f
--- /dev/null
+++ b/y2023/localizer/corrections_plotter.ts
@@ -0,0 +1,140 @@
+import {ByteBuffer} from 'flatbuffers';
+import {AosPlotter} from '../../aos/network/www/aos_plotter';
+import {MessageHandler, TimestampedMessage} from '../../aos/network/www/aos_plotter';
+import {BLUE, BROWN, CYAN, GREEN, PINK, RED, WHITE} from '../../aos/network/www/colors';
+import {Connection} from '../../aos/network/www/proxy';
+import {Table} from '../../aos/network/www/reflection';
+import {Schema} from 'flatbuffers_reflection/reflection_generated';
+import {Visualization, TargetEstimateDebug} from './visualization_generated';
+
+
+const TIME = AosPlotter.TIME;
+// magenta, yellow, cyan, black
+const PI_COLORS = [[255, 0, 255], [255, 255, 0], [0, 255, 255], [0, 0, 0]];
+
+class VisionMessageHandler extends MessageHandler {
+  constructor(private readonly schema: Schema) {
+    super(schema);
+  }
+
+  private readScalar(table: Table, fieldName: string): number|BigInt|null {
+    return this.parser.readScalar(table, fieldName);
+  }
+
+  addMessage(data: Uint8Array, time: number): void {
+    const message = Visualization.getRootAsVisualization(new ByteBuffer(data));
+    for (let ii = 0; ii < message.targetsLength(); ++ii) {
+      const target = message.targets(ii);
+      const time = Number(target.imageMonotonicTimestampNs()) * 1e-9;
+      if (time == 0) {
+        console.log('Dropping message without populated time?');
+        continue;
+      }
+      const table = Table.getNamedTable(
+          target.bb, this.schema, 'y2023.localizer.TargetEstimateDebug', target.bb_pos);
+      this.messages.push(new TimestampedMessage(table, time));
+    }
+  }
+}
+
+export function plotVision(conn: Connection, element: Element): void {
+  const aosPlotter = new AosPlotter(conn);
+
+  const targets = [];
+  for (const pi of ['pi1', 'pi2', 'pi3', 'pi4']) {
+    targets.push(aosPlotter.addRawMessageSource(
+        '/' + pi + '/camera', 'y2023.localizer.Visualization',
+        new VisionMessageHandler(
+            conn.getSchema('y2023.localizer.Visualization'))));
+  }
+  const localizerStatus = aosPlotter.addMessageSource(
+      '/localizer', 'y2023.localizer.Status');
+  const localizerOutput = aosPlotter.addMessageSource(
+      '/localizer', 'frc971.controls.LocalizerOutput');
+
+  const rejectionPlot = aosPlotter.addPlot(element);
+  rejectionPlot.plot.getAxisLabels().setTitle('Rejection Reasons');
+  rejectionPlot.plot.getAxisLabels().setXLabel(TIME);
+  rejectionPlot.plot.getAxisLabels().setYLabel('[bool, enum]');
+
+  rejectionPlot
+      .addMessageLine(localizerStatus, ['statistics[]', 'total_accepted'])
+      .setDrawLine(false)
+      .setColor(BLUE);
+  rejectionPlot
+      .addMessageLine(localizerStatus, ['statistics[]', 'total_candidates'])
+      .setDrawLine(false)
+      .setColor(RED);
+  for (let ii = 0; ii < targets.length; ++ii) {
+    rejectionPlot.addMessageLine(targets[ii], ['rejection_reason'])
+        .setDrawLine(false)
+        .setColor(PI_COLORS[ii])
+        .setLabel('pi' + (ii + 1));
+  }
+
+  const xPlot = aosPlotter.addPlot(element);
+  xPlot.plot.getAxisLabels().setTitle('X Position');
+  xPlot.plot.getAxisLabels().setXLabel(TIME);
+  xPlot.plot.getAxisLabels().setYLabel('[m]');
+
+  for (let ii = 0; ii < targets.length; ++ii) {
+    xPlot.addMessageLine(targets[ii], ['implied_robot_x'])
+        .setDrawLine(false)
+        .setColor(PI_COLORS[ii])
+        .setLabel('pi' + (ii + 1));
+  }
+  xPlot.addMessageLine(localizerOutput, ['x'])
+      .setDrawLine(false)
+      .setColor(BLUE);
+
+  const correctionXPlot = aosPlotter.addPlot(element);
+  correctionXPlot.plot.getAxisLabels().setTitle('X Corrections');
+  correctionXPlot.plot.getAxisLabels().setXLabel(TIME);
+  correctionXPlot.plot.getAxisLabels().setYLabel('[m]');
+
+  for (let ii = 0; ii < targets.length; ++ii) {
+    correctionXPlot.addMessageLine(targets[ii], ['correction_x'])
+        .setDrawLine(false)
+        .setColor(PI_COLORS[ii])
+        .setLabel('pi' + (ii + 1));
+  }
+
+  const yPlot = aosPlotter.addPlot(element);
+  yPlot.plot.getAxisLabels().setTitle('Y Position');
+  yPlot.plot.getAxisLabels().setXLabel(TIME);
+  yPlot.plot.getAxisLabels().setYLabel('[m]');
+
+  for (let ii = 0; ii < targets.length; ++ii) {
+    yPlot.addMessageLine(targets[ii], ['implied_robot_y'])
+        .setDrawLine(false)
+        .setColor(PI_COLORS[ii])
+        .setLabel('pi' + (ii + 1));
+  }
+  yPlot.addMessageLine(localizerOutput, ['y'])
+      .setDrawLine(false)
+      .setColor(BLUE);
+
+  const correctionYPlot = aosPlotter.addPlot(element);
+  correctionYPlot.plot.getAxisLabels().setTitle('Y Corrections');
+  correctionYPlot.plot.getAxisLabels().setXLabel(TIME);
+  correctionYPlot.plot.getAxisLabels().setYLabel('[m]');
+
+  for (let ii = 0; ii < targets.length; ++ii) {
+    correctionYPlot.addMessageLine(targets[ii], ['correction_y'])
+        .setDrawLine(false)
+        .setColor(PI_COLORS[ii])
+        .setLabel('pi' + (ii + 1));
+  }
+
+  const aprilTagPlot = aosPlotter.addPlot(element);
+  aprilTagPlot.plot.getAxisLabels().setTitle('April Tag IDs');
+  aprilTagPlot.plot.getAxisLabels().setXLabel(TIME);
+  aprilTagPlot.plot.getAxisLabels().setYLabel('[id]');
+
+  for (let ii = 0; ii < targets.length; ++ii) {
+    aprilTagPlot.addMessageLine(targets[ii], ['april_tag'])
+        .setDrawLine(false)
+        .setColor(PI_COLORS[ii])
+        .setLabel('pi' + (ii + 1));
+  }
+}