#include "ContigPath.h"
#include "ContigProperties.h"
#include "Graph/Assemble.h"
#include "Graph/ContigGraph.h"
#include "Graph/ContigGraphAlgorithms.h"
#include "Graph/DirectedGraph.h"
#include "Graph/GraphAlgorithms.h"
#include "Graph/GraphIO.h"
#include "Graph/GraphUtil.h"
#include "Uncompress.h"
#include "config.h"
#include <cassert>
#include <cstdlib>
#include <fstream>
#include <getopt.h>
#include <iostream>
#include <sstream>
#include <string>
#include <utility>

using namespace std;
using boost::tie;

#define PROGRAM "abyss-layout"

static const char VERSION_MESSAGE[] = PROGRAM " (" PACKAGE_NAME ") " VERSION "\n"
                                              "Written by Shaun Jackman.\n"
                                              "\n"
                                              "Copyright 2012 Shaun Jackman\n";

static const char USAGE_MESSAGE[] =
    "Usage: " PROGRAM " [OPTION]... OVERLAP\n"
    "Layout contigs using the sequence overlap graph.\n"
    "Output sequence paths.\n"
    "\n"
    " Arguments:\n"
    "\n"
    "  OVERLAP  the sequence overlap graph\n"
    "\n"
    " Options:\n"
    "\n"
    "  -s, --min-length=N    minimum sequence length [0]\n"
    "  -m, --min-overlap=N   minimum overlap [0]\n"
    "  -k, --kmer=N          length of a k-mer\n"
    "  -o, --out=FILE        write the paths to FILE\n"
    "  -g, --graph=FILE      write the graph to FILE\n"
    "      --tred            remove transitive edges\n"
    "      --no-tred         do not remove transitive edges [default]\n"
    "      --SS              expect contigs to be oriented correctly\n"
    "      --no-SS           no assumption about contig orientation [default]\n"
    "  -v, --verbose         display verbose output\n"
    "      --help            display this help and exit\n"
    "      --version         output version information and exit\n"
    "\n"
    "Report bugs to <" PACKAGE_BUGREPORT ">.\n";

namespace opt {
unsigned k; // used by ContigProperties

/** Minimum sequence length. */
static unsigned minLength;

/** Minimum overlap. */
static unsigned minOverlap;

/** Write the paths to this file. */
static string out;

/** Write the graph to this file. */
static string graphPath;

/** Remove transitive edges. */
static int tred;

/** Run a strand-specific RNA-Seq assembly. */
static int ss;

/** Verbose output. */
int verbose; // used by PopBubbles

/** Output format */
int format = DOT;
}

static const char shortopts[] = "g:k:m:o:s:v";

enum
{
	OPT_HELP = 1,
	OPT_VERSION
};

static const struct option longopts[] = { { "graph", required_argument, NULL, 'g' },
	                                      { "kmer", required_argument, NULL, 'k' },
	                                      { "min-overlap", required_argument, NULL, 'm' },
	                                      { "out", required_argument, NULL, 'o' },
	                                      { "min-length", required_argument, NULL, 's' },
	                                      { "tred", no_argument, &opt::tred, true },
	                                      { "no-tred", no_argument, &opt::tred, false },
	                                      { "SS", no_argument, &opt::ss, 1 },
	                                      { "no-SS", no_argument, &opt::ss, 0 },
	                                      { "verbose", no_argument, NULL, 'v' },
	                                      { "help", no_argument, NULL, OPT_HELP },
	                                      { "version", no_argument, NULL, OPT_VERSION },
	                                      { NULL, 0, NULL, 0 } };

/** An overlap graph. */
typedef DirectedGraph<Length, Distance> DG;
typedef ContigGraph<DG> Graph;

/** Remove short vertices. */
static void
filterVertices(Graph& g, unsigned minLength)
{
	typedef graph_traits<Graph> GTraits;
	typedef GTraits::vertex_descriptor V;
	typedef GTraits::vertex_iterator Vit;

	if (minLength == 0)
		return;

	// Remove short sequences.
	unsigned numRemoved = 0;
	std::pair<Vit, Vit> urange = vertices(g);
	for (Vit uit = urange.first; uit != urange.second; ++uit) {
		V u = *uit;
		if (g[u].length < minLength)
			clear_vertex(u, g);
		if (out_degree(u, g) == 0 && in_degree(u, g) == 0) {
			remove_vertex(u, g);
			numRemoved++;
		}
	}

	if (opt::verbose > 0) {
		cerr << "Ignored " << numRemoved << " sequences shorter than " << minLength << " bp.\n";
		printGraphStats(cerr, g);
	}
}

/** Return true if the edge is a small overlap. */
struct IsSmallOverlap
{
	IsSmallOverlap(Graph& g)
	  : m_g(g)
	{}
	bool operator()(graph_traits<Graph>::edge_descriptor e) const
	{
		int maxDistance = -opt::minOverlap;
		return m_g[e].distance > maxDistance;
	}
	const Graph& m_g;
};

/** Remove small overlaps. */
static void
filterEdges(Graph& g, unsigned minOverlap)
{
	if (minOverlap == 0)
		return;
	unsigned numBefore = num_edges(g);
	remove_edge_if(IsSmallOverlap(g), static_cast<DG&>(g));
	unsigned numRemoved = numBefore - num_edges(g);
	if (opt::verbose > 0) {
		cerr << "Removed " << numRemoved << " small overlaps.\n";
		printGraphStats(cerr, g);
	}
}

/** Read a graph from the specified file. */
static void
readGraph(const string& path, Graph& g)
{
	if (opt::verbose > 0)
		cerr << "Reading `" << path << "'...\n";
	ifstream fin(path.c_str());
	istream& in = path == "-" ? cin : fin;
	assert_good(in, path);
	in >> g;
	assert(in.eof());
	if (opt::verbose > 0)
		printGraphStats(cerr, g);
	g_contigNames.lock();
}

/** Return the length histogram. */
static Histogram
buildLengthHistogram(const Graph& g)
{
	typedef graph_traits<Graph>::vertex_descriptor V;
	typedef graph_traits<Graph>::vertex_iterator Vit;
	Histogram h;
	Vit uit, ulast;
	for (tie(uit, ulast) = vertices(g); uit != ulast; ++++uit) {
		V u = *uit;
		if (!get(vertex_removed, g, u))
			h.insert(g[u].length);
	}
	return h;
}

/** Run abyss-layout. */
int
main(int argc, char** argv)
{
	bool die = false;
	for (int c; (c = getopt_long(argc, argv, shortopts, longopts, NULL)) != -1;) {
		istringstream arg(optarg != NULL ? optarg : "");
		switch (c) {
		case '?':
			die = true;
			break;
		case 'k':
			arg >> opt::k;
			break;
		case 'g':
			arg >> opt::graphPath;
			break;
		case 'm':
			arg >> opt::minOverlap;
			break;
		case 'o':
			arg >> opt::out;
			break;
		case 's':
			arg >> opt::minLength;
			break;
		case 'v':
			opt::verbose++;
			break;
		case OPT_HELP:
			cout << USAGE_MESSAGE;
			exit(EXIT_SUCCESS);
		case OPT_VERSION:
			cout << VERSION_MESSAGE;
			exit(EXIT_SUCCESS);
		}
		if (optarg != NULL && !arg.eof()) {
			cerr << PROGRAM ": invalid option: `-" << (char)c << optarg << "'\n";
			exit(EXIT_FAILURE);
		}
	}

	if (argc - optind < 1) {
		cerr << PROGRAM ": missing arguments\n";
		die = true;
	}

	if (die) {
		cerr << "Try `" << PROGRAM << " --help' for more information.\n";
		exit(EXIT_FAILURE);
	}

	Graph g;
	if (optind < argc) {
		for (; optind < argc; optind++)
			readGraph(argv[optind], g);
	} else
		readGraph("-", g);

	// Remove short sequences.
	filterVertices(g, opt::minLength);

	// Remove small overlaps.
	filterEdges(g, opt::minOverlap);

	// Remove transitive edges.
	if (opt::tred) {
		unsigned numTransitive = remove_transitive_edges(g);
		if (opt::verbose > 0) {
			cerr << "Removed " << numTransitive << " transitive edges.\n";
			printGraphStats(cerr, g);
		}
	}

	/** A container of contig paths. */
	typedef vector<ContigPath> ContigPaths;

	// Assemble the paths.
	ContigPaths paths;
	assembleDFS(g, back_inserter(paths), opt::ss);
	sort(paths.begin(), paths.end());
	if (opt::verbose > 0) {
		unsigned n = 0;
		for (ContigPaths::const_iterator it = paths.begin(); it != paths.end(); ++it)
			n += it->size();
		cerr << "Assembled " << n << " sequences in " << paths.size() << " contigs.\n";
		printGraphStats(cerr, g);
	}

	// Output the paths.
	ofstream fout(opt::out.c_str());
	ostream& out = opt::out.empty() || opt::out == "-" ? cout : fout;
	assert_good(out, opt::out);
	g_contigNames.unlock();
	for (vector<ContigPath>::const_iterator it = paths.begin(); it != paths.end(); ++it)
		out << createContigName() << '\t' << *it << '\n';
	assert_good(out, opt::out);

	// Create the new vertices.
	for (vector<ContigPath>::const_iterator it = paths.begin(); it != paths.end(); ++it) {
		const ContigPath& path = *it;
		merge(g, path.begin(), path.end());
		remove_vertex_if(
		    g, path.begin(), path.end(), [](const ContigNode& c) { return !c.ambiguous(); });
	}
	if (opt::verbose > 0)
		printGraphStats(cerr, g);

	// Output the graph.
	if (!opt::graphPath.empty()) {
		ofstream out(opt::graphPath.c_str());
		assert_good(out, opt::graphPath);
		write_dot(out, g);
		assert_good(out, opt::graphPath);
	}

	// Print assembly contiguity statistics.
	if (opt::verbose > 0) {
		Histogram h = buildLengthHistogram(g);
		const unsigned STATS_MIN_LENGTH = 200; // bp
		printContiguityStats(cerr, h, STATS_MIN_LENGTH) << '\n';
	}

	return 0;
}