#include <doctest/doctest.h>

#include "fgc/ImuMount.h"

#include <cmath>

using namespace fgc;

namespace {

// A sample with distinguishable values in every slot, so a transform that
// touches the wrong field is visible.
ImuSample makeSample() {
    ImuSample s;
    s.valid          = true;
    s.roll_deg       = 0.f;
    s.pitch_deg      = 0.f;
    s.yaw_deg        = 0.f;
    s.acc[0]         = 1.f;
    s.acc[1]         = 2.f;
    s.acc[2]         = 3.f;
    s.gyr[0]         = 0.1f;
    s.gyr[1]         = 0.2f;
    s.gyr[2]         = 0.3f;
    s.mag[0]         = 0.4f;
    s.mag[1]         = 0.5f;
    s.mag[2]         = 0.6f;
    s.temp_c         = 24.5f;
    s.sample_counter = 4321;
    return s;
}

double norm(const float v[3]) {
    return std::sqrt(double(v[0]) * v[0] + double(v[1]) * v[1] + double(v[2]) * v[2]);
}

// Heading difference on the circle, so 359.9 and 0.1 are 0.2 apart, not 359.8.
double headingDiff(double a, double b) {
    double d = std::fmod(std::abs(a - b), 360.0);
    return d > 180.0 ? 360.0 - d : d;
}

}  // namespace

TEST_CASE("an identity mount changes nothing at all") {
    const ImuSample in = makeSample();
    const ImuSample out = applyMount(ImuMount{}, in);

    // Exact, not approximate: a rig with no mount configured must be bit-for-bit
    // what it was before this feature existed.
    CHECK(out.roll_deg == in.roll_deg);
    CHECK(out.pitch_deg == in.pitch_deg);
    CHECK(out.yaw_deg == in.yaw_deg);
    for (int i = 0; i < 3; ++i) {
        CHECK(out.acc[i] == in.acc[i]);
        CHECK(out.gyr[i] == in.gyr[i]);
        CHECK(out.mag[i] == in.mag[i]);
    }
    CHECK(ImuMount{}.identity());
    CHECK_FALSE(ImuMount{0.0, 0.0, 90.0}.identity());
}

TEST_CASE("the mount rotation is a rotation: orthonormal, determinant +1") {
    for (const ImuMount& m : {ImuMount{0, 0, 0}, ImuMount{180, 0, 0}, ImuMount{0, 0, -90},
                              ImuMount{30, -45, 120}, ImuMount{-90, 17, 240}}) {
        const Mat3 R = mountRotationZXY(m);

        // R^T R = I
        for (int i = 0; i < 3; ++i)
            for (int j = 0; j < 3; ++j) {
                const double dot = R[0][i] * R[0][j] + R[1][i] * R[1][j] + R[2][i] * R[2][j];
                CHECK(dot == doctest::Approx(i == j ? 1.0 : 0.0).epsilon(1e-12));
            }

        const double det = R[0][0] * (R[1][1] * R[2][2] - R[1][2] * R[2][1]) -
                           R[0][1] * (R[1][0] * R[2][2] - R[1][2] * R[2][0]) +
                           R[0][2] * (R[1][0] * R[2][1] - R[1][1] * R[2][0]);
        CHECK(det == doctest::Approx(1.0));   // +1, not -1: a rotation, never a reflection
    }
}

TEST_CASE("z = 90: the IMU's forward axis reads as the gimbal's left") {
    // The angles are what you apply TO THE IMU to bring it onto the gimbal. If it
    // takes a +90 turn about down to align it, the IMU is currently pointing 90
    // deg to the LEFT of the gimbal's nose - so what the sensor calls forward is
    // what the gimbal calls left (-Y).
    ImuSample s = makeSample();
    s.acc[0] = 1.f; s.acc[1] = 0.f; s.acc[2] = 0.f;      // sensor-frame forward

    const ImuSample out = applyMount(ImuMount{0, 0, 90}, s);
    CHECK(out.acc[0] == doctest::Approx(0.0).scale(1));
    CHECK(out.acc[1] == doctest::Approx(-1.0));
    CHECK(out.acc[2] == doctest::Approx(0.0).scale(1));
}

TEST_CASE("x = 180: an upside-down IMU has its gravity put back the right way up") {
    // Mounted inverted, the sensor's own Z points at the sky, so a stationary unit
    // reports gravity along -Z. In the gimbal frame it must come back as +Z (down).
    ImuSample s = makeSample();
    s.acc[0] = 0.f; s.acc[1] = 0.f; s.acc[2] = -9.81f;

    const ImuSample out = applyMount(ImuMount{180, 0, 0}, s);
    CHECK(out.acc[0] == doctest::Approx(0.0).scale(1));
    CHECK(out.acc[1] == doctest::Approx(0.0).scale(1));
    CHECK(out.acc[2] == doctest::Approx(9.81));
}

TEST_CASE("a pure yaw mount shifts the reported heading by exactly that much") {
    for (double sensor_yaw : {0.0, 45.0, 170.0, 300.0, 359.0}) {
        ImuSample s = makeSample();
        s.yaw_deg = static_cast<float>(sensor_yaw);

        const ImuSample out = applyMount(ImuMount{0, 0, 90}, s);
        CHECK(headingDiff(out.yaw_deg, sensor_yaw + 90.0) == doctest::Approx(0.0).scale(1));
        CHECK(out.roll_deg == doctest::Approx(0.0).scale(1));
        CHECK(out.pitch_deg == doctest::Approx(0.0).scale(1));
        // Stays in the device's own range: the mount rotates, it does not re-range.
        CHECK(out.yaw_deg >= -180.f);
        CHECK(out.yaw_deg <= 180.f);
    }
}

TEST_CASE("the composition order is Z-X-Y, not any other permutation") {
    // z=90 and x=180 together are what separate the orders: intrinsic Z-X-Y rolls
    // about the X axis the yaw has ALREADY carried around, which lands sensor-
    // forward somewhere a Y-X-Z composition never would. Without this, the order
    // could be silently swapped and every single-axis test above would still pass.
    ImuSample s = makeSample();
    s.acc[0] = 1.f; s.acc[1] = 0.f; s.acc[2] = 0.f;

    // R = Rz(90)Rx(180) = [[0,1,0],[1,0,0],[0,0,-1]], so R^T sends sensor-forward
    // to gimbal +Y (right).
    const ImuSample zxy = applyMount(ImuMount{180, 0, 90}, s);
    CHECK(zxy.acc[0] == doctest::Approx(0.0).scale(1));
    CHECK(zxy.acc[1] == doctest::Approx(1.0));
    CHECK(zxy.acc[2] == doctest::Approx(0.0).scale(1));

    // Feeding the same two angles to the other axes gives a different answer -
    // R = Rz(180)Rx(90) sends sensor-forward to gimbal -X instead - which is what
    // proves the order is load-bearing rather than incidental.
    const ImuSample swapped = applyMount(ImuMount{90, 0, 180}, s);
    CHECK(swapped.acc[0] == doctest::Approx(-1.0));
    CHECK(swapped.acc[1] == doctest::Approx(0.0).scale(1));
}

TEST_CASE("rotating a vector never changes its length") {
    // The cheapest guard against a malformed matrix: a rotation preserves
    // magnitude, a typo in a sine almost never does.
    const ImuSample in = makeSample();
    for (const ImuMount& m : {ImuMount{30, -45, 120}, ImuMount{180, 0, -90}, ImuMount{-12, 88, 7}}) {
        const ImuSample out = applyMount(m, in);
        CHECK(norm(out.acc) == doctest::Approx(norm(in.acc)));
        CHECK(norm(out.gyr) == doctest::Approx(norm(in.gyr)));
        CHECK(norm(out.mag) == doctest::Approx(norm(in.mag)));
    }
}

TEST_CASE("everything that is not a frame passes through untouched") {
    const ImuSample in  = makeSample();
    const ImuSample out = applyMount(ImuMount{30, -45, 120}, in);
    CHECK(out.temp_c == in.temp_c);
    CHECK(out.sample_counter == in.sample_counter);
    CHECK(out.valid == in.valid);
}

TEST_CASE("Euler round-trips through the rotation matrix") {
    struct { double roll, pitch, yaw; } cases[] = {
        {0, 0, 0}, {10, 20, 30}, {-45, 12, 200}, {179, -80, 359}, {-170, 0, 45},
    };
    for (const auto& c : cases) {
        const Mat3 R = rotationFromEulerZYX(c.roll, c.pitch, c.yaw);
        double roll = 0, pitch = 0, yaw = 0;
        eulerZYXFromRotation(R, roll, pitch, yaw);
        CHECK(roll == doctest::Approx(c.roll).epsilon(1e-9));
        CHECK(pitch == doctest::Approx(c.pitch).epsilon(1e-9));
        CHECK(headingDiff(yaw, c.yaw) == doctest::Approx(0.0).scale(1));
    }
}

TEST_CASE("applying a mount and then its inverse gives back the original") {
    const ImuMount m{30, -45, 120};
    // The inverse of Rz(z)Rx(x)Ry(y) is Ry(-y)Rx(-x)Rz(-z), which is NOT another
    // Z-X-Y mount - so undo it by rotating the matrix back rather than by
    // negating the angles, and check against the composed identity.
    const Mat3 R = mountRotationZXY(m);
    Mat3 Rt{};
    for (int i = 0; i < 3; ++i)
        for (int j = 0; j < 3; ++j) Rt[i][j] = R[j][i];

    // R^T R must be the identity, i.e. the round trip is exact by construction.
    for (int i = 0; i < 3; ++i)
        for (int j = 0; j < 3; ++j) {
            const double v = Rt[i][0] * R[0][j] + Rt[i][1] * R[1][j] + Rt[i][2] * R[2][j];
            CHECK(v == doctest::Approx(i == j ? 1.0 : 0.0).epsilon(1e-12));
        }

    // And a yaw mount really is undone by the opposite yaw mount.
    const ImuSample in = makeSample();
    const ImuSample there_and_back = applyMount(ImuMount{0, 0, -35}, applyMount(ImuMount{0, 0, 35}, in));
    for (int i = 0; i < 3; ++i) CHECK(there_and_back.acc[i] == doctest::Approx(in.acc[i]));
}

TEST_CASE("pitch at the poles produces angles, not NaN") {
    // Straight up and straight down are gimbal lock: roll and yaw stop being
    // separable. The result must still be finite - a NaN here would propagate
    // into the calibration fit and the drift test with no warning.
    for (double pitch : {90.0, -90.0, 89.9999999}) {
        ImuSample s = makeSample();
        s.pitch_deg = static_cast<float>(pitch);
        s.yaw_deg   = 137.f;

        const ImuSample out = applyMount(ImuMount{0, 0, 25}, s);
        CHECK(std::isfinite(out.roll_deg));
        CHECK(std::isfinite(out.pitch_deg));
        CHECK(std::isfinite(out.yaw_deg));
        CHECK(out.yaw_deg >= -180.f);
        CHECK(out.yaw_deg <= 180.f);
    }
}

// --- The earth reference ----------------------------------------------------

namespace {

Mat3 matMul(const Mat3& a, const Mat3& b) {
    Mat3 r{};
    for (int i = 0; i < 3; ++i)
        for (int j = 0; j < 3; ++j)
            r[i][j] = a[i][0] * b[0][j] + a[i][1] * b[1][j] + a[i][2] * b[2][j];
    return r;
}

// What a device with earth frame `ref` and body mount `m` would report, given the
// gimbal's true attitude in NED. The inverse of what applyFrames has to undo.
ImuSample deviceReports(ImuReference ref, const ImuMount& m, double roll, double pitch,
                        double yaw) {
    const Mat3 ned_from_local = nedFromReference(ref);
    Mat3 local_from_ned{};
    for (int i = 0; i < 3; ++i)
        for (int j = 0; j < 3; ++j) local_from_ned[i][j] = ned_from_local[j][i];

    const Mat3 ned_from_gimbal = rotationFromEulerZYX(roll, pitch, yaw);
    // R_gimbal<-sensor is the transpose of the mount (which is R_sensor<-gimbal).
    const Mat3 R = mountRotationZXY(m);
    Mat3 gimbal_from_sensor{};
    for (int i = 0; i < 3; ++i)
        for (int j = 0; j < 3; ++j) gimbal_from_sensor[i][j] = R[j][i];

    const Mat3 local_from_sensor =
        matMul(local_from_ned, matMul(ned_from_gimbal, gimbal_from_sensor));
    double r = 0, p = 0, y = 0;
    eulerZYXFromRotation(local_from_sensor, r, p, y);
    ImuSample s = makeSample();
    s.roll_deg = static_cast<float>(r);
    s.pitch_deg = static_cast<float>(p);
    s.yaw_deg = static_cast<float>(y);
    return s;
}

}  // namespace

TEST_CASE("a Z-up earth reference is what makes pitch fall as the nose rises") {
    // The bug this whole path exists for. An upside-down IMU on a NWU-referenced
    // device reports pitch with the wrong sign, and - the crucial part - the mount
    // alone cannot fix it, because a mount only ever shifts roll here.
    const ImuMount upside_down{180, 0, 0};

    for (double nose_up : {5.0, 10.0, 20.0}) {
        const ImuSample raw = deviceReports(ImuReference::Nwu, upside_down, 0.0, nose_up, 0.0);
        CHECK(raw.pitch_deg == doctest::Approx(-nose_up));         // falls, as observed
        CHECK(raw.roll_deg == doctest::Approx(0.0).scale(1));      // and reads level

        // Mount only: pitch is untouched, roll merely picks up 180.
        const ImuSample mount_only = applyMount(upside_down, raw);
        CHECK(mount_only.pitch_deg == doctest::Approx(-nose_up));
        CHECK(std::abs(mount_only.roll_deg) == doctest::Approx(180.0));

        // Earth reference + mount: correct.
        const ImuSample fixed = applyFrames(ImuReference::Nwu, upside_down, raw);
        CHECK(fixed.pitch_deg == doctest::Approx(nose_up));
        CHECK(fixed.roll_deg == doctest::Approx(0.0).scale(1));
        CHECK(fixed.yaw_deg == doctest::Approx(0.0).scale(1));
    }
}

TEST_CASE("applyFrames recovers the true attitude for every reference and mount") {
    struct { double roll, pitch, yaw; } poses[] = {
        {0, 0, 0}, {0, 10, 0}, {0, -25, 45}, {15, 0, 0}, {-30, 12, -120}, {0, 0, 179},
    };
    for (ImuReference ref : {ImuReference::Ned, ImuReference::Nwu, ImuReference::Enu}) {
        for (const ImuMount& m : {ImuMount{}, ImuMount{180, 0, 0}, ImuMount{0, 0, -90},
                                  ImuMount{90, 0, 180}}) {
            for (const auto& p : poses) {
                const ImuSample raw   = deviceReports(ref, m, p.roll, p.pitch, p.yaw);
                const ImuSample fixed = applyFrames(ref, m, raw);
                CHECK(fixed.roll_deg == doctest::Approx(p.roll).epsilon(1e-4));
                CHECK(fixed.pitch_deg == doctest::Approx(p.pitch).epsilon(1e-4));
                CHECK(headingDiff(fixed.yaw_deg, p.yaw) == doctest::Approx(0.0).scale(1));
            }
        }
    }
}

TEST_CASE("the earth reference leaves body-frame vectors alone") {
    // acc/gyr/mag say which way the sensor felt something, which has nothing to do
    // with what the orientation is measured against. Only the mount may touch them.
    const ImuSample in = makeSample();
    for (ImuReference ref : {ImuReference::Ned, ImuReference::Nwu, ImuReference::Enu}) {
        const ImuSample out = applyFrames(ref, ImuMount{}, in);
        for (int i = 0; i < 3; ++i) {
            CHECK(out.acc[i] == doctest::Approx(in.acc[i]));
            CHECK(out.gyr[i] == doctest::Approx(in.gyr[i]));
            CHECK(out.mag[i] == doctest::Approx(in.mag[i]));
        }
    }
}

TEST_CASE("NED plus an identity mount is an exact no-op, and the matrices are rotations") {
    const ImuSample in = makeSample();
    const ImuSample out = applyFrames(ImuReference::Ned, ImuMount{}, in);
    CHECK(out.roll_deg == in.roll_deg);
    CHECK(out.pitch_deg == in.pitch_deg);
    CHECK(out.yaw_deg == in.yaw_deg);

    for (ImuReference ref : {ImuReference::Ned, ImuReference::Nwu, ImuReference::Enu}) {
        const Mat3 R = nedFromReference(ref);
        const double det = R[0][0] * (R[1][1] * R[2][2] - R[1][2] * R[2][1]) -
                           R[0][1] * (R[1][0] * R[2][2] - R[1][2] * R[2][0]) +
                           R[0][2] * (R[1][0] * R[2][1] - R[1][1] * R[2][0]);
        CHECK(det == doctest::Approx(1.0));   // a rotation, never a reflection
    }
    CHECK(std::string(referenceName(ImuReference::Nwu)) == "NWU");
}

TEST_CASE("NWU and ENU differ only in yaw, by exactly 90 degrees") {
    // Which of the two the device uses cannot be told from roll or pitch - only a
    // compass check on yaw distinguishes them. That is worth pinning, because it
    // is the whole risk of picking the wrong one.
    for (double pitch : {0.0, 15.0, -20.0}) {
        const ImuSample raw = deviceReports(ImuReference::Nwu, ImuMount{180, 0, 0}, 0, pitch, 30);
        const ImuSample as_nwu = applyFrames(ImuReference::Nwu, ImuMount{180, 0, 0}, raw);
        const ImuSample as_enu = applyFrames(ImuReference::Enu, ImuMount{180, 0, 0}, raw);
        CHECK(as_enu.roll_deg == doctest::Approx(as_nwu.roll_deg));
        CHECK(as_enu.pitch_deg == doctest::Approx(as_nwu.pitch_deg));
        CHECK(headingDiff(as_enu.yaw_deg, as_nwu.yaw_deg + 90.0) == doctest::Approx(0.0).scale(1));
    }
}