#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "Sequence.h"
#include "SeqMap.h"

typedef struct {
  double dicodon[64][64];
  double codon[64];
  double total_dicodons;
} coding_info;


void usage (void);
void count (SeqMap map, Sequence seq, coding_info * coding,
	    coding_info * noncoding, int dir);
void create_table (coding_info * coding, coding_info * noncoding,
		   double fraction);
void bell_correction (double freq[64][64], coding_info *info);
void add_pseudocounts (coding_info * info);
void mackay_correction (double freq[64][64], coding_info * info); 
int mackay_N(coding_info *info, int j, int f);
int mackay_fmax(coding_info *info, int j);
double mackay_K(coding_info *info, double alpha);
double mackay_G(coding_info *info, int j, int fmax);
double mackay_H(coding_info *info, int j, int fmax);

int
main (int argc, char *argv[])
{
  FILE *seqfile = NULL, *cdsfile = NULL;
  int i, start, end;
  double fraction = 0;
  Sequence seq, antiSeq;
  SeqMap map;
  char *arg, locus[255];
  coding_info *coding, *noncoding;

  if (argc < 2)
    usage ();

  for (i = 1; i < argc; i++) {
    arg = argv[i];
    if (arg[0] == '-') {
      if (strstr (arg, "-fraction-coding=") != NULL) {
	arg = strtok (arg, "=");
	arg = strtok (NULL, "=");
	fraction = atof (arg) / 6;
      }
      else {
	fprintf (stderr, "%s: option %s invalid\n", argv[0], arg);
	exit (1);
      }
    }
    else if (cdsfile == NULL) {
      cdsfile = fopen (arg, "r");
      if (cdsfile == NULL) {
	fprintf (stderr, "%s: %s not found.\n", argv[0], arg);
	exit (1);
      }
    }
    else if (seqfile == NULL) {
      seqfile = fopen (arg, "r");
      if (seqfile == NULL) {
	fprintf (stderr, "%s: %s not found.\n", argv[0], arg);
	exit (1);
      }
    }
  }
  if ((seqfile == NULL) || (cdsfile == NULL))
    usage ();

  coding = SeqMapCalloc (1, sizeof (coding_info));
  noncoding = SeqMapCalloc (1, sizeof (coding_info));
  while (!feof (seqfile)) {
    seq = loadNextContig (seqfile);
    antiSeq = complement (seq);
    map = newSeqMap (seq->length);
    rewind (cdsfile);
    while (!feof (cdsfile)) {
      fscanf (cdsfile, "%s %d %d%*[^\n]\n", locus, &start, &end);
      if (strcmp (seq->header, locus) == 0) {
	fillSeqMap (map, start, end, 0);
      }
    }
    count (map, seq, coding, noncoding, 1);
    count (map, antiSeq, coding, noncoding, 2);
    freeSequence (seq);
    freeSequence (antiSeq);
    freeSeqMap (map);
  }
  create_table (coding, noncoding, fraction);
  fclose (seqfile);
  fclose (cdsfile);
  free (coding);
  free (noncoding);
  exit (0);
}

void
count (SeqMap map, Sequence seq, coding_info * coding,
       coding_info * noncoding, int dir)
{
  int i, cn1, cn2, mapVal;

  for (i = 0; i < seq->length - 5; i++) {
    cn1 = codonNumber (&seq->seq[i]);
    cn2 = codonNumber (&seq->seq[i + 3]);
    if ((cn1 < 0) || (cn2 < 0)) {
      continue;
    }
    mapVal = mapPos (map, i, dir);
    if ((mapVal == dir) || (mapVal == 3)) {
      coding->dicodon[cn1][cn2]++;
      coding->codon[cn1]++;
      coding->total_dicodons++;
    }
    else {
      noncoding->dicodon[cn1][cn2]++;
      noncoding->codon[cn1]++;
      noncoding->total_dicodons++;
    }
  }
}

/*
 * Input:
 *
 *  fraction = prior probability of a hexamer coding
 *  coding->dicodon[i][j] = number of hexamers ij explicitly called coding
 *  noncoding->dicodon[i][j] = number of hexamers ij not classified (noncoding)
 *  coding->codon[i] = num. of expected coding hexamers that start with i
 *  noncoding->codon[i] = num. of expected noncoding hexamers that start with i
 *  coding->total_dicodons = total num. of dicodons explicitly called coding
 *  noncoding->total_dicodons = total num. of dicodons not classified 
 *
 * Correct for unclassifed hexamers:
 *
 *  ratio = ratio of expected coding hexamers to called coding hexamers
 *  total_noncoding = total number of noncoding hexamers
 *  sum = coding->dicodon[i][j] + noncoding->dicodon[i][j]
 *  ratio = fraction * total_dicodons / total_called
 *      
 *  coding->dicodon[i][j] = ratio * coding->dicodon[i][j]
 *  noncoding->dicodon[i][j] = sum - noncoding->dicodon[i][j]
 *
 * Conditional probability ratios:
 *
 *      freq_cod[i][j] = probability of finding codon j given preceding i in
 *                       coding DNA
 *      freq_non[i][j] = probability of finding codon j given preceding i in 
 *                       noncoding DNA
 *
 */

void
create_table (coding_info * coding, coding_info * noncoding, double fraction)
{
  int i, j;
  char *codon1, *codon2;
  double ratio, freq;
  double freq_cod[64][64], freq_non[64][64], sum;
  double total_dicodons = coding->total_dicodons +
  noncoding->total_dicodons;

  if (fraction > 0) {
    ratio = fraction * (total_dicodons / coding->total_dicodons);
    fprintf (stderr, "Fraction coding %4.2f. Adjusted to %4.2f\n",
      6 * coding->total_dicodons / noncoding->total_dicodons, 6 * fraction);
  }
  else
    ratio = 1.0;

  coding->total_dicodons = 0;
  noncoding->total_dicodons = 0;

  for (i = 0; i < 64; i++) {
    for (j = 0; j < 64; j++) {
      sum = coding->dicodon[i][j] + noncoding->dicodon[i][j];
      coding->dicodon[i][j] = ratio * coding->dicodon[i][j];
      noncoding->dicodon[i][j] = sum - coding->dicodon[i][j];
      coding->codon[i] += coding->dicodon[i][j];
      noncoding->codon[i] += noncoding->dicodon[i][j];
      coding->total_dicodons += coding->dicodon[i][j];
      noncoding->total_dicodons += noncoding->dicodon[i][j];
    }
  }

  bell_correction (freq_cod, coding);
  bell_correction (freq_non, noncoding);

  for (i = 0; i < 64; i++) {
    codon1 = codonString (i);
    for (j = 0; j < 64; j++) {
      codon2 = codonString (j);
      freq = noncoding->dicodon[i][j] / noncoding->total_dicodons;
      fprintf (stdout, "%3s%3s %10.3e %10.3e\n", codon1, codon2,
	       log (freq_cod[i][j] / freq_non[i][j]), freq);
      free (codon2);
    }
    free (codon1);
  }
}


/* correct for low counts by the method of Bell, et al., 1990 */

void
bell_correction (double freq[64][64], coding_info * info)
{
  int i, j;
  double alpha, rawfreq, lambda;

  add_pseudocounts (info);
  alpha = info->total_dicodons / 4096;
  
  for (i = 0; i < 64; i++) {
    for (j = 0; j < 64; j++) {
      lambda = alpha / (info->codon[i] + alpha);
      rawfreq = info->dicodon[i][j] / info->codon[i];
      freq[i][j] = (lambda * info->codon[i] / info->total_dicodons)
	+ ((1 - lambda) * rawfreq);
    }
  }
}

#define PSEUDOCOUNT 0.5

/* if counts==0, add small pseudocount, not really for smoothing (because
   there are better ways of doing that, but simply to avoid dividing by
   zero in bell_correction, above */

void
add_pseudocounts (coding_info * info)
{
  int i, j;

  for (i = 0; i < 64; i++) {
    for (j = 0; j < 64; j++) {
      if (info->dicodon[i][j] == 0) {
	info->dicodon[i][j] += PSEUDOCOUNT;
	info->codon[i] += PSEUDOCOUNT;
	info->total_dicodons += PSEUDOCOUNT;
      }
    }
  }
}

/* correct for low counts by the method of MacKay and Peto, 1994 */

void
mackay_correction (double freq[64][64], coding_info * info) 
{
  int i,fmax;
  double v,k,g,h, sum=info->total_dicodons / 64,alpha=-1, u[64]={0};

  while((sum/alpha<.99) || (sum/alpha>1.01)) {
    alpha=sum;
    sum=0;
    for(i=0; i<64; i++) {
      fmax=mackay_fmax(info, i);
      v=mackay_N(info, i, 1);
      k=mackay_K(info, alpha);
      g=mackay_G(info, i, fmax);
      h=mackay_H(info, i, fmax); 
      u[i]=2*v/(k-g+(sqrt(pow(k-g,2)+(4*h*v))));
      sum+=u[i];
    }
  }
  for(i=0;i<64;i++) {
    printf("%4.2f %4.2f\n",u[i],info->codon[i]);
  }
  exit(1);
}

int mackay_N(coding_info *info, int j, int f) 
{
  int i, n=0;

  for(i=0;i<64;i++) {
    if (info->dicodon[i][j]>=f) n++;
  }
  return n;
}

int mackay_fmax(coding_info *info, int j) 
{
  int i, fmax=0;

  for(i=0;i<64;i++) {
    if (info->dicodon[i][j]>fmax) fmax=info->dicodon[i][j];
  }
  return fmax;
}

double mackay_K(coding_info *info, double alpha) 
{
  double k=0,mod;
  int i;

  for(i=0;i<64;i++) {
    mod=info->codon[i]+alpha;
    k+=log(mod/alpha)+0.5*(info->codon[i]/(alpha*mod));
  }
  return k;
}

double mackay_G(coding_info *info, int j, int fmax) 
{
  double g=0;
  int f;

  for(f=2;f<=fmax;f++) {
    g+=mackay_N(info, j, f)/(f-1);
  }
  return g;
}

double mackay_H(coding_info *info, int j, int fmax) 
{
  double h=0;
  int f;

  for(f=2;f<=fmax;f++) {
    h+=mackay_N(info, j, f)/pow(f-1,2);
  }
  return h;
}

void
usage (void)
{
  fprintf (stderr, "Usage: dicodontable [options] cds-file seq-file\n");
  fprintf (stderr, "Valid dicontable options:\n");
  fprintf (stderr, "\t-fraction-coding=fraction\n");
  exit (1);
}
