#include "Physics.h"

#include <glm/gtc/constants.hpp>
#include <glm/gtc/quaternion.hpp>

// ============================================================================
// ComputeInvInertiaTensor
// Returns zero matrix for static or rotation-locked bodies.
// ============================================================================
glm::mat3 Physics::ComputeInvInertiaTensor(const p* body)
{
    if (!body->velocity || body->lockRotation)
        return glm::mat3(0.0f);

    constexpr float k = 1.0f / 12.0f;
    const glm::vec3 s = body->pointer->t.GetScale();
    const glm::vec3 sq = s * s;

    glm::mat3 localInvI(0.0f);
    localInvI[0][0] = 1.0f / (k * body->mass * (sq.y + sq.z));
    localInvI[1][1] = 1.0f / (k * body->mass * (sq.x + sq.z));
    localInvI[2][2] = 1.0f / (k * body->mass * (sq.x + sq.y));

    const glm::mat3 R = glm::mat3(body->pointer->t.GetRotationMatrix());
    return R * localInvI * glm::transpose(R);
}

// ============================================================================
// UpdateSleep
// ============================================================================
bool Physics::UpdateSleep(p* body, float dt)
{
    if (!body->velocity || body->frozen) return true;
    if (body->collisionGroup != 0) return false;  // groups settle as a whole, a joint can't hold a sleeping body

    const float linSq = glm::dot(body->linearvelocity, body->linearvelocity);
    const float angSq = glm::dot(body->angularvelocity, body->angularvelocity);

    if (linSq < PHYSICS_SLEEP_LINEAR_SQ && angSq < PHYSICS_SLEEP_ANGULAR_SQ) {
        body->sleepTimer += dt;
        if (body->sleepTimer >= PHYSICS_SLEEP_TIME_THRESHOLD) {
            body->isSleeping = true;
            body->linearvelocity = glm::vec3(0.0f);
            body->angularvelocity = glm::vec3(0.0f);
            return true;
        }
    }
    else {
        body->sleepTimer = 0.0f;
        body->isSleeping = false;
    }
    return false;
}

// ============================================================================
// Step
// ============================================================================
void Physics::Step(float dt)
{
    // The damping constants are per 1/60 s step; rescale so drag doesn't depend on frame rate
    const float steps = dt * 60.0f;
    const float linearDamping = std::pow(PHYSICS_LINEAR_DAMPING, steps);
    const float angularDamping = std::pow(PHYSICS_ANGULAR_DAMPING, steps);

    // Integrate forces ? velocities
    for (p* body : objects) {
        // Nothing rotates until the end of the step, so everything until then can share this
        body->invInertia = ComputeInvInertiaTensor(body);

        if (!body->velocity || body->isSleeping || body->frozen) continue;

        body->force += body->mass * gravity;

        body->linearvelocity += (body->force / body->mass) * dt;

        if (!body->lockRotation)
            body->angularvelocity += (body->invInertia * body->torque) * dt;

        body->linearvelocity *= linearDamping;
        body->angularvelocity *= angularDamping;

        body->force = glm::vec3(0.0f);
        body->torque = glm::vec3(0.0f);
    }

    ResolveCollisions(dt);

    // Integrate velocities ? positions
    for (p* body : objects) {
        if (!body->velocity) continue;
        if (UpdateSleep(body, dt)) continue;

        if (body->maxSpeed > 0.0f && glm::length(body->linearvelocity) > body->maxSpeed)
            body->linearvelocity = glm::normalize(body->linearvelocity) * body->maxSpeed;
        if (body->maxAngularSpeed > 0.0f && glm::length(body->angularvelocity) > body->maxAngularSpeed)
            body->angularvelocity = glm::normalize(body->angularvelocity) * body->maxAngularSpeed;

        body->pointer->t.TranslateBy(body->linearvelocity * dt);

        if (!body->lockRotation) {
            body->pointer->t.RotateByQuaternion(glm::quat(body->angularvelocity * dt));
            if (body->collisionMesh) body->pointer->t.NormalizeRotation();
        }
    }
}

// ============================================================================
// ResolveCollisions
// ============================================================================
void Physics::ResolveCollisions(float dt)
{
    FindCollisions();
    for (Collision& col : collisions)
        Resolve(&col);

    if (joints.empty()) return;

    // The first contact pass nudges positions, so these have to be worked out after it
    for (BallJoint* joint : joints) {
        DampJoint(*joint, dt);
        PrepareJoint(*joint, dt);
    }

    // Joints and contacts fight each other, so iterate them together or limbs get pulled through the floor
    for (int i = 0; i < PHYSICS_JOINT_ITERATIONS; i++) {
        for (BallJoint* joint : joints)
            SolveJoint(*joint);
        for (Collision& col : collisions)
            Resolve(&col, false);
    }
}

// ============================================================================
// FindCollisions
// Sort and sweep: once sorted along x, a body can only touch the ones after it
// that start before it ends.
// ============================================================================
void Physics::FindCollisions()
{
    collisions.clear();
    broadPhase.clear();

    for (p* body : objects) {
        if (!body->collision) continue;
        BroadPhaseEntry entry;
        entry.body = body;
        entry.bounds = body->pointer->t.GetAABB();
        broadPhase.push_back(entry);
    }

    std::sort(broadPhase.begin(), broadPhase.end(), [](const BroadPhaseEntry& first, const BroadPhaseEntry& second) {
        return first.bounds.min.x < second.bounds.min.x;
    });

    for (int i = 0; i < (int)broadPhase.size(); i++) {
        const BroadPhaseEntry& first = broadPhase[i];
        p* a = first.body;

        for (int j = i + 1; j < (int)broadPhase.size(); j++) {
            const BroadPhaseEntry& second = broadPhase[j];
            if (second.bounds.min.x >= first.bounds.max.x) break;

            p* b = second.body;
            if (!a->velocity && !b->velocity)   continue;
            if (a->isSleeping && b->isSleeping) continue;
            if ((a->frozen || !a->velocity) && (b->frozen || !b->velocity)) continue;
            if (a->collisionGroup != 0 && a->collisionGroup == b->collisionGroup) continue;
            if (!BoundingBoxInBoundingBox(first.bounds, second.bounds)) continue;

            if (a->collisionMesh || b->collisionMesh) {
                CollideTriangles(a, b);
                continue;
            }

            std::optional<TInTInfo> e = TInT(a->pointer->t, b->pointer->t);
            if (e.has_value()) {
                Collision tp(a, b);
                tp.POI = e.value().POI;
                tp.CN = e.value().CN;
                tp.overlap = e.value().overlap;
                collisions.push_back(tp);
            }
        }
    }
}

// ============================================================================
// CollideTriangles
// Tests every triangle of one body against the other's BVH. Triangles have no
// volume, so depth comes from how far the intersecting triangle reaches behind
// the other triangle's outward facing plane.
// ============================================================================
static Triangle TransformTriangle(const Triangle& tri, const glm::mat4& matrix)
{
    return {
        glm::vec3(matrix * glm::vec4(tri.a, 1.0f)),
        glm::vec3(matrix * glm::vec4(tri.b, 1.0f)),
        glm::vec3(matrix * glm::vec4(tri.c, 1.0f))
    };
}

// Keeps the deepest contact, then whichever contacts spread the rest out the most
static void KeepSpreadContacts(std::vector<Collision>& contacts)
{
    if ((int)contacts.size() <= PHYSICS_MAX_MESH_CONTACTS)
        return;

    int deepest = 0;
    for (int i = 1; i < (int)contacts.size(); i++) {
        if (contacts[i].overlap > contacts[deepest].overlap)
            deepest = i;
    }
    std::swap(contacts[0], contacts[deepest]);

    for (int kept = 1; kept < PHYSICS_MAX_MESH_CONTACTS; kept++) {
        int farthest = kept;
        float farthestDistance = -1.0f;
        for (int i = kept; i < (int)contacts.size(); i++) {
            float nearest = FLT_MAX;
            for (int j = 0; j < kept; j++) {
                const glm::vec3 offset = contacts[i].POI - contacts[j].POI;
                nearest = glm::min(nearest, glm::dot(offset, offset));
            }
            if (nearest > farthestDistance) {
                farthestDistance = nearest;
                farthest = i;
            }
        }
        std::swap(contacts[kept], contacts[farthest]);
    }

    contacts.erase(contacts.begin() + PHYSICS_MAX_MESH_CONTACTS, contacts.end());
}

// Two meshes can tangle: a hand between two faces gets contacts pushing opposite ways, which cancel
// and pin it in place. So the whole pair shares one normal that always pushes them apart.
static void UnifyContactNormals(const p* a, const p* b, std::vector<Collision>& contacts)
{
    glm::vec3 summed(0.0f);
    float totalDepth = 0.0f;
    for (const Collision& contact : contacts) {
        summed += contact.CN * contact.overlap;
        totalDepth += contact.overlap;
    }

    // Faces that mostly disagree say nothing useful, so fall back to the line between the two bodies
    const glm::vec3 apart = a->pointer->t.GetTranslation() - b->pointer->t.GetTranslation();
    glm::vec3 normal = summed;
    if (glm::length(summed) < totalDepth * 0.3f)
        normal = apart;
    if (glm::length(normal) < 1e-6f)
        return;

    normal = glm::normalize(normal);
    if (glm::dot(normal, apart) < 0.0f)
        normal = -normal;

    for (Collision& contact : contacts) {
        contact.CN = normal;
        contact.frictionScale = PHYSICS_MESH_PAIR_FRICTION;
    }
}

void Physics::CollideTriangles(p* a, p* b)
{
    const CollisionMesh& meshA = a->collisionMesh ? *a->collisionMesh : CollisionMesh::UnitBox();
    const CollisionMesh& meshB = b->collisionMesh ? *b->collisionMesh : CollisionMesh::UnitBox();

    // Loop over the smaller mesh and search the larger one
    const bool queryIsA = meshA.GetTriangles().size() <= meshB.GetTriangles().size();
    const CollisionMesh& queryMesh = queryIsA ? meshA : meshB;
    const CollisionMesh& targetMesh = queryIsA ? meshB : meshA;

    const glm::mat4 queryToWorld = (queryIsA ? a : b)->pointer->t.GetMatrix();
    const glm::mat4 targetToWorld = (queryIsA ? b : a)->pointer->t.GetMatrix();
    const glm::mat4 queryToTarget = glm::inverse(targetToWorld) * queryToWorld;

    // Box faces are the reliable surface. Between two meshes, b's faces are used.
    const bool surfaceIsA = !a->collisionMesh && b->collisionMesh;
    const bool surfaceIsQuery = surfaceIsA == queryIsA;

    meshContacts.clear();

    for (const Triangle& queryTriangle : queryMesh.GetTriangles()) {
        const Triangle queryInTarget = TransformTriangle(queryTriangle, queryToTarget);

        meshCandidates.clear();
        targetMesh.Query(queryInTarget, meshCandidates);

        for (int index : meshCandidates) {
            const Triangle& targetTriangle = targetMesh.GetTriangles()[index];
            if (!TriangleInTriangle(queryInTarget, targetTriangle))
                continue;

            const Triangle queryWorld = TransformTriangle(queryTriangle, queryToWorld);
            const Triangle targetWorld = TransformTriangle(targetTriangle, targetToWorld);
            const Triangle& surface = surfaceIsQuery ? queryWorld : targetWorld;
            const Triangle& other = surfaceIsQuery ? targetWorld : queryWorld;

            const glm::vec3 faceNormal = glm::cross(surface.b - surface.a, surface.c - surface.a);
            const float faceLength = glm::length(faceNormal);
            if (faceLength < 1e-8f)
                continue;
            const glm::vec3 normal = faceNormal / faceLength;

            glm::vec3 deepestPoint = other.a;
            float deepestDistance = glm::dot(other.a - surface.a, normal);
            for (const glm::vec3& corner : { other.b, other.c }) {
                const float distance = glm::dot(corner - surface.a, normal);
                if (distance < deepestDistance) {
                    deepestDistance = distance;
                    deepestPoint = corner;
                }
            }
            if (deepestDistance >= 0.0f)
                continue;

            // CN points from b to a, the same way TInT does
            Collision contact(a, b);
            contact.POI = deepestPoint;
            contact.CN = surfaceIsA ? -normal : normal;
            contact.overlap = glm::min(-deepestDistance, PHYSICS_MAX_MESH_DEPTH);
            meshContacts.push_back(contact);
        }
    }

    if (meshContacts.empty())
        return;

    if (a->frozen) WakeGroup(a->collisionGroup);
    if (b->frozen) WakeGroup(b->collisionGroup);

    if (a->collisionMesh && b->collisionMesh)
        UnifyContactNormals(a, b, meshContacts);

    KeepSpreadContacts(meshContacts);
    for (Collision& contact : meshContacts) {
        contact.correctionWeight = 1.0f / (float)meshContacts.size();
        collisions.push_back(contact);
    }
}

void Physics::WakeGroup(int group)
{
    for (p* body : objects) {
        if (body->collisionGroup == group)
            body->frozen = false;
    }
}

// ============================================================================
// Joints
// Everything that only depends on where the bodies are gets worked out once in
// Prepare. Solve then just pushes velocities around, and runs many times a step.
// ============================================================================
static glm::mat3 SkewMatrix(const glm::vec3& v)
{
    // Skew * x = cross(v, x). glm matrices are built from columns.
    return glm::mat3(
        glm::vec3(0.0f, v.z, -v.y),
        glm::vec3(-v.z, 0.0f, v.x),
        glm::vec3(v.y, -v.x, 0.0f));
}

// Asks for the child's spin relative to the parent around axis to close error
static void AddAngularRow(BallJoint& joint, const glm::vec3& axis, float error, float dt, bool oneSided)
{
    AngularRow& row = joint.rows[joint.rowCount++];
    row.axis = axis;
    row.oneSided = oneSided;
    row.targetSpeed = glm::clamp(error * (PHYSICS_JOINT_BIAS / dt),
        -PHYSICS_JOINT_MAX_BIAS_SPEED, PHYSICS_JOINT_MAX_BIAS_SPEED);
    row.inverseDenom = 1.0f / (glm::dot(axis, joint.parent->invInertia * axis) + glm::dot(axis, joint.child->invInertia * axis));
}

void Physics::PrepareJoint(BallJoint& joint, float dt)
{
    p* a = joint.parent;
    p* b = joint.child;
    if (a->frozen && b->frozen)
        return;

    const glm::quat orientA = a->pointer->t.GetRotationQuaternion();
    const glm::quat orientB = b->pointer->t.GetRotationQuaternion();

    const glm::vec3 rA = orientA * joint.anchorParent;
    const glm::vec3 rB = orientB * joint.anchorChild;
    joint.worldAnchorParent = rA;
    joint.worldAnchorChild = rB;

    const glm::mat3 skewA = SkewMatrix(rA);
    const glm::mat3 skewB = SkewMatrix(rB);
    const glm::mat3 K = glm::mat3(a->GetInvMass() + b->GetInvMass())
        - skewA * a->invInertia * skewA - skewB * b->invInertia * skewB;
    joint.inverseK = glm::inverse(K);

    const glm::vec3 drift = (b->pointer->t.GetTranslation() + rB) - (a->pointer->t.GetTranslation() + rA);
    joint.bias = drift * (PHYSICS_JOINT_BIAS / dt);
    if (glm::length(joint.bias) > PHYSICS_JOINT_MAX_BIAS_SPEED)
        joint.bias = glm::normalize(joint.bias) * PHYSICS_JOINT_MAX_BIAS_SPEED;

    joint.rowCount = 0;

    // ---- Hinge ----
    if (joint.isHinge) {
        const glm::vec3 hingeA = orientA * joint.hingeAxisParent;
        const glm::vec3 hingeB = orientB * joint.hingeAxisChild;

        // Keep both hinge axes lined up, leaving only spin around the hinge free
        const glm::vec3 misalignment = glm::cross(hingeB, hingeA);
        const glm::vec3 side = glm::normalize(orientA * joint.hingeRefParent);
        const glm::vec3 otherSide = glm::cross(hingeA, side);
        AddAngularRow(joint, side, glm::dot(misalignment, side), dt, false);
        AddAngularRow(joint, otherSide, glm::dot(misalignment, otherSide), dt, false);

        const glm::vec3 refA = side;
        const glm::vec3 refB = orientB * joint.hingeRefChild;
        const float bend = std::atan2(glm::dot(glm::cross(refA, refB), hingeA), glm::dot(refA, refB));

        if (bend < joint.minBend)
            AddAngularRow(joint, hingeA, joint.minBend - bend, dt, true);
        else if (bend > joint.maxBend)
            AddAngularRow(joint, -hingeA, bend - joint.maxBend, dt, true);
        return;
    }

    // ---- Swing cone ----
    const glm::vec3 axisA = orientA * joint.axisParent;
    const glm::vec3 axisB = orientB * joint.axisChild;
    const float angle = glm::acos(glm::clamp(glm::dot(axisA, axisB), -1.0f, 1.0f));
    if (angle > joint.maxAngle) {
        const glm::vec3 swingAxis = glm::cross(axisB, axisA);
        const float swingLength = glm::length(swingAxis);
        if (swingLength > 1e-5f)
            AddAngularRow(joint, swingAxis / swingLength, angle - joint.maxAngle, dt, true);
    }

    // ---- Twist limit ----
    const glm::quat relative = glm::conjugate(orientA) * orientB;
    const float twistAlongAxis = glm::dot(glm::vec3(relative.x, relative.y, relative.z), joint.axisChild);
    float twist = 2.0f * std::atan2(twistAlongAxis, relative.w);
    const float pi = glm::pi<float>();
    if (twist > pi) twist -= 2.0f * pi;
    else if (twist < -pi) twist += 2.0f * pi;

    const glm::vec3 twistAxis = orientA * joint.axisChild;
    if (twist > joint.maxTwist)
        AddAngularRow(joint, -twistAxis, twist - joint.maxTwist, dt, true);
    else if (twist < -joint.maxTwist)
        AddAngularRow(joint, twistAxis, -joint.maxTwist - twist, dt, true);
}

void Physics::SolveJoint(BallJoint& joint)
{
    p* a = joint.parent;
    p* b = joint.child;
    if (a->frozen && b->frozen)
        return;

    const glm::vec3& rA = joint.worldAnchorParent;
    const glm::vec3& rB = joint.worldAnchorChild;

    const glm::vec3 relVel = b->linearvelocity + glm::cross(b->angularvelocity, rB)
        - a->linearvelocity - glm::cross(a->angularvelocity, rA);

    const glm::vec3 impulse = joint.inverseK * (-relVel - joint.bias);
    a->linearvelocity -= a->GetInvMass() * impulse;
    a->angularvelocity -= a->invInertia * glm::cross(rA, impulse);
    b->linearvelocity += b->GetInvMass() * impulse;
    b->angularvelocity += b->invInertia * glm::cross(rB, impulse);

    for (int i = 0; i < joint.rowCount; i++) {
        const AngularRow& row = joint.rows[i];
        const float speed = glm::dot(b->angularvelocity - a->angularvelocity, row.axis);
        const float rowImpulse = (row.targetSpeed - speed) * row.inverseDenom;
        if (row.oneSided && rowImpulse <= 0.0f)
            continue;

        a->angularvelocity -= a->invInertia * row.axis * rowImpulse;
        b->angularvelocity += b->invInertia * row.axis * rowImpulse;
    }
}

// Joints have no friction of their own, so without this limbs flail for a long time
void Physics::DampJoint(BallJoint& joint, float dt)
{
    p* a = joint.parent;
    p* b = joint.child;
    if (a->frozen && b->frozen)
        return;

    const glm::vec3 relative = b->angularvelocity - a->angularvelocity;
    const glm::vec3 removed = relative * (1.0f - std::exp(-PHYSICS_JOINT_DAMPING * dt)) * 0.5f;
    a->angularvelocity += removed;
    b->angularvelocity -= removed;
}

// ============================================================================
// Resolve
// ============================================================================
void Physics::Resolve(Collision* c, bool firstPass)
{
    p* a = c->ObjA;
    p* b = c->ObjB;

    a->WakeUp();
    b->WakeUp();

    const glm::vec3 n = c->CN;
    const glm::vec3 rA = c->POI - a->pointer->t.GetTranslation();
    const glm::vec3 rB = c->POI - b->pointer->t.GetTranslation();

    const glm::mat3& invIA = a->invInertia;
    const glm::mat3& invIB = b->invInertia;

    const float invMassA = a->GetInvMass();
    const float invMassB = b->GetInvMass();

    const glm::vec3 velA = a->linearvelocity + glm::cross(a->angularvelocity, rA);
    const glm::vec3 velB = b->linearvelocity + glm::cross(b->angularvelocity, rB);
    const glm::vec3 relVel = velA - velB;

    const float relVelN = glm::dot(relVel, n);
    if (relVelN > 0.0f) return;

    const float e = (!firstPass || glm::abs(relVelN) < PHYSICS_RESTITUTION_THRESHOLD)
        ? 0.0f
        : glm::min(a->restitution, b->restitution);

    auto angDenom = [](const glm::mat3& invI, const glm::vec3& r, const glm::vec3& d) {
        return glm::dot(glm::cross(invI * glm::cross(r, d), r), d);
        };

    // ---- Normal impulse ----
    const float denomN = invMassA + invMassB + angDenom(invIA, rA, n) + angDenom(invIB, rB, n);
    const float jN = -(1.0f + e) * relVelN / denomN;

    const glm::vec3 impulseN = jN * n;
    if (a->velocity) {
        a->linearvelocity += invMassA * impulseN;
        if (!a->lockRotation)
            a->angularvelocity += invIA * glm::cross(rA, impulseN);
    }
    if (b->velocity) {
        b->linearvelocity -= invMassB * impulseN;
        if (!b->lockRotation)
            b->angularvelocity -= invIB * glm::cross(rB, impulseN);
    }

    // ---- Friction impulse ----
    glm::vec3 tangent = relVel - relVelN * n;
    const float tangLen = glm::length(tangent);
    if (tangLen > 1e-6f) {
        tangent /= tangLen;

        const float denomT = invMassA + invMassB + angDenom(invIA, rA, tangent) + angDenom(invIB, rB, tangent);
        float jT = -glm::dot(relVel, tangent) / denomT;

        const float mu = (a->friction + b->friction) * 0.5f * c->frictionScale;
        jT = glm::clamp(jT, -mu * jN, mu * jN);

        const glm::vec3 impulseT = jT * tangent;
        if (a->velocity) {
            a->linearvelocity += invMassA * impulseT;
            if (!a->lockRotation)
                a->angularvelocity += invIA * glm::cross(rA, impulseT);
        }
        if (b->velocity) {
            b->linearvelocity -= invMassB * impulseT;
            if (!b->lockRotation)
                b->angularvelocity -= invIB * glm::cross(rB, impulseT);
        }
    }

    // ---- Baumgarte positional correction ----
    if (!firstPass)
        return;

    const float correction = glm::max(c->overlap - PHYSICS_BAUMGARTE_SLOP, 0.0f)
        / (invMassA + invMassB)
        * PHYSICS_BAUMGARTE_FACTOR
        * c->correctionWeight;
    const glm::vec3 corrVec = correction * n;
    if (a->velocity) a->pointer->t.TranslateBy(invMassA * corrVec);
    if (b->velocity) b->pointer->t.TranslateBy(-invMassB * corrVec);
}