/*
 * 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 "longest_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([0123])   P([234])   P([56])   P([67])
// P([01234567]) = --------- * -------- * ------- * -------
//    ---- --         1.0      P([23])      1.0     P([6])
//      --- --
//
///////////////////////////////////////////////////////////////////////////

probability
LongestRunsProbability::Probability(HDictionary* dictionary,
				    const Haplotype& h, size_type pos) const
{
    if (Verbose()) {
	std::cout << "\t\t\t"
		  << "Longest Runs ( " << h << ", " << pos << " )" << endl;
    }
    
    /*
     * Haplotype segments are identified as [i, j).
     */
    size_type i = 0;		// starting index of current search area
    size_type j;		// end index of current sub-segment
    size_type prev;		// end index of previous sub-segment
    size_type end;		// end index of current search area
    size_type size = h.size();

    /*
     * Invariants:	i <= j <= end <= size;
     *			i <= prev <= end <= 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;
    }
    prev = 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;
    j = end;

    /*
     * While we haven't looked at the whole haplotype h...
     */
    while (prev < size) {
	probability segment_prob = 1.0;
	size_type segment_begin = i;
	/*
	 * While we haven't looked at all of the current search area...
	 */
	while (prev < end) {
	    /*
	     * Find the longest fragment [i,j) in the current search area.
	     * Must end past previous fragment, though.
	     */
	    j = end;
	    probability frag_prob = 0;
	    Haplotype fragment;
	    while (j > prev) {
		fragment = h.Substr(i, j);
		if (Verbose()) {
		    std::cout << "\t\t\t"
			      << "Search: " << "[" << pos + i + 1
			      << ", " << pos + j + 1 << ")"
			      << " = " << fragment << endl;
		}
		frag_prob = dictionary->Frequency(fragment, pos + i);
		if (frag_prob > 0)
		    break;
		--j;
	    }
	    if (frag_prob > 0) {
		/*
		 * Either we've found a fragment already present in the
		 * dictionary...
		 */
		if (Verbose()) {
		    std::cout << "\t\t\t\t"
			      << "Found: " << "[" << pos + i + 1
			      << ", " << pos + j + 1 << ")"
			      << " = " << fragment << ": " << frag_prob
			      << endl;
		}
		/*
		 * We need to divide by the probability of the "overlapping"
		 * sub-fragment, or 1 if fragment doesn't overlap previous 
		 * fragment.
		 */
		probability overlap_prob = 1.0;
		if (i < prev) {
		    Haplotype overlap(h.Substr(i, prev));
		    overlap_prob = dictionary->Frequency(overlap, pos + i);
		    if (Verbose()) {
			std::cout << "\t\t\t\t"
				  << "Overlap: " << "[" << pos + i + 1
				  << ", " << pos + prev + 1 << ")"
				  << " = " << overlap << ": "
				  << overlap_prob << endl;
		    }
		}
		segment_prob *= frag_prob / overlap_prob;
		prev = j;
	    } else {
		/*
		 * ...or it wasn't there.  If i wasn't within the previous
		 * fragment, we're down to a single marker, which we
		 * estimate with probability 1/(NumberOfHaplotypes() + 1).
		 */
		if (i >= prev) {
		    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;
		    prev = i + 1;
		}
	    }
	    /*
	     * Advance the starting index of the next search to the
	     * next index.
	     */
	    ++i;
	}
	/*
	 * Record Pr(segment) = segment_prob in dictionary to speed
	 * future lookups.
	 */
	dictionary->AddFrequency(h.Substr(segment_begin, end),
				 pos + segment_begin, segment_prob);
	prob *= segment_prob;
	/*
	 * We've reached the end of the current search area.
	 * Advance i past all the NA's, recording NAProbability(),
	 * for each one passed.
	 */
	i = end;
	while ((i < size) && h[i].IsNA()) {
	    prob *= NAProbability();
	    ++i;
	}
	prev = 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;
	j = end;
    }
    /*
     * ...and we're done!
     */
    return prob;
}

