#include "core/BitArray.hpp"

#include <cassert>
#include <cstring>

#include "core/BitArray.hpp"
#include "core/ToString.hpp"
#include "core/Utility.hpp"

using Core::BitArray;

static constexpr size_t U64_BITS = 64;
static constexpr size_t DIVIDE_BITS = 6;

static u64 readBits(const u64* data, size_t index, u64 bits) {
    u64 dataIndexA = (index * bits) >> DIVIDE_BITS;
    u64 dataIndexB = ((index * bits) + (bits - 1lu)) >> DIVIDE_BITS;
    u64 shifts = (index * bits) & (U64_BITS - 1lu);
    if(dataIndexA == dataIndexB) {
        return (data[dataIndexA] >> shifts) & ((1lu << bits) - 1lu);
    }
    u64 bitsInA = U64_BITS - shifts;
    u64 r = (data[dataIndexA] >> shifts) & ((1lu << bitsInA) - 1lu);
    r |= (data[dataIndexB] & ((1lu << (bits - bitsInA)) - 1lu)) << bitsInA;
    return r;
}

static void setBits(u64* data, size_t index, size_t bits, u64 value) {
    u64 mask = (1lu << bits) - 1lu;
    value &= mask;
    u64 dataIndexA = (index * bits) >> DIVIDE_BITS;
    u64 dataIndexB = ((index * bits) + (bits - 1lu)) >> DIVIDE_BITS;
    u64 shifts = (index * bits) & (U64_BITS - 1lu);
    data[dataIndexA] &= ~(mask << shifts);
    data[dataIndexA] |= (value << shifts);
    if(dataIndexA != dataIndexB) {
        u64 leftBits = bits - (U64_BITS - shifts);
        data[dataIndexB] &= ~((1lu << leftBits) - 1lu);
        data[dataIndexB] |= (value >> (U64_BITS - shifts));
    }
}

static u64 roundUpDivide(u64 a, u64 b) {
    return a / b + ((a % b) != 0);
}

static size_t getArrayLength(size_t length, size_t bits) {
    return roundUpDivide(length * bits, U64_BITS);
}

BitArray::BitArray() : length(0), bits(0), data(nullptr) {
}

BitArray::BitArray(const BitArray& other) : BitArray() {
    resize(other.getLength(), other.getBits());
    size_t l = getLength();
    for(size_t i = 0; i < l; i++) {
        set(i, other.get(i));
    }
}

BitArray::BitArray(BitArray&& other) noexcept : BitArray() {
    swap(other);
}

BitArray::~BitArray() {
    coreDeleteN(data);
}

BitArray& BitArray::operator=(BitArray other) noexcept {
    swap(other);
    return *this;
}

BitArray& BitArray::set(size_t index, u64 value) {
    assert(data != nullptr);
    assert(index < length);
    setBits(data, index, bits, value);
    return *this;
}

u64 BitArray::get(size_t index) const {
    assert(data != nullptr);
    assert(index < length);
    return readBits(data, index, bits);
}

size_t BitArray::getLength() const {
    return length;
}

size_t BitArray::getBits() const {
    return bits;
}

size_t Core::BitArray::getInternalByteSize() const {
    if(getLength() <= 0 || getBits() <= 0) {
        return 0;
    }
    return getArrayLength(getLength(), getBits()) * sizeof(u64);
}

i64 BitArray::select(u64 index) const {
    if(index <= 0) {
        return -1;
    }
    u64 found = 0;
    size_t end = getArrayLength(length, bits);
    for(size_t i = 0; i < end; i++) {
        u64 ones = Core::popCount<u64, u64>(data[i]);
        found += ones;
        if(found >= index) {
            found -= ones;
            u64 a = i * U64_BITS - 1;
            u64 d = data[i];
            while(found < index) {
                found += d & 1;
                d >>= 1;
                a++;
            }
            return static_cast<i64>(a);
        }
    }
    return -1;
}

void BitArray::fill(u64 value) {
    size_t l = getLength();
    for(size_t i = 0; i < l; i++) {
        set(i, value);
    }
}

void BitArray::resize(size_t newLength, size_t newBits) {
    if(newLength == 0 || newBits == 0) {
        return;
    } else if(newBits > 64) {
        newBits = 64;
    }
    size_t arrayLength = getArrayLength(newLength, newBits);
    u64* newData = coreNewN(u64, arrayLength);
    memset(newData, 0, arrayLength * sizeof(u64));

    size_t end = Core::min(length, newLength);
    for(size_t i = 0; i < end; i++) {
        setBits(newData, i, newBits, get(i));
    }
    for(size_t i = end; i < newLength; i++) {
        setBits(newData, i, newBits, 0);
    }
    delete[] data;
    data = newData;
    length = newLength & 0x00FF'FFFF;
    bits = newBits & 0xFF;
}

size_t BitArray::toString(char* s, size_t n) const {
    size_t total = 0;
    addString("[", s, n, total);
    size_t l = length;
    if(l > 0) {
        l--;
        for(size_t i = 0; i < l; i++) {
            addString(get(i), s, n, total);
            addString(", ", s, n, total);
        }
        addString(get(l), s, n, total);
    }
    addString("]", s, n, total);
    return total;
}

void BitArray::swap(BitArray& other) {
    u64 l = length;
    u64 b = bits;
    length = other.length;
    bits = other.bits;
    other.length = l & 0x00FF'FFFF;
    other.bits = b & 0xFF;
    Core::swap(data, other.data);
}