#include "fixed.h"
#include "bits.h"
#include <assert.h>
#include <stdint.h>
#include <stdio.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};

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
  int leading_zeros = com_bits_needed(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
  int 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);
}