#include <cstdio>
#include <cstring>
#include <string>
#include <cstdlib>
#include <vector>
#include "Utils.h"
#include "LabelledCEdge.h"
#include <fstream>

#define MAX 1024

using namespace std;

struct NodeSeq {
  char *seq;
  int node;
  NodeSeq(char *initSeq, int initNode) : seq(initSeq), node(initNode) { }
  bool operator== (const NodeSeq& rhs) const { return (strcmp(seq, rhs.seq) == 0); }
  bool operator<  (const NodeSeq& rhs) const { return (strcmp(seq, rhs.seq) <  0); }
};

int count_nb_lines( FILE* file )
{
  int ch, number_of_lines = 0;

  while (EOF != (ch=getc(file)))
    if ('\n' == ch)
      number_of_lines++;
  
  // Set the cursor back to the begining of the file.
  rewind(file);
  
  // Don't care if the last line has a '\n' or not. We over-estimate it.
  return number_of_lines + 1; 
}


void readEdgeFile( char* nodes_fname, vector<LabelledCEdge>& edges )
{
  FILE *edge_file = open_file(nodes_fname);

  char* buffer = new char[100 * MAX];
  char* u = new char[MAX];
  char* v = new char[MAX];
  char* label = new char[MAX];
  
  edges.reserve(count_nb_lines(edge_file));

  while ( fgets(buffer, 100 * MAX, edge_file) != NULL )
  {
    char* p;

    // outgoing node
    p = strtok( buffer, "\t\n" );
    strcpy( u, p );

    // incoming node
    p = strtok( NULL, "\t\n" );
    strcpy( v, p );

    // edge label
    p = strtok( NULL, "\t\n" );
    strcpy(label, p);
  
    edges.push_back( LabelledCEdge( atoi(u), atoi(v), label ) );
  }
  
  sort( edges.begin(), edges.end() );
  
  delete [] buffer;
  delete [] u;
  delete [] v;
  delete [] label;
  
  fclose(edge_file);
}

// rev_seq needs to be global! let's make sure we allocate only once
char *rev_seq = NULL;
char *revComp(char *seq)
{
  int size = strlen(seq);
  if (rev_seq == NULL)
    rev_seq = new char[size + 1];
   
  for (int i = size - 1; i >= 0; i--)
    rev_seq[(size - 1) - i] = complement(seq[i]); 
  rev_seq[size] = '\0';

  return rev_seq;
}

int findNode(char *query, vector<NodeSeq>& nodes)
{
  vector<NodeSeq>::iterator low, low_r;
  char *query_r = revComp(query);

  // The value of "x" in NodeSeq(query, x) doesn't matter 
  low   = lower_bound(nodes.begin(), nodes.end(), NodeSeq(query, 0));
  low_r = lower_bound(nodes.begin(), nodes.end(), NodeSeq(query_r, 0));

  if (low != nodes.end() && *low == NodeSeq(query, 0))
    return low->node;
  if (low_r != nodes.end() && *low_r == NodeSeq(query_r, 0)) 
    return low_r->node;

  return -1;
}

void readAbundanceFile (char* abundanceFileName, std::vector<double> &abundances) {
  auto abundanceFile = std::ifstream(abundanceFileName);

  if (!abundanceFile) {
    std::fprintf(stderr, "Error opening %s\n", abundanceFileName);
    exit(1);
  }

  double abundance = 0;

  while (abundanceFile >> abundance) {
    abundances.push_back(abundance);
  }
}

int findEdge(vector<LabelledCEdge>& allEdges, LabelledCEdge e)
{
  vector<LabelledCEdge>::iterator low; 

  low = lower_bound(allEdges.begin(), allEdges.end(), e);
  if (low != allEdges.end())
    return low - allEdges.begin();
  
  fprintf(stderr, "inconsistent graph!");
  return -1;
}

// Move this to CEdge
LabelledCEdge reverse(LabelledCEdge e)
{
  char *label = new char[3];
  
  label[0] = e.label[1] == 'F' ? 'R' : 'F';
  label[1] = e.label[0] == 'F' ? 'R' : 'F';
  label[2] = '\0';
  
  return LabelledCEdge(e.getSecond(), e.getFirst(), label);
}


void errorRemoval(vector<LabelledCEdge>& allEdges, int nbNodes, double cutoff, vector<bool>& removed, const vector<double> &abundances)
{
  int offset = 0;
  for ( int src = 0 ; src < nbNodes; src++ ) {
    // Count the number of edges (they are ordered in all_edges) of node src
    int size = 0;
    while (offset + size < (int) allEdges.size() && allEdges[offset + size].getFirst() == src)
      size++;

    // Compute the sum of coverage for each direction
    double sum_F = 0, sum_R = 0;
    for (int k = 0; k < size; k++) {
      if (!removed[offset + k]) {
        int target = allEdges[offset + k].getSecond();
        if (allEdges[offset + k].label[0] == 'F')
          sum_F += abundances[target];

        if (allEdges[offset + k].label[0] == 'R')
          sum_R += abundances[target];
      }
    }

    // Remove the edges with relative coverage below cutoff
    for (int k = 0; k < size; k++) {
      if (!removed[offset + k]) {
        int target = allEdges[offset + k].getSecond();
        double count = abundances[target];
        double ratio;

        ratio = count / sum_F;
        if (allEdges[offset + k].label[0] == 'F' && ratio < cutoff) {
          removed[offset + k] = true;
          removed[findEdge(allEdges, reverse(allEdges[offset + k]))] = true;
        }

        ratio = count / sum_R;
        if (allEdges[offset + k].label[0] == 'R' && ratio < cutoff) {
          removed[offset + k] = true;
          removed[findEdge(allEdges, reverse(allEdges[offset + k]))] = true;
        }
      }
    }

    offset += size;
  }
}

int main( int argc, char** argv )
{
  string base_name = "graph";

  if ( argc < 4 )
  { 
    fprintf( stderr, "Wrong number of arguments!\n" );
    fprintf( stderr, "Usage: ./error_removal edge_file unitig_abundance_file cutoff [base_name]\n" );
    return 0;
  }
  
  if ( argc == 5 )
  {
    base_name = argv[4];
  }
  
  double cutoff = atof( argv[3] );

  // Read edge file
  vector<LabelledCEdge> allEdges;
  readEdgeFile(argv[1], allEdges );

  // Read abundance file
  vector<double> abundances;
  abundances.reserve(allEdges.size()); //over-estimating, but it is ok
  readAbundanceFile(argv[2], abundances); //read the abundances

  vector<bool> removed(allEdges.size(), false);
  errorRemoval(allEdges, (int)abundances.size(), cutoff, removed, abundances);
  
  int nb_removed = 0;
  FILE *output = fopen((base_name + ".edges").c_str(), "w");
  FILE *removed_out = fopen((base_name + "_removed.edges").c_str(), "w");

  for (int i = 0 ; i < (int)allEdges.size(); i++)
    if (!removed[i])
      fprintf(output, "%d\t%d\t%s\n", allEdges[i].getFirst(), allEdges[i].getSecond(), allEdges[i].label);
    else
    {
      fprintf(removed_out, "%d\t%d\t%s\n", allEdges[i].getFirst(), allEdges[i].getSecond(), allEdges[i].label);
      nb_removed++;
    }

  fclose(output);
  fclose(removed_out);
  
  fprintf(stdout, "%d out of %d edges removed\n", nb_removed, (int)allEdges.size());
 
  return 0;
}