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

     Copyright : All right reserved by ICT

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

#include <spdlog/spdlog.h>

#include <cstdio>

#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"

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 offset = 0; offset < readLength; ++offset) {
            // if (read.rid == 46114) {
            //     fprintf(gf[3], "%d %d\n", offset, info.skips[offset] ? 1 : 0);
            // }
            if (!info.skips[offset]) {
                // if (true){ // 不跳过当前位置
                for (int idx = 0; idx < validEventTypes.size(); ++idx) {
                    // 获取四个值，readgroup / qualityscore / context / cycle
                    EventTypeValue& event = validEventTypes[idx];
                    vector<int>& covariatesAtOffset = readCovars[event.index][offset];
                    uint8_t qual = info.getQual(event, offset);
                    double isError = info.getErrorFraction(event, offset);

                    int readGroup = covariatesAtOffset[ReadGroupCovariate::index];
                    int baseQuality = covariatesAtOffset[BaseQualityCovariate::index];

                    // 处理base quality score协变量
                    // RecalUtils::IncrementDatum3keys(qualityScoreTable, qual, isError, readGroup, baseQuality, event.index);
                    //if (read.rid == 46114) {
                        // fprintf(gf[3], "%d %d %f\n", offset, baseQuality, isError);
                    //}
                    qualityScoreTable[readGroup][baseQuality][event.index].increment(1, isError, baseQuality);

                    auto& d = qualityScoreTable[readGroup][baseQuality][event.index];
                    // spdlog::info("isError {} : {}, mis {}, obs {}", isError, info.snp_errs[offset], d.numMismatches, d.numObservations);

                    // 处理context covariate
                    int contextCovariate = covariatesAtOffset[ContextCovariate::index];
                    if (contextCovariate >= 0)
                        contextTable[readGroup][baseQuality][contextCovariate][event.index].increment(1, isError, baseQuality);
                    // RecalUtils::IncrementDatum4keys(nsgv::gRecalTables.contextTable, qual, isError, readGroup, baseQuality, contextCovariate,
                    //                                 event.index);
                    // 处理cycle covariate
                    int cycleCovariate = covariatesAtOffset[CycleCovariate::index];
                    if (cycleCovariate >= 0)
                        cycleTable[readGroup][baseQuality][cycleCovariate][event.index].increment(1, isError, baseQuality);
                    //     RecalUtils::IncrementDatum4keys(nsgv::gRecalTables.cycleTable, qual, isError, readGroup, baseQuality, cycleCovariate,
                    //     event.index);
                }
            }
        }
        // fprintf(gf[3], "%ld %d %ld\n", read.rid, read.read_len, read.start_pos+1);
        // for (auto& arr1 : qualityScoreTable.data) {
        //     for (size_t si = 0; si < arr1.size(); ++si) {
        //         for (auto &val : arr1[si]) {
        //             if (val.numObservations > 0)
        //                 fprintf(gf[3], "%ld %f %f ", val.numObservations, val.getNumMismatches(), val.reportedQuality);
        //         }
        //     }
        // }
        // fprintf(gf[3], "\n");

        // fprintf(gf[3], "%ld %d %ld\n", read.rid, read.read_len, read.start_pos+1);
        // for (auto& arr1 : contextTable.data) {
        //     for (size_t si = 0; si < arr1.size(); ++si) {
        //         for (auto &arr2 : arr1[si]) {
        //             for (auto& val : arr2) {
        //                 if (val.numObservations > 0)
        //                     fprintf(gf[3], "%ld %f %f ", val.numObservations, val.getNumMismatches(), val.reportedQuality);
        //             }
        //         }
        //     }
        // }
        // fprintf(gf[3], "\n");
    }

    // 输出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\n", REPORT_HEADER_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[k1],
                    ReportUtil::ToString(EventType::EVENTS[k2].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[k1],
                    ReportUtil::ToString(k2),
                    ReportUtil::ToString(EventType::EVENTS[k3].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[k1],
                    ReportUtil::ToString(k2),
                    ReportUtil::ToString(ContextCovariate::ContextFromKey(k3)),
                    "Context",
                    ReportUtil::ToString(EventType::EVENTS[k4].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[k1],
                    ReportUtil::ToString(k2),
                    ReportUtil::ToString(CycleCovariate::CycleFromKey(k3)),
                    "Cycle",
                    ReportUtil::ToString(EventType::EVENTS[k4].representation),
                    ReportUtil::ToString(dat.getEmpiricalQuality(), 4),
                    ReportUtil::ToString(dat.getNumObservations()),
                    ReportUtil::ToString(dat.getNumMismatches(), 2)
                });
            }
        });
        table.write(fpout);
    }
};