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

#include <algorithm>
#include <array>
#include <cassert>
#include <cstdlib>
#include <iostream>
#include <fstream>
#include <graph>
#include <graphalg>
#include <strstream>
#include <iterator>
#include <deque>
#include <limits>
#include <cmdargs.h>

#include "commandline.h"
#include "dictionary.h"
#include "geometric.h"
#include "haplotype.h"
#include "path_criterion_factory.h"
#include "path_strategy_factory.h"
#include "probability.h"
#include "segment_factory.h"

void
ProfileInit(int argc, char* argv[],
	    int& verbose, unsigned int& disease, unsigned int& numODH,
	    unsigned int& lmask, unsigned int& rmask,
	    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)
{
    CommandLine::options_t options[] = {
	CommandLine::VERBOSE,
	CommandLine::SEGMENT,
	CommandLine::LEFT_MASK,
	CommandLine::RIGHT_MASK,
	CommandLine::SSM,
	CommandLine::IBD,
	CommandLine::P_MUTATION,
	CommandLine::P_RECOMBINATION,
	CommandLine::DIST_FNAME,
	CommandLine::STR_FNAME,
	CommandLine::RECOMB_FNAME,
	CommandLine::NO_RECOMB_SUM,
	CommandLine::MAX_LATENT_DEPTH,
	CommandLine::PATH_STRATEGY,
	CommandLine::DISEASE,
	CommandLine::DISEASED_FNAME,
	CommandLine::NUM_ODH,
	CommandLine::CRITERION,
	CommandLine::COMPONENTS_FNAME,
	CommandLine::ALLELE_FREQ,
	CommandLine::ROOT_PRIOR_FNAME,
	CommandLine::NO_CLADES,
	CommandLine::EXCLUDE_NORMALS,
	CommandLine::NUM_NORMALS,
	CommandLine::NORMAL_FNAME
    };

    CommandLine::options_t required[] = {
	CommandLine::DISEASE,
	CommandLine::DISEASED_FNAME,
	CommandLine::NORMAL_FNAME
    };
    
    CommandLine::Instance()->Parse(argc, argv,
				   begin_array(options), end_array(options),
				   begin_array(required), end_array(required));

    verbose = CommandLine::Instance()->Verbose();
    disease = CommandLine::Instance()->Disease();
    numODH = CommandLine::Instance()->NumOdh();
    diseasedFilename = CommandLine::Instance()->DiseasedFname();
    componentsFilename = CommandLine::Instance()->ComponentsFname();
    root_prior_fname = CommandLine::Instance()->RootPriorFname();
    no_clades = CommandLine::Instance()->NoClades();
    exclude_normals = CommandLine::Instance()->ExcludeNormals();
    
    if (verbose > 0)
	std::cout << "Initialization." << endl;
    
    /* Get length of haplotype */
    unsigned int hap_length = 0;
    {
	Haplotype h;
	std::ifstream istr(diseasedFilename.data());
	h.ReadDelimited(istr);
	hap_length = h.size();
    }
    
    lmask = CommandLine::Instance()->LeftMask();
    rmask = std::min(CommandLine::Instance()->RightMask(), hap_length);

    if (verbose > 0) {
	CommandLine::Instance()->Write(std::cout, 
				       begin_array(options),
				       end_array(options));
    }

    /*
     * Handle degenerate case.
     */
    if ((disease < lmask) || (disease > rmask)) {
	std::cout << "Tree Weight: "
		  << std::numeric_limits<log_probability>::max() << endl;
	exit(1);
    }

    segment = SegmentProbabilityFactory::Instance()->Make(
	CommandLine::Instance()->Segment(),
	(CommandLine::Instance()->Verbose() > 6));

    dictionary.reset(new HDictionary(segment.get(),
	(CommandLine::Instance()->Verbose() > 5)));
    {
	ifstream istr(CommandLine::Instance()->NormalFname().data());
	if (CommandLine::Instance()->AlleleFreq()) {
	    dictionary->ReadFrequencies(
		istr, CommandLine::Instance()->NumNormals());
	} else {
	    dictionary->ReadDelimited(istr, '\n',
				      CommandLine::Instance()->NumNormals());
	}
    }

    rules.reset(new Probability(Geometric(CommandLine::Instance()->Ssm()),
				dictionary.get(),
				CommandLine::Instance()->Ibd(),
				CommandLine::Instance()->PMutation(),
				CommandLine::Instance()->PRecombination(),
				hap_length,
				lmask,
				rmask,
				CommandLine::Instance()->StrFname(),
				CommandLine::Instance()->DistFname(),
				CommandLine::Instance()->RecombFname(),
				CommandLine::Instance()->NoRecombSum(),
				(CommandLine::Instance()->Verbose() > 4)));

    strategy = PathStrategyFactory::Instance()->Make(
	CommandLine::Instance()->PathStrategy(),
	CommandLine::Instance()->MaxLatentDepth(),
	(CommandLine::Instance()->Verbose() > 1));

    criterion = PathCriterionFactory::Instance()->Make(
	CommandLine::Instance()->Criterion());
}

std::istream&
operator>>(std::istream& istr, std::pair<Haplotype, probability>& x)
{
    istr >> x.second;
    x.first.ReadDelimited(istr);
    return istr;
}

bool
operator<(const Link& lhs, const Link& rhs)
{
    return lhs.priority < rhs.priority;
}

int
PathsIndex(int i, int j)
{
    assert(i > j);
    return i * (i-1) / 2 + j;
}

void
LinkWithPaths(const MutationClade& clade1, const MutationClade& clade2,
	      std::vector<std::pair<Path, Path> >& clade_paths,
	      unsigned int disease, Probability* rules, PathStrategy* strategy,
	      int verbose)
{
    Path best_path_12, best_path_21, cur_path;
    bool seen_12 = false, seen_21 = false;
    log_probability min_weight_12 = std::numeric_limits<log_probability>::max();
    log_probability min_weight_21 = std::numeric_limits<log_probability>::max();
    log_probability cur_weight = 0;
    
    MutationClade::const_node_iterator m, n;
    for (m = clade1.n_begin(); m != clade1.n_end(); ++m)
	for (n = clade2.n_begin(); n != clade2.n_end(); ++n) {
	    cur_path = strategy->MakePath(clade1.value(*m), clade2.value(*n),
					  disease, rules);
	    cur_weight = cur_path.Weight();
	    if ((cur_weight <= min_weight_12) || !seen_12) {
		min_weight_12 = cur_weight;
		best_path_12 = cur_path;
		seen_12 = true;
	    }
	    cur_path = strategy->MakePath(clade2.value(*n), clade1.value(*m),
					  disease, rules);
	    cur_weight = cur_path.Weight();
	    if ((cur_weight < min_weight_21) || !seen_21) {
		min_weight_21 = cur_weight;
		best_path_21 = cur_path;
		seen_21 = true;
	    }
	}
    if (verbose > 2) {
	std::cout << "\t" << "Linking: " << best_path_12.From() << endl;
	std::cout << "\t" << "   ----> " << best_path_12.To() << endl;
	std::cout << "\t" << "Linking: " << best_path_21.To() << endl;
	std::cout << "\t" << "   <---- " << best_path_21.From() << endl;
	std::cout << best_path_12 << endl;
	std::cout << best_path_21 << endl;
    }
    clade_paths.push_back(std::make_pair(best_path_12, best_path_21));
    if (verbose > 2)
	std::cout << endl;
}

void
ExplodeClade(MutationClade& exploded_clade,
	     const LinkClade& meta_clade,
	     const vector<MutationClade*>& clades, int verbose)
{
    if (verbose > 2)
	std::cout << "\t" << "Exploding clades:" << endl;
    for (LinkClade::const_node_iterator n = meta_clade.n_begin();
	 n != meta_clade.n_end(); ++n) {
	if (verbose > 2) {
	    std::cout << " " << meta_clade.value(*n);
	    std::cout.flush();
	}
	const MutationClade &g = *clades[meta_clade.value(*n)];
	exploded_clade.add_node_range(g, g.n_begin(), g.n_end());
    }
    if (verbose > 2)
	std::cout << endl << "\t" << "Adding paths:" << endl;
    for (LinkClade::const_edge_iterator e = meta_clade.e_begin();
	 e != meta_clade.e_end(); ++e) {
	if (verbose > 2) {
	    std::cout << "\t\t" << meta_clade.value(e->source())
		      << " --> " << meta_clade.value(e->target()) << endl;
	}
	AddPath(exploded_clade, meta_clade.weight(*e), verbose);
    }
    if (verbose > 2)
	std::cout << endl;
}

void
AddPath(MutationClade& tree, const Path& path, int verbose)
{
    typedef graph_traits<MutationClade>::node_type node_type;
    
    if (path.begin() == path.end())
	return;
    
    Path::const_iterator i(path.begin());
    node_type head_result(tree.find_node(i->haplotype));
    node_type prev;
    if (head_result != node_type()) {
	prev = head_result;
    } else {
	prev = tree.add_node(i->haplotype);
    }
    ++i;
    
    if (verbose > 3)
	std::cout << "Linking: " << tree.value(prev) << endl;
    while (i != path.end()) {
	node_type result(tree.find_node(i->haplotype));
	node_type curr;
	if (result != node_type()) {
	    curr = result;
	} else {
	    curr = tree.add_node(i->haplotype);
	}
	tree.add_edge(prev, curr, i->weight);
	if (verbose > 3) {
	    std::cout << "  -----> "
		      << i->haplotype << "\t" << i->weight << endl;
	}
	prev = curr;
	++i;
    }
    if (verbose > 3)
	std::cout << "End Linking." << endl;
}

static
graph_traits<MutationClade>::node_type
PruneLatentRoot(MutationClade& clade,
		graph_traits<MutationClade>::node_type root)
{
    graph_traits<MutationClade>::node_type true_root(root);
    graph_traits<MutationClade>::node_type test_root(root);
    
    while (test_root.outdeg() == 1 && clade.value(test_root).IsLatent()) {
	true_root = *(test_root.adj_nodes_begin());
	clade.erase(test_root);
	test_root = true_root;
    }
    return true_root;
}

///////////////////////////////////////////////////////////////////////////
//
// PruneLatentChildren()
//
///////////////////////////////////////////////////////////////////////////

static
void
PruneLatentChildren(MutationClade& clade,
		    graph_traits<MutationClade>::node_type& root)
{
    typedef graph_traits<MutationClade>::node_type node_type;

    if (root.outdeg() == 0)
	return;
    
    vector<node_type> neighbors;
    for (typename node_type::adj_nodes_iterator n = root.adj_nodes_begin();
	 n != root.adj_nodes_end(); ++n) {
	neighbors.push_back(*n);
    }
    for (vector<node_type>::iterator n = neighbors.begin();
	 n != neighbors.end(); ++n) {
	PruneLatentChildren(clade, *n);
	if ((n->outdeg() == 0) && clade.value(*n).IsLatent()) {
	    clade.erase(*n);
	}
    }
}

graph_traits<MutationClade>::node_type
PruneLatents(MutationClade& clade, graph_traits<MutationClade>::node_type root)
{
    PruneLatentChildren(clade, root);
    return PruneLatentRoot(clade, root);
}

void
PrintTreePathWeights(ostream& ostr,
		     const MutationClade& clade,
		     const graph_traits<MutationClade>::node_type& u,
		     graph_traits<MutationClade>::weight_type acc_weight)
{
    typedef graph_traits<MutationClade>::node_type node_type;
    typedef graph_traits<MutationClade>::weight_type weight_type;
    typedef typename node_type::adj_edges_iterator adj_edges_iterator;
    
    for (adj_edges_iterator e = u.adj_edges_begin();
	 e != u.adj_edges_end(); ++e) {
	acc_weight += clade.weight(*e);
	if (!clade.value(e->target()).IsLatent()) {
	    ostr << acc_weight << endl;
	    acc_weight = static_cast<weight_type>(0);
	}
	PrintTreePathWeights(ostr, clade, e->target(), acc_weight);
    }
}

void
PrintRootedTree(ostream& ostr, int component, unsigned int disease,
		const TreeProperties& tp)
{
    ostr << "===== Tree[" << component << "]:" << endl
	 << "Tree Disease: " << disease << endl
	 << "Tree Root: " << tp.tree->value(tp.root) << endl
	 << "Tree Root Count: " << tp.root_count << endl
	 << "Tree Weight: " << tp.weight << endl
	 << "Tree Overcount: " << tp.overcount << endl
	 << "Tree Subset Size: " << tp.subset_size << endl;
    WriteRootedTree(ostr, *tp.tree, tp.root);
    ostr << "Tree Topology: " << endl;
    WriteTreeTopology(ostr, *tp.tree, tp.root, "  ");
    ostr << endl;
}

auto_ptr<MutationClade>
FindBestRootedTree(const MutationClade& c,
		   const graph_traits<MutationClade>::node_type& root,
		   graph_traits<MutationClade>::node_type* tree_root,
		   graph_traits<MutationClade>::weight_type* tree_weight,
		   node_map<graph_traits<MutationClade>::weight_type>* tree_dist)
{
    typedef typename graph_traits<MutationClade>::node_type node_type;
    typedef typename graph_traits<MutationClade>::edge_type edge_type;
    typedef typename graph_traits<MutationClade>::weight_type weight_type;
    typedef typename MutationClade::const_node_iterator const_node_iterator;

    node_map<weight_type> clade_dist;
    node_map<node_type> node_translate(c.topology());
    std::list<edge_type> tree_list;
    IsObserved<MutationClade> visible(&c);
    DijkstraSubset(c, root, visible, &clade_dist,
		   std::back_inserter(tree_list));

    auto_ptr<MutationClade> tree(new MutationClade());
    if (tree_list.begin() == tree_list.end()) {
	node_translate[root] = tree->add_node(c.value(root));
    } else {
	tree->add_edge_range(c, tree_list.begin(), tree_list.end(),
			     &node_translate);
    }

    *tree_root = tree->find_node(c.value(root));
    *tree_weight = weight(*tree);
    tree_dist->init(c.topology(), numeric_limits<weight_type>::max());

    const_node_iterator f = c.n_begin();
    const_node_iterator end = c.n_end();
    while (f != end) {
	node_type tree_node = node_translate[*f];
	weight_type weight = clade_dist[*f];
	if (tree_node != node_type()) {
	    (*tree_dist)[tree_node]
		= (weight > numeric_limits<weight_type>::max()
		   ? numeric_limits<weight_type>::max()
		   : weight);
	}
	++f;
    }

    return tree;
}

graph_traits<MutationClade>::weight_type
StarCriterion(const MutationClade& tree,
	      const graph_traits<MutationClade>::node_type& tree_root,
	      const node_map<graph_traits<MutationClade>::weight_type>& tree_dist)
{
    typedef typename graph_traits<MutationClade>::node_type node_type;
    typedef typename graph_traits<MutationClade>::weight_type weight_type;
    typedef typename MutationClade::const_node_iterator const_node_iterator;
    
    weight_type acc_path_weight = 0;
    unsigned int num_observed = 0;

    const_node_iterator f = tree.n_begin();
    const_node_iterator end = tree.n_end();
    while (f != end) {
	if (!tree.value(*f).IsLatent()) {
	    acc_path_weight += tree_dist[*f];
	    ++num_observed;
	}
	++f;
    }
    if ((num_observed == 0) ||
	(acc_path_weight > numeric_limits<weight_type>::max())) {
	return numeric_limits<weight_type>::max();
    } else if (num_observed == 1) {
	return 0;
    } else {
	return acc_path_weight / (num_observed - 1);
    }
}

void
TotalPath(const MutationClade& tree,
	  const graph_traits<MutationClade>::node_type& root,
	  unsigned int root_dist, unsigned int* total_dist)
{
    typedef typename graph_traits<MutationClade>::node_type node_type;
    typedef typename node_type::adj_nodes_iterator adj_nodes_iterator;

    if (!tree.value(root).IsLatent())
	*total_dist += root_dist;

    adj_nodes_iterator f = root.adj_nodes_begin();
    adj_nodes_iterator end = root.adj_nodes_end();
    while (f != end) {
	TotalPath(tree, *f, root_dist + 1, total_dist);
	++f;
    }
}
