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[] = {
        0.312449302886265,
        0.04883417185589417,
        -0.12945104079301352
    };
    
    // 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) * 2 /*??*/);
    private final double SCALE_ACCL = (1/16.384);
 
    public IMUGyro(I2C gyro) 
    {
        mpu = gyro;
        reset();

        // mpu.write(CONFIG_ADDR, 0b000000);
        mpu.write(SMPLRT_DIV, 0b00000010);
        // mpu.write(GYRO_CONFIG, 0b00011000);
        // mpu.write(ACCL_CONFIG, 0b00000000);
    }
 
    public void reset()
    {
        mpu.write(PWR_MGMT_1, 0b000001000);

        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;
    }
 
    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;
    }
}