module bitfield

import math.bits

pub struct BitField {
mut:
    size  int
    field []u32
}

// helper functions
const slot_size = 32

// from_bytes converts a byte array into a bitfield.
// [0x0F, 0x01] => 0000 1111 0000 0001
// Each byte is bit-reversed via a 256-entry LUT (bits.reverse_8).
pub fn from_bytes(bytes []u8) BitField {
    mut output := new(bytes.len * 8)
    for i, b in bytes {
        output.field[i / 4] |= u32(bits.reverse_8(b)) << ((i % 4) * 8)
    }
    return output
}

// from_bytes_lowest_bits_first converts a byte array into a bitfield.
// For example: [0x0F, 0x01] => 1111 0000 1000 0000
pub fn from_bytes_lowest_bits_first(bytes []u8) BitField {
    mut output := new(bytes.len * 8)
    for i, b in bytes {
        output.field[i / 4] |= u32(b) << ((i % 4) * 8)
    }
    return output
}

// from_str converts a string of characters (`0` and `1`) to a bitfield.
// Any character different from `0` is treated as `1`.
pub fn from_str(str string) BitField {
    mut output := new(str.len)
    for i in 0 .. str.len {
        if str[i] != `0` {
            output.set_bit(i)
        }
    }
    return output
}

// str converts the bit array to a string of characters (`0` and `1`).
pub fn (bf BitField) str() string {
    mut output := []u8{len: bf.size}
    for i in 0 .. bf.size {
        output[i] = if bf.get_bit(i) == 1 { `1` } else { `0` }
    }
    return output.bytestr()
}

// new creates an empty bit array capable of storing `size` bits.
pub fn new(size int) BitField {
    output := BitField{
        size: size
        // field: *u32(calloc(zbitnslots(size) * slot_size / 8))
        field: []u32{len: zbitnslots(size)}
    }
    return output
}

// free frees the memory allocated for a bitfield.
@[unsafe]
pub fn (bf &BitField) free() {
    bf.field.free()
}

// get_bit returns the value (0 or 1) of bit number `bitnr` (count from 0).
@[inline]
pub fn (bf BitField) get_bit(bitnr int) int {
    if bitnr < 0 || bitnr >= bf.size {
        return 0
    }
    return int((bf.field[bitslot(bitnr)] >> (bitnr % slot_size)) & u32(1))
}

// set_bit sets bit number `bitnr` to 1 (count from 0).
@[inline]
pub fn (mut bf BitField) set_bit(bitnr int) {
    if bitnr < 0 || bitnr >= bf.size {
        return
    }
    bf.field[bitslot(bitnr)] |= bitmask(bitnr)
}

// clear_bit sets bit number `bitnr` to 0 (count from 0).
@[inline]
pub fn (mut bf BitField) clear_bit(bitnr int) {
    if bitnr < 0 || bitnr >= bf.size {
        return
    }
    bf.field[bitslot(bitnr)] &= ~bitmask(bitnr)
}

// extract returns the value converted from a slice of bit numbers from `start` to length of `len`.
// For example 0101 . extract(1, 2) => 0b10
pub fn (bf BitField) extract(start int, len int) u64 {
    // panic?
    if start < 0 {
        return 0
    }
    mut output := u64(0)
    for i in 0 .. len {
        output |= u64(bf.get_bit(start + len - i - 1)) << i
    }
    return output
}

// insert sets bit numbers from `start` to `len` length with the value converted from the number `_value`.
// For example 0000.insert(1, 2, 0b10) => 0100
pub fn (mut bf BitField) insert[T](start int, len int, _value T) {
    // panic?
    if start < 0 {
        return
    }
    mut value := _value
    for i in 0 .. len {
        pos := start + len - i - 1
        if value & 1 == 1 {
            bf.set_bit(pos)
        } else {
            bf.clear_bit(pos)
        }
        value >>= 1
    }
}

// extract_lowest_bits_first returns the value converted from a slice of bit numbers from `start` to length of `len`.
// For example 0101.extract_lowest_bits_first(1, 2) => 0b01
pub fn (bf BitField) extract_lowest_bits_first(start int, len int) u64 {
    // panic?
    if start < 0 {
        return 0
    }
    mut output := u64(0)
    for i in 0 .. len {
        output |= u64(bf.get_bit(start + i)) << i
    }
    return output
}

// insert_lowest_bits_first sets bit numbers from `start` to `len` length with the value converted from the number `_value`.
// For example 0000.insert_lowest_bits_first(1, 2, 0b10) => 0010
pub fn (mut bf BitField) insert_lowest_bits_first[T](start int, len int, _value T) {
    // panic?
    if start < 0 {
        return
    }
    mut value := _value
    for pos in start .. start + len {
        if value & 1 == 1 {
            bf.set_bit(pos)
        } else {
            bf.clear_bit(pos)
        }
        value >>= 1
    }
}

// set_all sets all bits in the bitfield to 1.
pub fn (mut bf BitField) set_all() {
    for i in 0 .. zbitnslots(bf.size) {
        bf.field[i] = u32(0xFFFF_FFFF)
    }
    bf.clear_tail()
}

// clear_all sets all bits in the bitfield to 0.
pub fn (mut bf BitField) clear_all() {
    for i in 0 .. zbitnslots(bf.size) {
        bf.field[i] = u32(0)
    }
}

// toggle_bit changes the value (from 0 to 1 or from 1 to 0) of bit number `bitnr`.
@[inline]
pub fn (mut bf BitField) toggle_bit(bitnr int) {
    if bitnr < 0 || bitnr >= bf.size {
        return
    }
    bf.field[bitslot(bitnr)] ^= bitmask(bitnr)
}

// set_if sets bit number `bitnr` to 1 (count from 0) if `cond` is true else clears the bit.
@[inline]
pub fn (mut bf BitField) set_if(cond bool, bitnr int) {
    if bitnr < 0 || bitnr >= bf.size {
        return
    }
    if cond {
        bf.field[bitslot(bitnr)] |= bitmask(bitnr)
    } else {
        bf.field[bitslot(bitnr)] &= ~bitmask(bitnr)
    }
}

// toggle_bits changes the value (from 0 to 1 or from 1 to 0) of bits.
// Example: mut bf := bitfield.new(10); bf.toggle_bits(1,3,5,7); assert bf.str() == '0101010100'
@[inline]
pub fn (mut bf BitField) toggle_bits(a ...int) {
    for bitnr in a {
        if bitnr < 0 || bitnr >= bf.size {
            continue
        }
        bf.field[bitslot(bitnr)] ^= bitmask(bitnr)
    }
}

// set_bits sets multiple bits in the bitfield to 1.
// Example: mut bf := bitfield.new(10); bf.set_bits(1,3,5,7); assert bf.str() == '0101010100'
@[inline]
pub fn (mut bf BitField) set_bits(a ...int) {
    for bitnr in a {
        if bitnr < 0 || bitnr >= bf.size {
            continue
        }
        bf.field[bitslot(bitnr)] |= bitmask(bitnr)
    }
}

// clear_bits sets multiple bits in the bitfield to 0.
// Example: mut bf := bitfield.from_str('1111111111111'); bf.clear_bits(1,2,5,6,7); assert bf.str() == '1001100011111'
@[inline]
pub fn (mut bf BitField) clear_bits(a ...int) {
    for bitnr in a {
        if bitnr < 0 || bitnr >= bf.size {
            continue
        }
        bf.field[bitslot(bitnr)] &= ~bitmask(bitnr)
    }
}

// has test if *at least one* of the bits is set.
// Example: mut bf := bitfield.from_str('111111100000000'); assert bf.has(1,3,5,7)
@[inline]
pub fn (bf BitField) has(a ...int) bool {
    for bitnr in a {
        if bitnr < 0 || bitnr >= bf.size {
            return false
        }
        if int((bf.field[bitslot(bitnr)] >> (bitnr % slot_size)) & u32(1)) == 1 {
            return true
        }
    }
    return false
}

// all test if *all* of the bits are set.
// Example: mut bf := bitfield.from_str('111111100000000'); assert !bf.all(1,3,5,7)
@[inline]
pub fn (bf BitField) all(a ...int) bool {
    for bitnr in a {
        if bitnr < 0 || bitnr >= bf.size {
            return false
        }
        if int((bf.field[bitslot(bitnr)] >> (bitnr % slot_size)) & u32(1)) == 0 {
            return false
        }
    }
    return true
}

// bf_and performs logical AND operation on every pair of bits from `input1` and `input2`.
// It returns the result as a new bitfield. If inputs differ in size, the tail of the longer one is ignored.
pub fn bf_and(input1 BitField, input2 BitField) BitField {
    size := min(input1.size, input2.size)
    bitnslots := zbitnslots(size)
    mut output := new(size)
    for i in 0 .. bitnslots {
        output.field[i] = input1.field[i] & input2.field[i]
    }
    output.clear_tail()
    return output
}

// bf_not toggles all bits in a bitfield and returns the result as a new bitfield.
pub fn bf_not(bf BitField) BitField {
    size := bf.size
    bitnslots := zbitnslots(size)
    mut output := new(size)
    for i in 0 .. bitnslots {
        output.field[i] = ~bf.field[i]
    }
    output.clear_tail()
    return output
}

// bf_or performs logical OR operation on every pair of bits from `input1` and `input2`.
// It returns the result as a new bitfield. If inputs differ in size, the tail of the longer one is ignored.
pub fn bf_or(input1 BitField, input2 BitField) BitField {
    size := min(input1.size, input2.size)
    bitnslots := zbitnslots(size)
    mut output := new(size)
    for i in 0 .. bitnslots {
        output.field[i] = input1.field[i] | input2.field[i]
    }
    output.clear_tail()
    return output
}

// bf_xor perform logical XOR operation on every pair of bits from `input1` and `input2`.
// It returns the result as a new bitfield. If inputs differ in size, the tail of the longer one is ignored.
pub fn bf_xor(input1 BitField, input2 BitField) BitField {
    size := min(input1.size, input2.size)
    bitnslots := zbitnslots(size)
    mut output := new(size)
    for i in 0 .. bitnslots {
        output.field[i] = input1.field[i] ^ input2.field[i]
    }
    output.clear_tail()
    return output
}

// join concatenates two bitfields and returns the result as a new bitfield.
pub fn join(input1 BitField, input2 BitField) BitField {
    output_size := input1.size + input2.size
    mut output := new(output_size)
    // copy the first input to output as is
    for i in 0 .. zbitnslots(input1.size) {
        output.field[i] = input1.field[i]
    }
    // find offset bit and offset slot
    offset_bit := input1.size % slot_size
    offset_slot := input1.size / slot_size
    for i in 0 .. zbitnslots(input2.size) {
        output.field[i + offset_slot] |= u32(input2.field[i] << u32(offset_bit))
    }
    // If offset_bit is not zero, additional operations are needed.
    // Number of iterations depends on the nr of slots in output. Two
    // options:
    // (a) nr of slots in output is the sum of inputs' slots. In this
    // case, the nr of bits in the last slot of output is less than the
    // nr of bits in the second input (i.e. ), OR
    // (b) nr of slots of output is the sum of inputs' slots less one
    // (i.e. less iterations needed). In this case, the nr of bits in
    // the last slot of output is greater than the nr of bits in the second
    // input.
    // If offset_bit is zero, no additional copies needed.
    if (output_size - 1) % slot_size < (input2.size - 1) % slot_size {
        for i in 0 .. zbitnslots(input2.size) {
            output.field[i + offset_slot + 1] |= u32(input2.field[i] >> u32(slot_size - offset_bit))
        }
    } else if (output_size - 1) % slot_size > (input2.size - 1) % slot_size {
        for i in 0 .. zbitnslots(input2.size) - 1 {
            output.field[i + offset_slot + 1] |= u32(input2.field[i] >> u32(slot_size - offset_bit))
        }
    }
    return output
}

// get_size returns the number of bits the array can hold.
@[inline]
pub fn (bf BitField) get_size() int {
    return bf.size
}

// clone creates a copy of a bit array.
pub fn (bf BitField) clone() BitField {
    bitnslots := zbitnslots(bf.size)
    mut output := new(bf.size)
    for i in 0 .. bitnslots {
        output.field[i] = bf.field[i]
    }
    return output
}

// == compares 2 bitfields, and returns true if they are equal.
pub fn (a BitField) == (b BitField) bool {
    if a.size != b.size {
        return false
    }
    for i in 0 .. zbitnslots(a.size) {
        if a.field[i] != b.field[i] {
            return false
        }
    }
    return true
}

// pop_count returns the number of set bits (ones) in the array.
pub fn (bf BitField) pop_count() int {
    mut count := 0
    for i in 0 .. zbitnslots(bf.size) {
        count += bits.ones_count_32(bf.field[i])
    }
    return count
}

// hamming computes the Hamming distance between two bit arrays.
@[inline]
pub fn hamming(input1 BitField, input2 BitField) int {
    input_xored := bf_xor(input1, input2)
    return input_xored.pop_count()
}

// pos checks if the bitfield contains a sub-array `needle`.
// It returns its position if it does, -1 if it does not, and -2 on error.
pub fn (bf BitField) pos(needle BitField) int {
    heystack_size := bf.size
    needle_size := needle.size
    diff := heystack_size - needle_size
    // needle longer than bitfield; return error code -2
    if diff < 0 {
        return -2
    }
    for i := 0; i <= diff; i++ {
        needle_candidate := bf.slice(i, needle_size + i)
        if needle_candidate == needle {
            // needle matches a section of the bitfield; return starting position of the section
            return i
        }
    }
    // nothing matched; return -1
    return -1
}

// slice returns a sub-array of bits between `_start` (included) and `_end` (excluded).
pub fn (bf BitField) slice(_start int, _end int) BitField {
    // boundary checks
    mut start := _start
    mut end := _end
    if end > bf.size {
        end = bf.size // or panic?
    }
    if start > end {
        start = end // or panic?
    }
    mut output := new(end - start)
    if end == start {
        // zero-length slice: nothing to copy; field is empty
        return output
    }
    start_offset := start % slot_size
    end_offset := (end - 1) % slot_size
    start_slot := start / slot_size
    end_slot := (end - 1) / slot_size
    output_slots := zbitnslots(end - start)
    if output_slots > 1 {
        if start_offset != 0 {
            for i in 0 .. output_slots - 1 {
                output.field[i] = u32(bf.field[start_slot + i] >> u32(start_offset))
                output.field[i] = output.field[i] | u32(bf.field[start_slot + i + 1] << u32(slot_size - start_offset))
            }
        } else {
            for i in 0 .. output_slots - 1 {
                output.field[i] = u32(bf.field[start_slot + i])
            }
        }
    }
    if start_offset > end_offset {
        output.field[(end - start - 1) / slot_size] = u32(bf.field[end_slot - 1] >> u32(start_offset))
        mut mask := u32((1 << (end_offset + 1)) - 1)
        mask = bf.field[end_slot] & mask
        mask = u32(mask << u32(slot_size - start_offset))
        output.field[(end - start - 1) / slot_size] |= mask
    } else if start_offset == 0 {
        mut mask := u32(0)
        if end_offset == slot_size - 1 {
            mask = u32(-1)
        } else {
            mask = u32(u32(1) << u32(end_offset + 1))
            mask = mask - u32(1)
        }
        output.field[(end - start - 1) / slot_size] = (bf.field[end_slot] & mask)
    } else {
        mut mask := u32(((1 << (end_offset - start_offset + 1)) - 1) << start_offset)
        mask = bf.field[end_slot] & mask
        mask = u32(mask >> u32(start_offset))
        output.field[(end - start - 1) / slot_size] |= mask
    }
    return output
}

// reverse reverses the order of bits in the bitfield (swap the first with the last, the second with the last but one and so on).
pub fn (bf BitField) reverse() BitField {
    size := bf.size
    bitnslots := zbitnslots(size)
    mut output := new(size)
    for i := 0; i < (bitnslots - 1); i++ {
        for j in 0 .. slot_size {
            if u32(bf.field[i] >> u32(j)) & u32(1) == u32(1) {
                output.set_bit(size - i * slot_size - j - 1)
            }
        }
    }
    bits_in_last_input_slot := (size - 1) % slot_size + 1
    for j in 0 .. bits_in_last_input_slot {
        if u32(bf.field[bitnslots - 1] >> u32(j)) & u32(1) == u32(1) {
            output.set_bit(bits_in_last_input_slot - j - 1)
        }
    }
    return output
}

// resize changes the size of the bit array to `new_size`.
pub fn (mut bf BitField) resize(new_size int) {
    new_bitnslots := zbitnslots(new_size)
    old_size := bf.size
    old_bitnslots := zbitnslots(old_size)
    mut field := []u32{len: new_bitnslots}
    for i := 0; i < old_bitnslots && i < new_bitnslots; i++ {
        field[i] = bf.field[i]
    }
    bf.field = field
    bf.size = new_size
    if new_size < old_size && new_size % slot_size != 0 {
        bf.clear_tail()
    }
}

// rotate performs a circular-shift on the bits by `offset` positions (move `offset` bit to 0, `offset+1` bit to 1, and so on).
pub fn (bf BitField) rotate(offset int) BitField {
    // This function "cuts" the bitfield into two and swaps the pieces.
    // If the offset is positive, the cutting point is counted from the
    // beginning of the bit array, otherwise from the end.
    size := bf.size
    if size == 0 {
        return bf
    }
    // removing extra rotations
    mut offset_internal := offset % size
    if offset_internal == 0 {
        // nothing to shift
        return bf
    }
    if offset_internal < 0 {
        offset_internal = offset_internal + size
    }
    first_chunk := bf.slice(0, offset_internal)
    second_chunk := bf.slice(offset_internal, size)
    output := join(second_chunk, first_chunk)
    return output
}

// shift_left shift-left the bits by `count` positions.
pub fn (bf BitField) shift_left(count int) BitField {
    size := bf.size
    if count <= 0 {
        return bf
    } else if count >= size {
        // return zeroes
        return new(size)
    }
    zeroes := new(count)
    return join(bf.slice(count, size), zeroes)
}

// shift_right shift-right the bits by `count` positions.
pub fn (bf BitField) shift_right(count int) BitField {
    size := bf.size
    if count <= 0 {
        return bf
    } else if count >= size {
        // return zeroes
        return new(size)
    }
    zeroes := new(count)
    return join(zeroes, bf.slice(0, size - count))
}

// Internal functions

// clear_tail clears the extra bits that are not part of the bitfield, but yet are allocated
@[inline]
fn (mut bf BitField) clear_tail() {
    tail := bf.size % slot_size
    if tail != 0 {
        // create a mask for the tail
        mask := u32((1 << tail) - 1)
        // clear the extra bits
        bf.field[zbitnslots(bf.size) - 1] = bf.field[zbitnslots(bf.size) - 1] & mask
    }
}

// bitmask is the bitmask needed to access a particular bit at offset bitnr
@[inline]
fn bitmask(bitnr int) u32 {
    return u32(u32(1) << u32(bitnr % slot_size))
}

// bitslot is the slot index (i.e. the integer) where a particular bit is located
@[inline]
fn bitslot(size int) int {
    return size / slot_size
}

// min returns the minimum of 2 integers; it is here to avoid importing math just for that
@[inline]
fn min(input1 int, input2 int) int {
    if input1 < input2 {
        return input1
    } else {
        return input2
    }
}

// zbitnslots returns the minimum number of whole integers, needed to represent a bitfield of size length
@[inline]
fn zbitnslots(length int) int {
    if length <= 0 {
        return 0
    }
    return (length - 1) / slot_size + 1
}