/*
     Description: Recalibrate计算过程的工具函数

     Copyright : All right reserved by ICT

     Author : Zhang Zhonghai
     Date : 2025/12/29
*/
#pragma once

#include "util/stable_array.h"
#include "util/sam_data.h"
#include "bqsr/aux_arg.h"
#include "util/bam_wrap.h"
#include "util/interval.h"
#include "util/vcf_parser.h"
#include "util/profiling.h"

struct RecalFuncs {
    static constexpr int MAX_SITES_INTERVAL = 100000;
    static constexpr uint8_t NO_BAQ_UNCERTAINTY = (uint8_t)'@';

    // 设置某个位置是indel
    static inline void updateIndel(StableArray<int>& isIndel, int index) {
        if (index >= 0 && index < isIndel.size()) {
            isIndel[index] = 1;
        }
    }

    // 计算该read的每个碱基位置是否是SNP或Indel
    static int calculateIsSNPOrIndel(AuxVar& aux, SamData& sd, StableArray<int>& isSNP, StableArray<int>& isIns, StableArray<int>& isDel) {
        isSNP.resize_fill(sd.read_len, 0);
        isIns.resize_fill(sd.read_len, 0);
        isDel.resize_fill(sd.read_len, 0);
        // 1. 读取参考基因组，先看看串行运行性能，稍后可以将读入ref和vcf合并起来做成一个并行流水线步骤
        Interval interval{sd.start_pos, sd.end_pos};  // 闭区间
        PROF_START(read_ref);
        read_ref_base(aux, interval.left, interval);
        PROF_GP_END(read_ref);
        string refBases(aux.ref_seq);

        // 2. 遍历cigar，计算每个碱基是否是SNP或Indel
        PROF_START(calc_snp);
        int readPos = 0, refPos = 0, nEvents = 0;
        for (int i = 0; i < sd.cigars.size(); ++i) {
            const char c = sd.cigars[i].op;
            int len = sd.cigars[i].len;
            if (c == 'M' || c == '=' || c == 'X') {
                for (int j = 0; j < len; ++j) {
                    // 按位置将read和ref碱基进行比较，不同则是snp，注意read起始位置要加上left_clip
                    int snpInt = sd.bases[readPos] == refBases[refPos] ? 0 : 1;
                    isSNP[readPos] = snpInt;
                    nEvents += snpInt;
                    readPos++;
                    refPos++;
                }
            } else if (c == 'D') {
                // 应该是在上一个消耗碱基的cigar的最后一个位置，标记Del
                int index = sd.bw->GetReadNegativeStrandFlag() ? readPos : readPos - 1;
                updateIndel(isDel, index);
                refPos += len;
            } else if (c == 'N') {
                refPos += len;
            } else if (c == 'I') {
                // 与Del不同，Ins应该是在下一个cigar开始的位置，标记Ins
                bool forwardStrandRead = !sd.bw->GetReadNegativeStrandFlag();
                if (forwardStrandRead) {
                    updateIndel(isIns, readPos - 1);
                }
                readPos += len;
                if (!forwardStrandRead) {
                    updateIndel(isIns, readPos);
                }
            } else if (c == 'S') {
                readPos += len;
            }
        }
        nEvents += std::accumulate(isIns.begin(), isIns.end(), 0) + std::accumulate(isDel.begin(), isDel.end(), 0);
        PROF_GP_END(calc_snp);
        return nEvents;
    }

    // 读取给定区间的reference
    static inline void read_ref_base(AuxVar& aux, int64_t cur_pos, Interval& interval) {
        if (aux.ref_seq != NULL)
            free(aux.ref_seq);
        int tid = BamWrap::bam_tid(cur_pos);
        const char* chr = sam_hdr_tid2name(aux.header, tid);
        int seq_begin = BamWrap::bam_pos(interval.left);  //- aux.REFERENCE_HALF_WINDOW_LENGTH;
        int seq_end = BamWrap::bam_pos(interval.right);   //+ aux.REFERENCE_HALF_WINDOW_LENGTH;
        aux.ref_seq = faidx_fetch_seq(aux.faidx, chr, seq_begin, seq_end, &aux.ref_len);
        // aux.offset = aux.REFERENCE_HALF_WINDOW_LENGTH;
    }

    // 获取一行字符串
    static void get_line_from_buf(char* buf, int64_t total, int64_t* cur, string* line) {
        line->clear();
        if (*cur >= total)
            return;
        char b;
        while (*cur < total && (b = buf[(*cur)++]) != '\n') {
            line->push_back(b);
        }
    }

    // 计算与read有交叉的已知位点信息， 应该要判断一下，是按照read的范围去读取vcf，还是按照一个batch read的范围去读取
    static void calculateKnownSites(SamData& sd, vector<VCFParser>& vcfs, sam_hdr_t* samHdr, StableArray<uint8_t>& knownSites) {
        BamWrap* bw = sd.bw;
        int tid = bw->contig_id();
        int64_t startPos = bw->start_pos();  // 闭区间
        int64_t endPos = bw->end_pos();      // 闭区间
        knownSites.resize_fill(sd.read_len, 0);

        // update vcfs
        for (auto& vcf : vcfs) {
            // 清理旧的interval
            while (!vcf.knownSites.empty()) {
                auto& intv = vcf.knownSites.front();
                // spdlog::info("intv bam {}, {}", intv.right, startPos);
                if (intv.right < startPos)
                    vcf.knownSites.pop_front();
                else
                    break;
            }
            if (!vcf.knownSites.empty() && vcf.knownSites.back().left > endPos)
                continue;

            // 读取新的interval
            int64_t fpos, flen;
            endPos = std::max(startPos + MAX_SITES_INTERVAL, endPos);
            Interval readIntv(startPos, endPos);
            vcf.index.SearchInterval(startPos, endPos, &fpos, &flen);
            if (flen > 0) {
                vcf.inStm.seekg(fpos, ios::beg);
                if (flen > vcf.bufLen) {
                    vcf.bufLen = flen;
                    vcf.buf = (char*)realloc(vcf.buf, flen);
                }
                char* buf = vcf.buf;
                vcf.inStm.read(buf, flen);
                string line;
                int64_t cur = 0;
                get_line_from_buf(buf, flen, &cur, &line);
                while (line.size() > 0) {
                    stringstream ss_line(line);
                    string stid;
                    int tid, pos;
                    int64_t locus;
                    string id, ref;
                    ss_line >> stid >> pos >> id >> ref;
                    tid = sam_hdr_name2tid(samHdr, stid.c_str());
                    int64_t varStart = BamWrap::bam_global_pos(tid, pos);
                    Interval varIntv(varStart, varStart + ref.size() - 1);
                    if (readIntv.overlaps(varIntv)) {
                        vcf.knownSites.push_back(Interval(tid, pos - 1, pos - 1 + ref.size() - 1));  // 闭区间
                    }
                    get_line_from_buf(buf, flen, &cur, &line);
                }
            }
        }
        for (auto& vcf : vcfs) {
            for (auto& intv : vcf.knownSites) {
                // knownSite is outside clipping window for the read, ignore
                if (intv.right < sd.softStart())
                    continue;
                if (intv.left > sd.softEnd())
                    break;
                // 计算对应的位点
                ReadIdxCigar rcStart = sd.getReadIndexForReferenceCoordinate(intv.left);
                int featureStartOnRead = rcStart.readIdx == SamData::READ_INDEX_NOT_FOUND ? 0 : rcStart.readIdx;
                if (rcStart.cigarOp == 'D') {
                    --featureStartOnRead;
                }
                ReadIdxCigar rcEnd = sd.getReadIndexForReferenceCoordinate(intv.right);
                int featureEndOnRead = rcEnd.readIdx == SamData::READ_INDEX_NOT_FOUND ? sd.read_len : rcEnd.readIdx;
                if (featureStartOnRead > sd.read_len) {
                    featureStartOnRead = featureEndOnRead = sd.read_len;
                }
                for (int i = max(0, featureStartOnRead); i < min(sd.read_len, featureEndOnRead + 1); ++i) {
                    knownSites[i] = true;
                }
            }
        }
    }

    // 应该是计算一段数据的平均值
    static void calculateAndStoreErrorsInBlock(int i, int blockStartIndex, StableArray<int>& errorArr, StableArray<double>& fracErrs) {
        int totalErrors = 0;
        for (int j = max(0, blockStartIndex - 1); j <= i; j++) {
            totalErrors += errorArr[j];
        }
        for (int j = max(0, blockStartIndex - 1); j <= i; j++) {
            fracErrs[j] = ((double)totalErrors) / ((double)(i - max(0, blockStartIndex - 1) + 1));
        }
    }

    // 应该是用来处理BAQ的，把不等于特定BAQ分数的碱基作为一段数据统一处理
    static void calculateFractionalErrorArray(StableArray<int>& errorArr, StableArray<uint8_t>& baqArr, StableArray<double>& fracErrs) {
        fracErrs.resize_fill(baqArr.size(), 0.0);
        // errorArray和baqArray必须长度相同
        const int BLOCK_START_UNSET = -1;
        bool inBlock = false;
        int blockStartIndex = BLOCK_START_UNSET;
        int i;
        for (i = 0; i < fracErrs.size(); ++i) {
            if (baqArr[i] == NO_BAQ_UNCERTAINTY) {  // 不等于NO_BAQ_UNCERTAINTY的碱基当成一段数据来处理
                if (!inBlock) {
                    fracErrs[i] = errorArr[i];
                } else {
                    calculateAndStoreErrorsInBlock(i, blockStartIndex, errorArr, fracErrs);
                    inBlock = false;                      // reset state variables
                    blockStartIndex = BLOCK_START_UNSET;  // reset state variables
                }
            } else {
                inBlock = true;
                if (blockStartIndex == BLOCK_START_UNSET) {
                    blockStartIndex = i;
                }
            }
        }
        // 处理最后一个block
        if (inBlock) {
            calculateAndStoreErrorsInBlock(i - 1, blockStartIndex, errorArr, fracErrs);
        }
    }
};