// clang-format off
// glad must be included before any other OpenGL libraries
#include <glad/gl.h>
// clang-format on

#include "simulator.h"

#include <iostream>

#include "Tracy.hpp"
#include "imgui.h"
#include "rbdlsim.h"
#include "srender.h"
#include "sthstry.h"
#include "utils.h"

using namespace std;
using namespace RBDLSim;

static World sWorld;
static SimShape sGroundShape;

typedef SimpleMath::Matrix<float, 3, 3> Matrix33f;
typedef SimpleMath::Matrix<float, 3, 1> Vector3f;
typedef SimpleMath::Matrix<float, 4, 4> Matrix44f;
typedef SimpleMath::Matrix<float, 4, 1> Vector4f;

static bool sIsPaused = true;
static bool nSteps = 0;
static double sSimTime = 0.;
static double sSimTimeAccumulator = 0.;
static double sSimTimeStep = 1.0e-2;
static sthstry* sStateHistory = nullptr;
static int sStateHistoryCurrent = 0;
static int sSolverNumIter = 20;

void simulator_init() {
  gLog("Initializing Simulator");
  sStateHistory = sthstry_create();

  sGroundShape.mType = SimShape::Plane;
  sGroundShape.pos.set(0., 0., 0.);
  sGroundShape.orientation.set(0., 0., 0., 1.);
  sGroundShape.scale.set(1.0, 1.0, 1.0);
  sGroundShape.restitution = 1.0;

  sWorld.mStaticShapes.push_back(sGroundShape);

  double restitution = 0.5;

  int num_bodies = 40;
  for (int i = 0; i < num_bodies; i++) {
    SimBody body;

    bool create_sphere = i %2;

    if (!create_sphere) {
      body = CreateBoxBody(
          1.,
          Vector3d(2., 1., 1.),
          restitution,
          Vector3d::Random() * 5.,
          Vector3d::Zero());
    } else {
      body = CreateSphereBody(
          1.,
          1.,
          restitution,
          Vector3d::Random() * 5.,
          Vector3d::Zero());
    }

    sWorld.mBodies.push_back(body);
  }

  for (int i = 0; i < sWorld.mBodies.size(); i++) {
    SimBody& body = sWorld.mBodies[i];
    sthstry_register(
        sStateHistory,
        body.q.data(),
        body.q.size() * sizeof(double));
    sthstry_register(
        sStateHistory,
        body.qdot.data(),
        body.qdot.size() * sizeof(double));
  }

  simulator_reset();
}

void simulator_reset() {
  // Reset all Quaternions:
  for (SimBody& body : sWorld.mBodies) {
    body.q.setZero();
    body.qdot.setZero();
    body.qddot.setZero();
    for (int i = 0; i < body.mModel.mBodies.size(); i++) {
      if (body.mModel.mJoints[i].mJointType
          == RigidBodyDynamics::JointTypeSpherical) {
        Quaternion orientation = Quaternion::Random().normalize();
        body.mModel.SetQuaternion(i, orientation, body.q);
      }
    }
  }

  for (int i = 0; i < sWorld.mBodies.size(); i++) {
    sWorld.mBodies[i].q.block(0, 0, 3, 1) =
        Vector3d::Random() * 2.5 + Vector3d(0., 5., 0.);
    sWorld.mBodies[i].q[2] = 0.;
  }

  sthstry_reset_storage(sStateHistory);
  sthstry_store(sStateHistory);

  sSimTime = 0.;
  sSimTimeAccumulator = 0.;
}

void simulator_gui() {
  ZoneScoped;

  sStateHistoryCurrent =
      std::min(sStateHistoryCurrent, sthstry_get_num_states(sStateHistory) - 1);

  if (ImGui::Button("Reset")) {
    simulator_reset();
  }
  ImGui::Checkbox("Paused", &sIsPaused);

  if (ImGui::Button("<")) {
    if (sStateHistoryCurrent > 0) {
      sStateHistoryCurrent--;
    }
    sthstry_restore(sStateHistory, sStateHistoryCurrent);
  }
  ImGui::SameLine();
  if (ImGui::Button(">")) {
    sStateHistoryCurrent++;

    if (sStateHistoryCurrent >= sthstry_get_num_states(sStateHistory)) {
      simulator_step(sSimTimeStep);
    } else {
      sthstry_restore(sStateHistory, sStateHistoryCurrent);
    }
  }
  ImGui::SameLine();
  if (ImGui::SliderInt(
          "Step",
          &sStateHistoryCurrent,
          0,
          sthstry_get_num_states(sStateHistory) - 1)) {
    sthstry_restore(sStateHistory, sStateHistoryCurrent);
  }
  sSimTime = sSimTimeStep * sStateHistoryCurrent;
  ImGui::Text("Time: %f", sSimTime);

	ImGui::SliderInt("Solver Steps", &sSolverNumIter, 0, 100);

  ImGui::Text("Ground Plane");
  Vector3f ground_pos = sGroundShape.pos;
  ImGui::DragFloat3("Position", ground_pos.data(), 0.1f, -5.0f, 5.0f);
  sGroundShape.pos = ground_pos;

  Vector4f orientation = sGroundShape.orientation;
  ImGui::DragFloat4("Normal", orientation.data(), 0.1f, -1.0f, 1.0f);
  orientation.normalize();
  sGroundShape.orientation
      .set(orientation[0], orientation[1], orientation[2], orientation[3]);

  ImGui::Text("Bodies");
  for (int i = 0; i < sWorld.mBodies.size(); i++) {
    ImGui::PushID(i);
    SimBody& body = sWorld.mBodies[i];
    Vector3f body_pos = body.q.block(0, 0, 3, 1);
    ImGui::DragFloat3("Pos", body_pos.data(), 0.01f, -5.0f, 5.0f);
    body.q.block(0, 0, 3, 1) = body_pos;

    Vector3f body_vel = body.qdot.block(0, 0, 3, 1);
    ImGui::DragFloat3("Vel", body_vel.data(), 0.01f, -20.0f, 20.0f);
    body.qdot.block(0, 0, 3, 1) = body_vel;

    Vector4f body_rot = body.mModel.GetQuaternion(2, body.q);
    ImGui::DragFloat4("Rot", body_rot.data(), 0.01f, -20.0f, 20.0f);

    ImGui::PopID();
    body.updateCollisionShapes();
  }
}

void simulator_update(double dt) {
  ZoneScoped;

  if (!sIsPaused) {
    sSimTimeAccumulator += dt;
    while (sSimTimeAccumulator > sSimTimeStep) {
      FrameMarkStart("SimStep");
      simulator_step(sSimTimeStep);
      sSimTimeAccumulator -= sSimTimeStep;
      FrameMarkEnd("SimStep");
    }
  }
}

void simulator_step(double dt) {
  ZoneScoped;

  sWorld.calcUnconstrainedVelUpdate(dt);
  sWorld.updateCollisionShapes();
  sWorld.detectCollisions();

  sWorld.resolveCollisions(dt, sSolverNumIter);
  sWorld.integrateWorld(dt);

  {
    ZoneScoped;
    sthstry_store(sStateHistory);
    sStateHistoryCurrent = sthstry_get_num_states(sStateHistory);
  }
  sSimTime += dt;
}

void simulator_draw(srcmdbuf* cmdbuf) {
  srcmd rcmd;
  srcmd_clear(&rcmd);

  // Ground Plane
  Vector3f ground_pos = sGroundShape.pos;

  rcmd.type = SRndrCmdTypeGrid;
  for (int i = -1; i <= 1; i++) {
    for (int j = -1; j <= 1; j++) {
      simd4x4f_translation(
          &rcmd.mat,
          ground_pos[0] + i * 10.0f,
          ground_pos[1],
          ground_pos[2] + j * 10.f);
      srcmdbuf_add(cmdbuf, &rcmd);
    }
  }

  // World Bodies
  for (int i = 0; i < sWorld.mBodies.size(); i++) {
    const SimBody& body = sWorld.mBodies[i];
    for (int j = 0; j < body.mCollisionShapes.size(); j++) {
      const SimBody::BodyCollisionInfo& cinfo = body.mCollisionShapes[j];
      switch (cinfo.second.mType) {
        case SimShape::Box:
          rcmd.type = SRndrCmdTypeCube;
          break;
        case SimShape::Sphere:
          rcmd.type = SRndrCmdTypeSphere;
          break;
        default:
          gLog("Error: cannot render shape of type %d", cinfo.second.mType);
      }

      simd4x4f trans = simd4x4f_create(
          // clang-format off
          simd4f_create(1.f, 0.f, 0.f, 0.f),
          simd4f_create(0.f, 1.f, 0.f, 0.f),
          simd4f_create(0.f, 0.f, 1.f, 0.f),
          simd4f_create(cinfo.second.pos[0], cinfo.second.pos[1], cinfo.second.pos[2], 1.f)
          // clang-format on
      );

      simd4x4f scale = simd4x4f_create(
          // clang-format off
          simd4f_create(cinfo.second.scale[0], 0.f, 0.f, 0.f),
          simd4f_create(0.f, cinfo.second.scale[1], 0.f, 0.f),
          simd4f_create(0.f, 0.f, cinfo.second.scale[2], 0.f),
          simd4f_create(0.f, 0.f, 0.f, 1.f)
          // clang-format on
      );

      Matrix3d R = cinfo.second.orientation.toMatrix();
      simd4x4f rot = simd4x4f_create(
          // clang-format off
          simd4f_create(R(0,0), R(0,1), R(0,2), 0.f),
          simd4f_create(R(1,0), R(1,1), R(1,2), 0.f),
          simd4f_create(R(2,0), R(2,1), R(2,2), 0.f),
          simd4f_create(0.f, 0.f, 0.f, 1.f)
          // clang-format on
      );

      simd4x4f rot_scale;
      simd4x4f_matrix_mul(&rot, &scale, &rot_scale);
      simd4x4f_matrix_mul(&trans, &rot_scale, &rcmd.mat);
      srcmdbuf_add(cmdbuf, &rcmd);
    }
  }
}