#include <assert.h>

#include <strings.h>

#include <htslib/vcf.h>

#include <htslib/vcfutils.h>

#include <htslib/hts_os.h>

#include <htslib/kbitset.h>

#include "bcftools.h"

#include "vcfbuf.h"

#include "rbuf.h"

typedef struct

{

double max[VCFBUF_LD_N];

int rand_missing, filter1;

}

ld_t;

typedef struct

{

bcf1_t *rec;

double af;

unsigned int af_set:1, filter:1, idx:30;

}

vcfrec_t;

#define PRUNE_MODE_MAX_AF 1

#define PRUNE_MODE_1ST 2

#define PRUNE_MODE_RAND 3

typedef struct

{

int max_sites, mvrec, mac, mfarr, mode;

int *ac, *idx;

float *farr;

char *af_tag;

vcfrec_t **vrec;

}

prune_t;

#define CLUSTER_MODE_PRUNE 1 // remove cluster

#define CLUSTER_MODE_SIZE 2 // make cluster size available via vcfbuf_get_val(buf,int,CLUSTER_SIZE);

typedef struct

{

int max_sites; // used with CLUSTER_PRUNE, removes cluster with more than this many sites within the window

int mode; // one of CLUSTER_MODE_PRUNE or CLUSTER_MODE_SIZE

int last; // the value of the currently removed element

int *size; // cluster size for this site

rbuf_t rbuf;

}

cluster_t;

#define MARK_OVERLAP 1

#define MARK_DUP 2

#define MARK_EXPR 3

#define MARK_MISSING_SCALAR 0 // actual value to use

#define MARK_MISSING_MAX_DP 1 // max overlap_t.value scaled by INFO/DP

typedef struct

{

double value; // the sort value

int rmme, idx; // mark for removal, index in vcfbuf_t.rbuf

int dp; // with MARK_MISSING_MAX_DP, INFO/DP is used extrapolate missing QUAL

kbitset_t *bset; // mark which records it overlaps with, given as 0-based indexes to vcfbuf_t.rbuf

bcf1_t *rec;

}

overlap_t;

typedef struct

{

// modes

int mode;

char *expr;

// sites marked according to expr, returned to the caller via vcfbuf_get()

rbuf_t rbuf;

uint8_t *mark;

int last;

// MARK_OVERLAP

int overlap_rid, overlap_end;

// MARK_EXPR

int nbuf;

overlap_t *buf, **buf_ptr;

int missing_expr; // the value to use when min(QUAL) encounters a missing value

float missing_value; // the default missing value

float max_qual; // with MARK_MISSING_MAX_DP

int max_qual_dp; //

int ntmpi; // temporary int array and the allocated memory

int32_t *tmpi;

}

mark_t;

struct _vcfbuf_t

};

vcfbuf_t *vcfbuf_init(bcf_hdr_t *hdr, int win)

{

vcfbuf_t *buf = (vcfbuf_t*) calloc(1,sizeof(vcfbuf_t));

buf->hdr = hdr;

buf->win = win;

buf->status = clean;

buf->mark.overlap_rid = -1;

int i;

for (i=0; i<VCFBUF_LD_N; i++) buf->ld.max[i] = HUGE_VAL;

rbuf_init(&buf->rbuf, 0);

rbuf_init(&buf->mark.rbuf, 0);

rbuf_init(&buf->cluster.rbuf, 0);

return buf;

}

void vcfbuf_destroy(vcfbuf_t *buf)

{

int i;

for (i=0; i<buf->rbuf.m; i++)

if ( buf->vcf[i].rec ) bcf_destroy(buf->vcf[i].rec);

free(buf->vcf);

free(buf->prune.farr);

free(buf->prune.vrec);

free(buf->prune.ac);

free(buf->prune.af_tag);

free(buf->prune.idx);

free(buf->mark.mark);

free(buf->mark.expr);

for (i=0; i<buf->mark.nbuf; i++) kbs_destroy(buf->mark.buf[i].bset);

free(buf->mark.buf);

free(buf->mark.buf_ptr);

free(buf->mark.tmpi);

free(buf->cluster.size);

free(buf);

}

int vcfbuf_set(vcfbuf_t *buf, vcfbuf_opt_t key, ...)

{

va_list args;

switch (key)

{

case LD_FILTER1:

va_start(args, key);

buf->ld.filter1 = va_arg(args,int);

va_end(args);

return 0;

case LD_RAND_MISSING:

va_start(args, key);

buf->ld.rand_missing = va_arg(args,int);

va_end(args);

return 0;

case LD_MAX_R2:

va_start(args, key);

buf->ld.max[VCFBUF_LD_IDX_R2] = va_arg(args,double);

va_end(args);

return 0;

case LD_MAX_LD:

va_start(args, key);

buf->ld.max[VCFBUF_LD_IDX_LD] = va_arg(args,double);

va_end(args);

return 0;

case LD_MAX_HD:

va_start(args, key);

buf->ld.max[VCFBUF_LD_IDX_HD] = va_arg(args,double);

va_end(args);

return 0;

case VCFBUF_DUMMY:

va_start(args, key);

buf->dummy = va_arg(args,int);

va_end(args);

return 0;

case CLUSTER_PRUNE:

va_start(args, key);

buf->cluster.max_sites = va_arg(args,int);

buf->cluster.mode = CLUSTER_MODE_PRUNE;

va_end(args);

return 0;

case CLUSTER_SIZE:

va_start(args, key);

buf->cluster.max_sites = va_arg(args,int);

buf->cluster.mode = CLUSTER_MODE_SIZE;

va_end(args);

return 0;

case PRUNE_NSITES:

va_start(args, key);

buf->prune.max_sites = va_arg(args,int);

if ( !buf->prune.mode ) buf->prune.mode = PRUNE_MODE_MAX_AF;

va_end(args);

return 0;

case PRUNE_NSITES_MODE:

va_start(args, key);

char *mode = va_arg(args,char*);

va_end(args);

if ( !strcasecmp(mode,"maxAF") ) buf->prune.mode = PRUNE_MODE_MAX_AF;

else if ( !strcasecmp(mode,"1st") ) buf->prune.mode = PRUNE_MODE_1ST;

else if ( !strcasecmp(mode,"rand") ) buf->prune.mode = PRUNE_MODE_RAND;

else error("The mode \"%s\" is not recognised\n",mode);

return 0;

case PRUNE_AF_TAG:

va_start(args, key);

buf->prune.af_tag = strdup(va_arg(args,char*));

va_end(args);

return 0;

case MARK:

va_start(args, key);

buf->mark.expr = strdup(va_arg(args,char*));

if ( !strcasecmp(buf->mark.expr,"overlap") ) buf->mark.mode = MARK_OVERLAP;

else if ( !strcasecmp(buf->mark.expr,"dup") ) buf->mark.mode = MARK_DUP;

else buf->mark.mode = MARK_EXPR;

va_end(args);

return 0;

case MARK_MISSING_EXPR:

va_start(args, key);

char *expr = va_arg(args,char*);

if ( !strcasecmp(expr,"0") )

{

buf->mark.missing_expr = MARK_MISSING_SCALAR;

buf->mark.missing_value = 0;

}

else if ( !strcasecmp(expr,"DP") )

{

if ( buf->mark.mode!=MARK_EXPR ) error("Only the combination of --mark 'min(QUAL)' with --missing DP is currently supported\n");

buf->mark.missing_expr = MARK_MISSING_MAX_DP;

}

else

error("todo: MARK_MISSING_EXPR=%s\n",expr);

va_end(args);

return 0;

}

return 0;

}

void *vcfbuf_get(vcfbuf_t *buf, vcfbuf_opt_t key, ...)

{

va_list args;

va_start(args, key);

if ( key==MARK )

return &buf->mark.last;

if ( key==CLUSTER_SIZE )

return &buf->cluster.last;

va_end(args);

return NULL;

}

int vcfbuf_nsites(vcfbuf_t *buf)

{

return buf->rbuf.n;

}

bcf1_t *vcfbuf_push(vcfbuf_t *buf, bcf1_t *rec)

{

// make sure the caller is using the buffer correctly and calls vcfbuf_flush()

// before placing next vcfbuf_push() call

assert(buf->status!=dirty);

if ( !buf->dummy ) buf->status = dirty;

rbuf_expand0(&buf->rbuf, vcfrec_t, buf->rbuf.n+1, buf->vcf);

int i = rbuf_append(&buf->rbuf);

if ( !buf->vcf[i].rec ) buf->vcf[i].rec = bcf_init1();

bcf1_t *ret = buf->vcf[i].rec;

buf->vcf[i].rec = rec;

buf->vcf[i].af_set = 0;

buf->vcf[i].filter = buf->ld.filter1;

buf->ld.filter1 = 0;

return ret;

}

bcf1_t *vcfbuf_peek(vcfbuf_t *buf, int idx)

{

buf->status = clean;

int i = rbuf_kth(&buf->rbuf, idx);

return i<0 ? NULL : buf->vcf[i].rec;

}

bcf1_t *vcfbuf_remove(vcfbuf_t *buf, int idx)

{

int i = rbuf_kth(&buf->rbuf, idx);

if ( i<0 ) return NULL;

bcf1_t *rec = buf->vcf[i].rec;

rbuf_remove_kth(&buf->rbuf, vcfrec_t, idx, buf->vcf);

return rec;

}

static int cmpvrec(const void *_a, const void *_b)

{

vcfrec_t *a = *((vcfrec_t**) _a);

vcfrec_t *b = *((vcfrec_t**) _b);

if ( a->af < b->af ) return -1;

if ( a->af == b->af ) return 0;

return 1;

}

static int cmpint_desc(const void *_a, const void *_b)

{

int a = *((int*)_a);

int b = *((int*)_b);

if ( a < b ) return 1;

if ( a == b ) return 0;

return -1;

}

static void _prune_sites(vcfbuf_t *buf, int flush_all)

{

int nbuf = flush_all ? buf->rbuf.n : buf->rbuf.n - 1;

int nprune = nbuf - buf->prune.max_sites;

int i,k,irec = 0;

if ( buf->prune.mode==PRUNE_MODE_1ST )

{

int eoff = flush_all ? 1 : 2;

for (i=0; i<nprune; i++)

rbuf_remove_kth(&buf->rbuf, vcfrec_t, buf->rbuf.n - eoff, buf->vcf);

return;

}

if ( buf->prune.mode==PRUNE_MODE_RAND )

{

int eoff = flush_all ? 0 : 1;

for (i=0; i<nprune; i++)

{

int j = (buf->rbuf.n - eoff) * hts_drand48();

rbuf_remove_kth(&buf->rbuf, vcfrec_t, j, buf->vcf);

}

return;

}

if ( nbuf > buf->prune.mvrec )

{

buf->prune.idx = (int*) realloc(buf->prune.idx, nbuf*sizeof(int));

buf->prune.vrec = (vcfrec_t**) realloc(buf->prune.vrec, nbuf*sizeof(vcfrec_t*));

buf->prune.mvrec = nbuf;

}

// set allele frequency and prepare buffer for sorting

for (i=-1; rbuf_next(&buf->rbuf,&i) && irec<nbuf; )

{

bcf1_t *line = buf->vcf[i].rec;

if ( line->n_allele > buf->prune.mac )

{

buf->prune.ac = (int*) realloc(buf->prune.ac, line->n_allele*sizeof(*buf->prune.ac));

buf->prune.mac = line->n_allele;

}

if ( !buf->vcf[i].af_set )

{

buf->vcf[i].af = 0;

if ( buf->prune.af_tag )

{

if ( bcf_get_info_float(buf->hdr,line,buf->prune.af_tag,&buf->prune.farr, &buf->prune.mfarr) > 0 ) buf->vcf[i].af = buf->prune.farr[0];

}

else if ( bcf_calc_ac(buf->hdr, line, buf->prune.ac, BCF_UN_INFO|BCF_UN_FMT) )

{

int ntot = 0;

for (k=0; k<line->n_allele; k++) ntot += buf->prune.ac[k];

buf->vcf[i].af = ntot ? (float)(ntot - buf->prune.ac[0])/ntot : 0;

}

buf->vcf[i].af_set = 1;

}

buf->vcf[i].idx = irec;

buf->prune.vrec[irec++] = &buf->vcf[i];

}

// sort by allele frequency, low AF will be removed preferentially

qsort(buf->prune.vrec, nbuf, sizeof(*buf->prune.vrec), cmpvrec);

// sort the rbuf indexes to be pruned descendently so that j-th rbuf index

// is removed before i-th index if i<j

for (i=0; i<nprune; i++)

buf->prune.idx[i] = buf->prune.vrec[i]->idx;

qsort(buf->prune.idx, nprune, sizeof(int), cmpint_desc);

for (i=0; i<nprune; i++)

rbuf_remove_kth(&buf->rbuf, vcfrec_t, buf->prune.idx[i], buf->vcf);

}

static int mark_dup_can_flush_(vcfbuf_t *buf, int flush_all)

{

int flush = flush_all;

mark_t *mark = &buf->mark;

if ( buf->status==dirty )

{

// a new site was just added by vcfbuf_push()

rbuf_expand0(&mark->rbuf, uint8_t, buf->rbuf.n, mark->mark);

int i = rbuf_append(&mark->rbuf);

mark->mark[i] = 0;

if ( buf->rbuf.n==1 ) goto flush;

// there is at least one previous site, check if it's a duplicate

int k1 = rbuf_kth(&buf->rbuf, -1);

int k2 = rbuf_kth(&buf->rbuf, -2);

vcfrec_t *rec1 = &buf->vcf[k1];

vcfrec_t *rec2 = &buf->vcf[k2];

int is_dup = 1;

if ( rec1->rec->rid!=rec2->rec->rid ) is_dup = 0;

else if ( rec1->rec->pos!=rec2->rec->pos ) is_dup = 0;

if ( is_dup )

{

// it is, mark the last two sites as duplicates

int k1 = rbuf_kth(&mark->rbuf, -1);

int k2 = rbuf_kth(&mark->rbuf, -2);

mark->mark[k1] = 1;

mark->mark[k2] = 1;

goto flush;

}

// the last site is not a duplicate with the previous, all sites but the last one can be flushed

flush = 1;

}

else if ( buf->rbuf.n > 1 ) flush = 1;

flush:

if ( !flush ) return 0;

int i = rbuf_shift(&mark->rbuf);

mark->last = mark->mark[i];

return 1;

}

static int mark_overlap_helper_(vcfbuf_t *buf, int flush_all)

{

if ( buf->status!=dirty ) return flush_all;

int flush = flush_all;

mark_t *mark = &buf->mark;

// a new site was just added by vcfbuf_push()

buf->status = clean;

rbuf_expand0(&mark->rbuf, uint8_t, buf->rbuf.n, mark->mark);

int i = rbuf_append(&mark->rbuf);

mark->mark[i] = 0;

// determine beg and end of the last record that was just added

i = rbuf_last(&buf->rbuf);

vcfrec_t *last = &buf->vcf[i];

if ( mark->overlap_rid != last->rec->rid ) mark->overlap_end = 0;

int beg_pos = last->rec->pos;

int end_pos = last->rec->pos + last->rec->rlen - 1;

// Assuming left-aligned indels. In case it is a deletion, the real variant

// starts one base after. If an insertion, the overlap with previous is zero

int imin = last->rec->rlen;

for (i=0; i<last->rec->n_allele; i++)

{

char *ref = last->rec->d.allele[0];

char *alt = last->rec->d.allele[i];

if ( *alt == '<' ) continue; // ignore symbolic alleles

while ( *ref && *alt && nt_to_upper(*ref)==nt_to_upper(*alt) ) { ref++; alt++; }

if ( imin > ref - last->rec->d.allele[0] ) imin = ref - last->rec->d.allele[0];

}

if ( beg_pos <= mark->overlap_end )

{

// the new site overlaps with the previous

beg_pos += imin;

if ( beg_pos > end_pos ) end_pos = beg_pos;

}

if ( buf->rbuf.n==1 )

{

mark->overlap_rid = last->rec->rid;

mark->overlap_end = end_pos;

return flush;

}

if ( beg_pos <= mark->overlap_end )

{

if ( mark->overlap_end < end_pos ) mark->overlap_end = end_pos;

int k1 = rbuf_kth(&mark->rbuf, -1);

int k2 = rbuf_kth(&mark->rbuf, -2);

mark->mark[k1] = 1;

mark->mark[k2] = 1;

}

else

{

if ( mark->overlap_end < end_pos ) mark->overlap_end = end_pos;

flush = 1;

}

return flush;

}

static int mark_overlap_can_flush_(vcfbuf_t *buf, int flush_all)

{

int flush = flush_all;

if ( buf->status==dirty ) flush = mark_overlap_helper_(buf,flush_all);

else if ( buf->rbuf.n > 1 ) flush = 1;

if ( !flush ) return 0;

mark_t *mark = &buf->mark;

int i = rbuf_shift(&mark->rbuf);

mark->last = mark->mark[i];

return 1;

}

static int records_overlap(bcf1_t *a, bcf1_t *b)

{

if ( a->rid != b->rid ) return 0;

if ( a->pos + a->rlen - 1 < b->pos ) return 0;

return 1;

}

static int cmp_overlap_ptr_asc(const void *aptr, const void *bptr)

{

overlap_t *a = *((overlap_t**)aptr);

overlap_t *b = *((overlap_t**)bptr);

if ( a->value < b->value ) return -1;

if ( a->value > b->value ) return 1;

return 0;

}

static void mark_expr_missing_reset_(vcfbuf_t *buf)

{

buf->mark.max_qual = 0;

buf->mark.max_qual_dp = 0;

}

static void mark_expr_missing_prep_(vcfbuf_t *buf, overlap_t *olap)

{

int nval = bcf_get_info_int32(buf->hdr,olap->rec,"DP",&buf->mark.tmpi,&buf->mark.ntmpi);

if ( nval!=1 ) return;

olap->dp = buf->mark.tmpi[0];

if ( bcf_float_is_missing(olap->rec->qual) ) return;

if ( buf->mark.max_qual < olap->rec->qual )

{

buf->mark.max_qual = olap->rec->qual;

buf->mark.max_qual_dp = olap->dp;

}

}

static void mark_expr_missing_set_(vcfbuf_t *buf, overlap_t *olap)

{

if ( !bcf_float_is_missing(olap->rec->qual) ) return;

if ( !buf->mark.max_qual_dp ) return;

// scale QUAL of the most confident variant in the overlap proportionally to the coverage

// and use that to prioritize the records

olap->value = buf->mark.max_qual * olap->dp / buf->mark.max_qual_dp;

}

static int mark_expr_can_flush_(vcfbuf_t *buf, int flush_all)

{

mark_t *mark = &buf->mark;

if ( strcasecmp("min(QUAL)",mark->expr) ) error("Todo; at this time only min(QUAL) is supported\n");

int flush = flush_all;

if ( buf->status==dirty )

{

flush = mark_overlap_helper_(buf,flush_all);

if ( !flush ) return 0;

if ( mark->missing_expr==MARK_MISSING_MAX_DP ) mark_expr_missing_reset_(buf);

// init overlaps, each overlap_t structure keeps a list of overlapping records, symmetrical

size_t nori = mark->nbuf;

hts_resize(overlap_t, buf->rbuf.n, &mark->nbuf, &mark->buf, HTS_RESIZE_CLEAR);

hts_resize(overlap_t*, buf->rbuf.n, &nori, &mark->buf_ptr, HTS_RESIZE_CLEAR);

int i;

for (i=0; i<buf->rbuf.n; i++)

{

overlap_t *oi = &mark->buf[i];

int j = rbuf_kth(&buf->rbuf, i);

assert(j>=0);

bcf1_t *rec = buf->vcf[j].rec;

assert(rec);

oi->rec = rec;

// todo: other than QUAL values

oi->value = bcf_float_is_missing(rec->qual) ? mark->missing_value : rec->qual;

if ( mark->missing_expr==MARK_MISSING_MAX_DP ) mark_expr_missing_prep_(buf,oi);

if ( oi->bset )

{

kbs_resize(&oi->bset,buf->rbuf.n);

kbs_clear(oi->bset);

}

else

oi->bset = kbs_init(buf->rbuf.n);

oi->idx = i;

mark->buf_ptr[i] = oi;

mark->mark[oi->idx] = 0;

}

int nolap = 0;

for (i=0; i<buf->rbuf.n; i++)

{

overlap_t *oi = &mark->buf[i];

if ( mark->missing_expr==MARK_MISSING_MAX_DP ) mark_expr_missing_set_(buf,oi);

int j;

for (j=i+1; j<buf->rbuf.n; j++)

{

overlap_t *oj = &mark->buf[j];

if ( !records_overlap(oi->rec,oj->rec) ) continue;

kbs_insert(oi->bset,j);

kbs_insert(oj->bset,i);

nolap++;

}

}

// sort according to the requested criteria, currently only min(QUAL)

qsort(mark->buf_ptr,buf->rbuf.n,sizeof(*mark->buf_ptr),cmp_overlap_ptr_asc); // todo: other than min()

// go through the list sorted by overlap_t.value, eg QUAL

for (i=0; nolap && i<buf->rbuf.n; i++)

{

kbitset_iter_t itr;

overlap_t *oi = mark->buf_ptr[i];

kbs_start(&itr);

int j;

while ((j=kbs_next(oi->bset, &itr)) >= 0)

{

kbs_delete(oi->bset,j);

assert(nolap);

assert(kbs_exists(mark->buf[j].bset,oi->idx));

kbs_delete(mark->buf[j].bset,oi->idx);

nolap--;

}

j = rbuf_kth(&mark->rbuf,oi->idx);

mark->mark[j] = 1;

}

}

else if ( buf->rbuf.n > 1 ) flush = 1;

if ( !flush ) return 0;

int i = rbuf_shift(&mark->rbuf);

mark->last = mark->mark[i];

return 1;

}

int cluster_can_flush_(vcfbuf_t *buf, int flush_all)

{

cluster_t *cluster = &buf->cluster;

if ( buf->status==dirty )

{

// a new site was just added by vcfbuf_push()

rbuf_expand0(&cluster->rbuf, int, buf->rbuf.n, cluster->size);

int i = rbuf_append(&cluster->rbuf);

cluster->size[i] = 0;

}

assert( cluster->rbuf.n==buf->rbuf.n );

// The following cases can occur:

// - if flush_all is set, then the entire buffer must be within the window

// - else if the last record can be on a different chr, then everything before can be flushed

// - else the last record can be either within or outside the window with respect to the first record

if ( buf->status==dirty )

{

int ib = 0;

while ( ib < buf->rbuf.n )

{

int b = rbuf_kth(&buf->rbuf, ib);

int ie = ib + 1;

while ( ie < buf->rbuf.n )

{

int e = rbuf_kth(&buf->rbuf, ie);

if ( buf->vcf[b].rec->rid != buf->vcf[e].rec->rid ) break;

if ( buf->vcf[e].rec->pos - buf->vcf[b].rec->pos + 1 > -buf->win ) break; // win is negative

ie++;

}

// now ie is just outside the window or beyond the last element of the window

// count the number of unfiltered sites that contribute to the cluster. Note this is inefficient,

// recalculating the same bits over and over, should be improved..

int ix, nbuf = 0;

for (ix=ib; ix<ie; ix++)

{

int x = rbuf_kth(&buf->rbuf, ix);

if ( buf->vcf[x].filter ) continue;

nbuf++;

}

for (ix=ib; ix<ie; ix++)

{

int x = rbuf_kth(&cluster->rbuf, ix);

if ( cluster->size[x] < nbuf ) cluster->size[x] = nbuf;

}

ib++;

}

buf->status = clean;

}

int b = rbuf_kth(&buf->rbuf, 0); // first

int e = rbuf_last(&buf->rbuf); // last

int can_flush = flush_all;

if ( buf->vcf[b].rec->rid != buf->vcf[e].rec->rid ) can_flush = 1;

if ( buf->vcf[e].rec->pos - buf->vcf[b].rec->pos + 1 > -buf->win ) can_flush = 1;

if ( !can_flush ) return 0;

if ( buf->cluster.mode==CLUSTER_MODE_PRUNE )

{

int flush = 0;

while ( buf->rbuf.n )

{

flush = 0;

int b = rbuf_kth(&buf->rbuf, 0);

int e = rbuf_kth(&buf->rbuf, -1);

if ( buf->vcf[b].filter )

{

// not to be pruned, not counted as part of the cluster

flush = 1;

break;

}

if ( flush_all ) flush = 1;

else if ( buf->vcf[b].rec->rid != buf->vcf[e].rec->rid ) flush = 1;

else if ( buf->vcf[e].rec->pos - buf->vcf[b].rec->pos + 1 > -buf->win ) flush = 1;

if ( !flush ) break;

b = rbuf_kth(&cluster->rbuf, 0);

if ( cluster->size[b] <= cluster->max_sites ) break; // not to be pruned

rbuf_remove_kth(&buf->rbuf, vcfrec_t, 0, buf->vcf);

rbuf_remove_kth(&cluster->rbuf, int, 0, cluster->size);

}

if ( !flush ) return 0;

}

if ( !cluster->rbuf.n ) return 0;

b = rbuf_shift(&cluster->rbuf);

cluster->last = cluster->size[b];

b = rbuf_kth(&buf->rbuf, 0);

if ( buf->vcf[b].filter ) cluster->last = 0;

return 1;

}

bcf1_t *vcfbuf_flush(vcfbuf_t *buf, int flush_all)

{

int i,j;

// nothing to do, no lines in the buffer

if ( buf->rbuf.n==0 ) return NULL;

// dummy mode, always flushing

if ( buf->dummy ) goto ret;

// either annotate or print clustered sites

if ( buf->cluster.mode )

{

if ( !cluster_can_flush_(buf,flush_all) ) return NULL;

goto ret;

}

// pruning mode

if ( buf->win )

{

int can_flush = flush_all;

i = rbuf_kth(&buf->rbuf, 0); // first

j = rbuf_last(&buf->rbuf); // last

if ( buf->vcf[i].rec->rid != buf->vcf[j].rec->rid ) can_flush = 1;

else if ( buf->win > 0 )

{

if ( buf->rbuf.n > buf->win ) can_flush = 1;

}

else if ( buf->win < 0 )

{

if ( !(buf->vcf[i].rec->pos - buf->vcf[j].rec->pos > buf->win) ) can_flush = 1;

}

buf->status = clean;

if ( !can_flush ) return NULL;

if ( buf->prune.max_sites && buf->prune.max_sites < buf->rbuf.n ) _prune_sites(buf, flush_all);

goto ret;

}

// overlaps and duplicates

if ( buf->mark.mode )

{

int can_flush = 0;

if ( buf->mark.mode==MARK_OVERLAP )

{

if ( mark_overlap_can_flush_(buf,flush_all) ) can_flush = 1;

}

else if ( buf->mark.mode==MARK_DUP )

{

if ( mark_dup_can_flush_(buf,flush_all) ) can_flush = 1;

}

if ( buf->mark.mode==MARK_EXPR )

{

if ( mark_expr_can_flush_(buf,flush_all) ) can_flush = 1;

}

buf->status = clean;

if ( !can_flush ) return NULL;

goto ret;

}

ret:

buf->status = clean;

i = rbuf_shift(&buf->rbuf);

return buf->vcf[i].rec;

}

static double _estimate_af(int8_t *ptr, int size, int nvals, int nsamples)

{

int i,j, nref = 0, nalt = 0;

for (i=0; i<nsamples; i++)

{

for (j=0; j<nvals; j++)

{

if ( bcf_gt_is_missing(ptr[j]) ) break;

if ( ptr[j]==bcf_int8_vector_end ) break;

if ( bcf_gt_allele(ptr[j]) ) nalt++;

else nref++;

}

ptr += size;

}

if ( nref+nalt == 0 ) return 0;

return (double)nalt/(nref+nalt);

}

static int _calc_r2_ld(vcfbuf_t *buf, bcf1_t *arec, bcf1_t *brec, vcfbuf_ld_t *ld)

{

if ( arec->n_sample!=brec->n_sample ) error("Different number of samples: %d vs %d\n",arec->n_sample,brec->n_sample);

assert( arec->n_sample );

int i,j,igt = bcf_hdr_id2int(buf->hdr, BCF_DT_ID, "GT");

bcf_unpack(arec, BCF_UN_FMT);

bcf_unpack(brec, BCF_UN_FMT);

bcf_fmt_t *afmt = NULL, *bfmt = NULL;

for (i=0; i<arec->n_fmt; i++)

if ( arec->d.fmt[i].id==igt ) { afmt = &arec->d.fmt[i]; break; }

if ( !afmt ) return -1; // no GT tag

for (i=0; i<brec->n_fmt; i++)

if ( brec->d.fmt[i].id==igt ) { bfmt = &brec->d.fmt[i]; break; }

if ( !bfmt ) return -1; // no GT tag

if ( afmt->n==0 ) return -1; // empty?!

if ( bfmt->n==0 ) return -1; // empty?!

if ( afmt->type!=BCF_BT_INT8 ) error("TODO: the GT fmt_type is not int8!\n");

if ( bfmt->type!=BCF_BT_INT8 ) error("TODO: the GT fmt_type is not int8!\n");

// Determine allele frequencies, this is to sample randomly missing genotypes

double aaf = 0, baf = 0;

if ( buf->ld.rand_missing )

{

aaf = _estimate_af((int8_t*)afmt->p, afmt->size, afmt->n, arec->n_sample);

baf = _estimate_af((int8_t*)bfmt->p, bfmt->size, bfmt->n, brec->n_sample);

}

// Calculate r2, lf, hd

double nhd[] = {0,0,0,0,0,0,0,0,0};

double ab = 0, aa = 0, bb = 0, a = 0, b = 0;

int nab = 0, ndiff = 0;

int an_tot = 0, bn_tot = 0;

for (i=0; i<arec->n_sample; i++)

{

int8_t *aptr = (int8_t*) (afmt->p + i*afmt->size);

int8_t *bptr = (int8_t*) (bfmt->p + i*bfmt->size);

int adsg = 0, bdsg = 0; // dosages (0,1,2) at sites (a,b)

int an = 0, bn = 0; // number of alleles at sites (a,b)

for (j=0; j<afmt->n; j++)

{

if ( aptr[j]==bcf_int8_vector_end ) break;

if ( bcf_gt_is_missing(aptr[j]) )

{

if ( !buf->ld.rand_missing ) break;

if ( hts_drand48() >= aaf ) adsg += 1;

}

else if ( bcf_gt_allele(aptr[j]) ) adsg += 1;

an++;

}

for (j=0; j<bfmt->n; j++)

{

if ( bptr[j]==bcf_int8_vector_end ) break;

if ( bcf_gt_is_missing(bptr[j]) )

{

if ( !buf->ld.rand_missing ) break;