/*
  Copyright (C) 2008-2009 Benjamin Redelings

  This file is part of BAli-Phy.

  BAli-Phy is free software; you can redistribute it and/or modify it under
  the terms of the GNU General Public License as published by the Free
  Software Foundation; either version 2, or (at your option) any later
  version.

  BAli-Phy is distributed in the hope that it will be useful, but WITHOUT ANY
  WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with BAli-Phy; see the file COPYING.  If not see
  <http://www.gnu.org/licenses/>.  */

#include "alignment/index-matrix.H"
#include <utility>
#include <fstream>
#include <string>
#include <cmath>
#include <list>
#include <numeric>
#include "util/myexception.H"
#include "optimize.H"
#include "findroot.H"
#include "alignment/alignment-util.H"
#include "alignment/load.H"
#include "distance-methods.H"
#include "util/rng.H"
#include "statistics.H"
#include "util/log-level.H"

#include <boost/graph/graph_traits.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/topological_sort.hpp>

using namespace boost;

typedef adjacency_list< vecS, vecS, bidirectionalS> Graph; 
typedef graph_traits<Graph>::vertex_descriptor Vertex;
typedef graph_traits<Graph>::edge_descriptor Edge_t;

#include <boost/program_options.hpp>

namespace po = boost::program_options;
using po::variables_map;

using namespace std;

void do_setup(const variables_map& args,vector<alignment>& alignments) 
{
    //------------ Try to load alignments -----------//
    int maxalignments = args["max-alignments"].as<int>();
    unsigned skip = args["skip"].as<unsigned>();

    // --------------------- try ---------------------- //
    if (log_verbose) std::cerr<<"alignment-max: Loading alignments...";
    list<alignment> As = load_alignments(std::cin,get_alphabet_name(args),skip,maxalignments);
    alignments.insert(alignments.begin(),As.begin(),As.end());
    if (log_verbose) std::cerr<<"done. ("<<alignments.size()<<" alignments)"<<std::endl;
    if (not alignments.size())
	throw myexception()<<"Alignment sample is empty.";
}


variables_map parse_cmd_line(int argc,char* argv[]) 
{ 
    using namespace po;

    // named options
    options_description all("Allowed options");
    all.add_options()
	("help,h", "Produce help message")
	("alphabet",value<string>(),"Specify the alphabet: DNA, RNA, Amino-Acids, Amino-Acids+stop, Triplets, Codons, or Codons+stop.")
	("skip,s",value<unsigned>()->default_value(0),"Number of alignment samples to skip")
	("max-alignments,m",value<int>()->default_value(1000),"Maximum number of alignments to analyze")
	("analysis",value<string>()->default_value("wsum"),"sum, wsum, wsum2, multiply")
	("out,o",value<string>()->default_value("-"),"Output file (defaults to stdout)")
	("out-probabilities,p",value<string>(),"Output file for column probabilities, if specified")
	("verbose,V","Output more log messages on stderr.")
	;

    variables_map args;     
    store(parse_command_line(argc, argv, all), args);
    notify(args);    

    if (args.count("help")) {
	cout<<"Construct a posterior decoding alignment to summarize an alignment sample.\n\n";
	cout<<"Usage: alignment-max [OPTIONS] < alignments-file\n\n";
	cout<<all<<"\n";
	exit(0);
    }

    if (args.count("verbose")) log_verbose = 1;

    return args;
}


/// Extract column @c from the index-matrix @M
vector<int> get_column(const matrix<int>& M, int c)
{
    vector<int> column(M.size2(),-1);

    for(int i=0;i<column.size();i++)
	if (M(c,i) >= 0)
	    column[i] = M(c,i);

    return column;
}

/// How many letters (e.g. non-gap) are there in column @v?
unsigned n_letters(const vector<int>& v)
{
    unsigned n=0;
    for(int i=0;i<v.size();i++)
	if (v[i]>=0)
	    n++;

    return n;
}

/// Augments a @column of residues with what has been @emitted so far.
struct emitted_column
{
    /// How many letters have been emitted in each sequence, after this column?
    vector<int> emitted;

    /// What are the indices of letters aligned in this column?
    vector<int> column;

    /// How many sequences are there?
    unsigned size() const {return emitted.size();}

    bool character(int i) const {return column[i] >= 0;}

    emitted_column(const vector<int>& e,const vector<int>& c)
	:emitted(e),column(c)
	{ 
	    assert(emitted.size() == column.size());
	}

    emitted_column(int n)
	:emitted(n,0),column(n,-1)
	{}
};

/// How many characters must be emitted by any column that comes immediately before c ?
vector<int> emitted_before(const emitted_column& c)
{
    vector<int> emitted = c.emitted;
    for(int i=0;i<c.size();i++)
	if (c.character(i))
	    emitted[i]--;

    return emitted;
}

bool get_emitted_column(emitted_column& C,const matrix<int>& m, int c)
{
    // the "emitted" value carries over from the previous iteration.
    C.column = get_column(m,c);

    // skip empty columns
    if (not n_letters(C.column))
	return false;
  
    // Update the "emitted" value of C
    // Check that sequential columns don't jump by more than one emitted character
    for(int i=0;i<C.size();i++)
	if (C.character(i)) {
	    assert(C.column[i] == C.emitted[i]);
	    C.emitted[i] = C.column[i]+1;
	}

    return true;
}


/// Gives -1:0:1 if the relationship of @c1 and @c2 is <:=:>
int get_column_order(const vector<int>& c1, const vector<int>& c2)
{
    if (c1.size() != c2.size()) {
	throw myexception()<<"comparing 2 columns: different sizes!";
	assert(c1.size() == c2.size());
    }

    for(int i=0;i<c1.size();i++) {
	if (c1[i] < c2[i])
	    return -1;
	else if (c1[i] > c2[i])
	    return 1;
    }

    return 0;
}

/// An ordering class for column indices
struct column_order
{
    bool operator()(const vector<int>& c1, const vector<int>& c2) const
	{
	    return (get_column_order(c1,c2) == -1);
	}
};

/// An ordering class for emitted_column structures
struct emitted_column_order
{
    bool operator()(const emitted_column& c1,const emitted_column& c2) const
	{
	    int o = get_column_order(c1.emitted, c2.emitted);

	    if (o == -1)
		return true;
	    else if (o== 1)
		return false;

	    o = get_column_order(c1.column, c2.column);
    
	    return (o == -1);
	}
};

/*
 * An emitted column specifies the character index (-1 for a gap) and the total number
   of emitted characters for each sequence.

   A "bare" column includes only the character indices, so it is possible for multiple
   emitted columns to map to each "bare" column.

   This probably makes sense for columns that consist only of leaf sequences, since those
   sequences are fixed.

   For ancestral sequences, knowing only the characters that we are aligned too might be...
   fine?  However, for columns that are ONLY ancestral and have gaps for all leaf characters,
*/

typedef map<emitted_column,int, emitted_column_order> emitted_column_map;
typedef map< vector<int>, int, column_order> column_map;
typedef map< vector<int>, vector<int>, column_order> emitted_map;

struct MPD
{
    const int N;

    vector<int> L;

    alignment A0;

    int n_samples;

    Graph g;
    Vertex vertex_start;
    Vertex vertex_end;

    /// map emitted columns -> x
    emitted_column_map emitted_columns;

    /// map bare columns    -> y
    column_map columns;

    /// map x -> y
    vector<int> emitted_to_bare;

    /// how many times did we see each bare column?
    vector<int> counts;

    /// emitted before column -> {x1...xn}
    emitted_map before;

    /// emitted after column -> {x1...xn}
    emitted_map after;

    /// map x -> &(C,x).  This is only valid after calling get_best_path( ), and before add_alignment( )
    vector<emitted_column_map::iterator> ec_from_x;

    void check_edges_go_forwards_only() const;

    emitted_column_map::iterator create_new_emitted_column(const emitted_column& C);

    void add_emitted_column(const emitted_column& C);

    void add_alignment(const alignment& A);

    vector<double> get_score(int) const;

    vector<int> get_best_path(const vector<double>& score);
  
    alignment get_best_alignment(int);
  
    int n_vertices() const {return emitted_to_bare.size();}

    vector<double> get_column_probabilities(const alignment&) const;

    MPD(const alignment&);
};


MPD::MPD(const alignment& A)
    :N( A.n_sequences() ), L( N ), A0(A), n_samples( 0 )
{
    //------------ Determine sequence lengths ----------//
    for(int i=0;i<L.size();i++)
	L[i] = A.seqlength(i);
    
    vertex_start = add_vertex(g); // add the start node
    emitted_to_bare.push_back(-1); // the start node doesn't correspond to a column
    int x_start = get(vertex_index, g, vertex_start);

    vector<int> nothing_emitted(N,0);
    // after x_start, nothing has been emitted
    after[nothing_emitted].push_back(x_start);

    vertex_end = add_vertex(g); // add the end node
    emitted_to_bare.push_back(-1); // the start node doesn't correspond to a column
    int x_end = get(vertex_index, g, vertex_end);

    vector<int> everything_emitted = L;
    // before x_end, everything has been emitted
    before[everything_emitted].push_back(x_end);
}

void MPD::check_edges_go_forwards_only() const
{
    emitted_column_order eco;

    graph_traits<Graph>::vertex_iterator vi, vi_end;
    for (tie(vi, vi_end) = ::vertices(g); vi != vi_end; ++vi) 
    {
	int index1 = get(vertex_index,g,*vi);
	graph_traits<Graph>::out_edge_iterator ei, eend;
	for(tie(ei,eend) = out_edges(*vi,g); ei != eend; ++ei)
	{ 
	    int index2 = get(vertex_index,g,target(*ei,g));
      
	    if (index1 >= ec_from_x.size() or index1 < 0)
		throw myexception()<<"Out of bounds...";
	    if (index2 >= ec_from_x.size() or index2 < 0)
		throw myexception()<<"Out of bounds...";
	    emitted_column_map::iterator ec1 = ec_from_x[index1];
	    emitted_column_map::iterator ec2 = ec_from_x[index2];
      
	    // skip the Start and End states
	    if (index1 == 0 or index2 == 1) continue;
      
	    // check that the columns all have the same size
	    if (ec1->first.size() != ec2->first.size() and log_verbose >=2) {
		cerr<<"alignment-max: index1 = "<<index1<<endl;
		cerr<<"alignment-max: index2 = "<<index2<<endl;
		cerr<<"alignment-max: ec1->first.size() = "<<ec1->first.size()<<endl;
		cerr<<"alignment-max: ec2->first.size() = "<<ec2->first.size()<<endl;
	    }
      
	    // check that if ec1 -> ec2, then also ec1 < ec2
	    if (not eco(ec1->first,ec2->first) and log_verbose >=2)
	    {
		cerr<<"alignment-max: ";
		for(int i=0;i<ec1->first.size();i++)
		    cerr<<ec1->first.emitted[i]<<" ";
		cerr<<endl;
		cerr<<"alignment-max: ";
		for(int i=0;i<ec1->first.size();i++)
		    cerr<<ec1->first.column[i]<<" ";
		cerr<<endl;
		cerr<<endl;
		cerr<<"alignment-max: ";
		for(int i=0;i<ec2->first.size();i++)
		    cerr<<ec2->first.emitted[i]<<" ";
		cerr<<endl;
		cerr<<"alignment-max: ";
		for(int i=0;i<ec2->first.size();i++)
		    cerr<<ec2->first.column[i]<<" ";
		cerr<<endl;
	    } 
      
	}
    }
}

emitted_column_map::iterator 
MPD::create_new_emitted_column(const emitted_column& C)
{
    Vertex v = add_vertex(g);
    int vi = get(vertex_index,g,v);
    assert(vi == emitted_columns.size()+2);

    // Add the mapping from C -> vi to the emitted_column map
    emitted_columns.insert(emitted_column_map::value_type(C,vi));
    emitted_column_map::iterator x_record = emitted_columns.find(C);
    assert(x_record != emitted_columns.end());

    // if this bare column has not been seen before
    column_map::iterator y_record = columns.find(C.column);
    if (y_record == columns.end()) 
    {
	// Add the mapping from C.column -> columns.size() to the bare column map
	columns.insert(column_map::value_type(C.column, columns.size()));
	y_record = columns.find(C.column);
	assert(y_record != columns.end());
    
	// This new bare column has no counts
	counts.push_back(0);
	assert(counts.size() == columns.size());
    }
  
    // map the emitted_column index (emitted.size()) to the bare column index (y_record->second)
    emitted_to_bare.push_back(y_record->second);
    assert(emitted_to_bare.size()-1 == vi);
  
    assert(emitted_columns.size()+2 == emitted_to_bare.size());

    return x_record;
}

void
MPD::add_emitted_column(const emitted_column& C)
{
    int x_current = -1;

    // Look up the column, creating a new index if necessary
    emitted_column_map::iterator x_record = emitted_columns.find(C);

    // if this emitted column has not been seen before
    if (x_record == emitted_columns.end()) 
    {
	x_record = create_new_emitted_column(C);
	int vi = x_record->second;
    
	// If this emitted column is new, then add edges to it.
	vector<int> e_before = emitted_before(C);
	vector<int> e_after = C.emitted;
	before[e_before].push_back(vi);
	after[e_after].push_back(vi);
    
	// Find the vertex index for the current column
	x_current = x_record->second;
    
	// Add edges TO this vertex
	const vector<int>& prev_vertex_indices = after[e_before];
	assert(prev_vertex_indices.size());
	for(const auto& x_prev: prev_vertex_indices)
	{
	    Vertex v1 = vertex(x_prev, g);
	    Vertex v2 = vertex(x_current, g);
	    ::add_edge(v1, v2, g);
	}
    
	// Add edges FROM this vertex
	const vector<int>& next_vertex_indices = before[e_after];
	//  assert(next_vertex_indices.size());
	for(const auto& x_next: next_vertex_indices)
	{
	    Vertex v1 = vertex(x_current, g);
	    Vertex v2 = vertex(x_next, g);
	    ::add_edge(v1, v2, g);
	}
    }
  
    x_current = x_record->second;
  
    // Increment column count
    ++counts[emitted_to_bare[x_current]];
}

void MPD::add_alignment(const alignment& A)
{
    check_same_sequence_lengths(L, A);

    emitted_column C(N);

    matrix<int> m = M(A);

    for(int c=0; c<m.size1(); c++)
    {
	// the "emitted" value carries over from the previous iteration.
	if (not get_emitted_column(C, m, c)) continue;
    
	add_emitted_column(C);
    }

    n_samples++;
}

vector<double> MPD::get_score(int type) const
{
    vector<double> score(counts.size());

    for(const auto& c: columns)
    {
	int n = n_letters(c.first);
	assert(n > 0);
    
	int i = c.second;
    
	score[i] = double(counts[i])/n_samples;
	if (type == 1)
	    score[i] *= n;
	else if (type == 2)
	    score[i] = log(score[i]);
        else if (type == 3)
            score[i] *= n*(n-1)/2;
    }

    return score;
}

vector<int>
MPD::get_best_path(const vector<double>& score)
{
    //----------------- construct a map from index -> &(EC,index) ---------------//
    ec_from_x = vector<emitted_column_map::iterator>(n_vertices(), emitted_columns.end());
    for(auto ec = emitted_columns.begin(); ec != emitted_columns.end(); ec++)
	ec_from_x[ec->second] = ec;
  
    if (log_verbose) {
	cerr<<"\nalignment-max: checking edges...\n";
	check_edges_go_forwards_only();
	cerr<<"alignment-max: done."<<endl;
    }

    //----------------- Forward Sums -------------------//
    vector<Vertex> sorted_vertices;
    topological_sort(g, std::back_inserter(sorted_vertices));
    std::reverse(sorted_vertices.begin(), sorted_vertices.end());

    vector<int> sorted_indices(sorted_vertices.size());
    for(int i=0;i<sorted_indices.size();i++)
	sorted_indices[i] = get(vertex_index,g,sorted_vertices[i]);
    assert(sorted_indices[0] == 0);
    assert(sorted_indices.back() == 1);

    vector<double> forward(n_vertices(), -1);
    vector<int> visited(n_vertices(), 0);
    vector<int> from(n_vertices(), -1);
    forward[0] = 0;
    visited[0] = 1;

    for(int i=1;i<n_vertices();i++)
    {
	int v2i = sorted_indices[i];
	assert(not visited[v2i]);

	int v2 = vertex(v2i, g);

	double best = 0;
	int argmax = -1;
	graph_traits<Graph>::in_edge_iterator e, end;
	for(tie(e,end) = in_edges(v2,g); e != end; ++e)
	{ 
	    Vertex v1 = source(*e,g);
	    int v1i = get(vertex_index,g,v1);
	    assert(visited[v1i]);

	    if (argmax == -1) {
		best = forward[v1i];
		argmax = v1i;
	    }
	    else {
		if (forward[v1i] > best) {
		    argmax = v1i;
		    best = forward[v1i];
		}
	    }
	}
	assert(argmax != -1);
	from[v2i] = argmax;
      
	forward[v2i] = best;
	if (emitted_to_bare[v2i] != -1)
	    forward[v2i] += score[emitted_to_bare[v2i]];
      
	visited[v2i] = 1;
    }

    assert(visited[1]);

    if (log_verbose) cerr<<"alignment-max: Best score is: "<<forward[1]<<endl;

    //----------------- Backward Path Selection -------------------//
    vector<int> path(1,1); // start with just the end state

    while(path.back() != 0) {
	int S = from[path.back()];
	path.push_back(S);
    }
    std::reverse(path.begin(),path.end());

    return path;
}

alignment MPD::get_best_alignment(int type)
{
    vector<double> score = get_score(type);
  
    //----------------- Forward Sums -------------------//
  
    vector<int> path = get_best_path(score);
  
    //---------------- Create alignment matrix -------------------//
    matrix<int> M(path.size()-2, L.size());
  
    for(int i=0;i<M.size1();i++) {
	int S = path[i+1];
	emitted_column_map::iterator ec = ec_from_x[S];
	for(int j=0;j<L.size();j++)
	    M(i,j) = (ec->first).column[j];
    }
  
    alignment amax = get_alignment(M, A0);
    return amax;
}

vector<double> MPD::get_column_probabilities(const alignment& A) const
{
    vector<double> column_pr(A.length(), 0);
      
    matrix<int> m = ::M(A);

    for(int c=0; c<A.length(); c++)
    {
	vector<int> column = get_column(m, c);
	column_map::const_iterator y_record = columns.find(column);

	if (y_record != columns.end())
	{
	    int y_index = y_record->second;
	    int count = counts[y_index];
	    column_pr[c] = double(count)/n_samples;
	}
    }

    return column_pr;
}

int main(int argc,char* argv[]) 
{ 
    try {
	//---------- Parse command line  -------//
	variables_map args = parse_cmd_line(argc,argv);

	int type = -1;
	string analysis = args["analysis"].as<string>();
	if (analysis == "sum")
	    type = 0;
	else if (analysis == "wsum")
	    type = 1;
	else if (analysis == "multiply")
	    type = 2;
        else if (analysis == "wsum2")
            type = 3;
	else
	    throw myexception()<<"I don't recognize analysis type '"<<analysis<<"'.";

	//------------ Load alignment and tree ----------//
	vector<alignment> alignments;

	do_setup(args,alignments);
	for(int i=0;i<alignments.size();i++)
	    alignments[i] = chop_internal(alignments[i]);

	if (not alignments.size())
	    throw myexception()<<"Didn't read any alignments!";      

	//--------- Construct alignment indexes ---------//
	MPD mpd( alignments[0] );

	for(int i=0;i<alignments.size();i++)
	    mpd.add_alignment( alignments[i] );

	alignment amax = mpd.get_best_alignment( type );
	amax = get_ordered_alignment(amax);

	//------------------ Write best alignment -------------------//
	string out = args["out"].as<string>();

	if (out == "-")
	    cout<<amax<<endl;
	else {
	    ofstream outfile(out.c_str());
	    if (not outfile)
		throw myexception()<<"Can't open '"<<out<<"' to write result!";

	    outfile<<amax<<endl;
	    outfile.close();
	}

	//------------------ Write column probabilites -------------------//
	if (args.count("out-probabilities")) 
	{
	    vector<double> column_probabilities = mpd.get_column_probabilities( amax );

	    string outp = args["out-probabilities"].as<string>();
	    ofstream outfile(outp.c_str());

	    if (not outfile)
		throw myexception()<<"Can't open '"<<out<<"' to write column probabilities!";

	    for(int c=0; c < column_probabilities.size(); c++)
		outfile<<column_probabilities[c]<<endl;

	    outfile.close();
	}
    }
    catch (std::exception& e) {
	std::cerr<<"alignment-max: Error! "<<e.what()<<endl;
	exit(1);
    }
    return 0;
}
