001/************************* PROJECT RON *************************/
002/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved. */
003/* Use of this source code is governed by an MIT-style license */
004/* that can be found in the repository LICENSE file.           */
005/***************************************************************/
006package com.stuypulse.robot.util.simulation;
007
008import static edu.wpi.first.units.Units.*;
009
010import com.stuypulse.robot.constants.Settings;
011import com.stuypulse.robot.subsystems.intake.Intake.IntakeState;
012import edu.wpi.first.units.measure.*;
013import edu.wpi.first.wpilibj.smartdashboard.Mechanism2d;
014import edu.wpi.first.wpilibj.smartdashboard.MechanismLigament2d;
015import edu.wpi.first.wpilibj.smartdashboard.MechanismObject2d;
016import edu.wpi.first.wpilibj.smartdashboard.MechanismRoot2d;
017import edu.wpi.first.wpilibj.smartdashboard.SmartDashboard;
018import edu.wpi.first.wpilibj.util.Color;
019import edu.wpi.first.wpilibj.util.Color8Bit;
020import java.util.Arrays;
021
022/**
023 * <h2>Combined robot visualizer for all of the subsystems in one canvas.</h2>
024 * 
025 * <p>
026 * This class employs the singleton pattern to ensure that only one visualizer
027 * exists,
028 * since only one robot instance should exist.
029 * </p>
030 */
031public class RobotVisualizer {
032    public static RobotVisualizer instance;
033
034    static {
035        instance = new RobotVisualizer();
036    }
037
038    public static RobotVisualizer getInstance() {
039        return instance;
040    }
041
042    private static final double CANVAS_WIDTH = 67, CANVAS_HEIGHT = 67;
043
044    /**
045     * Number of spokes for wheels
046     * 
047     * @see #createSpokes(MechanismObject2d, String, double, double, Color8Bit)
048     */
049    private static final int NUM_SPOKES = 5;
050    /**
051     * Color of spokes for wheels
052     * 
053     * @see #createSpokes(MechanismObject2d, String, double, double, Color8Bit)
054     */
055    private static Color8Bit SPOKE_COLOR;
056
057    /**
058     * Sendable main canvas object for every drawn mechanism
059     */
060    private final Mechanism2d canvas;
061
062    /**
063     * Root of the {@link #bumper red} bumpers to show the silhouette of the robot
064     */
065    private final MechanismRoot2d bumperRoot;
066    /**
067     * Drawn representation of red bumpers to show the silhouette of the robot
068     * 
069     * @see #bumperRoot
070     */
071    private final MechanismLigament2d bumper;
072
073    /**
074     * Root for the drawn representation of the feeder
075     * 
076     * The actual feeder is a belt, but spokes are easier.
077     * Additionally, the {@link #feederSpokes spokes} are placed where the handoff should be...
078     */
079    private final MechanismRoot2d feederRoot;
080
081    /**
082     * Spokes of the feeder
083     * 
084     * @see #feederRoot
085     */
086    private final MechanismLigament2d[] feederSpokes;
087
088    /**
089     * Root for the drawn representation of the intake
090     * 
091     * @see #intakePivot
092     * @see #intakeSpokes
093     */
094    private final MechanismRoot2d intakeRoot;
095    /**
096     * Pivoting arm of the intake.
097     * Parent of the {@link #intakeSpokes intake rollers}
098     */
099    private final MechanismLigament2d intakePivot;
100    /**
101     * Intake rollers, made into flattened array of every set of rollers for easy iteration
102     * 
103     * @see #RobotVisualizer()
104     */
105    private final MechanismLigament2d[] intakeSpokes;
106
107    /**
108     * Root for the drawn representation of the shooter
109     */
110    private final MechanismRoot2d shooterRoot;
111    /**
112     * Shooter rollers
113     * 
114     * @see #shooterRoot
115     */
116    private final MechanismLigament2d[] shooterSpokes;
117    /**
118     * <h4>Constructs the visualizer instance</h4>
119     * <p>Uses specific internal logic, viewing the source code is recommended to better understand how everything is setup.</p>
120     */
121    private RobotVisualizer() {
122        SPOKE_COLOR = new Color8Bit(Color.kWhite);
123        canvas = new Mechanism2d(CANVAS_WIDTH, CANVAS_HEIGHT);
124        // Silhouette
125        bumperRoot = canvas.getRoot("Bumper Root", 10, 5);
126        bumper = new MechanismLigament2d("Bumper", 50, 0, 90, new Color8Bit(Color.kRed));
127        bumperRoot.append(bumper);
128        // Feeder
129        feederRoot = canvas.getRoot("Feeder Root", 45, 15);
130        feederSpokes = createSpokes(feederRoot, "Feeder Spoke", 6.7, 2);
131        // Shooter
132        shooterRoot = canvas.getRoot("Shooter Root", 55, 45);
133        shooterSpokes = createSpokes(shooterRoot, "Shooter Spoke", 6.7, 2);
134        // Intake
135        intakeRoot = canvas.getRoot("Intake Root", 15, 10);
136        intakePivot = new MechanismLigament2d(
137                "Intake Arm",
138                9,
139                IntakeState.IDLE.getTargetAngle().in(Degrees),
140                4,
141                new Color8Bit(Color.kGray));
142        intakeRoot.append(intakePivot);
143        MechanismLigament2d intakeTopRollers = new MechanismLigament2d("Intake Top Rollers", 4, -20, 4,
144                new Color8Bit(Color.kGray));
145        MechanismLigament2d intakeMiddleRollers = new MechanismLigament2d("Intake Middle Rollers", 4, 150, 4,
146                new Color8Bit(Color.kGray));
147        MechanismLigament2d intakeBottomRollers = new MechanismLigament2d("Intake Bottom Rollers", 8, 90, 4,
148                new Color8Bit(Color.kGray));
149        intakeBottomRollers.append(intakeMiddleRollers);
150        intakeMiddleRollers.append(intakeTopRollers);
151        {
152            /* 
153             * Draw the other lines in the intake structure.
154             * We do this Within a separate scope using {} to ensure the _a and _b variables are temporary.
155             */
156            var _a = new MechanismLigament2d("_a", 8, 50, 4, new Color8Bit(Color.kGray));
157            _a.append(intakeBottomRollers);
158            var _b = new MechanismLigament2d("_b", 13.75, 95, 4, new Color8Bit(Color.kGray));
159            intakeBottomRollers.append(_b);
160            intakePivot.append(_a);
161        }
162        // Flatten all the sets of intake rollers into one big array
163        intakeSpokes = Arrays.stream(
164                new MechanismLigament2d[][] {
165                        createSpokes(intakeTopRollers, "Intake Top Spoke", .67, 2),
166                        createSpokes(intakeMiddleRollers, "Intake Middle Spoke", .67, 2),
167                        createSpokes(intakeBottomRollers, "Intake Bottom Spoke", .67, 2)
168                })
169                .flatMap(Arrays::stream)
170                .toArray(MechanismLigament2d[]::new);
171        // Publish the canvas to Smartdashboard
172        SmartDashboard.putData("Visualizers/Robot", canvas);
173    }
174
175    /**
176     * 
177     * <h4>Helper to create wheels</h4>
178     * 
179     * <p>Creates a simulacrum of a wheel or side view of a roller
180     * by making a bunch of equally-spaced, outwards-radiating lines that like are spokes of a wheel</p>
181     * @param target the object to append the spokes to
182     * @param name name of each spoke, with name mangling applied internally
183     * @param length length of the spokes
184     * @param width width of the spokes
185     * @param color color of the spokes
186     * @return An array of the generated spokes
187     */
188    private MechanismLigament2d[] createSpokes(
189            MechanismObject2d target,
190            String name,
191            double length,
192            double width,
193            Color8Bit color) {
194        MechanismLigament2d[] spokes = new MechanismLigament2d[NUM_SPOKES];
195        double spacing = 360 / NUM_SPOKES;
196        for (int i = 0; i < NUM_SPOKES; i++) {
197            MechanismLigament2d spoke = new MechanismLigament2d(name.trim() + " " + i, length, spacing * i, width, color);
198            target.append(spoke);
199            spokes[i] = spoke;
200        }
201        return spokes;
202    }
203
204    /**
205     * 
206     * <h4>Helper to create wheels</h4>
207     * 
208     * <p>Creates a simulacrum of a wheel or side view of a roller
209     * by making a bunch of equally-spaced, outwards-radiating lines that like are spokes of a wheel.
210     * Defaults to the statically specified color {@link #SPOKE_COLOR}</p>
211     * @param target the object to append the spokes to
212     * @param name name of each spoke, with name mangling applied internally
213     * @param length length of the spokes
214     * @param width width of the spokes
215     * @return An array of the generated spokes
216     */
217    private MechanismLigament2d[] createSpokes(
218            MechanismObject2d target,
219            String name,
220            double length,
221            double width) {
222        return createSpokes(target, name, length, width, SPOKE_COLOR);
223    }
224
225    /**
226     * <h4>Update method for the feeder</h4>
227     * 
228     * <p>To be called by subsystems using motor measurements</p>
229     * @param angularVelocity measured angular velocity of the feeder rollers
230     */
231    public void updateFeeder(AngularVelocity angularVelocity) {
232        double rot = angularVelocity.in(RPM) * 6 * Settings.DT.in(Seconds);
233        for (MechanismLigament2d spoke : feederSpokes)
234            spoke.setAngle(spoke.getAngle() + rot);
235    }
236
237    /**
238     * <h4>Update method for the shooter</h4>
239     * 
240     * <p>To be called by subsystems using motor measurements</p>
241     * @param angularVelocity measured angular velocity of the shooter rollers
242     */
243    public void updateShooter(AngularVelocity angularVelocity) {
244        double rot = angularVelocity.in(RPM) * 6 * Settings.DT.in(Seconds);
245        for (MechanismLigament2d spoke : shooterSpokes)
246            spoke.setAngle(spoke.getAngle() + rot);
247    }
248
249    /**
250     * <h4>Update method for the intake</h4>
251     * 
252     * <p>To be called by subsystems using motor measurements</p>
253     * @param pivotAngle measured position of the intake pivot
254     * @param angularVelocity measured angular velocity of the intake rollers
255     */
256    public void updateIntake(Angle pivotAngle, AngularVelocity angularVelocity) {
257        intakePivot.setAngle(pivotAngle.in(Degrees) + 102); // counteract the weird zeroing
258        double rot = angularVelocity.in(RPM) * Settings.DT.in(Seconds);
259        for (MechanismLigament2d spoke : intakeSpokes)
260            spoke.setAngle(spoke.getAngle() + rot);
261    }
262
263    /**
264     * <p>Publish the canvas to {@link SmartDashboard}</p>
265     */
266    public void update() {
267        SmartDashboard.putData("Visualizers/Robot", canvas);
268    }
269}