/*
     Description: vcf index utils

     Copyright : All right reserved by ICT

     Author : Zhang Zhonghai
     Date : 2019/11/23
*/

#include "linear_index.h"
#include "bam_wrap.h"
#include "binary_util.h"

/*
 * LinearIndex 类
 */
// 搜索区间参数对应的数据区域，闭区间
void LinearIndex::SearchInterval(int64_t start, int64_t end, int64_t* file_pos, int64_t* content_len) {
    int stid, spos, etid, epos, stid_origin, spos_origin, etid_origin, epos_origin;
    stid = stid_origin = BamWrap::bam_tid(start);
    spos = spos_origin = BamWrap::bam_pos(start);
    etid = etid_origin = BamWrap::bam_tid(end);
    epos = epos_origin = BamWrap::bam_pos(end);

    while (stid < idx_.size() && idx_[stid].size() == 0) ++stid;
    while (etid >= 0 && idx_[etid].size() == 0) --etid;
    if (stid > etid || stid >= idx_.size()) {
        *content_len = 0;
        return;
    }
    int sb_idx, eb_idx;
    if (stid == stid_origin) {
        sb_idx = spos / idx_[stid].binWidth;
        if (sb_idx >= idx_[stid].size()) {  // 开始区域没有对应的block
            sb_idx = idx_[stid].size() - 1;
            Block& sb = idx_[stid][sb_idx];
            *file_pos = sb.startPosition + sb.size;
        } else {
            Block& sb = idx_[stid][sb_idx];
            *file_pos = sb.startPosition;
        }
    } else {
        sb_idx = 0;
        Block& sb = idx_[stid][sb_idx];
        *file_pos = sb.startPosition;
    }
    if (etid == etid_origin) {
        eb_idx = epos / idx_[etid].binWidth;
        if (eb_idx >= idx_[etid].size()) {
            eb_idx = idx_[etid].size() - 1;
        }
    } else {
        eb_idx = idx_[etid].size() - 1;
    }
    Block& eb = idx_[etid][eb_idx];
    if (eb.startPosition + eb.size > *file_pos)
        *content_len = eb.startPosition + eb.size - *file_pos;
    else
        *content_len = 0;
}

// 读入index文件信息
bool LinearIndex::ReadIndex(const string& idx_fn) {
    if (bam_hdr_ == NULL)
        return false;
    ifstream ifs(idx_fn, ios::in | ios::binary);
    // 如果idx文件不存在，则创建 todo
    if (!ifs.good())
        return false;
    // 全部读进来再处理
    uint64_t fsize = 0;
    uint64_t cur = 0;
    ifs.seekg(0, ios::end);
    fsize = (uint64_t)ifs.tellg();
    ifs.seekg(0, ios::beg);
    char* buf = (char*)malloc(fsize + 10);
    if (!ifs.read(buf, fsize)) {
        free(buf);
        ifs.close();
        return false;
    }
    ifs.close();
    int int_val, version, flags;
    uint64_t long_val;
    string str_val;

    // read index header
    BinaryUtil::ReadInt(buf, fsize, &cur, &int_val);
    // cout << "magic: " << int_val << endl;
    BinaryUtil::ReadInt(buf, fsize, &cur, &int_val);
    if (int_val != INDEX_TYPE) {
        free(buf);
        ifs.close();
        return false;
    }
    // cout << "type: " << int_val << endl;
    BinaryUtil::ReadInt(buf, fsize, &cur, &version);
    // cout << "version: " << version << endl;
    BinaryUtil::ReadStr(buf, fsize, &cur, &str_val);
    // cout << "path: " << str_val << endl;
    BinaryUtil::ReadLong(buf, fsize, &cur, &vcf_fsize);  // vcf file size
    // cout << "file size: " << long_val << endl;
    BinaryUtil::ReadLong(buf, fsize, &cur, &long_val);  // test number
    BinaryUtil::ReadStr(buf, fsize, &cur, &str_val);    // md5
    BinaryUtil::ReadInt(buf, fsize, &cur, &flags);

    if (version < 3 && (flags & 0x8000) == 0x8000) {
        BinaryUtil::ReadInt(buf, fsize, &cur, &int_val);
        if (int_val < 0)
            return false;
        for (int i = 0; i < int_val; ++i) {
            BinaryUtil::ReadStr(buf, fsize, &cur, &str_val);
            BinaryUtil::ReadInt(buf, fsize, &cur, &int_val);
        }
    }
    if (version >= 3) {
        int nProperties;
        BinaryUtil::ReadInt(buf, fsize, &cur, &nProperties);
        string key, val;
        while (nProperties-- > 0) {
            BinaryUtil::ReadStr(buf, fsize, &cur, &key);
            BinaryUtil::ReadStr(buf, fsize, &cur, &val);
            properties_[key] = val;
        }
    }
    // read index content
    int n_chrom;
    BinaryUtil::ReadInt(buf, fsize, &cur, &n_chrom);
    int last_tid = 0;
    string last_name;
    int bam_tid = 0;  // tid从0开始，索引的染色体也要补充一个数据, 使得idx_ vector与tid一一对应
    while (n_chrom-- > 0) {
        ChrIndex ci;
        BinaryUtil::ReadStr(buf, fsize, &cur, &ci.name);
        BinaryUtil::ReadInt(buf, fsize, &cur, &ci.binWidth);
        int nBins;
        BinaryUtil::ReadInt(buf, fsize, &cur, &nBins);
        BinaryUtil::ReadInt(buf, fsize, &cur, &ci.longestFeature);
        BinaryUtil::ReadInt(buf, fsize, &cur, &int_val);
        BinaryUtil::ReadInt(buf, fsize, &cur, &ci.nFeatures);
        ci.blocks.reserve(nBins);
        int tid = sam_hdr_name2tid(bam_hdr_, ci.name.c_str());
        if (tid < 0) {
            if (ci.name == last_name) {
                tid = last_tid;
            } else {
                tid = ++last_tid;
                last_name = ci.name;
            }
        }
        while (bam_tid + 1 < tid) {
            idx_.push_back(ChrIndex());  // 补充空染色体
            ++bam_tid;
        }
        bam_tid = tid;
        ci.tid = tid;
        uint64_t pos;
        uint64_t nextPos;
        BinaryUtil::ReadLong(buf, fsize, &cur, &pos);
        for (int i = 0; i < nBins; ++i) {
            BinaryUtil::ReadLong(buf, fsize, &cur, &nextPos);
            ci.blocks.push_back(Block(pos, nextPos - pos));
            pos = nextPos;
        }
        idx_.push_back(ci);
    }
    sort(idx_.begin(), idx_.end());
    ifs.close();
    free(buf);
    return true;
}

/*
 * LinearIndex::ChrIndex 类
 */
void LinearIndex::ChrIndex::write(ofstream& out) const {
    BinaryUtil::WriteStr(out, name);
    BinaryUtil::WriteInt(out, binWidth);
    BinaryUtil::WriteInt(out, blocks.size());
    BinaryUtil::WriteInt(out, longestFeature);
    BinaryUtil::WriteInt(out, 0);
    BinaryUtil::WriteInt(out, nFeatures);
    uint64_t pos = 0, si = 0;
    for (int i = 0; i < blocks.size(); ++i) {
        pos = blocks[i].startPosition;
        si = blocks[i].size;
        BinaryUtil::WriteLong(out, pos);
    }
    BinaryUtil::WriteLong(out, pos + si);
}

/*
 * LinearIndexCreator类
 */
// 初始化索引文件头部信息
void LinearIndexCreator::InitHeaderDict(bam_hdr_t* hdr) {
    for (int i = 0; i < hdr->n_targets; ++i) {
        v_contig_name_.push_back(hdr->target_name[i]);
        contig_name_to_id_[hdr->target_name[i]] = i;
        contig_len_[hdr->target_name[i]] = hdr->target_len[i];
        ++n_properties_;
    }
}

// 添加一条vcf记录
void LinearIndexCreator::AddFeature(const Feature& ft, uint64_t f_pos) {  // f_pos是vcf文件当前正要写入的位置
    if (idx_.size() == 0 || idx_.back().tid != ft.tid) {
        if (idx_.size() > 0) {
            for (int i = 0; i < blocks_.size() - 1; ++i) {
                blocks_[i].size = blocks_[i + 1].startPosition - blocks_[i].startPosition;
                idx_.back().blocks.push_back(blocks_[i]);
            }
            blocks_.back().size = f_pos - blocks_.back().startPosition;
            idx_.back().blocks.push_back(blocks_.back());
        }
        idx_.push_back(LinearIndex::ChrIndex(v_contig_name_[ft.tid], ft.tid, bin_width_));
        blocks_.clear();
        blocks_.push_back(Block(f_pos, 0));
        longest_feature_ = 0;
    }
    while (ft.start > blocks_.size() * bin_width_) blocks_.push_back(Block(f_pos, 0));
    if (ft.FeatureLen() > longest_feature_) {
        longest_feature_ = ft.FeatureLen();
        idx_.back().longestFeature = longest_feature_;
    }
    ++idx_.back().nFeatures;
    ++FEATURE_COUNT_;
    all_feature_len += ft.FeatureLen();
}

// 添加所有记录完毕
void LinearIndexCreator::FinalizeIndex(uint64_t f_pos) {
    if (f_pos == 0)
        error("[%s] Error: finalize index failed\n", __func__);
    if (blocks_.size() > 0) {
        for (int i = 0; i < blocks_.size() - 1; ++i) {
            blocks_[i].size = blocks_[i + 1].startPosition - blocks_[i].startPosition;
            idx_.back().blocks.push_back(blocks_[i]);
        }
        blocks_.back().size = f_pos - blocks_.back().startPosition;
        idx_.back().blocks.push_back(blocks_.back());
        blocks_.clear();
    }
    FEATURE_LENGTH_MEAN_ = (double)all_feature_len / (double)FEATURE_COUNT_;
    index_file_size_ = f_pos;
}

// 写入index文件
void LinearIndexCreator::WriteIndex(const string& out_fn) {
    ofstream out;
    out.open(out_fn, ios::out | ios::binary);

    BinaryUtil::WriteInt(out, MAGIC_NUMBER);
    BinaryUtil::WriteInt(out, LinearIndex::INDEX_TYPE);
    BinaryUtil::WriteInt(out, LinearIndex::INDEX_VERSION);
    BinaryUtil::WriteStr(out, out_fn);
    BinaryUtil::WriteLong(out, index_file_size_);
    BinaryUtil::WriteLong(out, -1);
    BinaryUtil::WriteStr(out, "");
    BinaryUtil::WriteInt(out, flags_);
    n_properties_ += 4;
    BinaryUtil::WriteInt(out, n_properties_);
    ostringstream oss;
    oss << FEATURE_LENGTH_MEAN_;
    BinaryUtil::WriteStr(out, "FEATURE_LENGTH_MEAN");
    BinaryUtil::WriteStr(out, oss.str());
    oss.str("");
    oss << FEATURE_LENGTH_STD_DEV_;
    BinaryUtil::WriteStr(out, "FEATURE_LENGTH_STD_DEV");
    BinaryUtil::WriteStr(out, oss.str());
    oss.str("");
    oss << MEAN_FEATURE_VARIANCE_;
    BinaryUtil::WriteStr(out, "MEAN_FEATURE_VARIANCE");
    BinaryUtil::WriteStr(out, oss.str());
    oss.str("");
    oss << FEATURE_COUNT_;
    BinaryUtil::WriteStr(out, "FEATURE_COUNT");
    BinaryUtil::WriteStr(out, oss.str());
    for (int i = 0; i < v_contig_name_.size(); ++i) {
        const string& cn = v_contig_name_[i];
        oss.str("");
        oss << contig_len_[cn];
        BinaryUtil::WriteStr(out, "DICT:" + cn);
        BinaryUtil::WriteStr(out, oss.str());
    }
    BinaryUtil::WriteInt(out, idx_.size());
    for (int i = 0; i < idx_.size(); ++i) {
        const LinearIndex::ChrIndex& ci = idx_[i];
        ci.write(out);
    }
    out.close();
}
