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

     Copyright : All right reserved by ICT

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

#include <spdlog/spdlog.h>
#include <numeric>

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

struct RecalFuncs {
    //static constexpr int MAX_SITES_INTERVAL = 100000;
    static constexpr int MAX_SITES_INTERVAL = 21500; // 经验值，这个数读取vcf和计算的性能最好
    static constexpr uint8_t NO_BAQ_UNCERTAINTY = (uint8_t)'@';

    // 过滤掉bqsr过程不符合要求的bam数据
    static bool bqsrReadFilterOut(const bam1_t* b) {
        // 过滤掉unmapped的read
        if (b->core.qual == 0)  // mapping quality 0
            return true;
        if (b->core.qual == 255)  // mapping quality not available
            return true;
        if (b->core.flag & BAM_FUNMAP || b->core.tid == -1 || b->core.pos == -1) {  // unmapped
            return true;
        }
        if (b->core.flag & BAM_FSECONDARY) {  // secondary alignment
            return true;
        }
        if (b->core.flag & BAM_FDUP) {  // secondary alignment
            return true;
        }
        if (b->core.flag & BAM_FQCFAIL) {  // Not passing quality controls
            return true;
        }
        return false;
    }

    // 过滤掉apply bqsr过程不符合要求的bam数据
    static bool applyBqsrReadFilterOut(const bam1_t* b) {
        // 好像4.6.2的GATK版本的welformed filter啥也没过滤
        // b->core.flag & BAM_FUNMAP ||
        if ((b->core.tid == -1 || b->core.pos == -1) && !(b->core.flag & BAM_FMUNMAP)) {  // unmapped
            return true;
        }
        return false;
    }

    // 设置某个位置是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, int thid) {
        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(TP_read_ref);
        read_ref_base(aux, interval.left, interval);
        PROF_TP_END(TP_read_ref);
        const char *refBases = aux.ref_seq;

        // 2. 遍历cigar，计算每个碱基是否是SNP或Indel
        PROF_START(TP_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 = BaseUtils::basesAreEqual(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_TP_END(TP_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) {
        if (*cur >= total) {
            (*cur)++;
            return;
        }
        while (*cur < total && buf[(*cur)++] != '\n');
    }

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

        // update vcfs
        PROF_START(TP_read_vcf);
        for (auto& vcf : vcfs) {
            if (!vcf.knownSites.empty() && vcf.knownSites.back().left > endPos) {// 此时vcf的区域包含bam，不需要读取
                // 清理旧的interval，只有此时才有清理的必要
                while (!vcf.knownSites.empty()) {
                    auto& intv = vcf.knownSites.front();
                    if (intv.right < startPos)
                        vcf.knownSites.pop_front();
                    else
                        break;
                }
                continue;
            }
            vcf.knownSites.clear(); // 清空，因为后面会读入覆盖bam的所有vcf位点
            // 读取新的interval
            int64_t fpos, flen;
            endPos = std::max(startPos + sitesStride, endPos);
            Interval readIntv(startPos, endPos);
            vcf.index.SearchInterval(startPos, endPos, &fpos, &flen);
            // fprintf(gf[thid * 2 + idx], "%s %d %ld %ld %ld\n", bam_get_qname(sd.bw->b), sd.bw->b->core.flag, sd.rid, 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);
                int64_t cur = 0;
                int64_t start = 0;
                get_line_from_buf(buf, flen, &cur);
                while (cur > start + 1) {
                    const string line(buf + start, buf + cur - 1);
                    // spdlog::info("s: {}, e: {}, line: {}", start, cur, line);
                    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 - 1);  // vcf中的pos是1-based，程序中计算是0-based
                    Interval varIntv(varStart, varStart + ref.size() - 1);
                    if (varIntv.overlaps(readIntv)) {
                        vcf.knownSites.push_back(varIntv);  // 闭区间
                    }
                    start = cur;
                    get_line_from_buf(buf, flen, &cur);
                }
            }
        }
        PROF_TP_END(TP_read_vcf);
        // fprintf(gf[0], "%s %d %ld ", bam_get_qname(sd.bw->b), sd.bw->b->core.flag, sd.rid);
        PROF_START(TP_calc_skips);
        for (auto& vcf : vcfs) {
            for (auto& intv : vcf.knownSites) {
                // knownSite is outside clipping window for the read, ignore
#if 1
                if (intv.right < sd.softStart())
                    continue;
                if (intv.left > sd.softEnd())
                    break;
#else
                if (intv.right < sd.softStart() || intv.left > sd.softEnd())
                    continue;
#endif
                //fprintf(gf[0], "%ld %ld ", intv.left, intv.right);

                // 计算对应的位点
                ReadIdxCigar rcStart = sd.getReadIndexForReferenceCoordinate(intv.contigLeft());
                int featureStartOnRead = rcStart.readIdx == SamData::READ_INDEX_NOT_FOUND ? 0 : rcStart.readIdx;
                if (rcStart.cigarOp == 'D') {
                    --featureStartOnRead;
                }
                ReadIdxCigar rcEnd = sd.getReadIndexForReferenceCoordinate(intv.contigRight());
                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;
                }
            }
        }
        PROF_TP_END(TP_calc_skips);
        //fprintf(gf[0], "\n");
    }

    // 应该是计算一段数据的平均值
    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);
        }
    }


    // for apply bqsr
    /**
     * Quality score recalibration algorithm works as follows:
     * - Start with the (approximate, or "estimated") reported quality score. (Approximation occurs when marginalizing/collapsing
     * over the reported qualities for each read group).
     * - Compute (retrieve?) the empirical quality score using the per-read group datum (i.e. counts). Call it y_1.
     * - Use y_1 just computed as the prior for the empirical quality score for the datum for the 2-tuple ( read group, quality score). Call it y_2.
     * - Use y_2 as the prior to compute the empirical quality for the 3-tuple ( read-group, quality-score, special covariate ). Call it y_3 for the
     * context covariate. Similarly define y_4 for the cycle covariate. Let d_3 = y_3 - y_2, d_4 = y_4 - y_2.
     * - (final recalibrated score) = y_2 + d_3 + d_4 = y_3 + y_4 - y_2.
     *
     * @param priorQualityScore the prior quality score (in log space). It is either the "estimated" or collapsed reported quality score
     *                          for the read group, or the constant prior if given. This value has type double because of the "combine" (or collapse)
     * operation that collapses the quality scores represented within the same read group.
     * @param readGroupDatum the RecalDatum object for a particular read group at hand. May be null.
     * @param qualityScoreDatum the RecalDatum object for a particular (read group, reported quality) tuple at hand. May be null.
     * @param specialCovariateDatums the array of RecalDatum objects for the non-required covariates (cycle and context covariates by default).
     * @return
     */
    static double hierarchicalBayesianQualityEstimate(double priorQualityScore, RecalDatum& readGroupDatum, RecalDatum& qualityScoreDatum,
                                               RecalDatum& contextDatum, RecalDatum& cycleDatum) {
        double empiricalQualityForReadGroup =
            readGroupDatum.getNumObservations() == 0 ? priorQualityScore : readGroupDatum.getEmpiricalQuality(priorQualityScore);
        double posteriorEmpiricalQualityForReportedQuality = qualityScoreDatum.getNumObservations() == 0
                                                                 ? empiricalQualityForReadGroup
                                                                 : qualityScoreDatum.getEmpiricalQuality(empiricalQualityForReadGroup);

        double deltaSpecialCovariates = 0.0;
        // At this point we stop being iterative; the special covariates (context and cycle by default) are treated differently.
        if (contextDatum.getNumObservations() > 0) {
            // TODO: the prior is ignored if the empirical quality for the datum is already cached.
            deltaSpecialCovariates +=
                contextDatum.getEmpiricalQuality(posteriorEmpiricalQualityForReportedQuality) - posteriorEmpiricalQualityForReportedQuality;
        }
        if (cycleDatum.getNumObservations() > 0) {
            // TODO: the prior is ignored if the empirical quality for the datum is already cached.
            deltaSpecialCovariates +=
                cycleDatum.getEmpiricalQuality(posteriorEmpiricalQualityForReportedQuality) - posteriorEmpiricalQualityForReportedQuality;
        }

        return posteriorEmpiricalQualityForReportedQuality + deltaSpecialCovariates;
    }

    // recalibrated quality is bound between 1 and MAX_QUAL
    static uint8_t getBoundedIntegerQual(double recalibratedQualDouble) {
        return QualityUtils::boundQual(MathUtils::fastRound(recalibratedQualDouble), RecalDatum::MAX_RECALIBRATED_Q_SCORE);
    }
};