#include "fgc/ImuMount.h"

#include <algorithm>
#include <cmath>

namespace fgc {

namespace {

constexpr double kDeg2Rad = M_PI / 180.0;
constexpr double kRad2Deg = 180.0 / M_PI;

// Below this the mount is not worth applying: 1e-9 deg is far under any angle an
// operator can measure, and treating it as zero keeps the untouched-sample
// guarantee exact rather than approximate.
constexpr double kAngleEpsilon = 1e-9;

Mat3 rotX(double rad) {
    const double c = std::cos(rad), s = std::sin(rad);
    return {{{1, 0, 0}, {0, c, -s}, {0, s, c}}};
}

Mat3 rotY(double rad) {
    const double c = std::cos(rad), s = std::sin(rad);
    return {{{c, 0, s}, {0, 1, 0}, {-s, 0, c}}};
}

Mat3 rotZ(double rad) {
    const double c = std::cos(rad), s = std::sin(rad);
    return {{{c, -s, 0}, {s, c, 0}, {0, 0, 1}}};
}

Mat3 mul(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;
}

// v' = R^T v, i.e. re-express a vector's components in the frame R carries the
// original onto. Written out rather than transposing then multiplying, since the
// transpose only ever appears here.
void applyTranspose(const Mat3& R, const float in[3], float out[3]) {
    double v[3] = {in[0], in[1], in[2]};
    for (int i = 0; i < 3; ++i)
        out[i] = static_cast<float>(R[0][i] * v[0] + R[1][i] * v[1] + R[2][i] * v[2]);
}

}  // namespace

bool ImuMount::identity() const {
    return std::abs(x_deg) < kAngleEpsilon && std::abs(y_deg) < kAngleEpsilon &&
           std::abs(z_deg) < kAngleEpsilon;
}

Mat3 mountRotationZXY(const ImuMount& m) {
    return mul(mul(rotZ(m.z_deg * kDeg2Rad), rotX(m.x_deg * kDeg2Rad)), rotY(m.y_deg * kDeg2Rad));
}

Mat3 rotationFromEulerZYX(double roll_deg, double pitch_deg, double yaw_deg) {
    return mul(mul(rotZ(yaw_deg * kDeg2Rad), rotY(pitch_deg * kDeg2Rad)), rotX(roll_deg * kDeg2Rad));
}

void eulerZYXFromRotation(const Mat3& R, double& roll_deg, double& pitch_deg, double& yaw_deg) {
    // For R = Rz(yaw) Ry(pitch) Rx(roll), R[2][0] is exactly -sin(pitch).
    // Repeated multiplication can push it a hair past +/-1, and asin() of that is
    // NaN - which would then propagate silently through every consumer.
    const double sin_pitch = std::clamp(-R[2][0], -1.0, 1.0);
    const double pitch     = std::asin(sin_pitch);
    const double cos_pitch = std::cos(pitch);

    double roll, yaw;
    if (std::abs(cos_pitch) < 1e-9) {
        // Gimbal lock: pointing straight up or straight down, roll and yaw are the
        // same rotation and only their sum is observable. Attribute all of it to
        // yaw, which is the axis the operator is steering.
        roll = 0.0;
        yaw  = std::atan2(-R[0][1], R[1][1]);
    } else {
        roll = std::atan2(R[2][1], R[2][2]);
        yaw  = std::atan2(R[1][0], R[0][0]);
    }

    // atan2/asin already land in -180..180 and -90..90, which is exactly the
    // range the MTi reports, so nothing is re-ranged on the way out.
    roll_deg  = roll * kRad2Deg;
    pitch_deg = pitch * kRad2Deg;
    yaw_deg   = yaw * kRad2Deg;
}

Mat3 nedFromReference(ImuReference r) {
    switch (r) {
        // Already NED: nothing to do.
        case ImuReference::Ned: return {{{1, 0, 0}, {0, 1, 0}, {0, 0, 1}}};
        // NWU -> NED: North stays, West becomes East, Up becomes Down. That is a
        // 180 deg turn about North, i.e. Rx(180).
        case ImuReference::Nwu: return {{{1, 0, 0}, {0, -1, 0}, {0, 0, -1}}};
        // ENU -> NED: North is ENU's Y, East is ENU's X, Down is -Up. A 180 deg
        // turn about the horizontal axis halfway between North and East.
        case ImuReference::Enu: return {{{0, 1, 0}, {1, 0, 0}, {0, 0, -1}}};
    }
    return {{{1, 0, 0}, {0, 1, 0}, {0, 0, 1}}};
}

const char* referenceName(ImuReference r) {
    switch (r) {
        case ImuReference::Ned: return "NED";
        case ImuReference::Nwu: return "NWU";
        case ImuReference::Enu: return "ENU";
    }
    return "?";
}

ImuSample applyFrames(ImuReference ref, const ImuMount& m, const ImuSample& s) {
    if (ref == ImuReference::Ned && m.identity()) return s;

    const Mat3 R = mountRotationZXY(m);

    ImuSample out = s;   // temp, counter and valid are frame-independent
    // acc/gyr/mag are body-frame directions, so only the mount touches them. The
    // earth reference is about what the ORIENTATION is measured against and has
    // no bearing on which way the sensor felt an acceleration.
    applyTranspose(R, s.acc, out.acc);
    applyTranspose(R, s.gyr, out.gyr);
    applyTranspose(R, s.mag, out.mag);

    // The device reports local_earth<-sensor. Left-multiplying re-references the
    // earth frame to NED; right-multiplying by the mount swaps the body frame for
    // the gimbal's. They are independent and both are needed - which is exactly
    // why a mount alone could never fix a Z-up reference.
    const Mat3 local_from_sensor = rotationFromEulerZYX(s.roll_deg, s.pitch_deg, s.yaw_deg);
    const Mat3 ned_from_gimbal =
        mul(mul(nedFromReference(ref), local_from_sensor), R);

    double roll = 0, pitch = 0, yaw = 0;
    eulerZYXFromRotation(ned_from_gimbal, roll, pitch, yaw);
    out.roll_deg  = static_cast<float>(roll);
    out.pitch_deg = static_cast<float>(pitch);
    out.yaw_deg   = static_cast<float>(yaw);
    return out;
}

ImuSample applyMount(const ImuMount& m, const ImuSample& s) {
    return applyFrames(ImuReference::Ned, m, s);
}

}  // namespace fgc