1 files changed, 192 insertions, 0 deletions
diff --git a/src/main/java/frc/robot/util/IMUGyro.java b/src/main/java/frc/robot/util/IMUGyro.java
new file mode 100644
index 0000000..1ea1983
--- /dev/null
+++ b/src/main/java/frc/robot/util/IMUGyro.java
@@ -0,0 +1,192 @@
+package frc.robot.util;
+ 
+import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.wpilibj.I2C;
+import edu.wpi.first.wpilibj.Timer;
+ 
+public class IMUGyro 
+{
+    public I2C mpu;
+ 
+    public double pitch;
+    public double roll;
+    public double yaw;
+ 
+    private double gyroOffsets[] = new double[3];
+    
+    // for deltatime
+    private double prev_time = Timer.getFPGATimestamp();
+ 
+    private final int PWR_MGMT_1 = 0x6B;
+    private final int CONFIG_ADDR = 0x1A;
+    private final int SMPLRT_DIV = 0x19;
+ 
+    private final int GYRO_ADDR = 0x43;
+    private final int GYRO_CONFIG = 0x1B;
+ 
+    private final int ACCL_ADDR = 0x3B;
+    private final int ACCL_CONFIG = 0x1C;
+ 
+    private final double SCALE_GYRO = (Math.PI/(180*16.384));
+    private final double SCALE_ACCL = (1/16.384);
+ 
+    public IMUGyro(I2C gyro) 
+    {
+        mpu = gyro;
+        reset();
+
+        mpu.write(CONFIG_ADDR, 0b00000110);
+        mpu.write(SMPLRT_DIV, 0b00010011);
+        mpu.write(GYRO_CONFIG, 0b00001000);
+        mpu.write(ACCL_CONFIG, 0b00000000);
+    }
+ 
+    public void reset()
+    {
+        mpu.write(PWR_MGMT_1, 0b10001000);
+
+        pitch = 0;
+        roll = 0;
+        yaw = 0;
+    }
+ 
+    public void update()
+    {
+        double cur_time = Timer.getFPGATimestamp();
+        double dt = cur_time - prev_time;
+        prev_time = cur_time;
+
+        // double[] gRates = gyroRates();
+
+        // roll = gRates[0] * dt;
+        // pitch = gRates[1] * dt;
+        // yaw = gRates[2] * dt;
+
+        double[] gRates = gyroRates();
+        double[] aRates = acclRates();
+
+        mahony(30, 0, gRates, aRates, dt);
+        double[] angles = fromQuaternion(q);
+
+        roll = angles[0]; 
+        pitch = angles[1]; 
+        yaw = angles[2];
+    }
+ 
+    public Rotation2d toRotation2d()
+    {
+        return new Rotation2d(yaw % 360);
+    }
+ 
+    public void calibrateGyro(int iter)
+    {
+        gyroOffsets = new double[3];
+
+        double x = 0, y = 0, z = 0;
+        for(int i = 0; i < iter; i++) 
+        {
+            double[] rates = gyroRates();
+            x += rates[0];
+            y += rates[1];
+            z += rates[2];
+        }
+ 
+        gyroOffsets[0] = x/iter;
+        gyroOffsets[1] = y/iter;
+        gyroOffsets[2] = z/iter;
+    }
+ 
+    private double[] gyroRates()
+    {
+        return scaledRates(GYRO_ADDR, SCALE_GYRO);
+    }
+ 
+    private double[] acclRates()
+    {
+        return scaledRates(ACCL_ADDR, SCALE_ACCL);
+    }
+ 
+    private double[] scaledRates(int addr, double scale)
+    {
+        byte[] buffer = new byte[6];
+        mpu.read(addr, 6, buffer);
+        double dx = decode(buffer[0], buffer[1]) * scale - gyroOffsets[0];
+        double dy = decode(buffer[2], buffer[3]) * scale - gyroOffsets[1];
+        double dz = decode(buffer[4], buffer[5]) * scale - gyroOffsets[2];
+        return new double[] {dx, dy, dz};
+    }
+ 
+    private double decode(byte first, byte second) 
+    {
+        return ((first << 8) | second);
+    }
+
+    // mahony quaternion
+    private double[] q = {1.0, 0.0, 0.0, 0.0};
+    // mahony integral term
+    private double[] i = new double[3];
+
+    private double[] fromQuaternion(double[] q)
+    {
+        return new double[] {
+            Math.atan2((q[0] * q[1] + q[2] * q[3]), 0.5 - (q[1] * q[1] + q[2] * q[2])),
+            Math.asin(2.0 * (q[0] * q[2] - q[1] * q[3])),
+            -Math.atan2((q[1] * q[2] + q[0] * q[3]), 0.5 - (q[2] * q[2] + q[3] * q[3]))
+        };
+    }
+
+    private void mahony(double Ki, double Kp, double[] gRates, double[] aRates, double dt)
+    {
+        double gx = gRates[0], gy = gRates[1], gz = gRates[2];
+        double ax = aRates[0], ay = aRates[1], az = aRates[2];
+
+        double tmp = ax * ax + ay * ay + az * az;
+        if(tmp > 0.0) 
+        {
+            tmp = Math.sqrt(tmp);
+            ax /= tmp; ay /= tmp; az /= tmp;
+
+            // estimated direction of gravity in the body frame
+            double vx = q[1] * q[3] - q[0] * q[2];
+            double vy = q[0] * q[1] + q[2] * q[3];
+            double vz = q[0] * q[0] - 0.5f + q[3] * q[3];
+
+            // error is cross product
+            double ex = (ay * vz - az * vy);
+            double ey = (az * vx - ax * vz);
+            double ez = (ax * vy - ay * vx);
+
+            // integral feedback
+            if(Ki > 0.0) {
+                i[0] += Ki * ex * dt;
+                i[1] += Ki * ey * dt;
+                i[2] += Ki * ez * dt;
+                gx += i[0];  
+                gy += i[1];
+                gz += i[2];
+            }
+
+            gx += Kp * ex;
+            gy += Kp * ey;
+            gz += Kp * ez;            
+        }
+
+        // integrate quaternion
+        dt /= 2; // ??
+        gx *= dt; gy *= dt; gz *= dt;
+        
+        double qa = q[0], qb = q[1], qc = q[2], qd = q[3];
+        q[0] += (-qb * gx - qc * gy - qd * gz);
+        q[1] += (qa * gx + qc * gz - qd * gy);
+        q[2] += (qa * gy - qb * gz + qd * gx);
+        q[3] += (qa * gz + qb * gy - qc * gx);
+
+        // renormalize
+        double qnorm_factor = Math.sqrt(q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]);
+        q[0] /= qnorm_factor;
+        q[1] /= qnorm_factor;
+        q[2] /= qnorm_factor;
+        q[3] /= qnorm_factor;
+    }
+}
+