/*
     Description: bqsr计算过程的辅助类

     Copyright : All right reserved by ICT

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

#include <spdlog/spdlog.h>

#include <cstdio>
#include <iostream>
#include <fstream>

#include "bqsr_args.h"
#include "nested_array.h"
#include "quant_info.h"
#include "recal_datum.h"
#include "recal_tables.h"
#include "util/report_table.h"
#include "util/utils.h"
#include "covariate.h"
#include "read_recal_info.h"
#include "util/debug.h"

struct RecalUtils {

    static constexpr int EMPIRICAL_QUAL_DECIMAL_PLACES = 4;
    static constexpr int REPORTED_QUALITY_DECIMAL_PLACES = 4;
    static constexpr int NUMBER_ERRORS_DECIMAL_PLACES = 2;

    static vector<EventTypeValue> validEventTypes; // 真正需要计算的event types

    static void initEventTypes(bool computeIndelBQSRTables) {
        validEventTypes.clear();
        validEventTypes.push_back(EventType::BASE_SUBSTITUTION);
        if (computeIndelBQSRTables) {
            validEventTypes.push_back(EventType::BASE_INSERTION);
            validEventTypes.push_back(EventType::BASE_DELETION);
        }
    }

    /**
     * Update the recalibration statistics using the information in recalInfo.
     *
     * Implementation detail: we only populate the quality score table and the optional tables.
     * The read group table will be populated later by collapsing the quality score table.
     *
     * @param recalInfo data structure holding information about the recalibration values for a single read
     */
    static void updateRecalTablesForRead(ReadRecalInfo& info, RecalTables& recalTables) {
        SamData& read = info.read;
        PerReadCovariateMatrix& readCovars = info.covariates;
        Array3D<RecalDatum>& qualityScoreTable = recalTables.qualityScoreTable;
        Array4D<RecalDatum>& contextTable = recalTables.contextTable;
        Array4D<RecalDatum>& cycleTable = recalTables.cycleTable;

        int readLength = read.read_len;
        for (int idx = 0; idx < validEventTypes.size(); ++idx) {
            for (int offset = 0; offset < readLength; ++offset) {
                if (!info.skips[offset]) {
                    // 获取四个值，readgroup / qualityscore / context / cycle
                    EventTypeValue& event = validEventTypes[idx];
                    CovariateValues& cv = readCovars[event.index][offset];
                    uint8_t qual = info.getQual(event, offset);
                    double isError = info.getErrorFraction(event, offset);
                    // fprintf(gf[4], "%d %d %f\n", offset, qual, isError);

                    // 处理quality score covariate
                    qualityScoreTable(event.index, cv.readGroup, cv.baseQuality).increment(1, isError, cv.baseQuality);

                    // 处理context covariate
                    if (cv.context >= 0)
                        contextTable(event.index, cv.readGroup, cv.baseQuality, cv.context).increment(1, isError, cv.baseQuality);

                    // 处理cycle covariate
                    if (cv.cycle >= 0)
                        cycleTable(event.index, cv.readGroup, cv.baseQuality, cv.cycle).increment(1, isError, cv.baseQuality);
                }
            }
        }
    }

    // 输出bqsr报告
    static void outputRecalibrationReport(const BQSRArg& RAC, const QuantizationInfo& quantInfo, const RecalTables& recalTables) {
        spdlog::info("output tables");
        // 输出文件
        FILE* fpout = fopen(RAC.OUTPUT_FILE.c_str(), "w");
        if (!fpout) {
            fprintf(stderr, "Failed to open file %s.\n", RAC.OUTPUT_FILE.c_str());
            exit(1);
        }
        // 输出version信息
        fprintf(fpout, "%s:%s\n", REPORT_HEADER_VERSION, REPORT_HEADER_MINOR_VERSION);
        // 输出参数信息
        outputArgsTable(RAC, fpout);
        // 输出量化质量分数信息
        outputQuantTable(quantInfo, fpout);
        // 输出协变量信息
        outputCovariateTable(recalTables, fpout);

        fclose(fpout);
    }

    // 输出运行时使用的参数信息
    static void outputArgsTable(const BQSRArg & p, FILE * fpout) {
        ReportTable table("Arguments", "Recalibration argument collection values used in this run");
        table.addColumn({"Argument", "%s"});
        table.addColumn({"Value", ""});
        // 添加行数据
        table.addRowData({"covariate", "ReadGroupCovariate,QualityScoreCovariate,ContextCovariate,CycleCovariate"});
        table.addRowData({"no_standard_covs", ReportUtil::ToString(p.DO_NOT_USE_STANDARD_COVARIATES)});
        table.addRowData({"run_without_dbsnp", ReportUtil::ToString(p.RUN_WITHOUT_DBSNP)});
        table.addRowData({"solid_recal_mode", ReportUtil::ToString(p.SOLID_RECAL_MODE)});
        table.addRowData({"solid_nocall_strategy", ReportUtil::ToString(p.SOLID_NOCALL_STRATEGY)});
        table.addRowData({"mismatches_context_size", ReportUtil::ToString(p.MISMATCHES_CONTEXT_SIZE)});
        table.addRowData({"indels_context_size", ReportUtil::ToString(p.INDELS_CONTEXT_SIZE)});
        table.addRowData({"mismatches_default_quality", ReportUtil::ToString(p.MISMATCHES_DEFAULT_QUALITY)});
        table.addRowData({"deletions_default_quality", ReportUtil::ToString(p.DELETIONS_DEFAULT_QUALITY)});
        table.addRowData({"insertions_default_quality", ReportUtil::ToString(p.INSERTIONS_DEFAULT_QUALITY)});
        table.addRowData({"maximum_cycle_value", ReportUtil::ToString(p.MAXIMUM_CYCLE_VALUE)});
        table.addRowData({"low_quality_tail", ReportUtil::ToString(p.LOW_QUAL_TAIL)});
        table.addRowData({"default_platform", ReportUtil::ToString(p.DEFAULT_PLATFORM)});
        table.addRowData({"force_platform", ReportUtil::ToString(p.FORCE_PLATFORM)});
        table.addRowData({"quantizing_levels", ReportUtil::ToString(p.QUANTIZING_LEVELS)});
        table.addRowData({"recalibration_report", ReportUtil::ToString(p.existingRecalibrationReport)});
        table.addRowData({"binary_tag_name", ReportUtil::ToString(p.BINARY_TAG_NAME)});
        // 输出到文件
        table.write(fpout);
    }

    // 输出量化质量分数table
    static void outputQuantTable(const QuantizationInfo& q, FILE* fpout) {
        ReportTable table("Quantized", "Quality quantization map");
        table.addColumn({"QualityScore", "%d"});
        table.addColumn({"Count", "%d"});
        table.addColumn({"QuantizedScore", "%d"});

        for (int qual = 0; qual <= QualityUtils::MAX_SAM_QUAL_SCORE; ++qual) {
            table.addRowData({ReportUtil::ToString(qual), ReportUtil::ToString(q.empiricalQualCounts[qual]), ReportUtil::ToString(q.quantizedQuals[qual])});
        }

        table.write(fpout);
    }

    // 输出协变量table
    static void outputCovariateTable(const RecalTables &r, FILE* fpout) {
        // 1. read group covariates
        outputReadGroupTable(r, fpout);
        // 2. quality score covariates
        outputQualityScoreTable(r, fpout);
        // 3. context and cycle covariates
        outputContextCycleTable(r, fpout);
    }

    // read group table
    static void outputReadGroupTable(const RecalTables& r, FILE* fpout) {
        ReportTable table("RecalTable0");
        table.addColumn({"ReadGroup", "%s"});
        table.addColumn({"EventType", "%s"});
        table.addColumn({"EmpiricalQuality", "%.4f"});
        table.addColumn({"EstimatedQReported", "%.4f"});
        table.addColumn({"Observations", "%d"});
        table.addColumn({"Errors", "%.2f"});

        _Foreach2DK(r.readGroupTable, datum, {
            RecalDatum &dat = const_cast<RecalDatum&>(datum);
            if (dat.getNumObservations() > 0) {
                table.addRowData({
                    ReadGroupCovariate::IdToRg[k2],
                    ReportUtil::ToString(EventType::EVENTS[k1].representation),
                    ReportUtil::ToString(dat.getEmpiricalQuality(), 4),
                    ReportUtil::ToString(dat.getReportedQuality(), 4),
                    ReportUtil::ToString(dat.getNumObservations()),
                    ReportUtil::ToString(dat.getNumMismatches(), 2)
                });
            }
        });
        table.write(fpout);
    }

    // quality table
    static void outputQualityScoreTable(const RecalTables& r, FILE* fpout) {
        ReportTable table("RecalTable1");
        table.addColumn({"ReadGroup", "%s"});
        table.addColumn({"QualityScore", "%d"});
        table.addColumn({"EventType", "%s"});
        table.addColumn({"EmpiricalQuality", "%.4f"});
        table.addColumn({"Observations", "%d"});
        table.addColumn({"Errors", "%.2f"});

        _Foreach3DK(r.qualityScoreTable, datum, {
            RecalDatum &dat = const_cast<RecalDatum&>(datum);
            if (dat.getNumObservations() > 0) {
                table.addRowData({
                    ReadGroupCovariate::IdToRg[k2],
                    ReportUtil::ToString(k3),
                    ReportUtil::ToString(EventType::EVENTS[k1].representation),
                    ReportUtil::ToString(dat.getEmpiricalQuality(), 4),
                    ReportUtil::ToString(dat.getNumObservations()),
                    ReportUtil::ToString(dat.getNumMismatches(), 2)
                });
            }
        });
        table.write(fpout);
    }

    static void outputContextCycleTable(const RecalTables& r, FILE* fpout) {
        ReportTable table("RecalTable2");
        table.addColumn({"ReadGroup", "%s"});
        table.addColumn({"QualityScore", "%d"});
        table.addColumn({"CovariateValue", "%s"});
        table.addColumn({"CovariateName", "%s"});
        table.addColumn({"EventType", "%s"});
        table.addColumn({"EmpiricalQuality", "%.4f"});
        table.addColumn({"Observations", "%d"});
        table.addColumn({"Errors", "%.2f"});

        _Foreach4DK(r.contextTable, datum, {
            RecalDatum &dat = const_cast<RecalDatum&>(datum);
            if (dat.getNumObservations() > 0) {
                table.addRowData({
                    ReadGroupCovariate::IdToRg[k2],
                    ReportUtil::ToString(k3),
                    ReportUtil::ToString(ContextCovariate::ContextFromKey(k4)),
                    "Context",
                    ReportUtil::ToString(EventType::EVENTS[k1].representation),
                    ReportUtil::ToString(dat.getEmpiricalQuality(), 4),
                    ReportUtil::ToString(dat.getNumObservations()),
                    ReportUtil::ToString(dat.getNumMismatches(), 2)
                });
            }
        });

        _Foreach4DK(r.cycleTable, datum, {
            RecalDatum &dat = const_cast<RecalDatum&>(datum);
            if (dat.getNumObservations() > 0) {
                table.addRowData({
                    ReadGroupCovariate::IdToRg[k2],
                    ReportUtil::ToString(k3),
                    ReportUtil::ToString(CycleCovariate::CycleFromKey(k4)),
                    "Cycle",
                    ReportUtil::ToString(EventType::EVENTS[k1].representation),
                    ReportUtil::ToString(dat.getEmpiricalQuality(), 4),
                    ReportUtil::ToString(dat.getNumObservations()),
                    ReportUtil::ToString(dat.getNumMismatches(), 2)
                });
            }
        });
        table.write(fpout);
    }

    // 解析bqsr table文件
    static void readRecalTables(const string& recalTableFn, BQSRArg& RAC, vector<uint8_t>& quantizedScore, RecalTables& recalTables) {
        std::ifstream infs(recalTableFn, ifstream::in);
        string line;
        // 读取文件头
        std::getline(infs, line);
        spdlog::info("header: {}", line);
        if (line.find(REPORT_HEADER_VERSION) != 0) {
            spdlog::error("BQSR table version not supported! {} {}", line, __LINE__);
            exit(1);
        }

        // 读取参数表
        readArgsTable(infs, RAC);

        readQuantTable(infs, quantizedScore);

        readReadGroupTable(infs, recalTables);

        readQualityScoreTable(infs, recalTables);

        readContextCycleTable(infs, recalTables);

        infs.close();
    }

    template<typename Arr>
    static void split(const string& str, Arr& dat) {
        dat.clear();
        std::istringstream iss(str);
        string item;
        while (iss >> item) {
            dat.push_back(item);
        }
    }

    template <typename Arr>
    static void split(const string& str, const char delimiter, Arr& dat) {
        dat.clear();
        std::istringstream iss(str);
        string token;
        while (std::getline(iss, token, delimiter)) {
            dat.push_back(token);
        }
    }

    // 读取参数表
    static void readArgsTable(std::ifstream& infs, BQSRArg& p) {
        string line;
        StableArray<string> dat;
        std::getline(infs, line);
        while (line.empty() && std::getline(infs, line));  // 跳过空行
        split(line, ':', dat);
        int argRowCnt = std::stoi(dat[3]);
        std::getline(infs, line); // 忽略表格说明的一行
        std::getline(infs, line); // 忽略表头column说明
        for (int i = 0; i < argRowCnt; ++i) {
            std::getline(infs, line);
            split(line, dat);
            if ("covariate" == dat[0]) {
            } else if ("no_standard_covs" == dat[0]) {
                p.DO_NOT_USE_STANDARD_COVARIATES = ReportUtil::ToBool(dat[1]);
            } else if ("run_without_dbsnp" == dat[0]) {
                p.RUN_WITHOUT_DBSNP = ReportUtil::ToBool(dat[1]);
            } else if ("solid_recal_mode" == dat[0]) {
                p.SOLID_RECAL_MODE = ReportUtil::ParseString(dat[1]);
            } else if ("solid_nocall_strategy" == dat[0]) {
                p.SOLID_NOCALL_STRATEGY = ReportUtil::ParseString(dat[1]);
            } else if ("mismatches_context_size" == dat[0]) {
                p.MISMATCHES_CONTEXT_SIZE = ReportUtil::ToInt(dat[1]);
            } else if ("indels_context_size" == dat[0]) {
                p.INDELS_CONTEXT_SIZE = ReportUtil::ToInt(dat[1]);
            } else if ("mismatches_default_quality" == dat[0]) {
                p.MISMATCHES_DEFAULT_QUALITY = ReportUtil::ToInt(dat[1]);
                // spdlog::info("int8_t : {}", (int)p.MISMATCHES_DEFAULT_QUALITY);
            } else if ("deletions_default_quality" == dat[0]) {
                p.DELETIONS_DEFAULT_QUALITY = ReportUtil::ToInt(dat[1]);
            } else if ("insertions_default_quality" == dat[0]) {
                p.INSERTIONS_DEFAULT_QUALITY = ReportUtil::ToInt(dat[1]);
            } else if ("maximum_cycle_value" == dat[0]) {
                p.MAXIMUM_CYCLE_VALUE = ReportUtil::ToInt(dat[1]);
            } else if ("low_quality_tail" == dat[0]) {
                p.LOW_QUAL_TAIL = ReportUtil::ToInt(dat[1]);
            } else if ("default_platform" == dat[0]) {
                p.DEFAULT_PLATFORM = ReportUtil::ParseString(dat[1]);
            } else if ("force_platform" == dat[0]) {
                p.FORCE_PLATFORM = ReportUtil::ParseString(dat[1]);
            } else if ("quantizing_levels" == dat[0]) {
                p.QUANTIZING_LEVELS = ReportUtil::ToInt(dat[1]);
            } else if ("recalibration_report" == dat[0]) {
                p.existingRecalibrationReport = ReportUtil::ParseString(dat[1]);
            } else if ("binary_tag_name" == dat[0]) {
                p.BINARY_TAG_NAME = ReportUtil::ParseString(dat[1]);
                // spdlog::info("BINARY_TAG_NAME : {}", p.BINARY_TAG_NAME);
            } else {
                spdlog::error("unknown argument : {} {}", dat[0], dat[1]);
            }
        }
    }

    // 读取量化质量分数信息表
    static void readQuantTable(std::ifstream& infs, vector<uint8_t>& quantizedScore) {
        string line;
        StableArray<string> dat;
        std::getline(infs, line);
        while (line.empty() && std::getline(infs, line)); // 跳过空行
        split(line, ':', dat);
        int argRowCnt = std::stoi(dat[3]);
        std::getline(infs, line);  // 忽略表格说明的一行
        std::getline(infs, line);  // 忽略表头column说明
        quantizedScore.resize(argRowCnt);
        for (int i = 0; i < argRowCnt; ++i) {
            std::getline(infs, line);
            // spdlog::info("{}", line);
            split(line, dat);
            quantizedScore[ReportUtil::ToInt(dat[0])] = ReportUtil::ToInt(dat[2]);
            // spdlog::info("{}, {}", ReportUtil::ToInt(dat[0]), quantizedScore[ReportUtil::ToInt(dat[0])]);
        }
    }

    // 读取read group表
    static void readReadGroupTable(std::ifstream& infs, RecalTables& recalTables) {
        string line;
        StableArray<string> dat;
        std::getline(infs, line);
        while (line.empty() && std::getline(infs, line));  // 跳过空行
        split(line, ':', dat);
        int argRowCnt = std::stoi(dat[3]);
        std::getline(infs, line);  // 忽略表格说明的一行
        std::getline(infs, line);  // 忽略表头column说明
        vector<vector<string>> dats(argRowCnt);
        int rgId = 0;
        for (int i = 0; i < argRowCnt; ++i) {
            std::getline(infs, line);
            split(line, dats[i]);
            if (ReadGroupCovariate::RgToId.find(dats[i][0]) == ReadGroupCovariate::RgToId.end()) {
                ReadGroupCovariate::RgToId[dats[i][0]] = rgId;
                ReadGroupCovariate::IdToRg[rgId] = dats[i][0];
                ++rgId;
            }
        }
        recalTables.initReadGroupTable(rgId);
        // spdlog::info("read group num: {}, {}", rgId, ReadGroupCovariate::IdToRg[0]);
        for (int i = 0; i < dats.size(); ++i) {
            int k2 = ReadGroupCovariate::RgToId[dats[i][0]];
            int k1 = EventType::GetIndexForEventRep(dats[i][1][0]);
            auto& datum = recalTables.readGroupTable(k1, k2);
            // datum.setEmpiricalQuality(ReportUtil::ToDouble(dats[i][2]));  // 这个数值在GATK4里是需要重新计算的
            datum.setEmpiricalQuality(RecalDatum::UNINITIALIZED_EMPIRICAL_QUALITY);
            datum.setReportedQuality(ReportUtil::ToDouble(dats[i][3]));
            datum.setNumObservations(ReportUtil::ToUint64(dats[i][4]));
            datum.setNumMismatches(ReportUtil::ToDouble(dats[i][5]));
        }
    }

    // 读取质量分数表
    static void readQualityScoreTable(std::ifstream& infs, RecalTables& recalTables) {
        string line;
        StableArray<string> dat;
        std::getline(infs, line);
        while (line.empty() && std::getline(infs, line));  // 跳过空行
        split(line, ':', dat);
        int argRowCnt = std::stoi(dat[3]);
        std::getline(infs, line);  // 忽略表格说明的一行
        std::getline(infs, line);  // 忽略表头column说明
        recalTables.initQualityScoreTable();
        int rgId = 0;
        for (int i = 0; i < argRowCnt; ++i) {
            std::getline(infs, line);
            split(line, dat);
            int k2 = ReadGroupCovariate::RgToId[dat[0]];
            int k3 = ReportUtil::ToInt(dat[1]);
            int k1 = EventType::GetIndexForEventRep(dat[2][0]);
            auto& datum = recalTables.qualityScoreTable(k1, k2, k3);
            // datum.setEmpiricalQuality(ReportUtil::ToDouble(dat[3]));  // 这个数值在GATK4里是需要重新计算的
            datum.setEmpiricalQuality(RecalDatum::UNINITIALIZED_EMPIRICAL_QUALITY);
            datum.setReportedQuality((double)k3);
            datum.setNumObservations(ReportUtil::ToUint64(dat[4]));
            datum.setNumMismatches(ReportUtil::ToDouble(dat[5]));
        }
    }

    // 读取context和cycle表
    static void readContextCycleTable(std::ifstream& infs, RecalTables& recalTables) {
        string line;
        StableArray<string> dat;
        std::getline(infs, line);
        while (line.empty() && std::getline(infs, line));  // 跳过空行
        split(line, ':', dat);
        int argRowCnt = std::stoi(dat[3]);
        std::getline(infs, line);  // 忽略表格说明的一行
        std::getline(infs, line);  // 忽略表头column说明
        recalTables.initContextTable();
        recalTables.initCycleTable();
        for (int i = 0; i < argRowCnt; ++i) {
            std::getline(infs, line);
            split(line, dat);
            int k2 = ReadGroupCovariate::RgToId[dat[0]];
            int k3 = ReportUtil::ToInt(dat[1]);
            int k4 = 0;
            int k1 = EventType::GetIndexForEventRep(dat[4][0]);
            RecalDatum* dp;
            if (dat[3] == "Cycle") {
                k4 = CycleCovariate::KeyFromCycle(ReportUtil::ToInt(dat[2]));
                dp = &recalTables.cycleTable(k1, k2, k3, k4);
            } else if (dat[3] == "Context") {
                k4 = ContextCovariate::KeyFromContext(dat[2]);
                dp = &recalTables.contextTable(k1, k2, k3, k4);
            } else {
                spdlog::error("Unknown CovariateName {}", dat[3]);
            }
            // dp->setEmpiricalQuality(ReportUtil::ToDouble(dat[5])); // 这个数值在GATK4里是需要重新计算的
            dp->setEmpiricalQuality(RecalDatum::UNINITIALIZED_EMPIRICAL_QUALITY);  // 这个数值在GATK4里是需要重新计算的
            dp->setReportedQuality((double)k3);
            dp->setNumObservations(ReportUtil::ToUint64(dat[6]));
            dp->setNumMismatches(ReportUtil::ToDouble(dat[7]));
        }
    }
};