/*
 * 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 <cmath>
#include <fstream>
#include <graph>
#include <graphalg>
#include <iostream>
#include <limits>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <utility>
#include <vector>

#include <GTL_connectivity.h>

#include "commandline.h"
#include "dictionary.h"
#include "haplotype.h"
#include "path_criterion.h"
#include "path_strategy_factory.h"
#include "probability.h"
#include "profile.h"
#include "segment_factory.h"
#include "types.h"

using namespace std;

class HDictionary;
class PathStrategy;

int
main(int argc, char* argv[])
{
    int			verbose;
    unsigned int	disease;
    unsigned int	numODH;
    unsigned int	left_mask, right_mask;
    string		diseasedFilename;
    string		componentsFilename;
    string		root_prior_fname;
    bool		no_clades;
    bool		exclude_normals;

    auto_ptr<SegmentProbability> segment;
    auto_ptr<HDictionary> dictionary;
    auto_ptr<Probability> rules;
    auto_ptr<PathStrategy> strategy;
    auto_ptr<PathCriterion> criterion;
    
    ProfileInit(argc, argv,
		verbose, disease, numODH,
		left_mask, right_mask,
		diseasedFilename,
		componentsFilename, root_prior_fname, no_clades,
		exclude_normals, 
		segment, dictionary, rules, strategy, criterion);
    
    // Read in all of the observed disease haplotypes here.
    //
    if (verbose > 0)
	cout << "Reading in diseased haplotypes." << endl;
    vector<Haplotype>	observed;
    multiset<Haplotype>	all_observed;
    {
	ifstream observed_istr(diseasedFilename.data());
	Haplotype h;
	while (h.ReadDelimited(observed_istr)) {
	    if (!exclude_normals ||
		(exclude_normals && !(h.Properties() & Haplotype::NORMAL))) {
		h.SetProperties(h.Properties() | Haplotype::DISEASED);
		h.Mask(rules->LeftMask(), rules->RightMask());
		if (all_observed.find(h) == all_observed.end())
		    observed.push_back(h);
		all_observed.insert(h);
	    }
	}
    }
    if (verbose > 1) {
	for (unsigned int i = 0; i < observed.size(); ++i)
	    cout << "\t" << observed[i] << endl;
    }
    map<Haplotype, probability> root_prior;
    {
	if (!root_prior_fname.empty()) {
	    ifstream istr(root_prior_fname.data());
	    copy(istream_iterator<pair<Haplotype, probability> >(istr),
		 istream_iterator<pair<Haplotype, probability> >(),
		 inserter(root_prior, root_prior.end()));
	} else {
	    /*
	     * XXX - This is a stupid prior!
	     */
	    for (unsigned int i = 0; i < observed.size(); ++i)
		root_prior[observed[i]] = all_observed.count(observed[i]);
	}
	probability sum = 0;
	map<Haplotype, probability>::iterator f = root_prior.begin();
	map<Haplotype, probability>::iterator end = root_prior.end();
	while (f != end) {
	    sum += f->second;
	    ++f;
	}
	if (verbose > 0)
	    std::cout << "Pr(S is root):" << endl;
	f = root_prior.begin();
	while (f != end) {
	    f->second /= sum;
	    if (verbose > 0) {
		std::cout << "\t" << f->second
			  << " :\t" << f->first << endl;
	    }
	    ++f;
	}
	if (verbose > 0)
	    std::cout << "End Pr(S is root)." << endl << endl;
    }
    
    // === Step 1:
    //
    // Partition observed haplotypes into the fewest possible mutation
    // clades, where each clade is composed of haplotypes which pairwise
    // differ from some other haplotype at only one locus.
    //
    // X = [x0] U [x1] U ... U [xn],
    // x ~ y if x,y differ at exactly one locus.
    //
    if (verbose > 0) {
	cout << "Step 1:  "
	     << "Partition observed haplotypes into mutation clades."
	     << endl;
    }
    vector<MutationClade*> mutation_clades;
    int M;
    if (no_clades) {
	M = observed.size();
	for (unsigned int i = 0; i < M; ++i) {
	    mutation_clades.push_back( new MutationClade() );
	    mutation_clades.back()->add_node(observed[i]);
	}
    } else {
	MutationClade distance_graph;
	for (unsigned int i = 0; i < observed.size(); ++i)
	    distance_graph.add_node(observed[i]);
	MutationClade::node_iterator m, n;
	for (m = distance_graph.n_begin(); m != distance_graph.n_end(); ++m) {
	    for (n = distance_graph.n_begin(); n != m; ++n) {
		graph_traits<MutationClade>::value_type m_value
		    = distance_graph.value(*m);
		graph_traits<MutationClade>::value_type n_value
		    = distance_graph.value(*n);
		
		if (rules->PossibleMutation(m_value, n_value)) {
		    distance_graph.add_edge(*m, *n,
					    rules->log_P(m_value, n_value,
							 disease));
		    distance_graph.add_edge(*n, *m,
					    rules->log_P(n_value, m_value,
							 disease));
		}
	    }
	}
	connectivity comp_dfs;
	comp_dfs.scan_whole_graph(true);
	comp_dfs.calc_comp_num(true);
	comp_dfs.run(distance_graph.topology());
	M = comp_dfs.number_of_components();

	for (int i = 0; i < M; ++i) {
	    mutation_clades.push_back( new MutationClade() );
	    mutation_clades.back()->add_node_range(distance_graph,
						   comp_dfs.component_begin(i),
						   comp_dfs.component_end(i));
	}
    }
    if (verbose > 1) {
	cout << "Begin Mutation Clades:" << endl;
	for ( unsigned int i = 0; i < mutation_clades.size(); ++i )
	    cout << "===== Clade[" << i << "] =====" << endl
		 << *mutation_clades[i] << endl;
	cout << "End Mutation Clades." << endl;
    }

    // === Step 2:
    //
    // For each pair of clades, compute a pair of paths between them
    // and insert them into a priority queue.
    //
    if (verbose > 0)
	cout << "Step 2:  Compute pairwise clade distances." << endl;
    
    vector<pair<Path, Path> > clade_paths;
    priority_queue<Link, vector<Link>, greater<Link> > link_pq;
    {
	clade_paths.reserve(PathsIndex(M, M-1) + 1);
	
	// For each pair of cladograms find the pair of haplotypes, one
	// from each cladogram, which minimizes the distance between the
	// clades.
	//
	for (int i = 0; i < M; ++i) {
	    for (int j = 0; j < i; ++j) {
		if (verbose > 1) {
		    cout << "===== " << "[" << i << "], [" << j << "]:"
			 << endl;
		}
		LinkWithPaths(*mutation_clades[i], *mutation_clades[j],
			      clade_paths, disease,
			      rules.get(), strategy.get(), verbose);
	    }
	}
	if (verbose > 1)
	    cout << "Begin Priority Queue:" << endl;
	for (int i = 0; i < M; ++i) {
	    for (int j = 0; j < i; ++j) {
		log_probability	priority
		    = criterion->compute(clade_paths[PathsIndex(i, j)].first,
					 clade_paths[PathsIndex(i, j)].second);
		link_pq.push(Link(i, j, priority));
		if (verbose > 1) {
		    cout << "\t"
			 << "[" << i << "]"
			 << " <---(" << priority << ")---> "
			 << "[" << j << "]"
			 << endl;
		}
	    }
	}
	if (verbose > 1)
	    cout << "End Priority Queue." << endl;
    }
    
    // === Step 3:
    //
    // Repeatedly choose the cladogram pair which yields the highest
    // value of the chosen criterion function and link them together.
    // Stop when numODH connected components are formed, or when all of
    // the cladograms are linked together.
    //
    if (verbose > 0)
	cout << "Step 3:  Link mutation clades in pairwise fashion." << endl;

    if (verbose > 1) {
	cout << "\t" << "Stop at " << numODH << " connected components."
	     << endl;
    }
    
    vector<LinkClade*> component_clades;
    {
	LinkClade* link_clade = new LinkClade();
	vector<LinkClade::node>	link_clade_nodes(M);
	
	for (int i = 0; i < M; ++i)
	    link_clade_nodes[i] = link_clade->add_node(i);
    
	vector<int> compnum_of_clade(M);
	set<int> used_compnums;
	
	if (verbose > 1)
	    cout << "\t" << "Adding links:" << endl;
	for (int i = 0; i < M; ++i) {
	    compnum_of_clade[i] = i;
	    used_compnums.insert( i );
	}
	if (!link_pq.empty()) {
	    Link link = link_pq.top();
	    while (used_compnums.size() > numODH && !link_pq.empty()) {
		int comp_from(compnum_of_clade[link.from_clade]);
		int comp_to(compnum_of_clade[link.to_clade]);
		
		if (verbose > 1)
		    cout << "\t\t" << "Current lowest criterion: "
			      << link.priority << endl;
		int index = PathsIndex(link.from_clade, link.to_clade);
		link_clade->add_edge(link_clade_nodes[link.from_clade],
				     link_clade_nodes[link.to_clade],
				     clade_paths[index].first);
		link_clade->add_edge(link_clade_nodes[link.to_clade],
				     link_clade_nodes[link.from_clade],
				     clade_paths[index].second);
		if (comp_from != comp_to) {
		    if (verbose > 1)
			cout << "\t\t"
			     << "Merging [" << comp_to << "] into ["
			     << comp_from << "]." << endl
			     << "\t\t" << "Components:";
		    for (int i = 0; i < M; ++i) {
			if (compnum_of_clade[i] == comp_to)
			    compnum_of_clade[i] = comp_from;
			if (verbose > 1)
			    cout << " " << compnum_of_clade[i];
		    }
		    if (verbose > 1)
			cout << endl;
		    used_compnums.erase(comp_to);
		    if (verbose > 1) {
			cout << "\t\t" << "Used:";
			for (set<int>::iterator i = used_compnums.begin();
			     i != used_compnums.end(); ++i)
			    cout << " " << *i;
			cout << endl;
		    }
		}
		link_pq.pop();
		link = link_pq.top();
	    }
	    if (verbose > 1) {
		cout << "Number of connected components: "
			  << used_compnums.size()
			  << endl;
	    }
	}
	if (verbose > 1)
	    cout << "End Adding links." << endl;
	
	if (verbose > 1)
	    cout << "Sort into components" << endl;
	map<int, list<LinkClade::node> > components;
	for (int i = 0; i < M; ++i) {
	    components[compnum_of_clade[i]].push_back(link_clade_nodes[i]);
	    if (verbose > 1)
		cout << "\t" << "Mutation clade " << i
			  << " is in component [" << compnum_of_clade[i] << "]"
			  << endl;
	}
	if (verbose > 1)
	    cout << "End Sort into components" << endl << endl;
	component_clades.reserve(components.size());
	{
	    map<int, list<LinkClade::node> >::const_iterator c;
	    for (c = components.begin(); c != components.end(); ++c) {
		if (verbose > 1) {
		    cout << "\t" << "Adding component [" << c->first << "]"
			 << endl;
		}
		component_clades.push_back(new stl_graph<int, Path>());
		component_clades.back()->add_node_range(*link_clade,
							c->second.begin(),
							c->second.end());
	    }
	}
    }
    if (verbose > 1) {
	cout << "Connected components:" << endl;
	for (unsigned int i = 0; i < component_clades.size(); ++i) {
	    cout << "============================" << endl
		 << *component_clades[i] << endl;
	}
	cout << "============================" << endl
	     << "End connected components:" << endl;
    }
    
    // === Step 4:
    //
    // Find the best tree which links all of the observed haplotypes
    // together.
    //
    if (verbose > 0)
	cout << "Step 4:  Find best trees." << endl;

    vector<TreeProperties>	trees;
    {
	trees.reserve(component_clades.size());

	ofstream components_ostr;
	if (!componentsFilename.empty()) {
	    if (verbose > 1) {
		cout << "Writing rooted trees per component to "
		     << componentsFilename << "." << endl;
	    }
	    components_ostr.open(componentsFilename.data());
	}

	for (unsigned int i = 0; i < component_clades.size(); ++i) {
	    if (!componentsFilename.empty()) {
		components_ostr << "==== Component " << i
				<< " ===================="
				<< "====================================="
				<< endl;
	    }
	    MutationClade exploded_clade;
	    ExplodeClade(exploded_clade, *component_clades[i],
			 mutation_clades, verbose);
	    if (verbose > 1) {
		cout << "Exploded clade:" << endl
		     << exploded_clade << endl;
	    }

	    MutationClade::const_node_iterator root = exploded_clade.n_begin();
	    MutationClade::const_node_iterator root_end = exploded_clade.n_end();

	    while (root != root_end) {
		/* Only allow observed haplotypes to be roots. */
		if ((exploded_clade.value(*root).IsLatent()) ||
		    (root_prior[exploded_clade.value(*root)] == 0))
		    break;

		if (verbose > 1) {
		    cout << "\t" << "Possible root "
				 << exploded_clade.value(*root)
				 << ":" << "\t";
		    cout.flush();
		}

		TreeProperties tp;
		node_map<log_probability> tree_dist;
		auto_ptr<MutationClade> tree(
		     FindBestRootedTree(exploded_clade, *root,
					&tp.root, &tp.weight, &tree_dist));
		tp.tree = tree.get();
		tp.root_count
		    = all_observed.count(exploded_clade.value(*root));
		tp.weight -= logP(root_prior[tree->value(tp.root)]);
		if (tp.weight > numeric_limits<log_probability>::max())
		    tp.weight = numeric_limits<log_probability>::max();

		unsigned int total_dist = 0;
		TotalPath(*tp.tree, tp.root, 0, &total_dist);
		tp.overcount = total_dist - (observed.size() - 1);

		tp.subset_size = 0;
		MutationClade::const_node_iterator f = tree->n_begin();
		MutationClade::const_node_iterator end = tree->n_end();
		while (f != end) {
		    tp.subset_size += all_observed.count(tree->value(*f));
		    ++f;
		}

		if (verbose > 1) 
		    cout << tp.weight << " (" << tp.overcount << ")"<< endl;
		if (!componentsFilename.empty())
		    PrintRootedTree(components_ostr, i, disease, tp);
		++root;
	    }
	}
    }

    return 0;
}
