#include <stdio.h>
#include <string.h>
#include <string>
#include <vector>
#include <map>
#include "LabelledCEdge.h"

#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; 
}

FILE* open_file( char* filename )
{
  FILE* file = fopen( filename, "r" );
  
  if ( file == NULL )
  {
    fprintf( stderr, "Problem opening %s!\n", filename );
    exit( EXIT_FAILURE );
  }
  
  return file;
}

static char complement(char b)
{
  switch(b)
    {
    case 'A': return 'T';
    case 'T': return 'A';
    case 'G': return 'C';
    case 'C': return 'G';
      
    case 'a': return 't';
    case 't': return 'a';
    case 'g': return 'c';
    case 'c': return 'g';
    
    case 'N': return 'N';
    case '*': return '*';  
    }
  return '?';
}

string reverse_complement(string seq)
{
  string s(seq.begin(),seq.end());
  string::iterator pos;
  for (pos = s.begin(); pos != s.end(); ++pos) {
    // cout << *pos;
  }
  // cout << endl;
  
  reverse(s.begin(), s.end());
  
  for(pos=s.begin();pos!=s.end();++pos)
    *pos=complement(*pos);
  
  return s;
}

void readNodeFile(char* nodes_fname, vector<NodeSeq>& nodesF, vector<NodeSeq>& nodesR, int k_val )
{
  FILE* node_file = open_file(nodes_fname);
  char* buffer = new char[100 * MAX];
  char* seq;
  
  nodesF.reserve(count_nb_lines(node_file));
  nodesR.reserve(count_nb_lines(node_file));
  
  while ( fgets(buffer, 100 * MAX, node_file) != NULL )
  {
    char* p;

    if (strlen(buffer) == 100 * MAX)
    {  
      p = strtok(buffer, "\t\n");
      fprintf(stdout, "ERROR: node %s with sequence larger than %d!", p, 100 * MAX);
      exit(0);
    }
      
    // Node label, should be contiguous
    p = strtok( buffer, "\t\n" );
    int node = atoi(p);
          
    // Node seq
    p = strtok( NULL, "\t\n"  );

    // We need the prefix of size k
    seq = new char[k_val + 1];
    strncpy( seq, p, k_val); seq[k_val] = '\0';
    nodesF.push_back( NodeSeq(seq, node) );

    // And the prefix of size k in the reverse complement
    seq = new char[k_val + 1];
    string rev_comp = reverse_complement(string(p));
    strncpy( seq, rev_comp.c_str(), k_val); seq[k_val] = '\0';
    nodesR.push_back( NodeSeq(seq, node) );

    //printf("%d %s %s\n", node, p, rev_comp.c_str());
  }
  
  delete [] buffer;
  fclose(node_file);
}

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 readCounts(char* counts_fname, map<string,int>& counts)
{
  FILE *count_file = open_file(counts_fname);
  char* buffer = new char[100 * MAX];
       
  while ( fgets(buffer, 100 * MAX, count_file) != NULL )
  {
    char* p;
     
    // Sequence
    p = strtok( buffer, "\t\n " );
    string kmer = p;
   
    // Count 
    p = strtok( NULL, "\t\n "  );
    counts[kmer] = atoi(p); 
  }
  
  delete [] buffer;
    
  fclose(count_file);
}

int kcounts(string query, map<string,int> counts)
{
  if (counts.find(query) != counts.end())
    return counts[query];
  
  if (counts.find(reverse_complement(query)) != counts.end())
    return counts[reverse_complement(query)];
  //  fprintf(stderr, "k-mer not found!\n");
  return 0;
}

char ACTG[4] = {'A','C','T','G'};

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);
}

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);
}

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;
}

int main( int argc, char** argv )
{
  if ( argc < 6 )
  { 
    fprintf( stderr, "Wrong number of arguments!\n" );
    fprintf( stderr, "Usage: ./check_error_removal removed.edges nodes_file k_value k_mer_count_file cutoff\n" );
    return 0;
  }
  
  int k_value = atoi( argv[3] );
  double cutoff = atof( argv[5] );

  // Read node file
  vector<NodeSeq> nodesF, nodesR;
  readNodeFile(argv[2], nodesF, nodesR, k_value );

  // Read count file
  map<string,int> counts;
  readCounts(argv[4], counts);

  // printf("\n");
  // vector<NodeSeq>::iterator it;
  // for (it = nodesF.begin(); it != nodesF.end(); it++)
  //   printf("%d %s %d\n", it->node, it->seq, kcounts(string(it->seq), counts));

  // printf("\n");

  // for (it = nodesR.begin(); it != nodesR.end(); it++)
  //   printf("%d %s %d\n", it->node, it->seq, kcounts(string(it->seq), counts));

  vector<LabelledCEdge> allEdges;
  readEdgeFile( argv[1], allEdges );
  vector<bool> verified(allEdges.size(), false);

  for (int i = 0; i < (int)allEdges.size(); i++)
    if (!verified[i])
    {
      int u =  allEdges[i].getFirst();
      int v = allEdges[i].getSecond();
      string label = allEdges[i].label;
      
      string kmer, pref;
      if (label[1] == 'F') 
	kmer = nodesF[v].seq;
      else 
	kmer = nodesR[v].seq;
      pref = kmer.substr(0,k_value -1);
      
      int sum = 0;
      for (int j = 0; j < 4; j++)
	sum += kcounts(string(pref+ACTG[j]), counts);
      
      double ratio = (double)kcounts(kmer, counts) / (double)sum;
      //fprintf(stderr, "edge %d -> %d (%s) ratio = %lf\n", u, v, label.c_str(), ratio);
      if (ratio < cutoff)
      {
	verified[i] = true;
	verified[findEdge(allEdges, reverse(allEdges[i]))] = true;
      }

    }

  bool flag = true;
  for (int i = 0; i < (int)allEdges.size(); i++)
    if (!verified[i])
    {
      int u = allEdges[i].getFirst();
      int v = allEdges[i].getSecond();
      string label = allEdges[i].label;
      fprintf(stderr, "edge %d -> %d (%s) should not be removed\n", u, v, label.c_str());
      flag = false;
    }

  if (flag)
    fprintf(stderr, "Removed edges are correct!\n");

  return 0;
}
