/*
 * 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 <graph>
#include <graphalg>
#include <fstream>
#include <iostream>
#include <map>
#include <string>
#include <vector>
#include <cmdargs.h>

#include "haplotype.h"
#include "marker.h"
#include "types.h"

typedef stl_graph<Haplotype, log_probability> Tree;
typedef typename graph_traits<Tree>::node_type node_type;
typedef typename graph_traits<Tree>::edge_type edge_type;
typedef typename graph_traits<Tree>::weight_type weight_type;
typedef typename node_type::adj_edges_iterator adj_edges_iterator;
    
typedef Haplotype::size_type size_type;

enum color_t { SMALL_MUTATION, LARGE_MUTATION, RECOMBINATION, NUM_COLORS };
const char* color_array[NUM_COLORS];
std::vector<unsigned int> mut_jump;

node_type	ReadTree(std::istream&, Tree&, unsigned int&, weight_type&);
void		WriteHaplotype(std::ostream&, int, int, weight_type,
			       const Haplotype&, color_t,
			       size_type, size_type);
void		WriteDifferences(std::ostream&, unsigned int,
				 const Haplotype&, const Haplotype&,
				 int, int, weight_type);

CmdArgBool document('d', "document",
		    "Set this option to produce a document suitable "
		    "for processing by the LaTeX program. "
		    "Default: (unset)");
CmdArgStr font_size('f', "font-size", "fontsize",
		    "LaTeX fontsize; see Section 5.8.1 in Lamport's "
		    "LaTeX User's Guide and Reference Manual. "
		    "Default: (none)");
CmdArgStr str_fname('j', "str-fname", "filename",
		    "File containing sizes of the repeated "
		    "units of the STR marker at each locus. "
		    "Default: (none)");
CmdArgStr small_mutation('s', "small-mutation", "color",
			 "Color for small mutations.  For more colors, "
			 "see colordvi.tex. "
			 "Default: Blue");
CmdArgStr large_mutation('l', "large-mutation", "color",
			 "Color for large mutations. "
			 "Default: Purple");
CmdArgStr recombination('r', "recombination", "color",
			"Color for recombinations. "
			"Default: Red");
CmdArgStr tree_fname("filename",
		     "File containing a tree, like the output of the "
		     "profile program.");

int
main(int argc, char* argv[])
{
    CmdLine cmd(argv[0], &document, &font_size, &str_fname,
		&small_mutation, &large_mutation, &recombination,
		&tree_fname, NULL);
    CmdArgvIter arg_iter(--argc, ++argv);

    document = false;
    font_size = 0;
    str_fname = 0;
    small_mutation = "Blue";
    large_mutation = "Purple";
    recombination = "Red";

    cmd.parse(arg_iter);

    color_array[SMALL_MUTATION] = small_mutation;
    color_array[LARGE_MUTATION] = large_mutation;
    color_array[RECOMBINATION] = recombination;

    Tree tree;
    unsigned int disease;
    weight_type weight;
    node_type root;
    
    if (string(tree_fname) == string("-")) {
	root = ReadTree(std::cin, tree, disease, weight);
    } else {
	std::ifstream tree_istr(tree_fname);
	root = ReadTree(tree_istr, tree, disease, weight);
    }

    const Haplotype& root_hap(tree.value(root));

    if (str_fname != 0) {
	std::ifstream str_istr(str_fname);
	std::copy(std::istream_iterator<unsigned int>(str_istr),
		  std::istream_iterator<unsigned int>(),
		  std::back_inserter(mut_jump));
	str_istr.close();
    } else {
	std::fill_n(std::back_inserter(mut_jump), root_hap.size(), 1);
    }

    if (document) {
	std::cout << "\\documentclass{article}" << endl
		  << "\\usepackage{fullpage}" << endl
		  << "\\usepackage{colordvi}" << endl
		  << endl
		  << "\\begin{document}" << endl
		  << "\\begin{flushright}" << endl
		  << "  {\\bf " << tree_fname << " \\\\" << endl
		  << "  Disease: " << disease << " \\\\" << endl
		  << "  Weight: " << weight << " }" << endl
		  << "\\end{flushright}" << endl
		  << endl
		  << "\\begin{table}[h]" << endl;
    }

    std::cout << "\\begin{center}" << endl;
    if (font_size != 0)
	std::cout << "{\\" << font_size << endl;
    std::cout << "\\begin{tabular}{rr|r|";
    for (size_type i = 0; i < root_hap.size(); ++i)
	std::cout << "r";
    std::cout << "ll}" << endl;

    node_map<int> ids(tree.topology());
    int index = 1;
    
    // parent of 0 indicates node is root
    ids[root] = index++;
    WriteHaplotype(std::cout, ids[root], 0, 0, root_hap, RECOMBINATION,
		   root_hap.size(), root_hap.size());

    adj_edges_iterator adj = root.adj_edges_begin();
    adj_edges_iterator adj_end = root.adj_edges_end();
    std::stack<edge_type> qu;
    while (adj != adj_end)
	qu.push(*adj++);
    while (!qu.empty()) {
	edge_type e = qu.top();
	node_type u = e.target();
	qu.pop();
	
	ids[u] = index++;
	WriteDifferences(std::cout, disease,
			 tree.value(e.source()), tree.value(u),
			 ids[u], ids[e.source()], tree.weight(e));
	adj = u.adj_edges_begin();
	adj_end = u.adj_edges_end();
	while (adj != adj_end)
	    qu.push(*adj++);
    }
    
    std::cout << "\\end{tabular}" << endl;
    if (font_size != 0)
	std::cout << "}" << endl;
    std::cout << "\\end{center}" << endl;

    if (document) {
	std::cout << "\\end{table}" << endl
		  << "\\end{document}" << endl;
    }
}

node_type
ReadTree(std::istream& istr, Tree& tree, unsigned int& disease,
	 weight_type& weight)
{
    node_type root;
    std::map<int, node_type> id_map;
    std::map<int, int> ancestor_map;
    char c;
    std::string word;

    istr >> word;		// potential "Tree"
    while (!istr.eof()) {
	while (word != "Tree")
	    istr >> word;
	istr >> word;		// potential "Tree Disease:"
	if (word == "Disease:") {
	    unsigned int d;
	    istr >> d;
	    disease = d;
	    break;
	}
    }
    
    istr >> word;		// potential "Tree"
    while (!istr.eof()) {
	while (word != "Tree")
	    istr >> word;
	istr >> word;		// potential "Tree Weight:"
	if (word == "Weight:") {
	    weight_type w;
	    istr >> w;
	    weight = w;
	    break;
	}
    }
    
    istr >> c;			// potential tree open brace
    while (c != '{')
	istr >> c;
    istr >> c;			// potential tree close brace
    while (c != '}') {
	istr.putback(c);
	
	int id, ancestor;
	weight_type weight;
	Haplotype h;
	
	istr >> id >> ancestor >> weight >> h;
	id_map[id] = tree.add_node(h);
	ancestor_map[id] = ancestor;
	if (ancestor == 0) {
	    root = id_map[id];
	} else {
	    tree.add_edge(id_map[ancestor], id_map[id], weight);
	}
	istr >> c;		// potential tree close brace
    }
    
    return root;
}

void
WriteHaplotype(std::ostream& ostr, int id, int ancestor, weight_type weight,
	       const Haplotype& h, color_t color,
	       size_type color_start, size_type color_end)
{
    ostr << "  " << id << " & " << ancestor << " & " << weight << " & ";
    size_type start = 0;
    while (start != color_start)
	ostr << h[start++] << " & ";
    while (start != color_end)
	ostr << "\\" << color_array[color] << "{" << h[start++] << "} & ";
    while (start != h.size())
	ostr << h[start++] << " & ";

    Haplotype::property_t properties = h.Properties();
    if (properties & Haplotype::LATENT)
	ostr << "[L]";
    ostr << " & ";
    if (properties & Haplotype::NORMAL)
	ostr << "[N]";
    ostr << "\\\\" << endl;
}

void
WriteDifferences(std::ostream& ostr, unsigned int disease,
		 const Haplotype& s, const Haplotype& t,
		 int id, int ancestor, weight_type weight)
{
    unsigned int delta = 0;
    unsigned int jump = 1;
    unsigned int diffs = 0;

    size_type left = disease;
    for (size_type i = disease; i != 0; --i) {
	if (s[i-1] != t[i-1]) {
	    ++diffs;
	    if (diffs > 1)
		break;
	    left = i;
	}
    }
    switch (diffs) {
    case 0:
	break;
    case 1:
	delta = Marker::diff(s[left - 1], t[left - 1]);
	jump = mut_jump[left - 1];
	if ((delta % jump == 0) && (delta / jump == 1)) {
	    WriteHaplotype(ostr, id, ancestor, weight, t, SMALL_MUTATION,
			   left - 1, left);
	} else {
	    WriteHaplotype(ostr, id, ancestor, weight, t, LARGE_MUTATION,
			   left - 1, left);
	}
	return;
    case 2:
	WriteHaplotype(ostr, id, ancestor, weight, t, RECOMBINATION, 0, left);
	return;
    }

    diffs = 0;
    size_type right = disease;
    for (size_type i = disease; i != t.size(); ++i) {
	if (s[i] != t[i]) {
	    ++diffs;
	    if (diffs > 1)
		break;
	    right = i;
	}
    }
    switch (diffs) {
    case 0:
	break;
    case 1:
	delta = Marker::diff(s[right], t[right]);
	jump = mut_jump[right];
	if ((delta % jump == 0) && (delta / jump == 1)) {
	    WriteHaplotype(ostr, id, ancestor, weight, t, SMALL_MUTATION,
			   right, right + 1);
	} else {
	    WriteHaplotype(ostr, id, ancestor, weight, t, LARGE_MUTATION,
			   right, right + 1);
	}
	return;
    case 2:
	WriteHaplotype(ostr, id, ancestor, weight, t, RECOMBINATION,
		       right, t.size());
	return;
    }
}
