#include "cell.h"
#include "prosperon.h"

#include <math.h>
#include <stdlib.h>
#include "HandmadeMath.h"
#include "cell.h"


// Utility function to get number from array index
static double js_getnum_uint32(JSContext *js, JSValue v, unsigned int i)
{
  JSValue val = JS_GetPropertyNumber(js,v,i);
  double ret = js2number(js, val);
  JS_FreeValue(js,val);
  return ret;
}

// Convert JS array to float array
static float *js2floats(JSContext *js, JSValue v, size_t *len)
{
  *len = JS_ArrayLength(js,v);
  float *arr = malloc(sizeof(float)* *len);
  for (int i = 0; i < *len; i++)
    arr[i] = js_getnum_uint32(js,v,i);
  return arr;
}

// Convert float array to JS array
static JSValue floats2array(JSContext *js, float *vals, size_t len) {
  JSValue arr = JS_NewArray(js);
  for (size_t i = 0; i < len; i++) {
    JS_SetPropertyNumber(js, arr, i, number2js(js, vals[i]));
  }
  return arr;
}

// Calculate vector length
static double arr_vec_length(JSContext *js,JSValue v)
{
  int len = JS_ArrayLength(js,v);
  switch(len) {
    case 2: return HMM_LenV2(js2vec2(js,v));
    case 3: return HMM_LenV3(js2vec3(js,v));
    case 4: return HMM_LenV4(js2vec4(js,v));
  }

  double sum = 0;
  for (int i = 0; i < len; i++)
    sum += pow(js_getnum_uint32(js, v, i), 2);

  return sqrt(sum);
}


// GCD helper function
static int gcd(int a, int b) {
    if (b == 0)
        return a;
    return gcd(b, a % b);
}

// MATH FUNCTIONS

// Rotate a 2D point (or array of length 2) by the given angle (in turns) around an optional pivot.
JSC_CCALL(math_rotate,
  HMM_Vec2 vec = js2vec2(js,argv[0]);
  double angle = js2angle(js, argv[1]);
  HMM_Vec2 pivot = JS_IsNull(argv[2]) ? v2zero : js2vec2(js,argv[2]);  
  vec = HMM_SubV2(vec,pivot);

  float r = HMM_LenV2(vec);
  angle += atan2f(vec.y, vec.x);

  vec.x = r * cosf(angle);
  vec.y = r * sinf(angle);

  vec = HMM_AddV2(vec,pivot);

  return vec22js(js, vec);
)

// Return a normalized copy of the given numeric array. For 2D/3D/4D or arbitrary length.
JSC_CCALL(math_norm,
  int len = JS_ArrayLength(js,argv[0]);

  switch(len) {
    case 2: return vec22js(js,HMM_NormV2(js2vec2(js,argv[0])));
    case 3: return vec32js(js, HMM_NormV3(js2vec3(js,argv[0])));
    case 4: return vec42js(js,HMM_NormV4(js2vec4(js,argv[0])));
  }

  double length = arr_vec_length(js,argv[0]);
  JSValue newarr = JS_NewArray(js);

  for (int i = 0; i < len; i++)
    JS_SetPropertyNumber(js, newarr, i, number2js(js,js_getnum_uint32(js, argv[0],i)/length));

  ret = newarr;
)

// Compute the angle between two vectors (2D/3D/4D).
JSC_CCALL(math_angle_between,
  int len = JS_ArrayLength(js,argv[0]);
  switch(len) {
    case 2: return angle2js(js,HMM_AngleV2(js2vec2(js,argv[0]), js2vec2(js,argv[1])));
    case 3: return angle2js(js,HMM_AngleV3(js2vec3(js,argv[0]), js2vec3(js,argv[1])));
    case 4: return angle2js(js,HMM_AngleV4(js2vec4(js,argv[0]), js2vec4(js,argv[1])));
  }
  return JS_ThrowReferenceError(js, "Input array must have a length between 2 and 4.");
)

// Linear interpolation between two numbers: lerp(a, b, t).
JSC_CCALL(math_lerp,
  double s = js2number(js,argv[0]);
  double f = js2number(js,argv[1]);
  double t = js2number(js,argv[2]);
  
  ret = number2js(js,(f-s)*t+s);
)

// Compute the greatest common divisor of two integers.
JSC_CCALL(math_gcd, ret = number2js(js,gcd(js2number(js,argv[0]), js2number(js,argv[1]))); )

// Compute the least common multiple of two integers.
JSC_CCALL(math_lcm,
  double a = js2number(js,argv[0]);
  double b = js2number(js,argv[1]);
  ret = number2js(js,(a*b)/gcd(a,b));
)

// Clamp a number between low and high. clamp(value, low, high).
JSC_CCALL(math_clamp,
  double x = js2number(js,argv[0]);
  double l = js2number(js,argv[1]);
  double h = js2number(js,argv[2]);
  return number2js(js,x > h ? h : x < l ? l : x);
)

// Compute the signed distance between two angles in 'turn' units, e.g. 0..1 range.
JSC_CCALL(math_angledist,
  double a1 = js2number(js,argv[0]);
  double a2 = js2number(js,argv[1]);
  a1 = fmod(a1,1);
  a2 = fmod(a2,1);
  double dist = a2-a1;
  if (dist == 0) return number2js(js,dist);
  if (dist > 0) {
    if (dist > 0.5) return number2js(js,dist-1);
    return number2js(js,dist);
  }
  
  if (dist < -0.5) return number2js(js,dist+1);
  
  return number2js(js,dist);
)

// Apply a random +/- percentage noise to a number. Example: jitter(100, 0.05) -> ~95..105.
JSC_CCALL(math_jitter,
  double n = js2number(js,argv[0]);
  double pct = js2number(js,argv[1]);
  return JS_NULL;
  
//  return number2js(js,n + (rand_range(js,-pct,pct)*n));
)

// Compute the arithmetic mean of an array of numbers.
JSC_CCALL(math_mean,
  double len =  JS_ArrayLength(js,argv[0]);
  double sum = 0;
  for (int i = 0; i < len; i++)
    sum += js_getnum_uint32(js, argv[0], i);
    
  return number2js(js,sum/len);
)

// Sum all elements of an array of numbers.
JSC_CCALL(math_sum,
  double sum = 0.0;
  int len = JS_ArrayLength(js,argv[0]);
  for (int i = 0; i < len; i++)
    sum += js_getnum_uint32(js, argv[0], i);
  
  return number2js(js,sum);
)

// Compute standard deviation of an array of numbers.
JSC_CCALL(math_sigma,
  int len = JS_ArrayLength(js,argv[0]);
  double sum = 0;
  for (int i = 0; i < len; i++)
    sum += js_getnum_uint32(js, argv[0], i);
  
  double mean = sum/(double)len;
  sum = 0;
  for (int i = 0; i < len; i++)
    sum += pow(js_getnum_uint32(js, argv[0], i) - mean, 2);
  
  return number2js(js,sqrt(sum/len));
)

// Compute the median of an array of numbers.
JSC_CCALL(math_median,
  int len = JS_ArrayLength(js,argv[0]);
  double vals[len];
  for (int i = 0; i < len; i++)
    vals[i] = js_getnum_uint32(js, argv[0], i);
  
  // Simple bubble sort for median calculation
  for (int i = 0; i < len-1; i++) {
    for (int j = 0; j < len-i-1; j++) {
      if (vals[j] > vals[j+1]) {
        double temp = vals[j];
        vals[j] = vals[j+1];
        vals[j+1] = temp;
      }
    }
  }
  
  if (len % 2 == 0)
    return number2js(js,(vals[len/2-1] + vals[len/2]) / 2.0);
  else
    return number2js(js,vals[len/2]);
)

// Return the length of a vector (i.e. sqrt of sum of squares).
JSC_CCALL(math_length, return number2js(js,arr_vec_length(js,argv[0])); )

// Return an array of points from a start to an end, spaced out by a certain distance.
JSC_CCALL(math_from_to,
  double start = js2number(js,argv[0]);
  double end = js2number(js,argv[1]);
  double step = js2number(js,argv[2]);
  int inclusive = JS_ToBool(js,argv[3]);
  int arr = JS_ToBool(js,argv[4]);
  
  JSValue jsarr = JS_NewArray(js);
  int i = 0;
  for (double val = start; val <= end; val += step) {
    if (val == end && !inclusive) break;
    JS_SetPropertyNumber(js, jsarr, i++, number2js(js, val));
  }
  
  return jsarr;
)

// Compute the dot product between two numeric arrays, returning a scalar. Extra elements are ignored.
JSC_CCALL(math_dot,
  size_t alen, blen;
  float *a = js2floats(js,argv[0], &alen);
  float *b = js2floats(js,argv[1], &blen);
  float dot = 0;
  size_t len = alen < blen? alen : blen;
  for (size_t i = 0; i < len; i++)
    dot += a[i] * b[i];
    
  free(a);
  free(b);
  return number2js(js,dot);
)

// Project one vector onto another, returning a new array of the same dimension.
JSC_CCALL(math_project,
  size_t alen, blen;
  float *a = js2floats(js, argv[0], &alen);
  float *b = js2floats(js, argv[1], &blen);

  if (!a || !b) {
    free(a);
    free(b);
    return JS_NULL;
  }

  // We'll work up to the smaller length
  size_t len = (alen < blen) ? alen : blen;
  if (len == 0) {
    free(a);
    free(b);
    return JS_NULL;
  }

  // Compute dot products: a·b and b·b
  float ab = 0, bb = 0;
  for (size_t i = 0; i < len; i++) {
    ab += a[i] * b[i];
    bb += b[i] * b[i];
  }

  // Build the result array
  float *proj = (float*)malloc(sizeof(float) * len);
  if (!proj) {
    free(a);
    free(b);
    return JS_EXCEPTION; // or some error
  }

  float scale = (bb != 0.0f) ? (ab / bb) : 0.0f;
  for (size_t i = 0; i < len; i++)
    proj[i] = scale * b[i];

  ret = floats2array(js, proj, len);

  free(a);
  free(b);
  free(proj);
)

// Compute the midpoint of two arrays of numbers. Only the first two entries are used if 2D is intended.
JSC_CCALL(math_midpoint,
  size_t alen, blen;
  float *a = js2floats(js, argv[0], &alen);
  float *b = js2floats(js, argv[1], &blen);

  if (!a || !b) {
    free(a);
    free(b);
    return JS_NULL;
  }

  size_t len = (alen < blen) ? alen : blen;
  if (len == 0) {
    free(a);
    free(b);
    return JS_NULL;
  }

  float *m = (float*)malloc(sizeof(float) * len);
  if (!m) {
    free(a);
    free(b);
    return JS_EXCEPTION;
  }

  for (size_t i = 0; i < len; i++)
    m[i] = (a[i] + b[i]) * 0.5f;

  ret = floats2array(js, m, len);

  free(a);
  free(b);
  free(m);
)

// Reflect a vector across a plane normal. Both arguments must be numeric arrays.
JSC_CCALL(math_reflect,
  size_t alen, blen;
  float *a = js2floats(js, argv[0], &alen);
  float *b = js2floats(js, argv[1], &blen);
  
  if (!a || !b) {
    free(a);
    free(b);
    return JS_NULL;
  }
  
  size_t len = (alen < blen) ? alen : blen;
  if (len == 0) {
    free(a);
    free(b);
    return JS_NULL;
  }
  
  // Reflect vector a across normal b
  // reflection = a - 2 * (a·b) * b
  float ab = 0, bb = 0;
  for (size_t i = 0; i < len; i++) {
    ab += a[i] * b[i];
    bb += b[i] * b[i];
  }
  
  float *result = (float*)malloc(sizeof(float) * len);
  if (!result) {
    free(a);
    free(b);
    return JS_EXCEPTION;
  }
  
  float scale = (bb != 0.0f) ? (2.0f * ab / bb) : 0.0f;
  for (size_t i = 0; i < len; i++)
    result[i] = a[i] - scale * b[i];
  
  ret = floats2array(js, result, len);
  
  free(a);
  free(b);
  free(result);
)

// Compute the normalized direction vector from the first array to the second.
JSC_CCALL(math_direction,
  size_t alen, blen;
  float *a = js2floats(js, argv[0], &alen);
  float *b = js2floats(js, argv[1], &blen);
  
  if (!a || !b) {
    free(a);
    free(b);
    return JS_NULL;
  }
  
  size_t len = (alen < blen) ? alen : blen;
  if (len == 0) {
    free(a);
    free(b);
    return JS_NULL;
  }
  
  // Direction vector from a to b (normalized)
  float *dir = (float*)malloc(sizeof(float) * len);
  if (!dir) {
    free(a);
    free(b);
    return JS_EXCEPTION;
  }
  
  float mag = 0;
  for (size_t i = 0; i < len; i++) {
    dir[i] = b[i] - a[i];
    mag += dir[i] * dir[i];
  }
  
  mag = sqrtf(mag);
  if (mag > 0.0f) {
    for (size_t i = 0; i < len; i++)
      dir[i] /= mag;
  }
  
  ret = floats2array(js, dir, len);
  
  free(a);
  free(b);
  free(dir);
)

// Given a 2D vector, return its angle from the X-axis in radians or some chosen units.
JSC_CCALL(math_angle,
  size_t len;
  float *v = js2floats(js, argv[0], &len);
  
  if (!v || len < 2) {
    free(v);
    return JS_NULL;
  }
  
  // Return angle in radians for 2D vector
  ret = number2js(js, atan2f(v[1], v[0]));
  
  free(v);
)

// Compute the Euclidean distance between two numeric arrays of matching length.
JSC_CCALL(math_distance,
  size_t alen, blen;
  float *a = js2floats(js, argv[0], &alen);
  float *b = js2floats(js, argv[1], &blen);

  if (!a || !b) {
    free(a);
    free(b);
    return JS_NULL;
  }

  size_t len = (alen < blen) ? alen : blen;
  if (len == 0) {
    free(a);
    free(b);
    return JS_NULL;
  }

  float distSq = 0.0f;
  for (size_t i = 0; i < len; i++) {
    float diff = b[i] - a[i];
    distSq += diff * diff;
  }
  float dist = sqrtf(distSq);

  free(a);
  free(b);
  return number2js(js, dist);
)

static const JSCFunctionListEntry js_math_funcs[] = {
  MIST_FUNC_DEF(math, dot,2),
  MIST_FUNC_DEF(math, project,2),
  MIST_FUNC_DEF(math, rotate, 3),
  MIST_FUNC_DEF(math, midpoint, 2),
  MIST_FUNC_DEF(math, reflect, 2),
  MIST_FUNC_DEF(math, distance, 2),
  MIST_FUNC_DEF(math, direction, 2),
  MIST_FUNC_DEF(math, angle, 1),
  MIST_FUNC_DEF(math, norm, 1),
  MIST_FUNC_DEF(math, angle_between, 2),
  MIST_FUNC_DEF(math, lerp, 3),
  MIST_FUNC_DEF(math, gcd, 2),
  MIST_FUNC_DEF(math, lcm, 2),
  MIST_FUNC_DEF(math, clamp, 3),
  MIST_FUNC_DEF(math, angledist, 2),
  MIST_FUNC_DEF(math, jitter, 2),
  MIST_FUNC_DEF(math, mean, 1),
  MIST_FUNC_DEF(math, sum, 1),
  MIST_FUNC_DEF(math, sigma, 1),
  MIST_FUNC_DEF(math, median, 1),
  MIST_FUNC_DEF(math, length, 1),
  MIST_FUNC_DEF(math, from_to, 5),
};

CELL_USE_FUNCS(js_math_funcs)