/* The MIT License

   Copyright (c) 2025-     Zhonghai Zhang (ICT)

   Permission is hereby granted, free of charge, to any person obtaining
   a copy of this software and associated documentation files (the
   "Software"), to deal in the Software without restriction, including
   without limitation the rights to use, copy, modify, merge, publish,
   distribute, sublicense, and/or sell copies of the Software, and to
   permit persons to whom the Software is furnished to do so, subject to
   the following conditions:

   The above copyright notice and this permission notice shall be
   included in all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
   MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
   NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
   BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
   ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
   CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
   SOFTWARE.
*/

#include <assert.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include <zlib.h>

#include "bwt.h"
#include "hyb_idx.h"
#include "kvec.h"
#include "rle.h"
#include "rope.h"
#include "utils.h"

#define KMER_START(tid, max_kmer, num_thread) ((uint64_t)(max_kmer) * (uint64_t)(tid) / (uint64_t)(num_thread))
#define KMER_END(tid, max_kmer, num_thread) ((uint64_t)(max_kmer) * (uint64_t)((tid) + 1) / (uint64_t)(num_thread))

///////////////////////////

typedef struct {
    uint64_t size;
    uint64_t cap;
    uint8_t* addr;
} BlockData;

#define _bd_init(bd)   \
    do {               \
        (bd).size = 0; \
        (bd).cap = 0;  \
        (bd).addr = 0; \
    } while (0)

#define _bd_claim(bd, add_size)                                 \
    do {                                                        \
        if ((bd).size + add_size >= (bd).cap) {                 \
            (bd).cap += add_size + 4L * 1024 * 1024;            \
            (bd).addr = (uint8_t*)realloc((bd).addr, (bd).cap); \
        }                                                       \
    } while (0)

#define _bd_addr(bd) ((bd).addr + (bd).size)

#define _bd_destory(bd)  \
    do {                 \
        free((bd).addr); \
    } while (0)

typedef struct {
    uint32_t idx;
    uint32_t cnt;
    uint32_t bytes;
} idx_cnt_t;

// 构建hyb-index时候用到的节点数据结构
typedef struct _Node {
    uint8_t header;     // 1字节头部信息 [0]
    uint8_t hit_bytes;  // 用多少bytes表示hits
    uint8_t off_bytes;  // 用多少bytes表示子节点数据偏移量

    uint8_t hits_neq;        // 和子节点hits数量不相等
    uint8_t hits_neq1;       // 子节点和孙子节点hits数量不相等
    uint8_t max_depth_leaf;  // 是否达到最大深度的叶节点
    uint8_t normal_leaf;     // 正常的叶子节点，hits=1

    uint16_t seq_len;          // 单一分支碱基序列长度，如果分支，那就是1
    uint16_t depth;            // 当前节点对应的位置，0对应着16（不包含16-kmer）
    uint16_t min_child_depth;  // 最近到hit=1的深度
    uint16_t max_child_depth;  // 最远到hit=1或者到达最大深度时的深度

    uint32_t hits;                // 匹配个数
    uint32_t child_num;           // 有多少分支子节点, 第一层子节点个数
    uint32_t num_offspring;       // 所有后代节点个数
    uint32_t min_offspring_hits;  // 所有后代中，最后最少的hits
    uint32_t node_bytes;          // 该节点包含所有后代所占内存
    uint32_t child_bytes;         // 子节点所占内存

    uint8_t type;      // 当前节点类型
    uint64_t mark;     // 与type对应，4bit，16bit，64bit
    uint64_t sa_pos;   // 对应的sa起始位置
    uint8_t seq[320];  // 保存所有单一分支的碱基序列

    uint8_t kmer;            // 当前节点对应的kmer，root节点是uint64，用这个不准确

    struct _Node* children[64];  // 子节点
    struct _Node* parent;        // 父节点
} Node;

// 并行构建hyb-index的数据结构
typedef struct {
    pthread_t ptid;       // 线程ID
    int tid;              // 线程id
    uint64_t kmer_start;  // [) 左闭右开
    uint64_t kmer_end;
    uint64_t bytes;         // 当前线程所有kmer数据占用的内存字节数
    uint8_t* kmer_data;     // kmer数据
    uint64_t prev_bd_size;  // 前一个kmer对应的index数据所占内存大小，用来调整多线程kmer数据对应的地址偏移量的
    BlockData idata;        // 二级索引数据
    bwt_t* bwt;             // bwt索引
} PartKmer;

typedef kvec_t(PartKmer) PartKmerVec;

// 当hit变成1后，继续扩展节点，为了保存2-kmer或者3-kmer对应的匹配信息
typedef struct {
    Node* node;
    uint64_t mark;
} MarkNode;

/////////////////

// 正向搜索kmer（hyb的16-kmer）在bwt中的区间信息
bwtintv_t bwt_search_kmer(const bwt_t* bwt, const uint8_t* q) {
    bwtintv_t ik = {0}, ok[4] = {0};
    int i, x, c;
    uint32_t qbit = 0;
    x = HASH_KMER_LEN;
    for (i = 0; i < x; ++i) {
        qbit |= q[i] << ((HASH_KMER_LEN - 1 - i) << 1);
    }
    bwt_kmer_get(&bwt->kmer_hash, &ik, qbit, x - 1);
    if (ik.x[2] < 1)
        return ik;
    for (i = x; i < HYB_KMER_LEN; ++i) {
        c = 3 - q[i];
        bwt_extend(bwt, &ik, ok, 0);
        if (ok[c].x[2] < 1)
            return ok[c];
        ik = ok[c];
    }
    return ik;
}

// 根据子节点更新父节点信息
static void update_parent(Node* n, Node* c, uint64_t k, int update_max) {
    // 创建完子节点后，更新一些数据
    n->children[n->child_num++] = c;
    c->kmer = k;
    c->parent = n;
    n->mark |= (1ULL << k);  // 设置当前节点的mark
    n->min_offspring_hits = MIN(n->min_offspring_hits, c->min_offspring_hits);
    n->min_child_depth = MIN(n->min_child_depth, c->min_child_depth + 1);
    if (update_max)
        n->max_child_depth = MAX(n->max_child_depth, c->max_child_depth + 1);
    n->num_offspring += c->num_offspring + 1;
    if (c->hits_neq)
        n->hits_neq1 = 1;  // 子节点和孙子节点hits数量不相等
}

// 清除父节点的一些信息, 不释放内存空间
static void clear_parent(Node* n) {
    uint32_t i = 0;
    n->max_child_depth = 0;
    for (; i < n->child_num; ++i) n->children[i] = NULL;  // 清除子节点
    n->child_num = 0;
    n->mark = 0;
    n->num_offspring = 0;
    n->hits_neq1 = 0;
}

// 更新节点信息
static void update_node_info(PartKmer* pk, Node* n, bwtintv_t* ik, int depth, int max_depth, int update_min) {
    n->depth = depth;                  // 当前节点对应的深度
    n->min_offspring_hits = ik->x[2];  // 所有子节点中最小的hits数量
    n->sa_pos = ik->x[0];              // bwt_get_sa(pk->fmt->sa, ik->x[0]);  // 对应的sa中第一个位置
    n->hits = ik->x[2];                // 匹配数量
    if (update_min)
        n->min_child_depth = max_depth;  // 初始化为当前节点所能达到的最大深度（人为设定的阈值）
}

// 普通的leaf节点扩展，扩展到第2层和第3层
static void extend_leaf_node(PartKmer* pk, Node* n, bwtintv_t* ik, int depth, int max_depth) {
    n->normal_leaf = 1;
    update_node_info(pk, n, ik, depth, max_depth, 0);
    if (max_depth == 0) {
        return;
    }
    uint8_t k;
    int child_hits = 0;
    bwtintv_t ok[4] = {0};
    bwt_extend(pk->bwt, ik, ok, 0);
    for (k = 0; k < 4; ++k) {
        bwtintv_t iv = ok[3 - k];
        child_hits += iv.x[2];
        if (iv.x[2] > 0) {  // hits必须大于0，即有匹配的位置，才创建子节点
            Node* child = (Node*)calloc(1, sizeof(Node));
            child->seq[child->seq_len++] = k;  // 保存序列
            // 递归构建子节点
            extend_leaf_node(pk, child, &iv, depth + 1, max_depth - 1);
            // 创建完子节点后，更新一些数据
            update_parent(n, child, k, 0);  // 更新父节点信息
            break;
        }
    }
    if (child_hits != (int)n->hits) {
        n->hits_neq = 1;  // 当前碱基有一个匹配到了ref的结尾，导致节点hits和子节点hits总和不等
    }
}

// 递归构建hyb-tree
// n: 父节点
// depth: kmer之后该碱基在第几个（位置/深度）
// max_depth: 遍历到最大深度停止(最长有多少碱基, 从当前节点算起)
static void build_hyb_tree(PartKmer* pk, Node* n, bwtintv_t* ik, int depth, int max_depth) {
    /* 递归终止条件 */
    if (ik->x[2] == 1) {  // 如果hits为0或者1，则不再需要子节点
        n->min_child_depth = 0;
        n->normal_leaf = 1;                     // 正常的叶子节点，hits=1
        extend_leaf_node(pk, n, ik, depth, 2);  // 继续扩展两个碱基
        return;
    }
    // 给节点n赋值
    update_node_info(pk, n, ik, depth, max_depth, 1);
    if (max_depth == 0) {  // 达到最大深度，不再进行扩展
        n->min_child_depth = 0;
        n->max_depth_leaf = 1;  // 达到最大深度的叶节点
        return;
    }
    int child_hits = 0;  // 所有子节点hits数量之和
    uint8_t k;           // 碱基
    bwtintv_t ok[4] = {0};
    bwt_extend(pk->bwt, ik, ok, 0);
    for (k = 0; k < 4; ++k) {  // 1-kmer
        bwtintv_t iv = ok[3 - k];
        child_hits += iv.x[2];
        if (iv.x[2] > 0) {  // hits必须大于0，即有匹配的位置，才创建子节点
            Node* child = (Node*)calloc(1, sizeof(Node));
            child->seq[child->seq_len++] = k;
            // 递归构建子节点
            build_hyb_tree(pk, child, &iv, depth + 1, max_depth - 1);
            // 创建完子节点后，更新一些数据
            update_parent(n, child, k, 1);  // 更新父节点信息
        }
    }
    // 子节点hits之和应该和父节点hits数量相等，否则，就是当前父节点有一个匹配到了ref的结尾
    if (child_hits != (int)n->hits) {
        n->hits_neq = 1;  // 当前碱基有一个匹配到了ref的结尾，导致节点hits和子节点hits总和不等
    }
}

// 将ptr的子节点赋值给n，并释放ptr节点
static void assign_children(Node* n, Node* ptr) {
    clear_parent(n);
    n->seq[n->seq_len++] = ptr->seq[0];  // 保存序列
    uint32_t i = 0;
    for (; i < ptr->child_num; ++i) {
        update_parent(n, ptr->children[i], ptr->children[i]->seq[0], 1);
    }
    free(ptr);  // 释放ptr节点
}

// 处理path node
static void handle_path_node(Node** n) {
    Node* cur = *n;
    Node* ptr = cur->children[0];                            // 只有一个子节点
    while (ptr->child_num == 1 && cur->hits == ptr->hits) {  // 继续向下遍历，直到没有子节点, 或父子hits不一致
        cur->seq[cur->seq_len++] = ptr->seq[0];              // 保存序列
        Node* next = ptr->children[0];                       // 获取下一个子节点
        free(ptr);                                           // 释放当前节点
        ptr = next;
    }
    if (ptr->child_num == 1) {  // hits不一致
        clear_parent(cur);
        update_parent(cur, ptr, ptr->seq[0], 1);
        cur->hits_neq = 1;                       // 继承子节点的hits_neq
        ptr->seq[ptr->seq_len++] = ptr->seq[0];  // 多保存一次该节点的碱基，因为他也是path节点
        *n = ptr;
        handle_path_node(n);  // 递归处理下一个path节点
    } else {
        assign_children(cur, ptr);  // 将ptr的子节点赋值给n
    }
}

// dfs遍历当前节点的子节点，扩展到2或3层，然后去掉中间层节点，将后续节点直接与父节点连接, layer = 2 or 3
static void shrink_node_layer(PartKmer* pk, Node* n, int layer) {
    kvec_t(MarkNode) layer_nodes_stack;
    kv_init(layer_nodes_stack);  // dfs遍历的stack
    kvec_t(MarkNode) kmer_stack;
    kv_init(kmer_stack);  // 装所有的kmer
    kvec_t(Node*) to_release;
    kv_init(to_release);  // 装所有的待删除的node

    int i;
    // init stack
    for (i = (int)n->child_num - 1; i >= 0; --i) {
        MarkNode ln;
        ln.node = n->children[i];
        ln.mark = 1;  // 第一层子节点
        kv_push(MarkNode, layer_nodes_stack, ln);
    }
    // dfs 遍历
    while (kv_size(layer_nodes_stack) > 0) {
        MarkNode ln = kv_pop(layer_nodes_stack);
        if ((int)ln.mark == layer) {  // 是第3层的节点，创建对应的kmer
            uint8_t k = 0;
            if (layer == 2)
                k = ln.node->parent->seq[0] << 2 | ln.node->seq[0];
            else if (layer == 3)
                k = ln.node->parent->parent->seq[0] << 4 | ln.node->parent->seq[0] << 2 | ln.node->seq[0];
            MarkNode kmer_node;
            kmer_node.node = ln.node;
            kmer_node.mark = k;                        // 保存kmer
            kv_push(MarkNode, kmer_stack, kmer_node);  // 保存kmer
        } else {
            kv_push(Node*, to_release, ln.node);  // 保存当前节点，待释放
            for (i = (int)ln.node->child_num - 1; i >= 0; --i) {
                MarkNode lnn;
                lnn.node = ln.node->children[i];
                lnn.mark = ln.mark + 1;
                kv_push(MarkNode, layer_nodes_stack, lnn);
            }
        }
    }
    // 处理所有kmer
    clear_parent(n);
    int child_hits = 0;
    for (i = 0; i < (int)kv_size(kmer_stack); ++i) {
        MarkNode kmer_node = kv_A(kmer_stack, i);
        child_hits += kmer_node.node->hits;                   // 统计子节点的hits数量
        update_parent(n, kmer_node.node, kmer_node.mark, 1);  // 更新父节点信息
    }
    if (child_hits != (int)n->hits) {
        n->hits_neq = 1;  // 当前碱基有一个匹配到了ref的结尾，导致节点hits和子节点hits总和不等
    }
    for (i = 0; i < (int)kv_size(to_release); ++i) {
        free(kv_A(to_release, i));
    }
    kv_destroy(layer_nodes_stack);
    kv_destroy(kmer_stack);
    kv_destroy(to_release);
}

// 构建完hybrid tree之后，从根节点开始遍历，设置节点的类型，构建path节点，3，2，1 kmer节点
static void update_hyb_tree(PartKmer* pk, Node* n) {
    if (n->child_num == 1 && !n->normal_leaf && !n->max_depth_leaf && !(n->max_child_depth <= 3)) {  // path类型
        handle_path_node(&n);
    }
    uint32_t i = 0;
#define __recursive_update_hyb_tree(pk, n)   \
    for (i = 0; i < n->child_num; ++i) {     \
        update_hyb_tree(pk, n->children[i]); \
    }
    // 计算子节点的idata地址，和父节点的match_cnt
    // 递归终点
    if (n->max_child_depth == 5) {
        if (n->hits_neq) {
            n->type = HYB_BP_1;
        } else {
            n->type = HYB_BP_2;
            shrink_node_layer(pk, n, 2);
        }
        __recursive_update_hyb_tree(pk, n);
    } else if (n->max_child_depth == 4) {
        n->type = HYB_BP_1;
        __recursive_update_hyb_tree(pk, n);
    } else if (n->max_child_depth == 3) {
        if (n->hits_neq) {
            n->type = HYB_BP_1;
        } else if (n->hits_neq1) {
            n->type = HYB_BP_2;
            shrink_node_layer(pk, n, 2);
        } else {
            n->type = HYB_BP_3;
            shrink_node_layer(pk, n, 3);
        }
    } else if (n->max_child_depth == 2) {
        if (n->hits_neq) {
            n->type = HYB_BP_1;
        } else {
            n->type = HYB_BP_2;
            shrink_node_layer(pk, n, 2);
        }
    } else if (n->max_child_depth == 1) {
        n->type = HYB_BP_1;
    } else {
        if (n->hits_neq) {
            n->type = HYB_BP_1;
        } else if (n->hits_neq1) {
            n->type = HYB_BP_2;
            shrink_node_layer(pk, n, 2);
        } else {
            n->type = HYB_BP_3;
            shrink_node_layer(pk, n, 3);
        }
        __recursive_update_hyb_tree(pk, n);
    }
}

// 根据构建的radix tree，创建hyb-index
static void calc_node_bytes(Node* n) {
    // 递归终点
    if (n->normal_leaf || n->max_depth_leaf) {  // 叶子节点没有子节点
        n->node_bytes = 0;
        return;
    }
    // 先处理子节点
    uint32_t i = 0;
    for (; i < n->child_num; ++i) {
        calc_node_bytes(n->children[i]);               // 递归构建子节点
        n->child_bytes += n->children[i]->node_bytes;  // 累加子节点所占内存
    }
    // 处理当前节点
    n->node_bytes = n->child_bytes;
    if (n->depth == 0) n->node_bytes += 5;  // sa_pos占用5字节
    
    // 计算偏移地址字节
    int off_bytes = 0;
    if (n->child_bytes == 0) {  // 没有子节点
        off_bytes = 0;
    } else if (n->child_bytes < 256) {
        off_bytes = 1;  // 子节点数据偏移量小于256, 1B就行
    } else if (n->child_bytes < 65536) {
        off_bytes = 2;  // 2B就行
    } else if (n->child_bytes < 16777216) {
        off_bytes = 3;  // 3B就行
    } else if (n->child_bytes < UINT32_MAX) {
        off_bytes = 4;  // 使用4B
    } else {
        fprintf(stderr, "too many child bytes: %d\n", n->child_bytes);
        exit(1);
    }
    n->off_bytes = off_bytes;  // 保存子节点数据偏移量字节数

    // 计算hits字节
    int hits_bytes = 0;
    if (n->hits < 256) {
        hits_bytes = 1;  // hits小于256, 占用1字节
    } else if (n->hits < 65536) {
        hits_bytes = 2;  // hits小于65536, 占用2字节
    } else if (n->hits < 16777216) {
        hits_bytes = 3;  // hits小于16777216, 占用3字节
    } else if (n->child_bytes < UINT32_MAX) {
        hits_bytes = 4;  // 使用4B
    } else {
        fprintf(stderr, "too many hits: %d\n", n->hits);
        exit(1);
    }
    n->hit_bytes = hits_bytes;  // 保存hits字节数

    if (n->depth == 0 && n->hits > HYB_HIT_THRESH) {
         n->node_bytes += hits_bytes;  // 根节点需要保存hits
    }

    if (n->seq_len > 1) {    // path节点
        n->node_bytes += 1;  // header, path节点相当于两个节点，一个path，一个普通的
        n->node_bytes += 1;  // 结合header 记录碱基path长度，代替mark
        n->node_bytes += ((n->seq_len - 1) * 2 + 7) / 8;  // 只有根节点保存本节点的碱基，其他的节点不保存seq[0]，因为seq[0]是子节点的碱基
    }
    // 有可能一个path下边还是一个path，因为可能父子hits不一致
    if (n->child_num > 1 || n->seq_len == 1) {  // 有多个子节点，或者当前节点不是path节点
        n->node_bytes += 1;                     // header_bits占用1字节
        // 叶节点只用bit-map来计算hits偏移就行了, 不用保存hits start信息
        int all_child_is_leaf = n->hits == n->child_num || (n->hits_neq && n->hits == n->child_num + 1);
        if (!all_child_is_leaf) {
            n->node_bytes += (hits_bytes + off_bytes) * (n->child_num - 1);  // 不保存第一个子节点offset，因为是0
            // for test，首个节点不保存hits start和offset，因为hits start是0或者1，offset是0，这样index-data减少了1GB，总共20GB
            // 这样不能根据概率调整内存位置了
        }
        if (n->type == HYB_BP_1) {         // 1-kmer节点
            n->node_bytes += 1;            // 1字节的mark
        } else if (n->type == HYB_BP_2) {  // 2-kmer节点
            n->node_bytes += 2;            // 2字节的mark
        } else if (n->type == HYB_BP_3) {  // 3-kmer节点
            n->node_bytes += 8;            // 8字节的mark
        } else {
            fprintf(stderr, "Unknown node type: %d\n", (int)n->type);
            exit(1);
        }
    }
}

// 释放所有子节点和当前节点
static void release_node(Node* p) {
    uint32_t k = 0;
    for (k = 0; k < p->child_num; k++) {
        release_node(p->children[k]);
    }
    free(p);
}

void write_seq_bits(uint8_t* bs, int seq_len, BlockData* idata) {
    if (seq_len > 0) {
        uint32_t num_bytes = ((seq_len << 1) + 7) >> 3;
        uint8_t* bit_seq = (uint8_t*)_bd_addr(*idata);
        uint8_t bybp = 0;
        int i = 0, j = 0, idx = 0;
        for (; i + 3 < seq_len; i += 4) {
            bybp = bs[i] | bs[i + 1] << 2 | bs[i + 2] << 4 | bs[i + 3] << 6;
            *(uint8_t*)&bit_seq[idx++] = bybp;
        }
        bybp = 0;
        for (; i < seq_len; ++i, ++j) {
            bybp |= bs[i] << j * 2;
        }
        *(uint8_t*)&bit_seq[idx++] = bybp;
        idata->size += num_bytes;  // 更新数据大小
    }
}

static void write_normal_node_data(Node* n, BlockData* idata, uint8_t write_sa) {
    uint8_t* ptr;
    int i, j;
    int all_child_is_leaf = 0;

    // 2bit子节点类型, 1bit父子hits是否一致，2bit hits字节数;
    uint8_t header = n->type << 6 | n->hits_neq << 5 | (n->hit_bytes - 1) << 3;
    // 叶节点只用bit-map来计算hits偏移就行了
    all_child_is_leaf = n->hits == n->child_num || (n->hits_neq && n->hits == n->child_num + 1);
    // 3bit子节点地址偏移字节数
    if (all_child_is_leaf) {
        header |= HYB_LEAF_NODE;  // 用off_bytes = 7来表示叶节点，不用读取hits和offset信息
    } else {
        header |= n->off_bytes;
//        assert(n->off_bytes != 0); // 这里会出错，为啥有这个断言？
    }

    // header
    idata->addr[idata->size++] = header;  // 写入header
    // sa_pos
    if (n->depth == 0 && write_sa) {
        ptr = (uint8_t*)&n->sa_pos;
        for (i = 0; i < 5; ++i) *(_bd_addr(*idata) + i) = *(ptr + i);
        idata->size += i;  // 写入sa_pos
    }
    // hits
    if (write_sa && n->depth == 0 && n->hits > HYB_HIT_THRESH) {
        ptr = (uint8_t*)&n->hits;
        for (i = 0; i < n->hit_bytes; ++i) *(_bd_addr(*idata) + i) = *(ptr + i);
        idata->size += i;  // 写入hits
    }

    // 处理mark
    if (n->type == HYB_BP_1) {
        idata->addr[idata->size++] = (uint8_t)(n->mark & 0xFF);  // 1-kmer节点
    } else if (n->type == HYB_BP_2) {
        ptr = (uint8_t*)&n->mark;
        for (i = 0; i < 2; ++i) *(_bd_addr(*idata) + i) = *(ptr + i);
        idata->size += i;  // 2-kmer节点
    } else if (n->type == HYB_BP_3) {
        ptr = (uint8_t*)&n->mark;
        for (i = 0; i < 8; ++i) *(_bd_addr(*idata) + i) = *(ptr + i);
        idata->size += i;  // 3-kmer节点
    }

    if (all_child_is_leaf) {
        return;  // 叶节点只用bit-map来计算hits start就行了
    }

    // 处理子节点
    uint32_t sa_start = 0;
    uint32_t offset = 0;  // 子节点数据偏移量
    // 剩下的节点
    for (j = 1; j < (int)n->child_num; ++j) {
        sa_start += n->children[j - 1]->hits;  // 累加前一个子节点的hits
        ptr = (uint8_t*)&sa_start;
        for (i = 0; i < n->hit_bytes; ++i) *(_bd_addr(*idata) + i) = *(ptr + i);
        idata->size += i;  // 写入sa_start
        if (n->children[j]->node_bytes == 0 && n->children[j]->hits > 1 && !n->children[j]->max_depth_leaf) {
            fprintf(stderr, "Error: no child! hits:%u, depth:%u, max-leaf:%u\n", n->children[j]->hits, n->children[j]->depth,
                    n->children[j]->max_depth_leaf);
        }
        offset += n->children[j - 1]->node_bytes;
        ptr = (uint8_t*)&offset;
        for (i = 0; i < n->off_bytes; ++i) *(_bd_addr(*idata) + i) = *(ptr + i);
        idata->size += i;  // 写入offset
    }
}

// 写入path node数据，一个path node包含path和normal两种数据
static void write_path_node_data(Node* n, BlockData* idata) {
    // 2bit子节点类型, 1bit父子hits是否一致，2bit hits字节数;
    uint8_t header = HYB_BP_PATH << 6 | n->hits_neq << 5 | (n->hit_bytes - 1) << 3;
    int path_len = n->seq_len - 1;        // path节点的长度
    header |= path_len >> 8;              // 此时，header最低位保存path_len的高8位（9bit足以表示path的长度）
    idata->addr[idata->size++] = header;  // 写入header
    int i;
    uint8_t* ptr;
    if (n->depth == 0) {
        ptr = (uint8_t*)&n->sa_pos;
        for (i = 0; i < 5; ++i) *(_bd_addr(*idata) + i) = *(ptr + i);
        idata->size += i;  // 写入sa_pos
    }
    // 处理hits
    if (n->depth == 0 && n->hits > HYB_HIT_THRESH) {
        ptr = (uint8_t*)&n->hits;
        for (i = 0; i < n->hit_bytes; ++i) *(_bd_addr(*idata) + i) = *(ptr + i);
        idata->size += i;  // 写入hits
    }
    // 处理mark
    idata->addr[idata->size++] = (uint8_t)(path_len & 0xFF);  // 写入mark
    // 写入seq
    write_seq_bits(n->seq + 1, path_len, idata);  // 跳过seq[0]，因为seq[0]是子节点的碱基
    // 继续把该节点当做正常节点处理
    if (n->child_num > 1)
        write_normal_node_data(n, idata, 0);  // 写入正常节点数据
}

/*
    header-1B: 2bit子节点类型，1bit父子hits是否一致，2bit hits字节，3bit子节点数据偏移量字节，
*/
static void build_index_from_node(Node* n, BlockData* idata) {
    if (n->normal_leaf || n->max_depth_leaf) {  // 叶子节点没有子节点
        return;
    }
    uint32_t i = 0;
    _bd_claim(*idata, n->node_bytes);  // 确保有足够的内存空间

    // 先处理当前节点
    // 1. 写入header
    // 2. 如果是根节点，写入sa_pos
    // 3. 写入hits
    // 4. 写入mark
    // 5. 写入子节点sa_pos偏移量和子节点数据偏移量
    uint32_t node_self_size = idata->size;
    if (n->seq_len > 1) {
        write_path_node_data(n, idata);  // 写入path节点数据
    } else {
        write_normal_node_data(n, idata, 1);  // 写入正常节点数据
    }
    node_self_size = idata->size - node_self_size;  // 当前节点数据大小
    if (node_self_size != n->node_bytes - n->child_bytes) {
        fprintf(stderr, "Error: node bytes not match: %u, %u, %u, %u, %u, %u, %u, %u, %u, %u\n", node_self_size,
                n->node_bytes - n->child_bytes, n->node_bytes, n->depth, n->max_child_depth, n->seq_len, n->hits, n->child_num,
                n->hits_neq, (int)n->type);
        Node* p = n;
        fprintf(stderr, "Info: father neq: %u, %u, %u\n", p->parent->hits_neq, p->parent->child_num, p->parent->hits);
        while (p != NULL) {
            fprintf(stderr, "Info: hits: %u, depth: %u, childs: %u, hits_neq: %u, seq_len: %u, max_depth: %u\n", p->hits,
                    p->depth, p->child_num, p->hits_neq, p->seq_len, p->max_child_depth);
            p = p->children[0];  // 只打印第一个子节点
        }
        p = n;
        while (p->depth > 0) {
            fprintf(stderr, "Info: hits: %u, depth: %u, childs: %u, hits_neq: %u, seq_len: %u, max_depth: %u\n", p->hits,
                    p->depth, p->child_num, p->hits_neq, p->seq_len, p->max_child_depth);
            p = p->parent;  // 向上打印父节点
        }
        fprintf(stderr, "Info: hits: %u, depth: %u, childs: %u, hits_neq: %u, seq_len: %u, max_depth: %u\n", p->hits, p->depth,
                p->child_num, p->hits_neq, p->seq_len, p->max_child_depth);

        exit(0);
    }

    for (i = 0; i < n->child_num; ++i) {               // 这里需要按照匹配概率来修改
        build_index_from_node(n->children[i], idata);  // 递归构建子节点
    }
}

// 多线程创建hybrid index
static void* thread_build_hyb_index(void* arg) {
    PartKmer* p = (PartKmer*)arg;
    uint8_t kmer[HYB_KMER_LEN + 1];
    uint64_t all_kmer_num = p->kmer_end - p->kmer_start;
    uint64_t percent_mod = all_kmer_num / 100;
    int percent = 0;
    if (percent_mod == 0)
        percent_mod = 1;

    uint8_t type_hits = 0;
    uint64_t ref_pos = 0;
    uint64_t sa_forward = 0;
    uint64_t idata_offset = 0;
    uint64_t packed_data = 0;
    uint8_t* ptr = (uint8_t*)&packed_data;
    uint64_t i, j, kmer_size;
    fprintf(stderr, "[thread %d] start building hybrid index for kmer range [%lu, %lu)\n", p->tid, p->kmer_start, p->kmer_end);
    for (i = p->kmer_start, kmer_size = 0; i < p->kmer_end; ++i, ++kmer_size) {
        uint8_t* kd = p->kmer_data + i * HYB_KMER_DATA_BYTES;
        for (j = 0; j < HYB_KMER_LEN; ++j) {
            kmer[j] = (i >> (j << 1)) & 3;
        }
        bwtintv_t iv = bwt_search_kmer(p->bwt, kmer);
        const uint32_t khits = iv.x[2]; // 有多少个匹配的位置
        if (khits == 0) {
            // no extra bytes
            packed_data = 0;
            for (int off = 0; off < 5; ++off) *(kd + off) = *(ptr + off);
        } else if (khits == 1) { // 只有一个匹配，直接保存对应的参考基因位置
            // no extra bytes
            type_hits = 1;
            ref_pos = bwt_get_sa(p->bwt->byte_sa, iv.x[0]);  // 直接记录kmer对应的reference pos
            packed_data = (ref_pos << HYB_KMER_DATA_TYPE_BITS) | type_hits;
            // 这里需要注意，packed_data是小端模式的, 5个字节
            for (int off = 0; off < 5; ++off) *(kd + off) = *(ptr + off);
        } else if (khits == 2) { // 只有两个匹配，直接保存对应的suffix array位置
            // no extra bytes
            type_hits = 2;
            sa_forward = iv.x[0];  // 记录kmer对应的后缀数组的位置
            packed_data = (sa_forward << HYB_KMER_DATA_TYPE_BITS) | type_hits;
            for (int off = 0; off < 5; ++off) *(kd + off) = *(ptr + off);
        } else {
            Node* root = (Node*)calloc(1, sizeof(Node));
            root->seq[root->seq_len++] = kmer[HYB_KMER_LEN - 1];
            build_hyb_tree(p, root, &iv, 0, HYB_MAX_SEQ_LEN - HYB_KMER_LEN);  // 创建基数树
            update_hyb_tree(p, root);                                         // 更新hybrid tree
            idata_offset = p->idata.size;                                     // 指向当前要用的地址偏移量(起始位置)
            calc_node_bytes(root);                                            // 计算每个节点所占内存
            build_index_from_node(root, &p->idata);  // 构建hybrid index
            // 将type和地址偏移信息写入kmer
            if (khits <= HYB_HIT_THRESH) {  // hits数量小于等于30
                type_hits = khits;          // hits数量
            } else {
                type_hits = HYB_HIT_THRESH + 1;
            }
            packed_data = (idata_offset << HYB_KMER_DATA_TYPE_BITS) | type_hits;
            for (int off = 0; off < 5; ++off) *(kd + off) = *(ptr + off);
            p->bytes += root->node_bytes;  // 累加当前节点所占内存
            release_node(root);            // 释放所有子节点和当前节点
        }

        if (p->tid == 0 && kmer_size % percent_mod == 0) {
            fprintf(stderr, "[thread %d] %d %% done, bytes: %fM\n", p->tid, percent++, (double)p->bytes / 1024 / 1024);
        }
    }

    fprintf(stderr, "Finish-[thread %d] bytes: %fM\n", p->tid, (double)p->bytes / 1024 / 1024);
    return NULL;
}

// 多线程调整hybrid index中的block data偏移地址，整合成连续的一块地址空间
static void* thread_adjust_hyb_index(void* arg) {
    PartKmer* p = (PartKmer*)arg;
    uint64_t all_kmer_num = p->kmer_end - p->kmer_start;
    uint64_t percent_mod = all_kmer_num / 100;
    int percent = 0;
    if (percent_mod == 0) percent_mod = 1;
    uint8_t type_hits = 0;
    uint64_t bd_offset = 0;
    uint64_t packed_data = 0;
    uint8_t* ptr = (uint8_t*)&packed_data;
    uint64_t i, kmer_size;
    for (i = p->kmer_start, kmer_size = 0; i < p->kmer_end; ++i, ++kmer_size) {
        uint8_t* kd = p->kmer_data + i * HYB_KMER_DATA_TYPE_BITS;
        type_hits = *(kd)&HYB_KMER_DATA_TYPE_MASK;
        if (type_hits > 2) {  // block data index数据，需要改offset
            // bd_offset = ((*(uint64_t *)kd) >> kTypeBits) & p->KMER_MASK;  // 指向当前要用的地址偏移量
            bd_offset = ((*(uint64_t*)kd) & HYB_KMER_DATA_MASK) >> HYB_KMER_DATA_TYPE_BITS;  // 指向当前要用的地址偏移量
            bd_offset += p->prev_bd_size;
            packed_data = (bd_offset << HYB_KMER_DATA_TYPE_BITS) | type_hits;
            for (int off = 0; off < 5; ++off) *(kd + off) = *(ptr + off);
        }

        if (p->tid == 0 && kmer_size % percent_mod == 0) {
            fprintf(stderr, "Info: %d%% done adjust.\n", percent++);
        }
    }
    return NULL;
}

// 保存hybrid index
static void save_hybrid_idx(const char* prefix, PartKmerVec* pkArr) {
    // fprintf(stderr, "size: %ld\n", kv_size(*pkArr));
    // return;
    char fn[MAX_PATH];
    char fn2[MAX_PATH];

    sprintf(fn, "%s.hyb.kmer", prefix);
    sprintf(fn2, "%s.hyb.data", prefix);
        
    FILE* fpk = fopen(fn, "wb");
    FILE* fpbd = fopen(fn2, "wb");

    if (fpk == NULL) {
        fprintf(stderr, "Error: open file %s failed\n", fn);
        return;
    }
    if (fpbd == NULL) {
        fprintf(stderr, "Error: open file %s failed", fn2);
        return;
    }

    uint64_t max_kmer = (1ULL << (HYB_KMER_LEN << 1));
    err_fwrite(kv_A(*pkArr, 0).kmer_data, 1, max_kmer * HYB_KMER_DATA_BYTES, fpk);  // 小端
    int i;
    for (i = 0; i < kv_size(*pkArr); ++i) {
        err_fwrite(kv_A(*pkArr, i).idata.addr, 1, kv_A(*pkArr, i).idata.size, fpbd);
    }
    fclose(fpk);
    fclose(fpbd);
}

// 构建hybrid-index入口函数
int hyb_idx_build_and_dump(int num_threads, bwt_t* bwt, const char* idx_prefix) {
    PartKmerVec pk_arr;
    kv_init(pk_arr);
    kv_resize(PartKmer, pk_arr, num_threads);
    pk_arr.n = num_threads;

    kvec_t(pthread_t) tid_arr;
    kv_init(tid_arr);
    kv_resize(pthread_t, tid_arr, num_threads);
    tid_arr.n = num_threads;

    uint64_t max_kmer = (1ULL << (HYB_KMER_LEN << 1));
    uint8_t* kmer_data = (uint8_t*)malloc(max_kmer * HYB_KMER_DATA_BYTES);

    int i;
    for (i = 0; i < num_threads; ++i) {
        PartKmer* p = &kv_A(pk_arr, i);
        p->tid = i;
        p->kmer_start = KMER_START(i, max_kmer, num_threads);
        p->kmer_end = KMER_END(i, max_kmer, num_threads);
        p->bytes = 0;
        p->kmer_data = kmer_data;
        p->prev_bd_size = 0;
        _bd_init(p->idata);
        p->bwt = bwt;
        pthread_create(&kv_A(tid_arr, i), NULL, thread_build_hyb_index, p);
    }

    uint64_t all_bytes = 0;
    uint64_t prev_bd_size = 0;
    for (i = 0; i < num_threads; ++i) {
        pthread_join(kv_A(tid_arr, i), NULL);
        PartKmer* p = &kv_A(pk_arr, i);
        all_bytes += p->bytes;
        p->prev_bd_size = prev_bd_size;
        prev_bd_size += p->idata.size;
    }

    for (i = 0; i < num_threads; ++i) {
        pthread_create(&kv_A(tid_arr, i), NULL, thread_adjust_hyb_index, &kv_A(pk_arr, i));
    }
    for (int i = 0; i < num_threads; ++i) {
        pthread_join(kv_A(tid_arr, i), NULL);
    }

    fprintf(stderr, "Hybrid index total bytes: %fM\n", (double)all_bytes / 1024 / 1024);

    double time_sava_hybrid_idx = realtime();
    save_hybrid_idx(idx_prefix, &pk_arr);
    time_sava_hybrid_idx = realtime() - time_sava_hybrid_idx;

    fprintf(stderr, "Time save hybrid index: %fs\n", time_sava_hybrid_idx);

    kv_destroy(pk_arr);
    kv_destroy(tid_arr);

    return 0;
}