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

#include <cstring>
#include <fstream>
#include <iostream>
#include <string>
#include <strstream>
#include "marker.h"
#include "segment.h"

///////////////////////////////////////////////////////////////////////////
//
// HDictionary
//
// Constructor.
//
// Takes the name of a file containing delimited haplotypes and adds
// every substring of those haplotypes into the dictionary.
//
///////////////////////////////////////////////////////////////////////////

HDictionary::HDictionary(SegmentProbability* segment, bool verbose)
    : _segment(segment), _verbose(verbose), _num_haps(0), _frozen(false)
{}

void
HDictionary::ReadDelimited(istream& istr, char delimiter,
			   unsigned int read_max)
{
    Haplotype h;
    unsigned int read = 0;
    
    if (h.ReadDelimited(istr, Haplotype::NORMAL, delimiter)) {
	std::vector<HaploMap> new_dicts(h.size());
	_dictionaries.swap(new_dicts);
	if (Verbose()) {
	    std::cout << "Creating "
		      << _dictionaries.size() << " dictionaries." << endl;
	}
	if (Verbose())
	    std::cout << "\t" << h << endl;
	++read;
	AddHaplotype(h);
	while (read != read_max &&
	       h.ReadDelimited(istr, Haplotype::NORMAL, delimiter)) {
	    if (Verbose())
		std::cout << "\t" << h << endl;
	    ++read;
	    AddHaplotype(h);
	}
    }
    Freeze();
}

void
HDictionary::ReadFrequencies(istream& istr, size_type num_haps)
{
    if (num_haps == 0) {
	if (Verbose())
	    std::cout << "The number of haplotypes is 0." << endl;
	exit(1);
    }

    char s[80];
    Haplotype h(1);
    Marker marker;
    probability freq;
    typename Volumes::size_type v = 0;

    while (!istr.eof()) {
	_dictionaries.push_back(HaploMap());
	istr.getline(s, 80);
	while (std::strlen(s) > 0) {
	    istrstream is(s);
	    is >> marker >> freq;
	    h.Mutate(0, marker);
	    AddFrequency(h, v, freq);
	    istr.getline(s, 80);
	}
	++v;
    }
    _frozen = true;
    _num_haps = num_haps;
}

///////////////////////////////////////////////////////////////////////////
//
// probability	Frequency( const Haplotype &, size_type)
//
// Returns the frequency of the haplotype h starting at pos.
// Returns 0 (zero) if h doesn't exist.
//
///////////////////////////////////////////////////////////////////////////

probability
HDictionary::Frequency(const Haplotype& h, size_type pos) const
{
    HaploMap::const_iterator search_result = _dictionaries[pos].find(h);
    
    if (search_result != _dictionaries[pos].end()) {
	return search_result->second;
    } else {
	return 0;
    }
}

///////////////////////////////////////////////////////////////////////////
//
// probability  Probability ( const Haplotype & h, size_type pos )
//
// Returns the probability of haplotype segment h, given that it starts
// at position pos.
//
//////////////////////////////////////////////////////////////////////////

probability
HDictionary::Probability(const Haplotype& h, size_type pos)
{
    if (Verbose()) {
	std::cout << "\t\t\t"
		  << "HDictionary::Probabiity( "
		  << h << ", " << pos << " )" << endl;
    }
    
    probability prob;
    if (h.size() == 0) {
	prob = 1.0;	// 0-length segments have probability 1.
    } else {
	prob = _segment->Probability(this, h, pos);
    }
    return prob;
}

///////////////////////////////////////////////////////////////////////////
//
// void	Write( ostream & ostr ) const
//
// Print a representation of dict onto the ostream:
//	<HaploPos1, prob1>
//	<HaploPos2, prob2>
//	     ...
//	<HaploPosn, probn>
//
///////////////////////////////////////////////////////////////////////////

void
HDictionary::Write(ostream &ostr) const
{
    ostr << "Number of haplotypes: " << _num_haps << endl;
    ostr << "Frozen: " << _frozen << endl;
    
    for (size_type i = 0; i < _dictionaries.size(); ++i) {
	HaploMap::const_iterator f = _dictionaries[i].begin();
	HaploMap::const_iterator end = _dictionaries[i].end();
	while (f != end) {
	    ostr << "<"  << i
		 << ", " << f->first
		 << ", " << f->second
		 << ">"  << endl;
	    ++f;
	}
    }
}

///////////////////////////////////////////////////////////////////////////
//
// void	AddHaplotype( const Haplotype & )
//
// Adds every sub-haplotype of the haplotype into the dictionary,
// keeping track of the counts of each sub-haplotype.
//
///////////////////////////////////////////////////////////////////////////

void
HDictionary::AddHaplotype(const Haplotype &h)
{
    if (_frozen)
	return;
    for (Haplotype::size_type i = 0; i < h.size(); ++i) {
	if (!h[i].IsNA()) {
	    for (size_type j = i + 1; j <= h.size(); ++j) {
		Haplotype s(h.Substr(i, j));
		HaploMap& dict(_dictionaries[i]);
		HaploMap::iterator search_result = dict.find(s);
		if (search_result == dict.end()) {
		    dict[s] = 1.0;
		    if (Verbose())
			std::cout << "\t\t" << "New: " << s << endl;
		} else {
		    search_result->second += 1.0;
		    if (Verbose()) {
			std::cout << "\t\t" << "Old: "
				  << search_result->first
				  << " == " << search_result->second << endl;
		    }
		}
		if (h[j].IsNA())
		    break;
	    }
	}
    }
    ++_num_haps;
}

///////////////////////////////////////////////////////////////////////////
//
// void	Freeze()
//
// "Freezes" the dictionary to additions, and changes all of the counts
// stored into frequencies by dividing the counts by NumHaplotypes().
//
///////////////////////////////////////////////////////////////////////////

void
HDictionary::Freeze()
{
    if (_frozen)
	return;
    for (size_type i = 0; i < _dictionaries.size(); ++i) {
	for (HaploMap::iterator it = _dictionaries[i].begin();
	     it != _dictionaries[i].end(); ++it)
	    it->second /= _num_haps;
    }
    _frozen = true;
}
