main.cpp

#include <htslib/synced_bcf_reader.h>
#include <htslib/vcf.h>
#include <htslib/hts.h>
#include <iostream>
#include <string>
#include <vector>

#include "CLI11.hpp"

struct FilterPositions {
    int mapqual_idx;
    int orientation_idx;
    int strandbias_idx;
};

FilterPositions cache_filter_positions(bcf_hdr_t *hdr) {
    int mapqual_idx = -1;
    int orientation_idx = -1;
    int strandbias_idx = -1;
    
    for (int i = 0; i < hdr->n[BCF_DT_ID]; ++i) {
        const bcf_idinfo_t *idinfo = hdr->id[BCF_DT_ID][i].val;
        if (idinfo && idinfo->hrec[BCF_HL_FLT]) {
            const char *filter_name = hdr->id[BCF_DT_ID][i].key;
            if (strcmp(filter_name, "map_qual") == 0) {
                mapqual_idx = i;
            } else if (strcmp(filter_name, "orientation") == 0) {
                orientation_idx = i;
            } else if (strcmp(filter_name, "strand_bias") == 0) {
                strandbias_idx = i;
            }
        }
    }
    return {mapqual_idx, orientation_idx, strandbias_idx};
}

struct HeaderInfo {
    std::vector<std::string> sample_names;
    int number_of_samples;
    FilterPositions filter_positions;
};

std::vector<HeaderInfo> process_headers(bcf_srs_t *sr, const std::vector<std::string>& filenames) {
    int nfiles = filenames.size();
    std::vector<HeaderInfo> header_infos(nfiles);

    for (int i = 0; i < nfiles; ++i) {
        bcf_hdr_t *hdr = sr->readers[i].header;
        int n_samples = bcf_hdr_nsamples(hdr);
        header_infos[i].number_of_samples = n_samples;

        for (int j = 0; j < n_samples; ++j) {
            const char *sample_name = hdr->samples[j];
            header_infos[i].sample_names.push_back(sample_name);
        }

        auto positions = cache_filter_positions(hdr);
        header_infos[i].filter_positions = positions;
    }
    return header_infos;
}

class HtsBuffer {
private:
    int32_t *data = nullptr;
    int capacity = 0;
    int size = 0;

public:
    HtsBuffer() = default;
    HtsBuffer(const HtsBuffer &) = delete; // Disable copy constructor
    HtsBuffer &operator=(const HtsBuffer &) = delete; // Disable copy assignment

    HtsBuffer(HtsBuffer &&other) noexcept : data(other.data), capacity(other.capacity) {
        other.data = nullptr;
        other.capacity = 0;
    }

    HtsBuffer &operator=(HtsBuffer &&other) noexcept {
        if (this != &other) {
            if (data) {
                free(data);
            }
            data = other.data;
            capacity = other.capacity;
            other.data = nullptr;
            other.capacity = 0;
        }
        return *this;
    }
    
    ~HtsBuffer() {
        free(data);
    }

    int fetch(bcf_hdr_t *header, bcf1_t *rec) {
        int ret = bcf_get_format_int32(header, rec, "AD", &data, &capacity);
        if (ret < 0) {
            std::cerr << "WARN: Failed to get FORMAT/AD for record at position " << rec->pos + 1 << std::endl;
            return -1; // Return -1 on failure
        }
        size = ret; // Store the number of AD values fetched
        return size; // Return the number of AD values fetched
    }
    
    int get_size() const { return size; }
    int get_capacity() const { return capacity; }

    int32_t &operator[](int index) { return data[index]; }
    const int32_t &operator[](int index) const { return data[index]; }
};
  

int32_t get_alt_allele_depth(bcf_hdr_t *header, bcf1_t *reci, const HeaderInfo &header_info, HtsBuffer &ad_buffer) {
  
    // Extract the Allele Depth (AD) values - populates an array of int values for each allele in each sample
    int ret = ad_buffer.fetch(header, reci);

    if (ret < 0) {
        std::cerr << "WARN: Failed to get FORMAT/AD for file: " << std::endl;
        return bcf_int32_missing; // Return missing value if AD extraction fails
    }
    
    // Sanity check of the data extraction
    int nsamples = header_info.number_of_samples;
    if (ret != nsamples * reci->n_allele) { // Expected number of AD values is // n_samples * n_allele
        std::cerr << "WARN: Number of AD values (" << ret << ") does not match expected (" << nsamples * reci->n_allele << ")" << std::endl;
    }

    // Indexing into the AD array to get the last sample's ALT allele depth
    int target_allele = 1; // Want the ALT allele of a biallelic site
    int target_sample = nsamples - 1; // Want the last sample
    int index = target_sample * reci->n_allele + target_allele;

    if (index < 0 || index > ret - 1) {
        std::cerr << "WARN: Expected index (" << index << ") is out of bounds for AD values array of size " << ret << std::endl;
    }

    if (index != ret - 1) {
        std::cerr << "WARN: Expected index (" << index << ") does not match the last index (" << ret - 1 << ") for AD values array of size " << ret << std::endl;
    }
                    
    return ad_buffer[ret - 1];
}

struct Config {
    std::string whitelist_vcf;
    std::vector<std::string> sample_vcfs;
    std::string pass_output_vcf = "-"; // Default to stdout
    std::string fail_output_vcf = "";
    std::string format = "z"; // Default output to vcf.gz
};

Config parse_args(int argc, char **argv) {
    CLI::App app{"VCF Flag based filtering"};
    Config config;

    app.add_option("-w,--whitelist", config.whitelist_vcf, "Whitelist VCF file")
        ->required()
        ->check(CLI::ExistingFile);

    app.add_option("inputs", config.sample_vcfs, "Input VCF files")
        ->required()
        ->check(CLI::ExistingFile)
        ->expected(1, -1); // At least one input file

    app.add_option("-o,--output-pass", config.pass_output_vcf,
                   "Output VCF file for passing variants (default: stdout)");

    app.add_option("-f,--output-fail", config.fail_output_vcf,
                   "Output VCF file for failing variants (optional)");

    app.add_option("-O,--output-type", config.format, "htslib output format (u|b|v|z, default: z)")
        ->check(CLI::IsMember({"u", "b", "v", "z"}));

    try {
        app.parse(argc, argv);
    } catch (const CLI::ParseError &e) {
        std::exit(app.exit(e));
    }
    
    return config;   
}

int main(int argc, char** argv) {
    Config config;
    try {
        config = parse_args(argc, argv);
    } catch (const std::exception &e) {
        std::cerr << "Error parsing command line arguments: " << e.what() << std::endl;
        return 1;
    }

    std::vector<std::string> filenames;
    filenames.push_back(config.whitelist_vcf);
    filenames.insert(filenames.end(), config.sample_vcfs.begin(), config.sample_vcfs.end());
    auto nfiles = filenames.size();

    std::unique_ptr<bcf_srs_t, decltype(&bcf_sr_destroy)> sr(bcf_sr_init(), bcf_sr_destroy);
    bcf_sr_set_opt(sr.get(), BCF_SR_REQUIRE_IDX);
    bcf_sr_set_opt(sr.get(), BCF_SR_PAIR_LOGIC, BCF_SR_PAIR_EXACT);

    for (const auto& filename : filenames) {
        if (!bcf_sr_add_reader(sr.get(), filename.c_str())) {
            std::cerr << "Failed to open VCF file: " << filename << std::endl;
            return 1;
        }
    }

    // Set up an output VCF to write to either a file or stdout, depending on the config
    std::unique_ptr<htsFile, decltype(&hts_close)> out_vcf(hts_open("-", "w"), hts_close);
    if (config.pass_output_vcf != "-") {
        out_vcf.reset(hts_open(config.pass_output_vcf.c_str(),
                               ("w" + config.format).c_str()));
    }
    if (!out_vcf) {
        std::cerr << "Failed to open output VCF file." << std::endl;
        return 1;
    }
    bcf_hdr_t *out_hdr = bcf_sr_get_header(sr.get(), 0); // Pointer is owned by bcf_srs_t
    if (!out_hdr) {
        std::cerr << "Failed to get header from the first VCF file." << std::endl;
        return 1;
    }

    if (bcf_hdr_write(out_vcf.get(), out_hdr) < 0) {
        std::cerr << "Error: Failed to write VCF header.\n";
        return 1;
    }

    // Set up an output VCF for failing variants if specified, with the file format specified in the config
    std::unique_ptr<htsFile, decltype(&hts_close)> fail_vcf(nullptr, hts_close);
    bool write_fail_vcf = !config.fail_output_vcf.empty();
    if (write_fail_vcf) {
        fail_vcf.reset(hts_open(config.fail_output_vcf.c_str(),
                                ("w" + config.format).c_str()));
        if (bcf_hdr_write(fail_vcf.get(), out_hdr) < 0) {
            std::cerr << "Error: Failed to write failing VCF header.\n";
            return 1;
        }
    }

    // PER FILE HEADER LOOP
    // Get the number of samples in each file
    // Get the FILTER positions of our favourite filters
    auto header_infos = process_headers(sr.get(), filenames);
    bcf_hdr_t *hdr = bcf_sr_get_header(sr.get(), 0);

    int processed = 0;
    HtsBuffer ad_buffer; // Buffer to hold AD values
    
    // LOOP OVER RECORDS
    while (bcf_sr_next_line(sr.get())) {
        if (bcf_sr_has_line(sr.get(), 0)) {
            processed++;
            bcf1_t *rec0 = bcf_sr_get_line(sr.get(), 0); // Allocation is owned by bcf_srs_t, so no need to free it manually
            // Assuming the input files are biallelic, so emit a warning when the are more than just a single REF and ALT allele
            if (rec0->n_allele > 2) {
                std::cerr << "FATAL ERROR: record in file " << filenames[0] << " has more than 2 alleles." << std::endl;
                auto hdr = bcf_sr_get_header(sr.get(), 0);
                auto chrom = bcf_hdr_id2name(hdr, rec0->rid);
                std::cerr << chrom << ":" << rec0->pos + 1 << " has " << rec0->n_allele << " alleles." << std::endl;
                return 1;
            }

            int considered = 0;
            int map_qual_count = 0;
            int orientation_count = 0;
            int strand_bias_count = 0;
            for (int i = 1; i < nfiles; ++i) {
                if (bcf_sr_has_line(sr.get(), i)) {
                    bcf1_t *reci = bcf_sr_get_line(sr.get(), i); // Pointer owned by bcf_srs_t

                    // Extract the Allele Depth (AD) values - populates an array of int values for each allele in each sample
                    int32_t last_alt_ad = get_alt_allele_depth(sr->readers[i].header, reci, header_infos[i], ad_buffer);
                    if (last_alt_ad == bcf_int32_missing) {
                        std::cerr << "WARN: Last AD value is missing for file: " << filenames[i] << std::endl;
                    }
                    if (last_alt_ad > 0) {
                        #ifdef DEBUG_BUILD
                        std::cerr << "DEBUG: Last AD value for file "
                                  << filenames[i] << " at position "
                                  << reci->pos + 1 << " is " << last_alt_ad
                                  << ".\n";
                        #endif
                        // Can extract flags for this row
                        considered++;
                        bcf_unpack(reci, BCF_UN_FLT);

                        int target_mq = header_infos[i].filter_positions.mapqual_idx;
                        int target_or = header_infos[i].filter_positions.orientation_idx;
                        int target_sb = header_infos[i].filter_positions.strandbias_idx;
                        
                        for (int j = 0; j < reci->d.n_flt; ++j) {
                            int filter_id = reci->d.flt[j];
                            if (filter_id == target_mq) {
                                #ifdef DEBUG_BUILD
                                std::cerr
                                    << "DEBUG: Found map_qual filter in file "
                                    << filenames[i] << " at position "
                                    << reci->pos + 1 << ".\n";
                                #endif
                                map_qual_count++;
                            } else if (filter_id == target_or) {
                                #ifdef DEBUG_BUILD
                                std::cerr << "DEBUG: Found orientation filter in file " << filenames[i] << " at position " << reci->pos + 1 << ".\n";
                                #endif
                                orientation_count++;
                            } else if (filter_id == target_sb) {
                                #ifdef DEBUG_BUILD
                                std::cerr << "DEBUG: Found strand_bias filter in file " << filenames[i] << " at position " << reci->pos + 1 << ".\n";
                                #endif
                                strand_bias_count++;
                            }
                        }
                    }
                    #ifdef DEBUG_BUILD
                    else {
                        std::cerr << "DEBUG: Last AD value for file " << filenames[i] << " at position " << reci->pos + 1 << " is " << last_alt_ad << ", not considered.\n";
                    }
                    #endif
                }
            }
            if (considered > 0) {
                float threshold = static_cast<float>(considered) / 2.0;
                #ifdef DEBUG_BUILD
                std::cerr << "DEBUG: Position " << rec0->pos + 1 << " considered in " << considered << " files. map_qual_count=" << map_qual_count
                          << ", orientation_count=" << orientation_count
                          << ", strand_bias_count=" << strand_bias_count
                          << ", threshold=" << threshold << std::endl;
                #endif
                bool map_qual_ok = map_qual_count < threshold;
                bool orientation_ok = orientation_count < threshold;
                bool strand_bias_ok = strand_bias_count < threshold;
                if (map_qual_ok && orientation_ok && strand_bias_ok) {
                    #ifdef DEBUG_BUILD
                    std::cerr << "DEBUG: Position " << rec0->pos + 1 << " passes all filters.\n";
                    #endif
                    // Write rec0 to output
                    if (bcf_write(out_vcf.get(), out_hdr, rec0) < 0) {
                        std::cerr << "Error: Failed to write VCF record.\n";
                        return 1;
                    }
                } else {
                    if (write_fail_vcf) {
                        if (bcf_write(fail_vcf.get(), out_hdr, rec0) < 0) {
                            std::cerr << "Error: Failed to write failing VCF record.\n";
                            return 1;
                        }
                    }
                    #ifdef DEBUG_BUILD
                    std::cerr << "DEBUG: FAIL: Position " << rec0->pos + 1 << " fails filters. map_qual_ok=" << map_qual_ok
                              << ", orientation_ok=" << orientation_ok
                              << ", strand_bias_ok=" << strand_bias_ok << std::endl;
                    #endif
                }
            }
        }
    }
    return 0;
}