/*
 * 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 "independent_segment.h"

#include <iostream>

#include "dictionary.h"
#include "haplotype.h"

///////////////////////////////////////////////////////////////////////////
//
// probability Probability ( HDictionary * dictionary, const Haplotype & h,
// 			     size_type pos )
//
// returns probability of segment h forming by pure recombination,
// given that it starts at position pos.
//
// E.g.:
//
// P([01234]) = P([0....])P([.1...])P([..2..])P([...3.])P([....4])
//
///////////////////////////////////////////////////////////////////////////

probability
IndependentProbability::Probability(HDictionary* dictionary,
				    const Haplotype& h, size_type pos) const
{
    if (Verbose()) {
	std::cout << "\t\t\t" << "Independent ( " << h << ", " << pos << " )"
		  << endl;
    }

    size_type i = 0;
    size_type end;
    size_type size = h.size();

    /*
     * Initialize.  Advance i past all the NA's, recording NAProbability()
     * for each one passed.
     */
    probability prob = 1.0;
    while ((i < size) && h[i].IsNA()) {
	prob *= NAProbability();
	++i;
    }
    /*
     * Set the current search to be the largest block of contiguous
     * non-NA markers, beginning at position i.
     */
    end = i;
    while ((end < size) && !h[end].IsNA())
	++end;

    /*
     * While we haven't looked at the whole haplotype h...
     */
    while (i < size) {
	probability segment_prob = 1.0;
	std::pair<size_type, Haplotype> key(
	    std::make_pair(pos, h.Substr(i, end)));

	FragmentMap::iterator result = prob_cache.find(key);
	if (result != prob_cache.end()) {
	    segment_prob = result->second;
	    if (Verbose()) {
		std::cout << "\t\t\t"
			  << "In cache:" << endl
			  << "\t\t\t\t"
			  << "Found: " << "[" << pos + 1
			  << ", " << pos + size + 1 << ")"
			  << " = " << h << ": "
			  << segment_prob << endl;
	    }
	    i = end;
	}

	/*
	 * While we haven't looked through the longest non-NA fragment...
	 */
	while (i < end) {
	    Haplotype fragment(h.Substr(i, i + 1));
	    if (Verbose()) {
		std::cout << "\t\t\t"
			  << "Search: " << "[" << pos + i + 1
			  << ", " << pos + i + 2 << ")"
			  << " = " << fragment << endl;
	    }
	    probability frag_prob = dictionary->Frequency(fragment, pos + i);
	    if (frag_prob > 0) {
		/*
		 * Either we found a fragment already present
		 * in the dictionary...
		 */
		if (Verbose()) {
		    std::cout << "\t\t\t\t"
			      << "Found: " << "[" << pos + i + 1
			      << ", " << pos + i + 2 << ")"
			      << " = " << fragment << ": "
			      << frag_prob << endl;
		}
		segment_prob *= frag_prob;
	    } else {
		/*
		 * ...or we didn't find it.  We estimate the marker
		 * probability with 1/(NumberOfHaplotypes() + 1).
		 */
		frag_prob = 1.0 / (dictionary->NumberOfHaplotypes() + 1);
		if (Verbose()) {
		    std::cout << "\t\t\t\t"
			      << "Not found: " << "[" << pos + i + 1 << ", "
			      << pos + i + 2 << ")"
			      << " = " << fragment << ": "
			      << frag_prob << endl;
		}
		segment_prob *= frag_prob;
	    }
	    /*
	     * Advance the starting index of the next fragment.
	     */
	    ++i;
	}
	/*
	 * Record Pr(segment) = segment_prob in cache to speed
	 * future lookups.
	 */
	if (result == prob_cache.end())
	    prob_cache.insert(result, std::make_pair(key, segment_prob));

	prob *= segment_prob;
	/*
	 * We've reached the end of the longest non-NA fragment.
	 * Advance i past all the NA's, recording NAProbability()
	 * for each one passed.
	 */
	while ((i < size) && h[i].IsNA()) {
	    prob *= NAProbability();
	    ++i;
	}
	/*
	 * Set the current search to be the largest block of contiguous
	 * non-NA markers, beginning at position i.
	 */
	end = i;
	while ((end < size) && !h[end].IsNA())
	    ++end;
    }
    return prob;
}
