Swerve odometry (maybe)

src/main/java/frc/robot/RobotState.java

@@ -1,7 +1,13 @@
 package frc.robot;
 
+import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.math.kinematics.SwerveDriveOdometry;
 import edu.wpi.first.math.kinematics.SwerveModulePosition;
+import edu.wpi.first.wpilibj.smartdashboard.Field2d;
+import edu.wpi.first.wpilibj.smartdashboard.SmartDashboard;
+import edu.wpi.first.wpilibj.smartdashboard.Field2d;
+import frc.robot.Constants.DriveConstants;
 
 public class RobotState {
 
@@ -13,12 +19,32 @@ public static RobotState getInstance() {
         return instance;
     }
 
-    public void setFieldToVehicle(){
+    private SwerveDriveOdometry odometry;
+    private Field2d field = new Field2d();
 
+    public void initializeOdometry(Rotation2d rotation, SwerveModulePosition[] modulePositions) {
+        odometry = new SwerveDriveOdometry(DriveConstants.kinematics, rotation, modulePositions);
     }
 
-    public void recordDriveObservations() {
+    public void initializeOdometry(Rotation2d rotation, SwerveModulePosition[] modulePositions, Pose2d pose) {
+        odometry = new SwerveDriveOdometry(DriveConstants.kinematics, rotation, modulePositions, pose);
 
     }
+
+    public void updateOdometry(Rotation2d rotation, SwerveModulePosition[] modulePositions) {
+        odometry.update(rotation, modulePositions);
+    }
+
+    public void resetOdometry(Rotation2d gyroAngle, SwerveModulePosition[] modulePositions, Pose2d pose) {
+        odometry.resetPosition(gyroAngle, modulePositions, pose);
+    }
+
+    public Pose2d getOdometryFieldToRobot() {
+        field.setRobotPose(odometry.getPoseMeters());
+        SmartDashboard.putData("Odometry Pose", field);
+
+        return odometry.getPoseMeters();
+    }
+
 
 }

src/main/java/frc/robot/subsystems/drive/DriveSubsystem.java

@@ -46,7 +46,7 @@ public DriveSubsystem() {
 
         //TODO: add swerve drive odometry
         setGoalChassisSpeeds(new ChassisSpeeds(0, 0, 0));
-        RobotState.getInstance();//.initializePoseEstimator(getRotation(), modulePositions);
+        RobotState.getInstance().initializeOdometry(getRotation(), getSwerveModulePositions());
     }
 
 
@@ -62,6 +62,8 @@ public void periodic(){
 
         }
 
+
+
         //estimation of position
         for (int i = 0; i < 4; i++) {
             driveModules[i].getModulePosition().distanceMeters = driveModules[i].getDrivePosition() * DriveConstants.wheelRadiusM;
@@ -72,8 +74,7 @@ public void periodic(){
             SmartDashboard.putNumber("Drive/modules/"+i+"/rotation stator current", driveModules[i].getRotationStatorCurrent());
         }
 
-        //TODO: add swerve odometry
-        RobotState.getInstance().recordDriveObservations(/*getRotation(), modulePositions*/);
+        RobotState.getInstance().updateOdometry(getRotation(), getSwerveModulePositions());
 
         SmartDashboard.putNumber("Drive/velocity magnitude", getChassisSpeeds().vxMetersPerSecond);
     }
@@ -126,7 +127,7 @@ public void setBabyMode(boolean baby) {
 
     //TODO: odometry
     public void setFieldToVehicle(Pose2d fieldToVehicle) {
-        RobotState.getInstance().setFieldToVehicle(/*getRotation(), modulePositions, fieldToVehicle*/);
+        RobotState.getInstance().resetOdometry(getRotation(), getSwerveModulePositions(), fieldToVehicle);
     }
 
     public void setVolts(double v) {
@@ -144,6 +145,16 @@ public SwerveModuleState[] getModuleStates() {
         return states;
     }
 
+    private SwerveModulePosition[] getSwerveModulePositions() {
+
+        SwerveModulePosition[] swerveModulePositions = new SwerveModulePosition[4];
+        for (int i = 0; i < 4; i++) {
+            swerveModulePositions[i] = driveModules[i].getModulePosition();
+        }
+        return swerveModulePositions;
+
+    }
+
     public ChassisSpeeds getChassisSpeeds() {
         return DriveConstants.kinematics.toChassisSpeeds(getModuleStates());
     }