#include <assert.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#define PDF_MAX_CAP       1024
#define PDF_MAX_ELM_CNT   2048
#define PDF_MAX_MIP_LVL   11

struct pdf_tree_elm {
  int val;
  float prob; 
}
struct pdf_tree {
  int vals[PDF_MAX_CAP];
  int act_cnt;
  float prob[PDF_MAX_ELM_CNT];
  int mip_off[PDF_MAX_MIP_LVL];
  int mip_siz[PDF_MAX_MIP_LVL];
  int cap;
  int mip_cnt;
};
static inline int
align_up(int v, int a) {
  return (v + a - 1) & ~(a - 1);
}
// Initialize the PdfTree
static void
pdf_init(struct pdf_tree *pdf, int cap) {
  assert(cap <= MAX_CAPACITY);
  pdf->act_cnt = 0;
  pdf->cap = cap;

  // Calculate mip levels
  int mip_cnt = 1;
  int mip_elm_cnt = cap;
  int elm_cnt = mip_elm_cnt;
  while (mip_elm_cnt > 1) {
    mip_elm_cnt = align_up(mip_elm_cnt, 2) / 2;
    elm_cnt += mip_elm_cnt;
    mip_cnt++;
  }
  assert(elm_cnt <= MAX_TOTAL_ELEMENTS);
  assert(mip_cnt <= MAX_MIP_LEVELS);
  pdf->mip_cnt = mip_cnt;

  memset(pdf->vals, 0, sizeof(pdf->vals));
  memset(pdf->prob, 0, sizeof(pdf->prob));
  memset(pdf->mip_off, 0, sizeof(pdf->mip_off));
  memset(pdf->mip_siz, 0, sizeof(pdf->mip_siz));
}
static void
pdf_fill(struct pdf_tree* pdf, const struct pdf_tree_elm *elm, int cnt) {
  assert(cnt <= pdf->cap);
  for (int i = 0; i < cnt; i++) {
    pdf->vals[i] = elm[i].value;
  }
  // Reset prob and mip data
  memset(pdf->prob, 0, sizeof(pdf->prob));
  memset(pdf->mip_off, 0, sizeof(pdf->mip_off));
  memset(pdf->mip_siz, 0, sizeof(pdf->mip_siz));

  // fill mip 0 (base level)
  pdf->act_cnt = 0;
  for (int i = 0; i < cnt; i++) {
    float prob = elm[i].prob > 0.0f ? elm[i].prob : 0.0f;
    pdf->prob[i] = prob;
    if (prob > 0.0f) {
      pdf->act_cnt++;
    }
  }
  pdf->mip_off[0] = 0;
  pdf->mip_siz[0] = cnt;

  // Build higher mip levels
  int mip_lvl = 1;
  int mip_elm_cnt = cnt;
  while (mip_elm_cnt > 1) {
    mip_elm_cnt = align_up(mip_elm_cnt, 2) / 2;
    pdf->mip_off[mip_lvl] = pdf->mip_off[mip_lvl - 1] + pdf->mip_siz[mip_lvl - 1];
    pdf->mip_siz[mip_lvl] = mip_elm_cnt;

    for (int e = 0; e < mip_elm_cnt; e++) {
      int child_idx = 2 * e;
      int child_prob_idx = pdf->mip_off[mip_lvl - 1] + child_idx;
      float prob_sum = pdf->prob[child_prob_idx];
      if (child_idx + 1 < pdf->mip_siz[mip_lvl - 1]) {
        prob_sum += pdf->prob[child_prob_idx + 1];
      }
      pdf->prob[pdf->mip_off[mip_lvl] + e] = prob_sum;
    }
    mip_lvl++;
  }
}
static void
pdf_del(struct pdf_tree* pdf, int index) {
  assert(pdf->act_cnt > 0);
  pdf->prob[index] = 0.0f;

  // Update higher mip levels
  int mip_lvl = 1;
  int e = index / 2;
  while (mip_lvl < pdf->mip_cnt) {
    int child_idx = 2 * e;
    int child_prob_idx = pdf->mip_off[mip_lvl - 1] + child_idx;
    float probabilitySum = pdf->prob[child_prob_idx];
    if (child_idx + 1 < pdf->mip_siz[mip_lvl - 1]) {
        probabilitySum += pdf->prob[child_prob_idx + 1];
    }
    pdf->prob[pdf->mip_off[mip_lvl] + e] = probabilitySum;
    mip_lvl++;
    e /= 2;
  }
  pdf->act_cnt--;
}
static inline float
getUniformFloat(float min, float max) {
  return min + (float)rand() / RAND_MAX * (max - min);
}
static int
pdf_gen(struct pdf_tree *tree) {
  assert(pdf->mip_cnt > 0);
  assert(pdf->act_cnt > 0);
  assert(pdf->mip_siz[pdf->mip_cnt - 1] == 1);

  float tot_prob = pdf->prob[pdf->mip_off[pdf->mip_cnt - 1]];
  float randomValue = getUniformFloat(0.0f, tot_prob);

  int e = 0;
  int mip_lvl = pdf->mip_cnt - 1;
  float cur_prob = randomValue;
  while (mip_lvl > 0) {
    mip_lvl--;

    int mip_off = pdf->mip_off[mip_lvl];
    int mip_siz = pdf->mip_siz[mip_lvl];
    int nxt_mip_e = 2 * e;
    assert(nxt_mip_e < mip_siz);

    float mip_prob = pdf->prob[mip_off + nxt_mip_e];
    if (cur_prob > mip_prob && nxt_mip_e + 1 < mip_siz) {
      nxt_mip_e++;
      cur_prob -= mip_prob;
    }
    e = nxt_mip_e;
  }
  pdf_del(pdf, e);
  return pdf->vals[e];
}

// Example usage
extern int
main(void) {
  struct pdf_tree tree;
  pdf_init(&tree, 4);
  Element elm[] = {
    {0, 1.0f},
    {1, 2.0f},
    {2, 3.0f},
    {3, 4.0f}
  };
  pdf_fill(&tree, elm, 4);

  srand(1234);  // Seed RNG for reproducibility
  for (int i = 0; i < 4; i++) {
    int value = pdf_gen(&tree);
    printf("Chosen value: %u\n", value);
  }
  return 0;
}