/*
 * Copyright (c) 1999 Johnny C. Lam.
 * All rights reserved.
 *
 * Permission to use, copy, modify, distribute and sell this software
 * and its documentation for any purpose is hereby granted without fee,
 * provided that the above copyright notice appear in all copies and
 * that both that copyright notice and this permission notice appear
 * in supporting documentation.  The author makes no representations
 * about the suitability of this software for any purpose.  It is
 * provided "as is" without express or implied warranty.
 */

#include "probability.h"

#include <algorithm>
#include <cmath>
#include <fstream>
#include <iostream>
#include <iterator>
#include <numeric>
#include <string>
#include <utility>
#include "dictionary.h"
#include "haplotype.h"

Probability::Probability(Geometric dist, HDictionary* dict, probability ibd,
			 probability p_mutation, probability p_recomb,
			 size_type hap_length,
			 size_type left_mask, size_type right_mask,
			 const string& jump_vector_fname,
			 const string& dist_fname,
			 const string& recomb_vector_fname,
			 bool no_recomb_sum,
			 bool verbose )
    : _distribution(dist),
      _dictionary(dict),
      _ibd(ibd),
      _verbose(verbose),
      _no_recomb_sum(no_recomb_sum),
      _left_mask(left_mask),
      _right_mask(right_mask),
      _Jump_at_k(hap_length),
      _P_mutation_at_k(hap_length),
      _P_recombination_at_k(hap_length - 1)
{
    if (verbose) {
	std::cout
	    << "Probability::Probability():" << endl
	    << "\t" << "Left mask:  "
	    << LeftMask() << endl
	    << "\t" << "Right mask: "
	    << RightMask() << endl;
    }
    {
	if (!jump_vector_fname.empty()) {
	    std::ifstream istr(jump_vector_fname.data());
	    copy_n(istream_iterator<unsigned int>(istr), hap_length,
		   _Jump_at_k.begin());
	} else {
	    fill(_Jump_at_k.begin(), _Jump_at_k.end(), 1);
	}
	if (verbose)
	    std::cout << "\t" << "Jumps:";
	typename JumpVector::iterator f = _Jump_at_k.begin();
	typename JumpVector::iterator end = _Jump_at_k.end();
	while (f != end) {
	    if (verbose)
		std::cout << " " << *f;
	    ++f;
	}
	if (verbose)
	    std::cout << endl;
    }
    probability rel_mut;
    {
	std::vector<probability> P_mutation_at_k;
	if (!dist_fname.empty()) {
	    std::ifstream istr(dist_fname.data());
	    std::copy_n(std::istream_iterator<probability>(istr), hap_length,
			std::back_inserter(P_mutation_at_k));
	} else {
	    std::fill_n(std::back_inserter(P_mutation_at_k), hap_length, 1);
	}
	probability region_sum
	    = std::accumulate(P_mutation_at_k.begin(),
			      P_mutation_at_k.end(),
			      static_cast<probability>(0));
	std::fill_n(P_mutation_at_k.begin(), LeftMask(), 0);
	std::fill(P_mutation_at_k.begin() + RightMask(),
		  P_mutation_at_k.end(), 0);
	probability masked_sum
	    = std::accumulate(P_mutation_at_k.begin(),
			      P_mutation_at_k.end(),
			      static_cast<probability>(0));
	std::copy(P_mutation_at_k.begin(), P_mutation_at_k.end(),
		  _P_mutation_at_k.begin());
	if (verbose)
	    std::cout << "\t" << "Pr(M = k):";
	typename ProbVector::iterator f = _P_mutation_at_k.begin();
	typename ProbVector::iterator end = _P_mutation_at_k.end();
	while (f != end) {
	    *f /= masked_sum;
	    if (verbose)
		std::cout << " " << *f;
	    ++f;
	}
	rel_mut = masked_sum / region_sum * p_mutation;
	if (verbose)
	    std::cout << endl;
    }
    probability rel_recomb;
    {
	std::vector<probability> P_recombination_at_k;
	if (!recomb_vector_fname.empty()) {
	    std::ifstream istr(recomb_vector_fname.data());
	    std::copy_n(std::istream_iterator<probability>(istr),
			hap_length - 1,
			std::back_inserter(P_recombination_at_k));
	} else {
	    std::fill_n(std::back_inserter(P_recombination_at_k),
			hap_length - 1, 1);
	}
	probability region_sum
	    = std::accumulate(P_recombination_at_k.begin(),
			      P_recombination_at_k.end(),
			      static_cast<probability>(0));
	std::fill_n(P_recombination_at_k.begin(), LeftMask(), 0);
	size_type rmask = std::min(RightMask() - 1,
				   P_recombination_at_k.size());
	std::fill(P_recombination_at_k.begin() + rmask,
		  P_recombination_at_k.end(), 0);
	probability masked_sum
	    = std::accumulate(P_recombination_at_k.begin(),
			      P_recombination_at_k.end(),
			      static_cast<probability>(0));
	std::copy(P_recombination_at_k.begin(), P_recombination_at_k.end(),
		  _P_recombination_at_k.begin());
	if (verbose)
	    std::cout << "\t" << "Pr(R = k):";
	typename ProbVector::iterator f = _P_recombination_at_k.begin();
	typename ProbVector::iterator end = _P_recombination_at_k.end();
	while (f != end) {
	    *f /= masked_sum;
	    if (verbose)
		std::cout << " " << *f;
	    ++f;
	}
	rel_recomb = masked_sum / region_sum * p_recomb;
	if (verbose)
	    std::cout << endl;
    }
    _gamma = rel_mut / (rel_mut + rel_recomb);
    if (verbose) {
	std::cout << "\t" << "P_mutation: " << rel_mut << endl
		  << "\t" << "P_recombination: " << rel_recomb << endl
		  << "\t" << "Gamma: " << _gamma << endl;
    }
}

///////////////////////////////////////////////////////////////////////////
//
// bool	PossibleMutation( const Haplotype &, const Haplotype & )
//
// Return true if s,t differ at 1 marker location.
// Return false otherwise.
//
///////////////////////////////////////////////////////////////////////////

bool
Probability::PossibleMutation(const Haplotype& s, const Haplotype& t) const
{
    typedef Haplotype::const_iterator iterator_type;

    iterator_type s_end = s.const_iterator_at(RightMask());
    std::pair<iterator_type, iterator_type> result
	= std::mismatch(s.const_iterator_at(LeftMask()), s_end,
			t.const_iterator_at(LeftMask()));
    if (result.first == s_end) {
	return false;
    } else {
	return std::equal(result.first + 1, s_end, result.second + 1);
    }
}

probability
Probability::P(const Haplotype &s, const Haplotype &t, size_type disease) const
{
    if (Verbose()) {
	std::cout << "\t\t" << "s: " << s << endl;
	std::cout << "\t\t" << "t: " << t << endl;
    }

    unsigned int delta = 0;

    /*
     * [left_mask, disease)
     */
    unsigned int left_diff = 0;
    size_type left_index = disease;
    for (size_type i = disease; i != LeftMask(); --i) {
	if (s[i-1] != t[i-1]) {
	    ++left_diff;
	    if (left_diff > 1)
		break;
	    left_index = i;
	    delta = Marker::diff(s[i-1], t[i-1]);
	}
    }

    /*
     * [disease, right_mask)
     */
    unsigned int right_diff = 0;
    size_type right_index = disease;
    for (size_type i = disease; i != RightMask(); ++i) {
	if (s[i] != t[i]) {
	    ++right_diff;
	    if (right_diff > 1)
		break;
	    right_index = i;
	    delta = Marker::diff(s[i], t[i]);
	}
    }

    unsigned int numdiff = left_diff + right_diff;
    if (numdiff == 0) {
	if (Verbose())
	    std::cout << "\t\t" << "Identical." << endl
		      << "\t\t" << "Probability: 1.0" << endl;
	return 1.0;
    }
    if ((left_diff > 0) && (right_diff > 0)) {
	if (Verbose()) {
	    std::cout << "\t\t" << "Differences on both sides." << endl
		      << "\t\t" << "Probability: 0.0" << endl;
	}
	return 0;
    }

    size_type mutation_index;
    if (left_diff > 0)
	mutation_index = left_index - 1;
    if (right_diff > 0)
	mutation_index = right_index;

    probability mutation;
    if ((numdiff == 1) && (delta % MutationJump(mutation_index) == 0)) {
	mutation = P_mutation(s, t, mutation_index);
    } else {
	mutation = 0;
    }
    if (Verbose())
	std::cout << "\t\t" << "Mutation: " << mutation << endl;

#if 0
    if (t.IsLatent() && (numdiff > 1)) {
	if (Verbose()) {
	    std::cout << "\t\t"
		      << "Recombinations to latents are forbidden."
		      << endl;
	}
	exit(1);
    }
#endif

    probability recombination = 0;
    if (left_diff > 0) {
	recombination = P_left_recombination(s, t, disease, left_index);
    } else /* right_diff > 0 */ {
	recombination = P_right_recombination(s, t, disease, right_index);
    }
    if (Verbose())
	std::cout << "\t\t" << "Recombination: " << recombination << endl;

    probability prob = _gamma * mutation + (1 - _gamma) * recombination;
    if (Verbose())
	std::cout << "\t\t" << "Probability: " << prob << endl;

    return prob;
}

///////////////////////////////////////////////////////////////////////////
//
// P(t|s, mutation) = P(t|s, mutation, mutation at k)*P(mutation at k)
//
// where the mutated marker position k is automatically found.  It is
// assumed that there is exactly one change between the two haplotypes.
//
///////////////////////////////////////////////////////////////////////////

probability
Probability::P_mutation(const Haplotype& s, const Haplotype& t) const
{
    for (size_type k = LeftMask(); k < RightMask(); ++k) {
	if (s[k] != t[k])
	    return P_mutation(s, t, k);
    }
    
    // Never gets here!
    if (Verbose())
	std::cerr << "P_mutation called with identical haplotypes" << endl;
    exit(1);
}

///////////////////////////////////////////////////////////////////////////
//
// P(t|s, recomb)
//
// where the correct direction of the recombination is automatically
// chosen.  It is assumed that there are differences on only one side
// of the disease.
//
///////////////////////////////////////////////////////////////////////////

probability
Probability::P_recombination(const Haplotype& s, const Haplotype& t,
			     size_type disease ) const
{
    /*
     * [left_mask, disease)
     */
    bool left = false;
    size_type left_index = disease;
    for (size_type k = disease; k != LeftMask(); --k) {
	if (s[k-1] != t[k-1]) {
	    left = true;
	    left_index = k;
	    break;
	}
    }
    /*
     * [disease, right_mask)
     */
    bool right = false;
    size_type right_index = disease;
    for (size_type k = disease; k < RightMask(); ++k) {
	if (s[k] != t[k]) {
	    right = true;
	    right_index = k;
	    break;
	}
    }

    if (left && right) {
	std::cout << "\t\t\t" << "No recombination possible." << endl;
	return 0;
    }
    if (left) {
	return P_left_recombination(s, t, disease, left_index);
    } else /* right */ {
	return P_right_recombination(s, t, disease, right_index);
    }

    // Never gets here!
    if (Verbose()) {
	std::cerr << "P_recombination called with identical haplotypes"
		  << endl;
    }
    exit(1);
}

///////////////////////////////////////////////////////////////////////////
//
// P(t|s, mutation) = P(t|s, mutation, mutation at k)*P(mutation at k)
//
///////////////////////////////////////////////////////////////////////////

probability
Probability::P_mutation(const Haplotype& s, const Haplotype& t,
			size_type k) const
{
    Marker a(s[k]);
    Marker b(t[k]);
    
    if (Verbose()) {
	if (Verbose()) {
	    std::cout << "\t\t\t" << "Mutation at " << k+1 << "." << endl;
	    std::cout << "\t\t\t" << "Mutation: "
		      << a << " --> " << b << "." << endl;
	}
    }

    unsigned int diff = 0;
    if (!a.IsNA() && !b.IsNA())
	diff = Marker::diff(a, b);
    assert(diff % MutationJump(k) == 0);
    unsigned int jump = diff / MutationJump(k);
    if (Verbose()) {
	std::cout << "\t\t\t"
		  << "Evaluate distribution at " << jump - 1 << "." << endl;
    }
    
    probability	mut_prob = _distribution(jump - 1);
    Haplotype	tprime(t.Substr(LeftMask(), RightMask()));
    probability	ibd_prob
	= _ibd
	/ (_ibd + _dictionary->Probability(tprime, LeftMask()) * (1 -_ibd));
    probability	at_k_prob = P_mutation_at_k(k);
    
    probability	prob = ibd_prob * mut_prob * at_k_prob;
    if (Verbose()) {
	std::cout << "\t\t\t" 
		  << ibd_prob << " * " << mut_prob << " * " << at_k_prob
		  << " = " << prob << endl;
    }
    return prob;
}

///////////////////////////////////////////////////////////////////////////
//
// Given that s = 0000000|000,
//            t = 1111100|000, (the segments before the gap)
//                       ^
//                       |--- disease
// computes:
// P(t|s, recomb)
//   = (Sum over k) P(t|s, t[1:k] is ibd, recomb, recomb at k)
//                  * P( t[1:k] is ibd | s) * Pr(recomb at k)
//
// where k goes from the gap after the first possible recombined piece to
// disease.  In the above example, disease = 7, and k runs from 5 to 7.
//
///////////////////////////////////////////////////////////////////////////

probability
Probability::P_left_recombination(const Haplotype& s, const Haplotype& t,
				  size_type disease, size_type index) const
{
    if ( Verbose() )
	std::cout << "\t\t\t" << "Left recombination." << endl;
    
    probability	prob = 0;
    size_type search_start = _no_recomb_sum ? index : disease;

    for (size_type end = search_start; end != LeftMask(); --end) {
	Haplotype h(t.Substr(LeftMask(), end));
	Haplotype match(t.Substr(end, RightMask()));
	probability ibd_prob
	    = _ibd
	    / (_ibd + _dictionary->Probability(match, end) * (1 -_ibd));
	probability seg_prob = _dictionary->Probability(h, LeftMask());
	probability at_k_prob = P_recombination_at_k(end);
	probability p = seg_prob * ibd_prob * at_k_prob;
	prob += p;
	if (Verbose()) {
	    std::cout << "\t\t\t"
		      << ibd_prob << " * " << seg_prob << " * " << at_k_prob
		      << " = " << p << endl;
	}
	if (s[end - 1] != t[end - 1])
	    break;
    }
    return prob;
}

///////////////////////////////////////////////////////////////////////////
//
// Given that s = 000|0000000,
//            t = 000|0011111, (the segments after the gap),
//                   ^
//                   |------- disease
// computes:
// P(t|s, recomb)
//   = (Sum over k) P(t|s, t[1:k] is ibd, recomb, recomb at k)
//                  * P( t[1:k] is ibd | s) * Pr(recomb at k)
//
// where k goes from disease to the first gap before the last possible
// recombined piece.  In the above example, disease = 3, and k runs from 3
// to 5.
//
///////////////////////////////////////////////////////////////////////////

probability
Probability::P_right_recombination(const Haplotype& s, const Haplotype& t,
				   size_type disease, size_type index) const
{
    if (Verbose())
	std::cout << "\t\t\t" << "Right recombination." << endl;

    probability	prob = 0;
    size_type search_start = _no_recomb_sum ? index : disease;
    for (size_type start = search_start; start < RightMask(); ++start) {
	Haplotype h(t.Substr(start, RightMask()));
	Haplotype match(t.Substr(LeftMask(), start));
	probability ibd_prob
	    = _ibd
	    / (_ibd + _dictionary->Probability(match, LeftMask()) * (1 -_ibd));
	probability seg_prob = _dictionary->Probability(h, start);
	probability at_k_prob = P_recombination_at_k(start);
	probability p = seg_prob * ibd_prob * at_k_prob;
	prob += p;
	if (Verbose()) {
	    std::cout << "\t\t\t"
		      << ibd_prob << " * " << seg_prob << " * " << at_k_prob
		      << " = " << p << endl;
	}
	if (s[start] != t[start])
	    break;
    }
    return prob;
}

/*
 * Return -log probabilities.
 */

log_probability
Probability::log_P(const Haplotype& s, const Haplotype& t,
		   size_type disease) const
{
    Triplet trip(s, t, disease);
    std::map<Triplet, log_probability>::iterator pos
	= log_prob_cache.find(trip);
    if (pos == log_prob_cache.end()) {
	probability p = P(s, t, disease);
	log_probability log_prob
	    = (p > 0) ? -logP(p) : std::numeric_limits<log_probability>::max();
	log_prob_cache.insert(pos, std::make_pair(trip, log_prob));
	return log_prob;
    } else {
	return pos->second;
    }
}
