001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.wpilibj;
006
007import edu.wpi.first.hal.IMUJNI;
008import edu.wpi.first.math.geometry.Quaternion;
009import edu.wpi.first.math.geometry.Rotation2d;
010import edu.wpi.first.math.geometry.Rotation3d;
011
012/** SystemCore onboard IMU. */
013public class OnboardIMU {
014  /** A mount orientation of SystemCore. */
015  public enum MountOrientation {
016    /** Flat (mounted parallel to the ground). */
017    kFlat,
018    /** Landscape (vertically mounted with long edge of SystemCore parallel to the ground). */
019    kLandscape,
020    /** Portrait (vertically mounted with the short edge of SystemCore parallel to the ground). */
021    kPortrait
022  }
023
024  /**
025   * Constructs a handle to the SystemCore onboard IMU.
026   *
027   * @param mountOrientation the mount orientation of SystemCore to determine yaw.
028   */
029  public OnboardIMU(MountOrientation mountOrientation) {
030    m_mountOrientation = mountOrientation;
031  }
032
033  /**
034   * Get the yaw value in radians.
035   *
036   * @return yaw value in radians
037   */
038  public double getYawRadians() {
039    return getYawNoOffset() - m_yawOffset;
040  }
041
042  /**
043   * Reset the current yaw value to 0. Future reads of the yaw value will be relative to the current
044   * orientation.
045   */
046  public void resetYaw() {
047    m_yawOffset = getYawNoOffset();
048  }
049
050  /**
051   * Get the yaw as a Rotation2d.
052   *
053   * @return yaw
054   */
055  public Rotation2d getRotation2d() {
056    return new Rotation2d(getYawRadians());
057  }
058
059  /**
060   * Get the 3D orientation as a Rotation3d.
061   *
062   * @return 3D orientation
063   */
064  public Rotation3d getRotation3d() {
065    return new Rotation3d(getQuaternion());
066  }
067
068  /**
069   * Get the 3D orientation as a Quaternion.
070   *
071   * @return 3D orientation
072   */
073  public Quaternion getQuaternion() {
074    double[] quatRaw = new double[4];
075    IMUJNI.getIMUQuaternion(quatRaw);
076    return new Quaternion(quatRaw[0], quatRaw[1], quatRaw[2], quatRaw[3]);
077  }
078
079  /**
080   * Get the angle about the X axis of the IMU in radians.
081   *
082   * @return angle about the X axis in radians
083   */
084  public double getAngleX() {
085    return getRawEulerAngles()[0];
086  }
087
088  /**
089   * Get the angle about the Y axis of the IMU in radians.
090   *
091   * @return angle about the Y axis in radians
092   */
093  public double getAngleY() {
094    return getRawEulerAngles()[1];
095  }
096
097  /**
098   * Get the angle about the Z axis of the IMU in radians.
099   *
100   * @return angle about the Z axis in radians
101   */
102  public double getAngleZ() {
103    return getRawEulerAngles()[2];
104  }
105
106  /**
107   * Get the angular rate about the X axis of the IMU in radians per second.
108   *
109   * @return angular rate about the X axis in radians per second
110   */
111  public double getGyroRateX() {
112    return getRawGyroRates()[0];
113  }
114
115  /**
116   * Get the angular rate about the Y axis of the IMU in radians per second.
117   *
118   * @return angular rate about the Y axis in radians per second
119   */
120  public double getGyroRateY() {
121    return getRawGyroRates()[1];
122  }
123
124  /**
125   * Get the angular rate about the Z axis of the IMU in radians per second.
126   *
127   * @return angular rate about the Z axis in radians per second
128   */
129  public double getGyroRateZ() {
130    return getRawGyroRates()[2];
131  }
132
133  /**
134   * Get the acceleration along the X axis of the IMU in meters per second squared.
135   *
136   * @return acceleration along the X axis in meters per second squared
137   */
138  public double getAccelX() {
139    return getRawAccels()[0];
140  }
141
142  /**
143   * Get the acceleration along the X axis of the IMU in meters per second squared.
144   *
145   * @return acceleration along the X axis in meters per second squared
146   */
147  public double getAccelY() {
148    return getRawAccels()[1];
149  }
150
151  /**
152   * Get the acceleration along the X axis of the IMU in meters per second squared.
153   *
154   * @return acceleration along the X axis in meters per second squared
155   */
156  public double getAccelZ() {
157    return getRawAccels()[2];
158  }
159
160  private double[] getRawEulerAngles() {
161    double[] anglesRaw = new double[3];
162    switch (m_mountOrientation) {
163      case kFlat -> IMUJNI.getIMUEulerAnglesFlat(anglesRaw);
164      case kLandscape -> IMUJNI.getIMUEulerAnglesLandscape(anglesRaw);
165      case kPortrait -> IMUJNI.getIMUEulerAnglesPortrait(anglesRaw);
166      default -> {
167        // NOP
168      }
169    }
170    return anglesRaw;
171  }
172
173  private double[] getRawGyroRates() {
174    double[] ratesRaw = new double[3];
175    IMUJNI.getIMUGyroRates(ratesRaw);
176    return ratesRaw;
177  }
178
179  private double[] getRawAccels() {
180    double[] accelsRaw = new double[3];
181    IMUJNI.getIMUAcceleration(accelsRaw);
182    return accelsRaw;
183  }
184
185  private double getYawNoOffset() {
186    return switch (m_mountOrientation) {
187      case kFlat -> IMUJNI.getIMUYawFlat();
188      case kLandscape -> IMUJNI.getIMUYawLandscape();
189      case kPortrait -> IMUJNI.getIMUYawPortrait();
190      default -> 0;
191    };
192  }
193
194  private final MountOrientation m_mountOrientation;
195  private double m_yawOffset;
196}