// -*- 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 "position_loghood_minfunc.hh"

#include "blt_util/log.hh"

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

#include <iostream>


const double onethird(1./3.);



///
/// likelihood function is P(obs|model)
/// where the model parameter is the reference allele frequency
///
/// this becomes:
/// sum{ s in seqs } { P(obs|s) * P(s|model) }
///
/// a direct expansion would look like:
/// L = P(obs=TC|true=AA) * P(true=AA|freq=.9)+
///     P(obs=TC|true=AC) * P(true=AC|freq=.9)+...
///
/// L = P(e1)/3*P(e2)/3*.1/3*.1/3 +
///     P(e1)/3*(1-P(e2))*.1/3*.9 +...
///
/// ...where P(e1) is the error prob for basecall 1 and P(e2) is the error
/// prob for basecall 2.
///
/// The above sequence can be factored into:
/// prod { pos in column_size } { P(A|pos,model)+P(C|pos,model)+... }
///
double
calc_pos_nonref_freq_loghood(const snp_pos_info& pi,
                             const double nonref_freq){

    if((nonref_freq < 0.) or (nonref_freq > 1.)){
        log_os << "ERROR:: invalid probability value: " << nonref_freq << "\n";
        exit(EXIT_FAILURE);
    }

    const unsigned n_calls(pi.calls.size());

    double loghood(0.);

    const uint8_t ref_id(base_to_id(pi.ref_base));

    for(unsigned i(0);i<n_calls;++i){
        const double eprob(pi.calls[i].error_prob());
        const bool is_ref_obs(ref_id == pi.calls[i].base_id);

        // each term is a sum of the probability of each of the four bases
        if(is_ref_obs){
            loghood += std::log((1-eprob)*(1-nonref_freq)+eprob*(nonref_freq*onethird));
        } else {
            loghood += std::log((eprob*onethird)*(1-nonref_freq)+(1-(eprob*onethird))*(nonref_freq*onethird));
        }
    }

    return loghood;
}



double
position_nonref_freq_loghood_minfunc::
val(const double nonref_freq) const {
    return -calc_pos_nonref_freq_loghood(_pi,arg_to_prob(nonref_freq));
}



// a sawtooth function to keep the argument in range, yet continous:
double
position_nonref_freq_loghood_minfunc::
arg_to_prob(const double arg){
    double prob(std::fabs(arg));
    const double pf(std::floor(prob));
    prob-=pf;
    if(1==static_cast<int>(pf)%2){ prob=1.-prob; }
    return prob;
}



/// same as above but allow separate frequencies for all nucleotides:
///
static
double
calc_pos_allele_distro_loghood(const snp_pos_info& pi,
                               const double* allele_distro,
                               const unsigned n_allele,
                               const unsigned* allele_map){

    // minimization parameters should already be normalized:
    for(unsigned i(0);i<n_allele;++i){
        if((allele_distro[i] < 0.) or (allele_distro[i] > 1.)){
            log_os << "ERROR:: invalid probability value: " << allele_distro[i] << "\n";
            exit(EXIT_FAILURE);
        }
    }

    const unsigned n_calls(pi.calls.size());

    double loghood(0.);

    for(unsigned i(0);i<n_calls;++i){
        const double eprob(pi.calls[i].error_prob());
        const unsigned idx(pi.calls[i].base_id);
        double site_prob(0.);
        for(unsigned n(0);n<n_allele;++n){
            const double nprob( (allele_map[n] == idx) ? (1.-eprob) : (eprob*onethird) );
            site_prob += nprob*allele_distro[n];
        }
        loghood += std::log(site_prob);
    }

    return loghood;
}


double
position_allele_distro_loghood_minfunc::
val(const double* allele_distro_in) {
    double allele_distro[N_BASE];
    arg_to_prob(allele_distro_in,allele_distro);
    return -calc_pos_allele_distro_loghood(_pi,allele_distro,_n_allele,_allele_map);
}



void
position_allele_distro_loghood_minfunc::
arg_to_prob(const double* allele_distro_in,
            double* allele_distro) {

    double sum(0.);
    for(unsigned i(0);i<_n_allele;++i){
        allele_distro[i] = std::fabs(allele_distro_in[i]);
        sum += allele_distro[i];
    }

    assert(sum>=0.);
    if(sum>0.){
        const double scale(1./sum);
        for(unsigned i(0);i<_n_allele;++i) allele_distro[i] *= scale;
    } else {
        static const double allele_expect(1./static_cast<double>(_n_allele));
        for(unsigned i(0);i<_n_allele;++i) allele_distro[i] = allele_expect;
    }
}