/*
    Fixed point arithmetic.
    https://github.com/howerj/q/blob/master/q.c
*/

#include "fixed.h"
#include <assert.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <tgmath.h>

// 4-bit LUT (16 entries) for the normalized range [0.5, 2.0)
// It stores the initial guess scaled to Q16.16.
static const uint32_t com_fixed_sqrt_lut[16] = {
    46340, 49547, 52521, 55314, 57954, 60464, 62862, 65161,
    67373, 69508, 71572, 73572, 75514, 77402, 79240, 81033};

/* Used for other quadrants as well as cosine eval, all from the same table. */
/* Only fractional part is present, reducing the cache footprint. */
const uint16_t com_fixed_sin_lut[128] = {
    0,     804,   1608,  2412,  3215,  4018,  4821,  5622,  6423,  7223,  8022,
    8819,  9616,  10410, 11204, 11995, 12785, 13573, 14359, 15142, 15923, 16702,
    17479, 18253, 19024, 19792, 20557, 21319, 22078, 22833, 23586, 24334, 25079,
    25820, 26557, 27291, 28020, 28745, 29465, 30181, 30893, 31600, 32302, 32999,
    33692, 34379, 35061, 35738, 36409, 37075, 37736, 38390, 39039, 39682, 40319,
    40950, 41575, 42194, 42806, 43412, 44011, 44603, 45189, 45768, 46340, 46906,
    47464, 48015, 48558, 49095, 49624, 50146, 50660, 51166, 51665, 52155, 52639,
    53114, 53581, 54040, 54491, 54933, 55368, 55794, 56212, 56621, 57022, 57414,
    57797, 58172, 58538, 58895, 59243, 59583, 59913, 60235, 60547, 60850, 61144,
    61429, 61705, 61971, 62228, 62475, 62714, 62942, 63162, 63371, 63571, 63762,
    63943, 64115, 64276, 64428, 64571, 64703, 64826, 64939, 65043, 65136, 65220,
    65294, 65358, 65412, 65457, 65491, 65516, 65531};

com_fixed_t com_fixed_sqrt(com_fixed_t a) {
  // 1. Handle sign and edge cases
  assert(a >= 0);
  if (a == 0)
    return 0;

  // 2. Normalize input to the range [0.5, 2.0) to maximize LUT precision
  // clz = count leading zeros. On modern hardware, use __builtin_clz
  uint32_t leading_zeros = __builtin_clz(a);

  // Calculate how much we need to shift to place the highest bit properly
  // We want the value to land squarely within an optimal window
  int32_t shift = (31 - leading_zeros) - COM_FIXED_FRACBITS;

  // Normalize shift to always be even so we can cleanly pull out 2^(shift/2)
  if (shift & 1)
    shift -= 1;

  uint32_t normalized_a;
  if (shift > 0) {
    normalized_a = a >> shift;
  } else {
    normalized_a = a << (-shift);
  }

  // 3. LUT Lookup using 4 MSBs of the normalized value
  // Extracted index corresponds to the interval [0.5, 2.0)
  uint32_t lut_index = (normalized_a >> (COM_FIXED_FRACBITS - 3)) & 0xF;
  uint64_t x = com_fixed_sqrt_lut[lut_index];

  // 4. Newton-Raphson Iterations: x = 0.5 * (x + normalized_a / x)
  // We upscale to 64-bit to prevent intermediate overflow during division
  x = (x + ((uint64_t)normalized_a << COM_FIXED_FRACBITS) / x) >>
      1; // Iteration 1
  x = (x + ((uint64_t)normalized_a << COM_FIXED_FRACBITS) / x) >>
      1; // Iteration 2

  // 5. Denormalize back to the target scale: result = x * 2^(shift / 2)
  int32_t final_shift = shift / 2;
  if (final_shift > 0) {
    return (int32_t)(x << final_shift);
  } else {
    return (int32_t)(x >> (-final_shift));
  }
}

void com_fixed_print(com_fixed_t a) {
  if (a < 0) {
    printf("-");
    a = -a;
  }
  int32_t int_part = a >> 16;
  int32_t frac_part = a & 0xFFFF;

  // Convert fractional 16-bit part to a decimal value (up to 4 decimal places)
  uint32_t decimal_val = (frac_part * 10000) >> 16;

  printf("%d.%04u\n", int_part, decimal_val);
}

void com_fixed_run_tests(void) {
  com_fixed_print(com_fixed_sqrt(COM_FIXED_FRACUNIT * 2));
  com_fixed_print(com_fixed_sqrt(COM_FIXED_FRACUNIT * 3));
  com_fixed_print(com_fixed_sqrt(COM_FIXED_FRACUNIT * 4));
  com_fixed_print(com_fixed_sqrt(COM_FIXED_FRACUNIT * 200));
  com_fixed_print(com_fixed_sqrt(COM_FIXED_FRACUNIT * 1000));

  com_fixed_print(COM_FIXED_PI);
  com_fixed_print(com_fixed_div(COM_FIXED_PI + 234, COM_FIXED_PI2) << 2);

  double max_sin_deviation = 0.0f;
  com_fixed_t sin_accumulator = -COM_FIXED_PI * 4;
  com_fixed_t sin_test_step = (COM_FIXED_PI << 1) >> 10;
  for (int i = 0; i < 1024 * 32; ++i) {
    com_fixed_t sin = com_fixed_sin(sin_accumulator);
    double deviation = fabs(sin(com_fixed_as_float(sin_accumulator)) -
                            com_fixed_as_float(sin));
    if (deviation > max_sin_deviation)
      max_sin_deviation = deviation;
    sin_accumulator += sin_test_step;
  }
  printf("max sin deviation: %f\n", max_sin_deviation);

  double max_cos_deviation = 0.0f;
  com_fixed_t cos_accumulator = -COM_FIXED_PI * 4;
  com_fixed_t cos_test_step = (COM_FIXED_PI << 1) >> 10;
  for (int i = 0; i < 1024 * 32; ++i) {
    com_fixed_t cos = com_fixed_cos(cos_accumulator);
    double deviation = fabs(cos(com_fixed_as_float(cos_accumulator)) -
                            com_fixed_as_float(cos));
    if (deviation > max_cos_deviation)
      max_cos_deviation = deviation;
    cos_accumulator += cos_test_step;
  }
  printf("max cos deviation: %f\n", max_cos_deviation);
}