shared.cpp

#include "shared.h"

#include <htslib/bgzf.h>  // for bgzf_close, bgzf_getline, bgzf_read, bgz...
#include <htslib/hts.h>   // for kstring_t, BGZF
#include <libgen.h>       // for basename
#include <string.h>       // for strlen, strtok, strdup, strcmp
#include <sys/stat.h>     // for stat, time_t
#include <unistd.h>       // for isatty
#include <zlib.h>         // for gzclose, gzgets, gzopen, Z_NULL, gzFile

#include <cstdio>   // for fprintf, stderr, snprintf, NULL, size_t
#include <cstdlib>  // for atoi, exit, malloc
#include <ctime>    // for time
#include <map>      // for operator!=, __map_iterator, operator==
#include <vector>   // for vector

static int decode_nt16_acgt(uint8_t nt16) {
    if (nt16 == 1) return 0; // A
    if (nt16 == 2) return 1; // C
    if (nt16 == 4) return 2; // G
    if (nt16 == 8) return 3; // T
    return -1;
}

double dust_score_nt16(const uint8_t *seq, int32_t l, int32_t window) {
    const int32_t WLEN = 3;
    const int32_t WTOT = 64;
    const int32_t WMASK = WTOT - 1;
    if (window < WLEN || window > 64)
        return -1.0;

    int32_t wCount[WTOT];
    int32_t wSeq[64];
    memset(wCount, 0, sizeof(wCount));
    memset(wSeq, 0, sizeof(wSeq));

    int64_t score = 0;
    int64_t maxScore = 0;
    int32_t t = 0;
    int32_t n = -WLEN;

    for (int32_t i = 0; i < l; ++i) {
        const int b = decode_nt16_acgt(bam_seqi(seq, i));
        if (b < 0)
            continue; // ignore N/ambiguous

        t = ((t << 2) | b) & WMASK;
        if (++n >= 0) {
            const int32_t k = n % window;
            if (n >= window) {
                const int32_t x = wSeq[k];
                if (wCount[x] > 0)
                    score -= --wCount[x];
                score += wCount[t]++;
                if (score > maxScore)
                    maxScore = score;
            } else {
                score += wCount[t]++;
            }
            wSeq[k] = t;
        }
    }

    if (n <= 0)
        return 0.0;
    if (n >= window)
        return (200.0 * (double)maxScore) / (window * (window - 1));
    return (200.0 * (double)score) / (n * (n + 1));
}


int file_is_older(const char *file1, const char *file2) {
    struct stat st1, st2;
    if (stat(file1, &st1) != 0) return -1;
    if (stat(file2, &st2) != 0) return -1;
    return st1.st_mtime < st2.st_mtime; // return 1 hvis file1 er ældre
}

BGZF *getbgzf(const char *str1, const char *mode, int nthreads) {
    BGZF *fp = NULL;
    fp = bgzf_open(str1, mode);
    fprintf(stderr, "\t-> opening file: \'%s\' mode: \'%s\'\n", str1, mode);
    if (fp == NULL) {
        fprintf(stderr, "\t-> Problem opening file: \"%s\"\n", str1);
        exit(1);
    }
    if (nthreads > 1) {
        fprintf(stderr, "\t-> Setting threads to: %d \n", nthreads);
        bgzf_mt(fp, nthreads, 64);
    }
    return fp;
}

BGZF *getbgzf2(const char *str1, const char *str2, const char *mode, int nthreads) {
    unsigned tmp_l = strlen(str1) + strlen(str2) + 1;  // +1 for '\0'

    char *tmp = (char*)malloc(tmp_l);
    if (!tmp) {
        fprintf(stderr, "\t-> Error: failed to allocate memory for tmp, will exit\n");
        exit(1);
    }

    snprintf(tmp, tmp_l, "%s%s", str1, str2);

    BGZF *ret = getbgzf(tmp, mode, nthreads);
    free(tmp);
    return ret;
}

BGZF *getbgzf3(const char *str1, const char *str2, const char *str3, const char *mode, int nthreads) {
    size_t tmp_l = strlen(str1) + strlen(str2) + strlen(str3) + 1;
    char *tmp = (char *)malloc(tmp_l);
    if (tmp == NULL) {
        fprintf(stderr, "\t-> Error: failed to allocate memory for tmp, will exit\n");
        exit(1);
    }
    snprintf(tmp, tmp_l, "%s%s%s", str1, str2, str3);
    BGZF *ret = getbgzf(tmp, mode, nthreads);
    free(tmp);
    return ret;
}
char *getfilename4(const char *str1, const char *str2, const char *str3, const char *str4) {
    unsigned tmp_l = strlen(str1) + strlen(str2) + strlen(str3) + strlen(str4) + 5;
    char *tmp=(char*)calloc(tmp_l,1);
    snprintf(tmp, tmp_l, "%s%s%s%s", str1, str2, str3, str4);
    return tmp;//getbgzf(tmp, mode, nthreads);
}


BGZF *getbgzf4(const char *str1, const char *str2, const char *str3, const char *str4, const char *mode, int nthreads) {
    size_t tmp_l = strlen(str1) + strlen(str2) + strlen(str3) + strlen(str4) + 1;

    char *tmp = (char *)malloc(tmp_l);
    if (tmp == NULL) {
        fprintf(stderr, "\t-> Error: failed to allocate memory for tmp, will exit\n");
        exit(1);
    }

    snprintf(tmp, tmp_l, "%s%s%s%s", str1, str2, str3, str4);

    BGZF *ret = getbgzf(tmp, mode, nthreads);
    free(tmp);
    return ret;
}

int fexists(const char *str) {  ///@param str Filename given as a string.
    fprintf(stderr, "\t-> Checking if exists: \'%s\'\n", str);
    struct stat buffer;
    return (stat(str, &buffer) == 0);  /// @return Function returns 1 if file exists.
}

int fexists2(const char *str1, const char *str2) {
    size_t tmp_l = strlen(str1) + strlen(str2) + 1;

    char *tmp = (char *)malloc(tmp_l);
    if (tmp == NULL) {
        fprintf(stderr, "\t-> Error: failed to allocate memory for tmp, will exit\n");
        exit(1);
    }

    snprintf(tmp, tmp_l, "%s%s", str1, str2);

    int ret = fexists(tmp);
    free(tmp);
    return ret;
}
size_t fsize(const char *fname) {
    struct stat st;
    stat(fname, &st);
    return st.st_size;
}

int fexists3(const char *str1, const char *str2, const char *str3) {
    size_t tmp_l = strlen(str1) + strlen(str2) + strlen(str3) + 1;

    char *tmp = (char *)malloc(tmp_l);
    if (tmp == NULL) {
        fprintf(stderr, "\t-> Error: failed to allocate memory for tmp, will exit\n");
        exit(1);
    }

    snprintf(tmp, tmp_l, "%s%s%s", str1, str2, str3);

    size_t fs = fsize(tmp);
    int ret = fexists(tmp) && fs > 0;

    free(tmp);
    return ret;
}

int fexists4(const char *str1, const char *str2, const char *str3, const char *str4) {
    size_t tmp_l = strlen(str1) + strlen(str2) + strlen(str3) + strlen(str4) + 1;

    char *tmp = (char *)malloc(tmp_l);
    if (tmp == NULL) {
        fprintf(stderr, "\t-> Error: failed to allocate memory for tmp, will exit\n");
        exit(1);
    }

    snprintf(tmp, tmp_l, "%s%s%s%s", str1, str2, str3, str4);

    size_t fs = fsize(tmp);
    int ret = fexists(tmp) && fs > 0;

    free(tmp);
    return ret;
}
// usefull little function to split
char *strpop(char **str, char split) {
    char *tok = *str;
    while (**str) {
        if (**str != split)
            (*str)++;
        else {
            **str = '\0';
            (*str)++;
            break;
        }
    }
    return tok;
}
// usefull little function to remove tab and newlines
void strip(char *line) {
    int at = 0;
    //  fprintf(stderr,"%s\n",line);
    for (size_t i = 0; i < strlen(line); i++)
        if (line[i] == '\t' || line[i] == '\n')
            continue;
        else
            line[at++] = line[i];
    //  fprintf(stderr,"at:%d line:%p\n",at,line);
    line[at] = '\0';
    // fprintf(stderr,"%s\n",line);
}

int2char parse_names(const char *fname) {
    gzFile gz = Z_NULL;
    gz = gzopen(fname, "rb");
    if (gz == Z_NULL) {
        fprintf(stderr, "\t-> Problems opening file: \'%s\'\n", fname);
        exit(1);
    }
    int2char name_map;
    char buf[4096];
    int at = 0;
    char **toks = new char *[5];

    while (gzgets(gz, buf, 4096)) {
        strip(buf);  // fprintf(stderr,"buf:%s\n",buf);
        char *saveptr = buf;
        toks[0] = strpop(&saveptr, '|');
        toks[1] = strpop(&saveptr, '|');
        toks[2] = strpop(&saveptr, '|');
        toks[3] = strpop(&saveptr, '|');
        for (int i = 0; 0 && i < 4; i++)
            fprintf(stderr, "%d):\'%s\'\n", i, toks[i]);

        int key = atoi(toks[0]);
        //    fprintf(stderr,"key:%d\n",key);
        if (toks[3] && strcmp(toks[3], "scientific name") == 0) {
            int2char::iterator it = name_map.find(key);

            if (it != name_map.end())
                fprintf(stderr, "\t->[%s] duplicate name(column1): %s\n", fname, toks[0]);
            else
                name_map[key] = strdup(toks[1]);
        }
        if (0 && at++ > 10)
            break;
    }
    //  int2char::iterator it = name_map.find(61564);  assert(it!=name_map.end());
    fprintf(stderr, "\t-> [%s] Number of unique names (column1): %lu with third column 'scientific name'\n", fname, name_map.size());
    gzclose(gz);
    delete[] toks;
    return name_map;
}
const char *find_nearest_non_clade_rank(int node, int2char &rank, int2int &parent) {
    int cur = node;

    while (1) {
        int2int::iterator pit = parent.find(cur);
        if (pit == parent.end())
            return NULL;

        int par = pit->second;
        if (par == cur)
            return NULL;

        int2char::iterator rit = rank.find(par);
        if (rit == rank.end())
            return NULL;

        if (strcmp(rit->second, "clade") != 0)
            return rit->second;

        cur = par;
    }
}

void parse_nodes(const char *fname, int2char &rank, int2int &parent, int2intvec &child, int dochild) {
    //  fprintf(stderr,"Parsing: %s\n",fname);
    gzFile gz = Z_NULL;
    gz = gzopen(fname, "rb");
    if (gz == Z_NULL) {
        fprintf(stderr, "\t-> Problems opening file: \'%s\'\n", fname);
        exit(1);
    }
    char buf[4096];
    char **toks = new char *[5];

    while (gzgets(gz, buf, 4096)) {
        strip(buf);  // fprintf(stderr,"buf:%s\n",buf);
        char *saveptr = buf;
        toks[0] = strpop(&saveptr, '|');
        toks[1] = strpop(&saveptr, '|');
        toks[2] = strpop(&saveptr, '|');
        for (int i = 0; 0 && i < 3; i++)
            fprintf(stderr, "%d):\'%s\'\n", i, toks[i]);

        int2int::iterator it = parent.find(atoi(toks[0]));
        if (it != parent.end())
            fprintf(stderr, "\t->[%s] duplicate name(column0): %s\n", fname, toks[0]);
        else {
            int key = atoi(toks[0]);
            int val = atoi(toks[1]);
            parent[key] = val;
            rank[key] = strdup(toks[2]);
            if (dochild) {
                if (key == val)  // catch 1 <-> 1
                    continue;
                int2intvec::iterator it2 = child.find(val);
                if (it2 == child.end()) {
                    std::vector<int> tmp;
                    tmp.push_back(key);
                    child[val] = tmp;
                } else
                    it2->second.push_back(key);
            }
        }
    }
    fprintf(stderr, "\t-> Number of unique names (column1): %lu from file: %s parent.size():%lu child.size():%lu\n", rank.size(), fname, parent.size(), child.size());
    // int2int::iterator it=parent.find(9532);
    // fprintf(stderr,"%d->%d\n",it->first,it->second);
    // int2intvec::iterator it=child.find(1);
    // fprintf(stderr,"%d->%lu\n",it->first,it->second.size());
    // exit(0);
    gzclose(gz);

     // fix clade -> nearest non-clade ancestor
    for (int2char::iterator it = rank.begin(); it != rank.end(); it++) {
        if (strcmp(it->second, "clade") == 0) {
            const char *newrank = find_nearest_non_clade_rank(it->first, rank, parent);
            if (newrank != NULL) {
                free(it->second);
                it->second = strdup(newrank);
            }
        }
    }
    delete[] toks;
}

// this generates a downtree, parent to childs taxid->vector<taxids>
void parse_nodes2(int2int &parent, int2intvec &child) {
    fprintf(stderr, "\t-> Generating reverse node table\n");
    for (int2int::iterator it = parent.begin(); it != parent.end(); it++) {
        int down = it->first;
        int up = it->second;
        //    fprintf(stdout,"%d\t%d\n",down,up);
        int2intvec::iterator it2 = child.find(up);
        if (it2 == child.end()) {
            std::vector<int> tmp;
            tmp.push_back(down);
            child[up] = tmp;
        } else
            it2->second.push_back(down);
    }
    fprintf(stderr, "\t-> Done generating reverse node table: contains: %lu\n", child.size());
    //    exit(0);
}

int SIG_COND = 1;
int2int *bamRefId2tax(bam_hdr_t *hdr, char *acc2taxfile, char *bamfile, int2int &errmap,
                      char *tempfolder, int usedump, char *filteredAcc2taxfile,
                      char2int *acc2taxidmap) {
  (void) errmap;
    fprintf(stderr, "\t-> Starting to extract (acc->taxid) from binary file: '%s'\n", acc2taxfile);
    fflush(stderr);

    int binExists = fexists4(tempfolder, basename(acc2taxfile), basename(bamfile), ".bin");
    int useBinary = 0;

    if (acc2taxidmap != NULL) {
        useBinary = 0;
    } else if (binExists && usedump) {
        useBinary = 1;
    } else if (!binExists && usedump) {
        useBinary = 0;
    } else if (binExists && !usedump) {
        useBinary = 1;
    } else {
        useBinary = 0;
    }

    fprintf(stderr, "\t-> Binary exists: %d, usedump: %d, acc2taxidmap: %p => useBinary: %d\n",
            binExists, usedump, acc2taxidmap, useBinary);

    time_t t = time(NULL);
    BGZF *fp = NULL;
    if (!useBinary && usedump)
        fp = getbgzf4(tempfolder, basename(acc2taxfile), basename(bamfile), ".bin", "wb", 4);
    else if (useBinary){
      char *binfile = getfilename4(tempfolder, basename(acc2taxfile), basename(bamfile), ".bin"); //not cleaned up
      
      int binIsOlder = 0;
      if (binExists && usedump) {
	binIsOlder = file_is_older(binfile, bamfile);
	if (binIsOlder == 1) {
	  fprintf(stderr, "\n\t-> WARNING: Binary cache file '%s' is older than BAM file '%s'.\n", binfile, bamfile);
	  fprintf(stderr, "\t-> This may indicate out-of-date accession-taxid mapping. Refusing to use it.\n");
	  exit(1);
	}
      }
      fp = getbgzf4(tempfolder, basename(acc2taxfile), basename(bamfile), ".bin", "rb", 4);
    }
    int2int *am = new int2int;

    gzFile filteredFile = Z_NULL;
    if (filteredAcc2taxfile != NULL && !useBinary) {
        filteredFile = gzopen(filteredAcc2taxfile, "wb");
        if (!filteredFile) {
            fprintf(stderr, "Error opening file '%s' for writing\n", filteredAcc2taxfile);
            exit(1);
        }
        gzprintf(filteredFile, "BAM_Reference\tTaxID\tBam_hdr_index\n");
    }

    if (!useBinary) {
        int at = 0;
        kstring_t *kstr = (kstring_t *)malloc(sizeof(kstring_t));
        kstr->l = kstr->m = 0;
        kstr->s = NULL;
        BGZF *fp2 = getbgzf(acc2taxfile, "rb", 2);
        bgzf_getline(fp2, '\n', kstr);  // skip header
        kstr->l = 0;
        int nprocs = 0;
        while (SIG_COND && bgzf_getline(fp2, '\n', kstr)) {
            if (kstr->l == 0) break;
            if (!((at++ % 100000)) && isatty(fileno(stderr)))
                fprintf(stderr, "\r\t-> At linenr: %d in '%s'      ", at, acc2taxfile);

            char *tok = strtok(kstr->s, "\t\n ");
            char *key = strtok(NULL, "\t\n ");
            tok = strtok(NULL, "\t\n ");
            int val = atoi(tok);
            if (!key) continue;

            if (acc2taxidmap != NULL) {
                if (!acc2taxidmap->insert(std::pair<char*, int>(strdup(key), val)).second) {
                    fprintf(stderr, "\t-> Problem inserting key: %s with value: %d\n", key, val);
                    exit(1);
                }
            }

            int valinbam = bam_name2id(hdr, key);
            if (valinbam == -1) continue;
            nprocs++;

            if (fp != NULL)
	      if (bgzf_write(fp, &valinbam, sizeof(int)) != sizeof(int) ||
		  bgzf_write(fp, &val, sizeof(int)) != sizeof(int)) {
		fprintf(stderr, "\t-> Error: failed to write expected number of bytes with bgzf_write, will exit\n");
		exit(1);
	      }

            if (am->find(valinbam) != am->end())
                fprintf(stderr, "\t-> Duplicate entries found '%s'\n", key);
            (*am)[valinbam] = val;

            if (filteredFile != Z_NULL)
                gzprintf(filteredFile, "%s\t%d\t%d\n", key, val, valinbam);
            kstr->l = 0;
        }
        bgzf_close(fp2);
        free(kstr->s);
        free(kstr);

        fprintf(stderr, "\t-> Number of items from acc2tax file that is relevant from the bamheader: %d\n", nprocs);
        fflush(stderr);
    } else {
        int valinbam, val;
        while (bgzf_read(fp, &valinbam, sizeof(int))) {
	  if (bgzf_read(fp, &val, sizeof(int)) != sizeof(int)) {
	    fprintf(stderr, "\t-> Error: failed to read expected number of bytes for int value with bgzf_read, will exit\n");
	    exit(1);
	  }
	  (*am)[valinbam] = val;
        }
    }

    if (filteredFile != Z_NULL) {
        gzclose(filteredFile);
        fprintf(stderr, "\n\t-> Filtered acc2taxfile saved to '%s'\n", filteredAcc2taxfile);
    }

    if (fp != NULL)
        bgzf_close(fp);

    fprintf(stderr, "\t-> Number of entries to use from accession to taxid: %lu, time taken: %.2f sec acc2taxidmap.size(): %lu\n",
            am->size(), (float)(time(NULL) - t), acc2taxidmap != NULL ? acc2taxidmap->size() : 0);
    return am;
}

queue *init_queue(size_t maxsize) {
    queue *ret = new queue;
    ret->l = 0;
    ret->m = maxsize;
    ret->ary = new bam1_t *[ret->m];
    for (size_t i = 0; i < ret->m; i++)
        ret->ary[i] = bam_init1();
    return ret;
}

// expand queue with 500 elements
void expand_queue(queue *ret) {
    bam1_t **newary = new bam1_t *[ret->m + 500];
    for (size_t i = 0; i < ret->l; i++)
        newary[i] = ret->ary[i];
    for (size_t i = ret->l; i < ret->m + 500; i++)
        newary[i] = bam_init1();
    delete[] ret->ary;
    ret->ary = newary;
    ret->m += 500;
}

void destroy_queue(queue *q) {
    for (size_t i = 0; i < q->m; i++)
        bam_destroy1(q->ary[i]);
    delete[] q->ary;
    delete q;
    q = NULL;
}