#include "bingo_multibit_tree.h"

using namespace bingo;

int MultibitTree::_compareBitWeights(_DistrWeight& bw1, _DistrWeight& bw2, void* context)
{
    if (bw1.weight > bw2.weight)
        return 1;
    else if (bw1.weight < bw2.weight)
        return -1;
    else
        return 0;
}

BingoPtr<MultibitTree::_MultibitNode> MultibitTree::_buildNode(Array<int>& fit_fp_indices, const Array<bool>& is_parrent_mb, int level)
{
    BingoPtr<_MultibitNode> node_ptr;
    node_ptr.allocate();
    _MultibitNode* node = new (node_ptr.ptr()) _MultibitNode();
    byte* fingerprints = _fingerprints_ptr.ptr();

    int fp_bitsize = _fp_size * 8;

    if (fit_fp_indices.size() <= 10 || level == _max_level)
    {
        node->fp_indices_array.allocate(fit_fp_indices.size());
        int* fp_indices = node->fp_indices_array.ptr();
        node->fp_indices_count = fit_fp_indices.size();

        for (int i = 0; i < fit_fp_indices.size(); i++)
            fp_indices[i] = fit_fp_indices[i];

        return node_ptr;
    }

    QS_DEF(Array<_DistrWeight>, bit_weights);
    bit_weights.clear_resize(fp_bitsize);
    bit_weights.zerofill();
    double bit_con = 1.0 / fit_fp_indices.size();

    for (int j = 0; j < fp_bitsize; j++)
        bit_weights[j].idx = j;

    for (int i = 0; i < fit_fp_indices.size(); i++)
    {
        for (int j = 0; j < fp_bitsize; j++)
        {
            bit_weights[j].weight += bitGetBit(fingerprints + fit_fp_indices[i] * _fp_size, j) * bit_con;
        }
    }

    bit_weights.qsort(_compareBitWeights, 0);

    QS_DEF(Array<int>, mb_indices);
    mb_indices.clear();

    double distr_coef = 1;
    for (int k = 0; k < bit_weights.size(); k++)
    {
        if (is_parrent_mb[bit_weights[k].idx])
            continue;
        distr_coef *= (1 - bit_weights[k].weight);
        mb_indices.push(bit_weights[k].idx);
        if (distr_coef < 0.5)
            break;
    }

    QS_DEF(Array<int>, left_fit);
    left_fit.clear();
    QS_DEF(Array<int>, right_fit);
    right_fit.clear();

    QS_DEF(Array<bool>, is_mb);
    is_mb.copy(is_parrent_mb);

    for (int i = 0; i < mb_indices.size(); i++)
    {
        is_mb[mb_indices[i]] = 1;
    }

    for (int i = 0; i < fit_fp_indices.size(); i++)
    {
        bool is_fit = true;
        for (int j = 0; j < mb_indices.size(); j++)
        {
            if (bitGetBit(fingerprints + fit_fp_indices[i] * _fp_size, mb_indices[j]))
            {
                is_fit = false;
                break;
            }
        }

        if (is_fit)
            right_fit.push(fit_fp_indices[i]);
        else
            left_fit.push(fit_fp_indices[i]);
    }

    if (left_fit.size() == 0 || right_fit.size() == 0)
    {
        node->fp_indices_array.allocate(fit_fp_indices.size());
        int* fp_indices = (int*)node->fp_indices_array.ptr();
        node->fp_indices_count = fit_fp_indices.size();

        for (int i = 0; i < fit_fp_indices.size(); i++)
            fp_indices[i] = fit_fp_indices[i];

        return node_ptr;
    }

    node->match_bits_array.allocate(mb_indices.size());
    _MatchBit* match_bits = node->match_bits_array.ptr();
    node->match_bits_count = mb_indices.size();

    for (int i = 0; i < mb_indices.size(); i++)
    {
        if (mb_indices[i] < 0 || mb_indices[i] > 511)
            i = i;

        match_bits[i] = _MatchBit(mb_indices[i], 0);
    }

    node->left = _buildNode(left_fit, is_mb, level + 1);
    node->right = _buildNode(right_fit, is_mb, level + 1);

    return node_ptr;
}

void MultibitTree::_build()
{
    QS_DEF(Array<int>, indices);
    indices.clear_resize(_fp_count);

    for (int i = 0; i < indices.size(); i++)
        indices[i] = i;

    QS_DEF(Array<bool>, is_mb);
    is_mb.clear_resize(_fp_size * 8);
    is_mb.zerofill();

    _tree_ptr = _buildNode(indices, is_mb, 0);
}

void MultibitTree::_findLinear(_MultibitNode* node, const byte* query, int query_bit_number, SimCoef& sim_coef, double min_coef, Array<SimResult>& sim_indices,
                               int fp_bit_number)
{
    profTimerStart(tmsl, "multibit_tree_search_linear");
    byte* fingerprints = _fingerprints_ptr.ptr();
    int* indices = _indices_ptr.ptr();

    int* fp_indices = node->fp_indices_array.ptr();

    for (int i = 0; i < node->fp_indices_count; i++)
    {
        const byte* fp = fingerprints + fp_indices[i] * _fp_size;
        int f_bit_number = bitGetOnesCount(fp, _fp_size);

        double coef = sim_coef.calcCoef(query, fp, query_bit_number, f_bit_number);
        if (coef < min_coef)
            continue;

        sim_indices.push(SimResult(indices[fp_indices[i]], (float)coef));
    }
}

void MultibitTree::_findSimilarInNode(BingoPtr<_MultibitNode> node_ptr, const byte* query, int query_bit_number, SimCoef& sim_coef, double min_coef,
                                      Array<SimResult>& sim_indices, int m01, int m10)
{
    if (node_ptr.isNull())
        return;

    _MultibitNode* node = node_ptr.ptr();

    if (node->fp_indices_count != 0)
    {
        if (_min_fp_bit_number == _max_fp_bit_number) // if fingerpint bits_count is fixed
            _findLinear(node, query, query_bit_number, sim_coef, min_coef, sim_indices, _min_fp_bit_number);
        else
            _findLinear(node, query, query_bit_number, sim_coef, min_coef, sim_indices);

        return;
    }

    _MatchBit* match_bits = node->match_bits_array.ptr();

    QS_DEF(Array<SimResult>, left_indices);
    left_indices.clear();
    QS_DEF(Array<SimResult>, right_indices);
    right_indices.clear();

    int right_m01 = m01, right_m10 = m10;
    for (int i = 0; i < node->match_bits_count; i++)
        if (match_bits[i].val == 0)
        {
            if (bitGetBit(query, match_bits[i].idx))
                right_m01++;
        }
        else if (!bitGetBit(query, match_bits[i].idx))
            right_m10++;

    double right_upper_bound = sim_coef.calcUpperBound(query_bit_number, _min_fp_bit_number, _max_fp_bit_number, right_m10, right_m01);

    if (!node->left.isNull())
        _findSimilarInNode(node->left, query, query_bit_number, sim_coef, min_coef, left_indices, m01, m10);
    if ((!node->left.isNull()) && right_upper_bound + EPSILON > min_coef)
        _findSimilarInNode(node->right, query, query_bit_number, sim_coef, min_coef, right_indices, right_m01, right_m10);

    for (int i = 0; i < left_indices.size(); i++)
        sim_indices.push(left_indices[i]);
    for (int i = 0; i < right_indices.size(); i++)
        sim_indices.push(right_indices[i]);
}

MultibitTree::MultibitTree(int fp_size) : _fp_size(fp_size)
{
    _tree_ptr.allocate();
    new (_tree_ptr.ptr()) _MultibitNode();
    _query_bit_number = -1;
    _max_level = 6;
}

void MultibitTree::build(BingoPtr<byte> fingerprints, BingoPtr<int> indices, int fp_count, int min_fp_bit_number, int max_fp_bit_number)
{
    _fingerprints_ptr = fingerprints;
    _indices_ptr = indices;

    _min_fp_bit_number = min_fp_bit_number;
    _max_fp_bit_number = max_fp_bit_number;

    _fp_count = fp_count;

    _build();
}

int MultibitTree::findSimilar(const byte* query, SimCoef& sim_coef, double min_coef, Array<SimResult>& sim_fp_indices)
{
    profTimerStart(tms, "multibit_tree_search");
    int query_bit_number = bitGetOnesCount(query, _fp_size);
    sim_fp_indices.clear();

    _findSimilarInNode(_tree_ptr, query, query_bit_number, sim_coef, min_coef, sim_fp_indices, 0, 0);

    return sim_fp_indices.size();
}