#include <iostream>
#include <string>
#include <stdio.h>
#include <math.h>
#include <vector>
#include <map>
#include <fstream>
#include <algorithm>
using namespace std;

struct block {
  string l[4], bcc;
  double rank;
};

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

bool comp_lower(block a, block b) 
{
  return a.l[3] < b.l[3];
}

bool comp_bcc(block a, block b) 
{
  return a.bcc < b.bcc;
}

block reverse_complement(block &b)
{
  block res;
  
  res.l[1] = reverse_complement(b.l[1]);
  res.l[3] = reverse_complement(b.l[3]);
  
  res.l[0] = b.l[0];
  res.l[2] = b.l[2];

  return res;
}

string extract_bcc(block &b)
{
  int sep = b.l[0].find_first_of("|");
   
  return b.l[0].substr(0, sep);
}

double extract_rank(block &b)
{
  int sep = b.l[0].find_last_of("|");
  
  string rank = b.l[0].substr(sep+1);
  sep = rank.find_first_of("_");
 
  float res = atof(rank.substr(sep+1).c_str());

  /* trick to test is res = nan */
  if (res != res)
    res = 0;
  
  return res;
}

void print_block(FILE *stream, block &b)
{
  for (int i = 0; i < 4; i++)
    fprintf(stream, "%s\n", b.l[i].c_str());
}

map<int, vector<int> > adj;  
vector<block> all_blocks;

int max_rank, max_ind;
void dfs(int node, map<int, bool> &visited)
{
  visited[node] = true;
  if (all_blocks[node].rank >= max_rank)
  {
    max_rank = all_blocks[node].rank;
    max_ind = node;
  }

  for (int i = 0; i < (int)adj[node].size(); i++)
  {
    int v = adj[node][i];
    if (!visited[v])
      dfs(v, visited);
  }
}

int nprinted = 0;
void print_clusters(FILE *stream)
{
  map<int, bool> visited;
  
  map<int, vector<int> >::iterator it;
  for (it = adj.begin(); it != adj.end(); it++)
    visited[it->first] = false;

  for (it = adj.begin(); it != adj.end(); it++)
    if (!visited[it->first])
    {
      max_rank = max_ind = -1;
      dfs(it->first, visited);
      if (max_ind != -1)
      {
	nprinted++;
	print_block(stream, all_blocks[it->first]);
      }
    }
}

int comp_base(const void *a, const void *b)
{
  return *(char *)a != *(char *)b;
}

int edit_distance(const void *s1, size_t l1, const void *s2, size_t l2, size_t nmemb, int (*comp)(const void*, const void*)) 
{
  unsigned int i;
  unsigned int j;
  size_t len = (l1 + 1) * (l2 + 1);
  char *p1;
  char *p2;
  unsigned int d1;
  unsigned int d2;
  unsigned int d3;
  unsigned int *d;
  unsigned int *dp;
  unsigned int res;
 
  if (l1 == 0)  
    return l2;
  else if (l2 == 0) 
    return l1;
   
  d = (unsigned int*)malloc(len * sizeof(unsigned int));
 
  *d = 0;
  for(i = 1, dp = d + l2 + 1; i < l1 + 1; ++i, dp += l2 + 1)
  {
    *dp = (unsigned) i;
  }
  for(j = 1, dp = d + 1; j < l2 + 1; ++j, ++dp)
  {
    *dp = (unsigned) j;
  }

  for(i = 1, p1 = (char*) s1, dp = d + l2 + 2; i < l1 + 1; ++i, p1 += nmemb, ++dp)
  {
    for(j = 1, p2 = (char*) s2; j < l2 + 1; ++j, p2 += nmemb, ++dp) 
    {
      if(!comp(p1, p2))
      {
        *dp = *(dp - l2 - 2);// same base
      }
      else 
      {
        d1 = *(dp - 1) + 1;// insertion
        d2 = *(dp - l2 - 1) + 1;// deletion
        d3 = *(dp - l2 - 2) + 1;// substitution
        *dp = std::min(d1, std::min(d2, d3));
      }
    }
  }
  res = *(dp - 2);
 
  dp = NULL;
  free(d);
  return res;
}

int min_dist = 1;
bool has_edge(int i, int j) 
{
  // lower path should be equal
  if (all_blocks[i].l[3] != all_blocks[j].l[3])
    return false;
  
  if (abs((int)all_blocks[i].l[1].size() - (int)all_blocks[j].l[1].size()) > min_dist)
    return false;

  int ed = edit_distance(all_blocks[i].l[1].c_str(), all_blocks[i].l[1].size(), all_blocks[j].l[1].c_str(), all_blocks[j].l[1].size(), sizeof(char), comp_base);

  return (ed <= min_dist);
}


int main(int argc, char **argv)
{
  ifstream bubbles_file;
  bubbles_file.open(argv[1]);
  
  if (argc < 3)
  {
    fprintf(stderr, "./cluster input.fasta output.fasta [min_dist]\n");
    exit(0);
  }
  if (argc == 4)
     min_dist = atoi(argv[3]);
   
  string line;
  int nlines = 0;
  block cur;

  while (bubbles_file.good())
  {
    getline(bubbles_file, line);
    nlines++;

    cur.l[(nlines-1) % 4] = line;
    if (nlines % 4 == 0)
    { 
      // convert the sequences to a canonical form, lexicografically smaller (forward and reverse) 
      if (comp_lower(reverse_complement(cur), cur))
      	cur = reverse_complement(cur);
      
      // extract the bcc number
      cur.bcc = extract_bcc(cur);
      cur.rank = extract_rank(cur);

      all_blocks.push_back(cur);
    }
  }
 
  FILE *output = fopen(argv[2], "w");
  if (output == NULL)
  {
    fprintf(stderr, "Problem opening %s\n", argv[2]);
    exit(0);
  }

  // sort by BCC, all bubbles from the same BCC are adjacent
  sort(all_blocks.begin(), all_blocks.end(), comp_bcc);
 
  for (int i = 0; i < (int)all_blocks.size(); i++)
  {
    map<int, vector<int> >::iterator it;
    adj[i] = vector<int>();
    for (it = adj.begin(); it != adj.end(); it++)
      if (has_edge(i, it->first))
      {
	adj[i].push_back(it->first);
	adj[it->first].push_back(i);
      }
  
    // finished the bubbles for a given BCC
    if ((i+1) == (int)all_blocks.size() || all_blocks[i].bcc != all_blocks[i+1].bcc)
    {
      print_clusters(output);
      map<int, vector<int> >().swap(adj);
    }
  }
  fclose(output);

  fprintf(stdout, "%d out of %d bubbles were outputed\n", nprinted, (int)all_blocks.size());

  //fprintf(stderr, "nprinted = %d, total = %d, max_nodes = %d, max_edges = %d\n", nprinted, (int)all_blocks.size(), max_sz, max_edges);

  bubbles_file.close();
  return 0;
}