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

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

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

int nucToNumber(const char& nuc) {
  switch (nuc) {
  case 'A': return 0;
  case 'C': return 1;
  case 'G': return 2;
  case 'T': return 3;
  default:  return 4;
  }
}

static double MinusPlogP(double value, double base) {
  return -1 * value * log(value)/log(base);
}

bool valid_nuc(char nuc)
{
  return nuc != 'N' && nuc != '*' && nuc != '?';
}

double entropy3(const string& window_nucs) 
{
  size_t window_length = window_nucs.length();
  vector<int> kmer_counts(444, 0);
  float subseqs = 0;
  for (size_t i = 0; i < window_length - 3; ++i) {
    char nuc1 = window_nucs.at(i);
    char nuc2 = window_nucs.at(i+1);
    char nuc3 = window_nucs.at(i+2);
    if (valid_nuc(nuc1) && valid_nuc(nuc2) && valid_nuc(nuc3)) {
      kmer_counts[nucToNumber(nuc1) + nucToNumber(nuc2)*10 + nucToNumber(nuc3)*10] += 1;
      subseqs+= 1;
    }
  }
  float entropy = 0, base = 64;
  
  if (window_length < 66)
    base = (float)window_length;

  for (size_t i = 0; i < kmer_counts.size(); ++i) {
    int count = kmer_counts.at(i);
    if (count !=0) {
      float p = static_cast<float>(count)/subseqs;
      entropy += MinusPlogP(p, base);
    }
  }
  return entropy;
}

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

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


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

  map<string, int> bcc_count;
  vector<block> all_blocks;

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

  int nprinted = 0;
  while (bubbles_file.good())
  {
    getline(bubbles_file, line);
    nlines++;

    cur.l[(nlines-1) % 4] = line;
    if (nlines % 4 == 0)
    {
      cur.bcc = extract_bcc(cur);
      all_blocks.push_back(cur);
      
      if (bcc_count.find(cur.bcc) == bcc_count.end())
	bcc_count[cur.bcc] = 0;
      bcc_count[cur.bcc] = bcc_count[cur.bcc] + 1;
    }
  }

  for (int i = 0; i < (int)all_blocks.size(); i++)
    if (entropy3(all_blocks[i].l[1]) > threshold || bcc_count[all_blocks[i].bcc] <= 10)
    {
      nprinted++;
      print_block(output, all_blocks[i]);
    } 
    else
    {
      print_block(filtered, all_blocks[i]);
    }
 
  fclose(output);
  fclose(filtered);
  bubbles_file.close();

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

  return 0;
}

