// This file is part of the 64k demo project.
// It tests the mathematical utility functions.
// Verifies vector operations, matrix transformations, and interpolation.

#include "util/mini_math.h"
#include <cassert>
#include <cmath>
#include <iostream>
#include <vector>

// Checks if two floats are approximately equal
bool near(float a, float b, float e = 0.001f) {
  return std::abs(a - b) < e;
}

// Generic test runner for any vector type (vec2, vec3, vec4)
template <typename T> void test_vector_ops(int n) {
  T a, b;
  // Set values
  for (int i = 0; i < n; ++i) {
    a[i] = (float)(i + 1);
    b[i] = 10.0f;
  }

  // Add
  T c = a + b;
  for (int i = 0; i < n; ++i)
    assert(near(c[i], (float)(i + 1) + 10.0f));

  // Scale
  T s = a * 2.0f;
  for (int i = 0; i < n; ++i)
    assert(near(s[i], (float)(i + 1) * 2.0f));

  // Dot Product
  // vec3(1,2,3) . vec3(1,2,3) = 1+4+9 = 14
  float expected_dot = 0;
  for (int i = 0; i < n; ++i)
    expected_dot += a[i] * a[i];
  assert(near(T::dot(a, a), expected_dot));

  // Norm (Length)
  assert(near(a.norm(), std::sqrt(expected_dot)));

  // Normalize
  T n_vec = a.normalize();
  assert(near(n_vec.norm(), 1.0f));

  // Normalize zero vector
  T zero_vec = T(); // Default construct to zero
  T norm_zero = zero_vec.normalize();
  for(int i = 0; i < n; ++i) assert(near(norm_zero[i], 0.0f));

  // Lerp
  T l = lerp(a, b, 0.3f);
  for (int i = 0; i < n; ++i)
    assert(near(l[i], .7 * (i + 1) + .3 * 10.0f));
}

// Specific test for padding alignment in vec3
void test_vec3_special() {
  std::cout << "Testing vec3 alignment..." << std::endl;
  // Verify sizeof is 16 bytes (4 floats) due to padding for WebGPU
  assert(sizeof(vec3) == 16);

  vec3 v(1, 0, 0);
  vec3 v2(0, 1, 0);

  // Cross Product
  vec3 c = vec3::cross(v, v2);
  assert(near(c.x, 0) && near(c.y, 0) && near(c.z, 1));
}

// Tests quaternion rotation, look_at, and slerp
void test_quat() {
  std::cout << "Testing Quat..." << std::endl;

  // Rotation (Rodrigues)
  vec3 v(1, 0, 0);
  quat q = quat::from_axis({0, 1, 0}, 1.5708f); // 90 deg Y
  vec3 r = q.rotate(v);
  assert(near(r.x, 0) && near(r.z, -1));

  // Rotation edge cases: 0 deg, 180 deg, zero vector
  quat zero_rot = quat::from_axis({1,0,0}, 0.0f);
  vec3 rotated_zero = zero_rot.rotate(v);
  assert(near(rotated_zero.x, 1.0f)); // Original vector

  quat half_pi_rot = quat::from_axis({0,1,0}, 3.14159f); // 180 deg Y
  vec3 rotated_half_pi = half_pi_rot.rotate(v);
  assert(near(rotated_half_pi.x, -1.0f)); // Rotated 180 deg around Y

  vec3 zero_vec(0,0,0);
  vec3 rotated_zero_vec = q.rotate(zero_vec);
  assert(near(rotated_zero_vec.x, 0.0f) && near(rotated_zero_vec.y, 0.0f) && near(rotated_zero_vec.z, 0.0f));

  // Look At
  // Looking from origin to +X, with +Y as up.
  // The local forward vector (0,0,-1) should be transformed to (1,0,0)
  quat l = quat::look_at({0, 0, 0}, {10, 0, 0}, {0, 1, 0});
  vec3 f = l.rotate({0, 0, -1});
  assert(near(f.x, 1.0f) && near(f.y, 0.0f) && near(f.z, 0.0f));

  // Slerp Midpoint
  quat q1(0, 0, 0, 1);
  quat q2 = quat::from_axis({0, 1, 0}, 1.5708f); // 90 deg
  quat mid = slerp(q1, q2, 0.5f);                // 45 deg
  assert(near(mid.y, 0.3826f));                  // sin(pi/8)

  // Slerp edge cases
  quat slerp_mid_edge = slerp(q1, q2, 0.0f);
  assert(near(slerp_mid_edge.w, q1.w) && near(slerp_mid_edge.x, q1.x) && near(slerp_mid_edge.y, q1.y) && near(slerp_mid_edge.z, q1.z));
  slerp_mid_edge = slerp(q1, q2, 1.0f);
  assert(near(slerp_mid_edge.w, q2.w) && near(slerp_mid_edge.x, q2.x) && near(slerp_mid_edge.y, q2.y) && near(slerp_mid_edge.z, q2.z));

  // FromTo
  quat from_to_test = quat::from_to({1,0,0}, {0,1,0}); // 90 deg rotation around Z
  vec3 rotated = from_to_test.rotate({1,0,0});
  assert(near(rotated.y, 1.0f));
}

// Tests WebGPU specific matrices
void test_matrices() {
  std::cout << "Testing Matrices..." << std::endl;
  float n = 0.1f, f = 100.0f;
  mat4 p = mat4::perspective(0.785f, 1.0f, n, f);

  // Check WebGPU Z-range [0, 1]
  // Z_ndc = (m10 * Z_view + m14) / -Z_view
  float z_near = (p.m[10] * -n + p.m[14]) / n;
  float z_far = (p.m[10] * -f + p.m[14]) / f;
  assert(near(z_near, 0.0f));
  assert(near(z_far, 1.0f));

  // Test mat4::look_at
  vec3 eye(0, 0, 5);
  vec3 target(0, 0, 0);
  vec3 up(0, 1, 0);
  mat4 view = mat4::look_at(eye, target, up);
  // Point (0,0,0) in world should be at (0,0,-5) in view space
  assert(near(view.m[14], -5.0f));

  // Test matrix multiplication
  mat4 t = mat4::translate({1, 2, 3});
  mat4 s = mat4::scale({2, 2, 2});
  mat4 ts = t * s; // Scale then Translate (if applied to vector on right: M*v)

  // v = (1,1,1,1) -> scale(2,2,2) -> (2,2,2,1) -> translate(1,2,3) -> (3,4,5,1)
  vec4 v(1, 1, 1, 1);
  vec4 res = ts * v;
  assert(near(res.x, 3.0f));
  assert(near(res.y, 4.0f));
  assert(near(res.z, 5.0f));

  // Test Rotation
  // Rotate 90 deg around Z. (1,0,0) -> (0,1,0)
  mat4 r = mat4::rotate({0, 0, 1}, 1.570796f);
  vec4 v_rot = r * vec4(1, 0, 0, 1);
  assert(near(v_rot.x, 0.0f));
  assert(near(v_rot.y, 1.0f));
}

// Tests easing curves
void test_ease() {
  std::cout << "Testing Easing..." << std::endl;
  // Boundary tests
  assert(near(ease::out_cubic(0.0f), 0.0f));
  assert(near(ease::out_cubic(1.0f), 1.0f));
  assert(near(ease::in_out_quad(0.0f), 0.0f));
  assert(near(ease::in_out_quad(1.0f), 1.0f));
  assert(near(ease::out_expo(0.0f), 0.0f));
  assert(near(ease::out_expo(1.0f), 1.0f));

  // Midpoint/Logic tests
  assert(ease::out_cubic(0.5f) > 0.5f); // Out curves should exceed linear value early
  assert(near(ease::in_out_quad(0.5f), 0.5f)); // Symmetric curves hit 0.5 at 0.5
  assert(ease::out_expo(0.5f) > 0.5f); // Exponential out should be above linear
}

// Tests spring solver
void test_spring() {
  std::cout << "Testing Spring..." << std::endl;
  float p = 0, v = 0;
  // Simulate approx 1 sec with 0.5s smooth time
  for (int i = 0; i < 60; ++i)
    spring::solve(p, v, 10.0f, 0.5f, 0.016f);
  assert(p > 8.5f); // Should be close to 10 after 1 sec

  // Test convergence over longer period
  p = 0; v = 0;
  for (int i = 0; i < 200; ++i)
    spring::solve(p, v, 10.0f, 0.5f, 0.016f);
  assert(near(p, 10.0f, 0.1f)); // Should be very close to target

  // Test vector spring
  vec3 vp(0, 0, 0), vv(0, 0, 0), vt(10, 0, 0);
  spring::solve(vp, vv, vt, 0.5f, 0.016f * 60.0f); // 1 huge step approx
  assert(vp.x > 1.0f); // Should have moved significantly
}

// Verifies that a matrix is approximately the identity matrix
void check_identity(const mat4& m) {
  for (int i = 0; i < 16; ++i) {
    float expected = (i % 5 == 0) ? 1.0f : 0.0f;
    if (!near(m.m[i], expected, 0.005f)) {
      std::cerr << "Matrix not Identity at index " << i << ": got " << m.m[i]
                << " expected " << expected << std::endl;
      assert(false);
    }
  }
}

// Tests matrix inversion and transposition correctness
void test_matrix_inversion() {
  std::cout << "Testing Matrix Inversion..." << std::endl;

  // 1. Identity
  mat4 id;
  check_identity(id.inverse());

  // 2. Translation
  mat4 t = mat4::translate({10.0f, -5.0f, 2.0f});
  mat4 t_inv = t.inverse();
  check_identity(t * t_inv);
  check_identity(t_inv * t);

  // 3. Scale (non-uniform)
  mat4 s = mat4::scale({2.0f, 0.5f, 4.0f});
  mat4 s_inv = s.inverse();
  check_identity(s * s_inv);

  // 4. Rotation
  mat4 r =
      mat4::rotate({1.0f, 2.0f, 3.0f}, 0.785f); // 45 deg around complex axis
  mat4 r_inv = r.inverse();
  check_identity(r * r_inv);

  // 5. Complex Transform (TRS)
  mat4 trs = t * (r * s);
  mat4 trs_inv = trs.inverse();
  check_identity(trs * trs_inv);
  check_identity(trs_inv * trs);

  // 6. Transposition
  std::cout << "Testing Matrix Transposition..." << std::endl;
  mat4 trs_t = mat4::transpose(trs);
  mat4 trs_tt = mat4::transpose(trs_t);
  for (int i = 0; i < 16; ++i) {
    assert(near(trs.m[i], trs_tt.m[i]));
  }

  // 7. Manual "stress" matrix (some small values, some large)
  mat4 stress;
  stress.m[0] = 1.0f;
  stress.m[5] = 2.0f;
  stress.m[10] = 3.0f;
  stress.m[15] = 4.0f;
  stress.m[12] = 100.0f;
  stress.m[1] = 0.5f;
  mat4 stress_inv = stress.inverse();
  check_identity(stress * stress_inv);

  // 8. Test Singular Matrix
  mat4 singular_scale;
  singular_scale.m[5] = 0.0f; // Scale Y by zero, making it singular
  mat4 singular_inv = singular_scale.inverse();
  // The inverse of a singular matrix should be the identity matrix as per the implementation
  check_identity(singular_inv);
}

int main() {
  std::cout << "Testing vec2..." << std::endl;
  test_vector_ops<vec2>(2);

  std::cout << "Testing vec3..." << std::endl;
  test_vector_ops<vec3>(3);
  test_vec3_special();

  std::cout << "Testing vec4..." << std::endl;
  test_vector_ops<vec4>(4);

  test_quat();
  test_matrices();
  test_matrix_inversion(); // New tests
  test_ease();
  test_spring();

  std::cout << "--- ALL TESTS PASSED ---" << std::endl;
  return 0;
}