/*  bam_sample.c -- group data by sample.

    Copyright (C) 2010, 2011 Broad Institute.
    Copyright (C) 2013, 2016-2022 Genome Research Ltd.

    Author: Heng Li <lh3@sanger.ac.uk>, Petr Danecek <pd3@sanger.ac.uk>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.  */

#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <htslib/hts.h>
#include <htslib/kstring.h>
#include <htslib/khash_str2int.h>
#include <khash_str2str.h>
#include "bam_sample.h"
#include "bcftools.h"


typedef struct
{
    char *fname;
    void *rg2idx;       // hash: read group name to BCF output sample index. Maintained by bsmpl_add_readgroup
    int default_idx;    // default BCF output sample index, set only when all readgroups are treated as one sample
}
file_t;

struct _bam_smpl_t
{
    kstring_t tmp;
    file_t *files;
    int ignore_rg, nsmpl, nfiles;
    char **smpl;        // list of BCF output sample names. Maintained by bsmpl_add_readgroup
    void *sample_list;  // hash: BAM input sample name to BCF output sample name. This is the -s/-S list
    int sample_logic;   // the -s/-S logic, 1: include, 0: exclude
    void *rg_list;      // hash: BAM/rg_id to sample name or */rg_id for global ids. This is the -G list
    int rg_logic;       // the -G logic, 1: include, 0: exclude
    void *name2idx;     // hash: BCF output sample name to BCF output sample index. Maintained by bsmpl_add_readgroup
};

bam_smpl_t *bam_smpl_init(void)
{
    bam_smpl_t *bsmpl;
    bsmpl = (bam_smpl_t*) calloc(1, sizeof(bam_smpl_t));
    bsmpl->name2idx = khash_str2int_init();
    return bsmpl;
}

void bam_smpl_destroy(bam_smpl_t *bsmpl)
{
    if ( !bsmpl ) return;
    if ( bsmpl->name2idx ) khash_str2int_destroy_free(bsmpl->name2idx);
    if ( bsmpl->sample_list ) khash_str2str_destroy_free_all(bsmpl->sample_list);
    if ( bsmpl->rg_list ) khash_str2str_destroy_free_all(bsmpl->rg_list);
    int i;
    for (i=0; i<bsmpl->nfiles; i++)
    {
        file_t *file = &bsmpl->files[i];
        if ( file->rg2idx ) khash_str2int_destroy_free(file->rg2idx);
        free(file->fname);
    }
    free(bsmpl->smpl);
    free(bsmpl->files);
    free(bsmpl->tmp.s);
    free(bsmpl);
}

void bam_smpl_ignore_readgroups(bam_smpl_t* bsmpl)
{
    bsmpl->ignore_rg = 1;
}

static void bsmpl_add_readgroup(bam_smpl_t *bsmpl, file_t *file, const char *rg_id, const char *smpl_name)
{
    int ismpl = -1;
    if ( smpl_name )
    {
        if ( khash_str2int_get(bsmpl->name2idx,smpl_name,&ismpl) < 0 )
        {
            // new sample
            bsmpl->nsmpl++;
            bsmpl->smpl = (char**) realloc(bsmpl->smpl,sizeof(char*)*bsmpl->nsmpl);
            bsmpl->smpl[bsmpl->nsmpl-1] = strdup(smpl_name);
            ismpl = khash_str2int_inc(bsmpl->name2idx,bsmpl->smpl[bsmpl->nsmpl-1]);
        }
    }
    if ( !strcmp("*",rg_id) )
    {
        // all read groups in the bam treated as the same sample
        file->default_idx = ismpl;
        return;
    }
    if ( !file->rg2idx ) file->rg2idx = khash_str2int_init();
    if ( khash_str2int_has_key(file->rg2idx,rg_id) ) return;    // duplicate @RG:ID
    khash_str2int_set(file->rg2idx, strdup(rg_id), ismpl);
}
static int bsmpl_keep_readgroup(bam_smpl_t *bsmpl, file_t *file, const char *rg_id, const char **smpl_name)
{
    char *rg_smpl = khash_str2str_get(bsmpl->rg_list,rg_id);    // unique read group present in one bam only
    if ( !rg_smpl )
    {
        // read group specific to this bam
        bsmpl->tmp.l = 0;
        ksprintf(&bsmpl->tmp,"%s\t%s",rg_id,file->fname);
        rg_smpl = khash_str2str_get(bsmpl->rg_list,bsmpl->tmp.s);
    }
    if ( !rg_smpl )
    {
        // any read group in this file?
        bsmpl->tmp.l = 0;
        ksprintf(&bsmpl->tmp,"*\t%s",file->fname);
        rg_smpl = khash_str2str_get(bsmpl->rg_list,bsmpl->tmp.s);
    }
    if ( !rg_smpl && bsmpl->rg_logic ) return 0;
    if ( rg_smpl && !bsmpl->rg_logic ) return 0;

    if ( rg_smpl && rg_smpl[0]!='\t' ) *smpl_name = rg_smpl;    // rename the sample
    return 1;
}

/*
    The logic of this function is a bit complicated because we want to work
    also with broken bams containing read groups that are not listed in the
    header. The desired behavior is as follows:
        - when -G is given, read groups which are not listed in the header must
          be given explicitly using the "?" symbol in -G.
          Otherwise:
        - if the bam has no header, all reads in the file are assigned to a
          single sample named after the file
        - if there is at least one sample defined in the header, reads with no
          read group id or with a read group id not listed in the header are
          assigned to the first sample encountered in the header
*/
int bam_smpl_add_bam(bam_smpl_t *bsmpl, char *bam_hdr, const char *fname)
{
    bsmpl->nfiles++;
    bsmpl->files = (file_t*) realloc(bsmpl->files,bsmpl->nfiles*sizeof(file_t));
    file_t *file = &bsmpl->files[bsmpl->nfiles-1];
    memset(file,0,sizeof(file_t));
    file->fname  = strdup(fname);
    file->default_idx = -1;

    if ( bsmpl->ignore_rg || !bam_hdr )
    {
        // The option --ignore-RG is set or there is no BAM header: use the file name as the sample name
        bsmpl_add_readgroup(bsmpl,file,"*",file->fname);
        return bsmpl->nfiles-1;
    }

    void *bam_smpls = khash_str2int_init();
    int first_smpl = -1, nskipped = 0;
    const char *p = bam_hdr, *q, *r;
    while (p != NULL && (q = strstr(p, "@RG")) != 0)
    {
        char *eol = strchr(q + 3, '\n');
        if (q > bam_hdr && *(q - 1) != '\n') { // @RG must be at start of line
            p = eol;
            continue;
        }
        p = q + 3;
        if ((q = strstr(p, "\tID:")) != 0) q += 4;
        if ((r = strstr(p, "\tSM:")) != 0) r += 4;
        if (r && q)
        {
            char *u, *v;
            int ioq, ior;
            for (u = (char*)q; *u && *u != '\t' && *u != '\n'; ++u);
            for (v = (char*)r; *v && *v != '\t' && *v != '\n'; ++v);
            ioq = *u; ior = *v; *u = *v = '\0';

            // q now points to a null terminated read group id
            // r points to a null terminated sample name
            if ( !strcmp("*",q) || !strcmp("?",q) )
                error("Error: the read group IDs \"*\" and \"?\" have a special meaning in the mpileup code. Please fix the code or the bam: %s\n", fname);

            int accept_rg = 1;
            if ( bsmpl->sample_list )
            {
                // restrict samples based on the -s/-S options
                char *name = khash_str2str_get(bsmpl->sample_list,r);
                if ( bsmpl->sample_logic==0 )
                    accept_rg = name ? 0 : 1;
                else if ( !name )
                    accept_rg = 0;
                else
                    r = name;
            }
            if ( accept_rg && bsmpl->rg_list )
            {
                // restrict readgroups based on the -G option, possibly renaming the sample
                accept_rg = bsmpl_keep_readgroup(bsmpl,file,q,&r);
            }
            if ( accept_rg )
                bsmpl_add_readgroup(bsmpl,file,q,r);
            else
            {
                bsmpl_add_readgroup(bsmpl,file,q,NULL); // ignore this RG but note that it was seen in the header
                nskipped++;
            }

            if ( first_smpl<0 )
                khash_str2int_get(bsmpl->name2idx,r,&first_smpl);
            if ( !khash_str2int_has_key(bam_smpls,r) )
                khash_str2int_inc(bam_smpls,strdup(r));

            *u = ioq; *v = ior;
        }
        else
            break;
        p = eol;
    }
    int nsmpls = khash_str2int_size(bam_smpls);
    khash_str2int_destroy_free(bam_smpls);

    const char *smpl_name = NULL;
    int accept_null_rg = 1;
    if ( bsmpl->rg_list && !bsmpl_keep_readgroup(bsmpl,file,"?",&smpl_name) ) accept_null_rg = 0;
    if ( bsmpl->sample_list && first_smpl==-1 ) accept_null_rg = 0;

    if ( !accept_null_rg && first_smpl==-1 )
    {
        // no suitable read group is available in this bam: ignore the whole file.
        free(file->fname);
        if ( file->rg2idx ) khash_str2int_destroy_free(file->rg2idx);
        bsmpl->nfiles--;
        return -1;
    }
    if ( !accept_null_rg ) return bsmpl->nfiles-1;
    if ( nsmpls==1 && !nskipped )
    {
        file->default_idx = first_smpl;
        return bsmpl->nfiles-1;
    }
    if ( !smpl_name ) smpl_name = first_smpl==-1 ? file->fname : bsmpl->smpl[first_smpl];

    bsmpl_add_readgroup(bsmpl,file,"?",smpl_name);
    return bsmpl->nfiles-1;
}

const char **bam_smpl_get_samples(bam_smpl_t *bsmpl, int *nsmpl)
{
    *nsmpl = bsmpl->nsmpl;
    return (const char**)bsmpl->smpl;
}

int bam_smpl_get_sample_id(bam_smpl_t *bsmpl, int bam_id, bam1_t *bam_rec)
{
    file_t *file = &bsmpl->files[bam_id];
    if ( file->default_idx >= 0 ) return file->default_idx;

    char *aux_rg = (char*) bam_aux_get(bam_rec, "RG");
    aux_rg = aux_rg ? aux_rg+1 : "?";

    int rg_id;
    if ( khash_str2int_get(file->rg2idx, aux_rg, &rg_id)==0 ) return rg_id;
    if ( khash_str2int_get(file->rg2idx, "?", &rg_id)==0 ) return rg_id;
    return -1;
}

int bam_smpl_add_samples(bam_smpl_t *bsmpl, char *list, int is_file)
{
    if ( list[0]!='^' ) bsmpl->sample_logic = 1;
    else list++;

    int i, nsamples = 0;
    char **samples = hts_readlist(list, is_file, &nsamples);
    if ( !samples || !nsamples ) return 0;

    kstring_t ori = {0,0,0};
    kstring_t ren = {0,0,0};

    bsmpl->sample_list = khash_str2str_init();
    for (i=0; i<nsamples; i++)
    {
        char *ptr = samples[i];
        ori.l = ren.l = 0;
        int escaped = 0;
        while ( *ptr )
        {
            if ( *ptr=='\\' && !escaped ) { escaped = 1; ptr++; continue; }
            if ( isspace(*ptr) && !escaped ) break;
            kputc(*ptr, &ori);
            escaped = 0;
            ptr++;
        }
        if ( *ptr )
        {
            while ( *ptr && isspace(*ptr) ) ptr++;
            while ( *ptr )
            {
                if ( *ptr=='\\' && !escaped ) { escaped = 1; ptr++; continue; }
                if ( isspace(*ptr) && !escaped ) break;
                kputc(*ptr, &ren);
                escaped = 0;
                ptr++;
            }
        }
        khash_str2str_set(bsmpl->sample_list,strdup(ori.s),strdup(ren.l?ren.s:ori.s));
        free(samples[i]);
    }
    free(samples);
    free(ori.s);
    free(ren.s);
    return nsamples;
}

int bam_smpl_add_readgroups(bam_smpl_t *bsmpl, char *list, int is_file)
{
    if ( list[0]!='^' ) bsmpl->rg_logic = 1;
    else list++;

    int i, nrows  = 0;
    char **rows = hts_readlist(list, is_file, &nrows);
    if ( !rows || !nrows ) return 0;

    kstring_t fld1 = {0,0,0};
    kstring_t fld2 = {0,0,0};
    kstring_t fld3 = {0,0,0};

    bsmpl->rg_list = khash_str2str_init();
    for (i=0; i<nrows; i++)
    {
        char *ptr = rows[i];
        fld1.l = fld2.l = fld3.l = 0;
        int escaped = 0;
        while ( *ptr )
        {
            if ( *ptr=='\\' && !escaped ) { escaped = 1; ptr++; continue; }
            if ( isspace(*ptr) && !escaped ) break;
            kputc(*ptr, &fld1);
            escaped = 0;
            ptr++;
        }
        if ( *ptr )
        {
            while ( *ptr && isspace(*ptr) ) ptr++;
            while ( *ptr )
            {
                if ( *ptr=='\\' && !escaped ) { escaped = 1; ptr++; continue; }
                if ( isspace(*ptr) && !escaped ) break;
                kputc(*ptr, &fld2);
                escaped = 0;
                ptr++;
            }
        }
        if ( *ptr )
        {
            while ( *ptr && isspace(*ptr) ) ptr++;
            while ( *ptr )
            {
                if ( *ptr=='\\' && !escaped ) { escaped = 1; ptr++; continue; }
                if ( isspace(*ptr) && !escaped ) break;
                kputc(*ptr, &fld3);
                escaped = 0;
                ptr++;
            }
        }
        if ( fld3.l )
        {
            // ID FILE SAMPLE
            kputc('\t',&fld1);
            kputs(fld2.s,&fld1);
            fld2.l = 0;
            kputs(fld3.s,&fld2);
        }
        // fld2.s now contains a new sample name. If NULL, use \t to keep the bam header name
        char *value = khash_str2str_get(bsmpl->rg_list,fld1.s);
        if ( !value )
            khash_str2str_set(bsmpl->rg_list,strdup(fld1.s),strdup(fld2.l?fld2.s:"\t"));
        else if ( strcmp(value,fld2.l?fld2.s:"\t") )
            error("Error: The read group \"%s\" was assigned to two different samples: \"%s\" and \"%s\"\n", fld1.s,value,fld2.l?fld2.s:"\t");
        free(rows[i]);
    }
    free(rows);
    free(fld1.s);
    free(fld2.s);
    free(fld3.s);
    return nrows;
}