#include "bitset.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include "util.h"
struct bitset {
  size_t nwords;     /* allocated */
  uint64_t *words;   /* array of 64-bit chunks */
};

#define BITWORD(i) ((i) >> 6)       /* divide by 64 */
#define BITMASK(i) (1ULL << ((i) & 63))

bitset_t *bitset_alloc(void) {
  bitset_t *bs = calloc(1, sizeof(bitset_t));
  if (bs == NULL) {
    die("Out of memory: cannot allocate bitset");
  }
  return bs;
}

void bitset_free(bitset_t *bs) {
  if (!bs) return;
  free(bs->words);
  free(bs);
}

void bitset_reserve(bitset_t *bs, size_t maxbit) {
  size_t need_words = BITWORD(maxbit) + 1;
  if (need_words <= bs->nwords)
    return;

  size_t new_words = bs->nwords ? bs->nwords : 1;
  while (new_words < need_words)
    new_words *= 2;

  uint64_t *new = calloc(new_words, sizeof(uint64_t));
  if (new == NULL) {
    die("Out of memory: cannot enlarge bitset");
  }
  if (bs->words) {
    memcpy(new, bs->words, bs->nwords * sizeof(uint64_t));
    free(bs->words);
  }
  bs->words = new;
  bs->nwords = new_words;
}

void bitset_set(bitset_t *bs, size_t bit) {
  bitset_reserve(bs, bit);
  bs->words[BITWORD(bit)] |= BITMASK(bit);
}

void bitset_clear(bitset_t *bs, size_t bit) {
  size_t w = BITWORD(bit);
  if (w >= bs->nwords)
    return;
  bs->words[w] &= ~BITMASK(bit);
}

bool bitset_get(const bitset_t *bs, size_t bit) {
  size_t w = BITWORD(bit);
  if (w >= bs->nwords)
    return false;
  return (bs->words[w] & BITMASK(bit)) != 0;
}

void bitset_reset(bitset_t *bs) {
  memset(bs->words, 0, bs->nwords * sizeof(uint64_t));
}

size_t bitset_count(const bitset_t *bs) {
  size_t count = 0;
  for (size_t i = 0; i < bs->nwords; i++)
    count += __builtin_popcountll(bs->words[i]);
  return count;
}


size_t bitset_max(const bitset_t *bs) {
  for (size_t w = bs->nwords; w > 0; w--) {
    uint64_t x = bs->words[w - 1];
    if (x != 0) {
      return ((w - 1) << 6)
        + (63 - __builtin_clzll(x));
    }
  }
  return 0;
}

void bitset_shift_up(bitset_t *bs, size_t n) {
  if (n == 0) return;

  size_t word_shift = n / 64;
  size_t bit_shift  = n % 64;

  size_t old_max = bitset_max(bs);
  bitset_reserve(bs, old_max + n - 1); // ensure enough space

  size_t old_nwords = bs->nwords;

  // shift words from end to start to avoid overwriting
  for (ssize_t i = old_nwords - 1; i >= 0; i--) {
    uint64_t w = bs->words[i];
    uint64_t upper = (bit_shift != 0 && i > 0) ? bs->words[i - 1] >> (64 - bit_shift) : 0;
    bs->words[i + word_shift] = (w << bit_shift) | upper;
  }

  // zero out lower words
  for (size_t i = 0; i < word_shift; i++)
    bs->words[i] = 0;
}

void bitset_shift_down(bitset_t *bs, size_t n) {
  if (n == 0) return;

  size_t word_shift = n / 64;
  size_t bit_shift  = n % 64;

  if (word_shift >= bs->nwords) {
    // all bits are discarded
    bitset_reset(bs);
    return;
  }

  // shift words from start to end
  for (size_t i = 0; i + word_shift < bs->nwords; i++) {
    uint64_t w = bs->words[i + word_shift];
    uint64_t lower = (bit_shift != 0 && i + word_shift + 1 < bs->nwords)
                        ? bs->words[i + word_shift + 1] << (64 - bit_shift)
                        : 0;
    bs->words[i] = (w >> bit_shift) | lower;
  }

  // zero out top words
  for (size_t i = bs->nwords - word_shift; i < bs->nwords; i++)
    bs->words[i] = 0;
}

void bitset_union_with_offset(bitset_t *bs1, const bitset_t *bs2, size_t skip) {
  size_t total_bits = bs2->nwords * 64;
  if (skip >= total_bits) return; // nothing to merge

  // ensure bs1 can hold all bits
  bitset_reserve(bs1, total_bits - skip - 1);

  size_t start_word = skip / 64;
  size_t start_bit  = skip % 64;

  uint64_t carry = 0;
  for (size_t i = start_word; i < bs2->nwords; i++) {
    uint64_t w = bs2->words[i];
    if (i == start_word && start_bit != 0) {
      // mask out the first 'start_bit' bits in the first word
      w &= ~((1ULL << start_bit) - 1);
    }

    uint64_t new_w = (w >> start_bit) | carry;
    carry = (start_bit != 0) ? (w << (64 - start_bit)) : 0;

    size_t dest_word = i - start_word;
    bs1->words[dest_word] |= new_w;
  }

  if (carry) {
    size_t dest_word = bs2->nwords - start_word;
    if (dest_word >= bs1->nwords) {
      bitset_reserve(bs1, dest_word * 64 + 63);
    }
    bs1->words[dest_word] |= carry;
  }
}

void bitset_union(bitset_t *bs1, const bitset_t *bs2) {
  if (bs2->nwords > bs1->nwords)
    bitset_reserve(bs1, (bs2->nwords * 64) - 1);

  for (size_t i = 0; i < bs2->nwords; i++)
    bs1->words[i] |= bs2->words[i];
}

void bitset_dump(const bitset_t *bs) {
  size_t start = 0;
  size_t bit;
  printf("{");
  bool isStart = 1;
  while ((bit = bitset_next(bs, &start)) != (size_t)-1) {
    if (!isStart)
      printf(", ");
    printf("%zu", bit);
    isStart = 0;
  }
  printf("}");
}

size_t bitset_next(const bitset_t *bs, size_t *start) {
  size_t i = *start;
  while (1) {
    size_t w = BITWORD(i);
    if (w >= bs->nwords) {
      *start = (size_t)-1;
      return (size_t)-1;
    }
    uint64_t word = bs->words[w];

    uint64_t mask = ~0ULL << (i & 63);
    uint64_t rem = word & mask;

    if (rem) {
      int bit = __builtin_ctzll(rem);
      size_t result = (w << 6) + bit;
      *start = result + 1;
      return result;
    }

    i = (w + 1) << 6;
  }
}