// -*- mode: c++; indent-tabs-mode: nil; -*-
//
// Copyright 2009 Illumina, Inc.
//
// This software is covered by the "Illumina Genome Analyzer Software
// License Agreement" and the "Illumina Source Code License Agreement",
// and certain third party copyright/licenses, and any user of this
// source file is bound by the terms therein (see accompanying files
// Illumina_Genome_Analyzer_Software_License_Agreement.pdf and
// Illumina_Source_Code_License_Agreement.pdf and third party
// copyright/license notices).
//
//

/// \file

/// \author Chris Saunders
///
#include "blt_common/position_snp_call_pprob_digt.hh"

#include "blt_util/log.hh"
#include "blt_util/seq_util.hh"

#include <cassert>
#include <cmath>
#include <cstdlib>

#include <iostream>
#include <iomanip>



/// find more accurate complement of posterior_probability:
static
double
prob_comp(const double* const digt_prob,
          const unsigned cgt){

    double val(0.);
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        if(gt == cgt) continue;
        val += digt_prob[gt];
    }
    return val;
}



///
/// The likelihood of an observed 'column' of observations for any
/// genotype: P(obs=ACT|genotype=AT) is a sum over all possible
/// columns:
///
/// sum { c in {A,C,G,T}**3 } { P(obs=ACT|col=c) * P(col=c|genotype=AT) }
///
/// Enumerating all {A,C,G,T}**col_size possible column values is
/// inefficent, so we only consider those columns for which
/// P(col|genotype) is nonzero.  For homozygous genotypes this always
/// leaves a single column; for heterozygous genotypes this is a sum
/// over all valid columns:
///
/// P(obs=ACT|genotype=AT)=
/// sum { c in {A,T}**3 } { P(obs=ACT|col=c) * P(col=c|genotype=AT }
///
/// that is...
///
/// P(obs|AAA)*P(AAA|genotype)+
/// P(obs|AAT)*P(AAT|genotype)+...etc
///
/// the first term is a product based on error probabilities:
/// P(obs=ACT|AAA) = (1-P(e1))*(P(e2)/3)*(P(e3)/3)
///
/// the second term is a product of genotype base frequencies:
/// P(AAA|genotype=AT) = .5**3
///
/// The likelihood calculated in the function below is an equivelent
/// factorization of the calculation described above.
///
void
position_snp_call_pprob_digt(const double theta,
                             const snp_pos_info& pi,
                             const std::vector<float>& dependent_eprob,
                             diploid_genotype& dgt,
                             const bool is_always_test){

    if(pi.ref_base=='N') return;

    const unsigned n_calls(pi.calls.size());
    dgt.ref_gt=base_to_id(pi.ref_base);

    // check that a non-reference call meeting quality criteria even exists:
    if(not is_always_test) {
        bool is_test(false);
        for(unsigned i(0);i<n_calls;++i){
            const uint8_t obs_id(pi.calls[i].base_id);
            assert(obs_id!=BASE_ID::ANY);
            if(dgt.ref_gt!=obs_id) {
                is_test=true;
                break;
            }
        }

        if (not is_test) return;
    }

    // get likelihood of each genotype
    double lhood[DIGT::SIZE];
    for(unsigned gt(0);gt<DIGT::SIZE;++gt) lhood[gt] = 0.;

    static const double one_third(1./3.);
    static const double log_one_third(std::log(one_third));
    static const double one_half(1./2.);
    static const double log_one_half(std::log(one_half));
    for(unsigned i(0);i<n_calls;++i){
        const double eprob(dependent_eprob[i]);
        const double ceprob(1.-pi.calls[i].error_prob());

        // precalculate the result for expect values of 0.0, 0.5 & 1.0
        double val[3];
        val[0] = std::log(eprob)+log_one_third;
        val[1] = std::log((ceprob)+((1.-ceprob)*one_third))+log_one_half;
        val[2] = std::log(ceprob);

        const uint8_t obs_id(pi.calls[i].base_id);
        for(unsigned gt(0);gt<DIGT::SIZE;++gt){
            lhood[gt] += val[DIGT::expect2(obs_id,gt)];
        }
    }

    //
    double prior[DIGT::SIZE];
    for(unsigned gt(0);gt<DIGT::SIZE;++gt) prior[gt] = 0.;

    double prior_sum(0.);
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        if(gt==dgt.ref_gt) continue;
        prior[gt]=(theta*one_third);
        if(DIGT::is_het(gt)){
            if(DIGT::expect(dgt.ref_gt,gt)<=0.) prior[gt]*=(theta);
        } else {
            prior[gt]*=.5;
        }
        prior_sum += prior[gt];
    }
    assert(prior_sum <= 1.);
    prior[dgt.ref_gt] = (1.-prior_sum);

    // mult by prior distro to get unnormalized pprob:
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        dgt.pprob[gt] = lhood[gt] + std::log(prior[gt]);
    }

    // scale and exp pprob values:
    dgt.max_gt=0;
    dgt.max2_gt=1;
    double max(dgt.pprob[dgt.max_gt]);
    double max2(dgt.pprob[dgt.max2_gt]);
    for(unsigned gt(1);gt<DIGT::SIZE;++gt){
        if(dgt.pprob[gt] > max){
            max2 = max;
            max = dgt.pprob[gt];
            dgt.max2_gt = dgt.max_gt;
            dgt.max_gt = gt;
        } else if(dgt.pprob[gt] > max2) {
            max2 = dgt.pprob[gt];
            dgt.max2_gt = gt;
        }
    }

    // note we could op at higher precision by keeping the
    // normalization in log space after sum is found, also, scaling
    // could (possibly?) be improved with the factor (the +1 on SIZE
    // is a safety fudge-factor for uniform distributions):
    //
    // static const double lnmax(std::log(std::numeric_limits<double>::max())-static_cast<double>(DIGT::SIZE+1));
    double sum(0.);
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        dgt.pprob[gt] = std::exp(dgt.pprob[gt]-max);
        sum += dgt.pprob[gt];
    }

    // normalize:
    sum = 1./sum;
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        dgt.pprob[gt] *= sum;
    }

    dgt.snp_qphred=error_prob_to_qphred(dgt.pprob[dgt.ref_gt]);
    dgt.max_gt_qphred=error_prob_to_qphred(prob_comp(dgt.pprob,dgt.max_gt));

    dgt.is_snp=(dgt.snp_qphred != 0);

    ///////////////////////////////////////////////////////////////
    // after calculating regular genotype probabilities, get
    // polymorphic site probabiities:
    //
    for(unsigned gt(0);gt<DIGT::SIZE;++gt) prior[gt] = 0.;

    prior_sum=0;
    const double theta2(theta);
    double ctheta2(1.-theta2);
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        if(gt==dgt.ref_gt) {
            prior[gt]=0.25*(ctheta2);
        } else if(DIGT::is_het(gt)){
            if(DIGT::expect(dgt.ref_gt,gt)<=0.) {
                prior[gt] = theta2*one_third;
            } else {
                prior[gt] = 0.5*one_third*ctheta2;
            }
        } else {
            prior[gt] = 0.25*one_third*ctheta2;
        }
        prior_sum += prior[gt];
    }
    assert(std::abs(1.-prior_sum) < 0.0001);

    // mult by prior distro to get unnormalized pprob:
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        dgt.pprob_poly[gt] = lhood[gt] + std::log(prior[gt]);
    }

    // scale and exp pprob values:
    dgt.max_gt_poly=0;
    max=dgt.pprob_poly[dgt.max_gt_poly];

    for(unsigned gt(1);gt<DIGT::SIZE;++gt){
        if(dgt.pprob_poly[gt] > max){
            max = dgt.pprob_poly[gt];
            dgt.max_gt_poly = gt;
        }
    }

    sum=0.;
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        dgt.pprob_poly[gt] = std::exp(dgt.pprob_poly[gt]-max);
        sum += dgt.pprob_poly[gt];
    }

    // normalize:
    sum = 1./sum;
    for(unsigned gt(0);gt<DIGT::SIZE;++gt){
        dgt.pprob_poly[gt] *= sum;
    }

    dgt.max_gt_poly_qphred=error_prob_to_qphred(prob_comp(dgt.pprob_poly,dgt.max_gt_poly));
}




std::ostream& operator<<(std::ostream& os,const diploid_genotype& dgt) {

    os << std::setprecision(10) << std::fixed;

    os << " Q(snp): " << dgt.snp_qphred
       << " max_gt: " << DIGT::label(dgt.max_gt)
       << " Q(max_gt): "  << dgt.max_gt_qphred
       << " max_gt|poly_site: " << dgt.max_gt_poly
       << " Q(max_gt|poly_site): " << dgt.max_gt_poly_qphred;

    os.unsetf(std::ios::fixed);

    return os;
}



void
write_diploid_genotype_allele(const blt_options& opt,
                              const snp_pos_info& pi,
                              const diploid_genotype& dgt,
                              std::ostream& os){

    os << std::setprecision(10) << std::fixed;

    os << dgt.snp_qphred
       << '\t' << DIGT::label(dgt.max_gt) 
       << '\t' << dgt.max_gt_qphred
       << '\t' << DIGT::label(dgt.max_gt_poly)
       << '\t' << dgt.max_gt_poly_qphred;

    if(opt.is_print_used_allele_counts) {
        unsigned base_count[N_BASE];
        for(unsigned i(0);i<N_BASE;++i) base_count[i] = 0;

        const unsigned n_calls(pi.calls.size());
        for(unsigned i(0);i<n_calls;++i){
            assert(pi.calls[i].base_id!=BASE_ID::ANY);
            base_count[pi.calls[i].base_id]++;
        }
        for(unsigned b(0);b<N_BASE;++b) {
            os << '\t' << base_count[b];
        }
    }

    if(opt.is_print_all_poly_gt) {
         for(unsigned gt(0);gt<DIGT::SIZE;++gt){
#if 1
             // print GT as prob:
             os << '\t' << dgt.pprob_poly[gt];
#else
             // print GT as qval:
             os << '\t' << error_prob_to_qphred(prob_comp(dgt.pprob_poly,gt));
#endif
         }
     }

     os.unsetf(std::ios::fixed);
}