#define PROG_NAME "dbd_tabulate"
#define PROG_DESC "computes {k}-tuple freqs and probs from training sets"
#define PROG_VERS "1.0"
/* Last edited on 2008-06-10 04:05:35 by stolfi */

#define dbd_tabulate_C_COPYRIGHT "Copyright © 2006,2008 by IC-UFF and IC-UNICAMP"

#define PROG_HELP \
  "  " PROG_NAME " \\\n" \
  "    {WINDOW_SIZE} \\\n" \
  "    {LAB_SET} \\\n" \
  "    {LAB_EQV} \\\n" \
  "    {BANK_FILE} \\\n" \
  "    > {OUT_FILE}"
  
#define PROG_INFO \
  "NAME\n" \
  "  " PROG_NAME " - " PROG_DESC "\n" \
  "\n" \
  "SYNOPSIS\n" \
  PROG_HELP "\n" \
  "\n" \
  "DESCRIPTION\n" \
  "  This program takes two input parameters, an integer {k=WINDOW_SIZE}" \
  " and the name of a file {BANK_FILE} that contains pairs of file" \
  " names \"{BAS_FILE_i} {LAB_FILE_i}\", one pair per line.\n" \
  "\n" \
  "INPUTS\n" \
  "  In each file pair, the first file {BAS_FILE_i} should contain the" \
  " nucleotide sequence of a DNA, encoded with the" \
  " letters [ATCG].  The second file {LAB_FILE_i} should" \
  " contain the corresponding labels, which may be any" \
  " of the letters in the string {LAB_SET}.\n" \
  "\n" \
  "  For users' convenience, the program automatically" \
  " converts lowercase nucleotide letters in the input files to upper" \
  " case, and replaces the RNA nucleotide letter 'U' by 'T'.  It also" \
  " maps each label letter listed in the {LAB_SET} string to the" \
  " corresponding letter in the {LAB_EQV} string.  These conversions are" \
  " performed before counting the events and pairs, and" \
  " affect the output file output file.\n" \
  "\n" \
  "OUTPUTS\n" \
  "  The program extracts from each pair of input files all the pairs {t,e}" \
  " where {t} is a /{k}-tuple/ --- a string of {k} consecutive base" \
  " letters --- and {e} is the corresponding /{k}-event/ --- the string" \
  " of {k} consecutive label letters with same indices." \
  " For each distinct pair {t,e} that occurs in the input files, the" \
  " program writes to standard output a line in the format" \
  " \"{e} {t} {FREQ} {PROB}\" where {FREQ} is the frequency (number" \
  " of occurrences) of that {k}-tuple with that" \
  " {k}-event; and {PROB} is the conditional probability {Pr(t|r)}" \
  " estimated from those frequencies.\n" \
  "OPTIONS\n" \
  "  None.\n" \
  "\n" \
  "DOCUMENTATION OPTIONS\n" \
  argparser_help_info_HELP_INFO "\n" \
  "\n" \
  "SEE ALSO\n" \
  "  dbd_predict(1), dbd_summary(1), dbd_paint(1).\n" \
  "\n" \
  "AUTHOR\n" \
  "  Created 2005-10-05 by Renatha Oliva Capua, IC-UFF.\n" \
  "\n" \
  "MODIFICATION HISTORY\n" \
  "  2008-06-06 J. Stolfi, IC-UNICAMP Rewritten for libs, style.\n" \
  "  2008-06-09 J. Stolfi, IC-UNICAMP Added {LAB_SET} arg.\n" \
  "\n" \
  "WARRANTY\n" \
  argparser_help_info_NO_WARRANTY "\n" \
  "\n" \
  "RIGHTS\n" \
  "  " dbd_tabulate_C_COPYRIGHT ".\n" \
  "\n" \
  argparser_help_info_STANDARD_RIGHTS

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include <jsfile.h>
#include <jsmath.h>
#include <affirm.h>
#include <fget.h>

#include <dbd_lib.h>

/* PROTOTYPES */

typedef struct Options_t
  { int windowSize;  /* Size of tuples and events to consider. */
    char *labSet;  /* List of valid input label letters. */
    char *labEqv;  /* List of output label letters. */
    char *bankName;  /* Name of file that describes the input data bank. */
  } Options_t;

Options_t *get_options(int argc, char **argv);
  /* Parses the command line, returns them as an {Options_t} record. */
  
void process_file_pair
  ( int k,
    char *basName, 
    char *labName,
    code_table_t *tbB,
    code_table_t *tbL, 
    int64_t *ctE,
    int64_t *ctP
  );
  /* Reads a string of bases from {basFile}, and a matching string of
    labels from {labFile}. Extracts all pairs {t,r} where {t} is a
    {k}-tuple (substring of {k} bases) and {r} is a {k}-event
    (substring of {k} labels). Maps {r} to an event index {ixE}, using
    the table {tbL}, and increments {ctE[ixE]}; then maps and the pair
    {t,r} to a pair index {ixP}, using the tables {tbB,tbL}, and
    increments {ctP[ixP]}.

    Ignores line breaks, blanks, comments, etc. Ignores any
    {k}-tuple/{k}-event pairs that contain invalid characters. */

void write_freqs_probs
  ( FILE *wr,
    int k,
    code_table_t *tbB,
    code_table_t *tbL,
    int64_t *ctP,
    int64_t *ctE
  );
  /* Writes to file {wr} the raw occurrence counts and the conditional
    probabilities {Pr(t|e)} for each event/tuple pair {(e,t)}. */

int main(int argc, char **argv);

/* IMPLEMENTATIONS */

int main(int argc, char **argv)
  {
    Options_t *o = get_options(argc, argv);
    
    int k = o->windowSize;        /* Window size: */
    
    /* Build the nucleotide and label encoding/decoding table: */
    char *basSet = "ATCGatcgUu";  /* Valid input nucleotide chars. */
    char *basEqv = "ATCGATCGTT";  /* Equivalent nucleotide chars. */
    code_table_t tbB = build_code_table(basSet, basEqv, strlen(basSet));
    code_table_t tbL = build_code_table(o->labSet, o->labEqv, strlen(o->labSet));
    
    /* Check whether we have enough memory: */
    double maxMem = 64.0e6;  /* Max reasonable table size in bytes. */
    double maxAlt = exp(log(maxMem)/k); /* Max alternatives per letter. */
    demand(tbB.n*tbL.n <= maxAlt, "tables would be too big");
    fprintf(stderr, "alphabet size: bases = %d  labels = %d\n", tbB.n, tbL.n);
    
    /* Total number of distinct events, tuples, pairs: */
    int64_t nT = ipow(tbB.n,k); 
    int64_t nE = ipow(tbL.n,k); 
    int64_t nP = nT*nE; 
    fprintf(stderr, "table size: tuples = %lld events = %lld  pairs = %lld\n", nT, nE, nP);
    
    /* Try allocating the count arrays: */
    int64_t *ctE = (int64_t *)notnull(malloc(nE*sizeof(int64_t)), "no mem for event counts"); 
    int64_t *ctP = (int64_t *)notnull(malloc(nP*sizeof(int64_t)), "no mem for pair counts"); 
    
    /* Clear all counts: */
    int64_t ix;
    for (ix = 0; ix < nP; ix++) { ctP[ix] = 0; }
    for (ix = 0; ix < nE; ix++) { ctE[ix] = 0; }
    
    /* Loop on file pairs: */
    FILE  *bankFile = open_read(o->bankName, TRUE);
    int nSeqs = 0; /* Number of biosequences (file pairs) processed. */
    while(TRUE)
      { 
        /* Skip initial spaces on line: */
        fget_skip_spaces(bankFile);
        
        /* Check for end-of-file: */
        if (feof(bankFile)) { break; }
        
        /* Omit blank lines: */
        if (fget_test_char(bankFile, '\n')) { continue; }
        
        /* Omit comment lines; */
        if (fget_test_char(bankFile, '>'))
          { int c = '>'; 
            while ((! feof(bankFile)) && (c != '\n')) { c = fgetc(bankFile); } 
            continue;
          }
        
        /* Get names of input files (bases and labels): */
        char *basName = fget_string(bankFile);
        fget_skip_spaces(bankFile);
        char *labName = fget_string(bankFile);
        fget_skip_spaces(bankFile);
        
        /* Process the two files: */
        process_file_pair(k, basName, labName, &tbB, &tbL, ctE, ctP);
        nSeqs++;
        
        /* Skip to next line: */
        fget_eol(bankFile);
      }
    fclose(bankFile);
    fprintf(stderr, "%d biosequences processed.\n", nSeqs);     

    write_freqs_probs(stdout, k, &tbB, &tbL, ctP, ctE);

    return(0);
  }

void process_file_pair
  ( int k,
    char *basName, 
    char *labName,
    code_table_t *tbB,
    code_table_t *tbL, 
    int64_t *ctE,
    int64_t *ctP
  )
  {
    FILE *basFile = open_read(basName, TRUE);
    FILE *labFile = open_read(labName, TRUE);

    /* Total number of distinct events, tuples, pairs: */
    int64_t nT = ipow(tbB->n,k); 
    int64_t nE = ipow(tbL->n,k); 

    int nBases = 0; /* Count of non-blank, non-comment chars read from each file. */
    
    char tuple[k];    /* Current {k}-tuple (circular queue). */
    char event[k];    /* Current {k}-event (circular queue). */
    bool_t bad[k];    /* {bad[j]} is TRUE iff {tuple[j]} or {event[j]} is invalid. */
    
    int nBad = 0;     /* Number of `bad' base/label pairs in {tuple}/{event}. */
    
    /* Tuple and event indices (computed assuming code 0 for bad slots). */
    int64_t ixT = 0;  /* Index of {tuple} in list of all {k}-tuples. */
    int64_t ixE = 0;  /* Index of {event} in list of all {k}-events. */

    /* Line numbers in data files: */
    int basLine = 1;
    int labLine = 1;
    
    /* Skip comment blocks: */
    skip_bio_file_header(basFile, &basLine);
    skip_bio_file_header(labFile, &labLine);

    /* Loop on character pairs: */
    int basChar = fgetc(basFile);
    int labChar = fgetc(labFile);
    while ((basChar != EOF) && (labChar != EOF))
      {
        /* Skip blanks and newlines: */
        if (basChar == '\n')
          { basLine++; basChar = fgetc(basFile); continue; }
        if (basChar == ' ')
          { basChar = fgetc(basFile); continue; }
        if (labChar == '\n')
          { labLine++; labChar = fgetc(labFile); continue;  }
        if (labChar == ' ')
          { labChar = fgetc(labFile); continue;  }
          
        /* Got the next pair of non-blank characters, process it: */

        /* Map nucleotide and label to numeric codes: */
        int basCode = tbB->num[(unsigned)basChar];
        int labCode = tbL->num[(unsigned)labChar];
 
        /* Compute slot {j} of current chars in {tuple} and {event}: */
        int j = nBases % k;

        /* Shift tuple and event indices, discarding oldest base/label pair: */
        ixT = (tbB->n*ixT) % nT;
        ixE = (tbL->n*ixE) % nE;

        /* Update {nBad} to account for discarded base/label pair: */
        if ((nBases >= k) && bad[j]) { nBad--; }

        /* Store characters in {tuple} and [event} queues: */
        tuple[j] = basChar;
        event[j] = labChar;
        bad[j] = (basCode < 0) | (basCode >= tbB->n) | (labCode < 0) | (labCode >= tbL->n);
        nBases++;

        if (bad[j]) 
          { nBad++; }
        else
          { /* Add base and label codes to the tuple and event indices: */
            ixT += basCode;
            ixE += labCode;
          }

        if ((nBases >= k) && (nBad == 0))
          { /* Count one more occurrence of pair and event: */
            int64_t ixP = ixE*nT + ixT;
            affirm((k <= 1) || (ixP != 0), "bug in index computation");
            ctP[ixP]++;
            ctE[ixE]++;
          }

        /* Get next characters from each file: */
        basChar = fgetc(basFile);
        labChar = fgetc(labFile);
      }
    fclose(basFile);
    fclose(labFile);
 }

void write_freqs_probs
  ( FILE *wr,
    int k,
    code_table_t *tbB,
    code_table_t *tbL,
    int64_t *ctP,
    int64_t *ctE
  )
  {
    /* Total number of distinct events, tuples, pairs: */
    int64_t nT = ipow(tbB->n,k); 
    int64_t nE = ipow(tbL->n,k); 
    int64_t nP = nT*nE; 
    
    /* Max number of entries to print the full table: */
    int nPfull = ipow(2*4,6); /* Allows {N,K}×{A,T,C,G} with {k=6}. */

    int64_t ixP = 0; 

    fprintf(wr, "> Frequencies and probabilities of %d-tuples and %d-events\n", k, k);
    fprintf(wr, "> Columns are: event, tuple, occurrences(tuple,event), cond-prob(tuple|event)\n");
    for(ixP = 0; ixP < nP; ixP++)
      {
        if ((nP <= nPfull) || (ctP[ixP] > 0))
          { 
            /* Split pair index {ixP} into components: */
            int64_t ixE = ixP / nT; /* {k}-event index. */
            int64_t ixT = ixP % nT; /* {k}-tuple index. */ 

            /* Write event: */
            int64_t d;
            d = ipow(tbL->n, k);
            int64_t jxE = ixE;
            int r;
            for (r = 0; r < k; r++)
              { d /= tbL->n; int labCode = jxE / d;  jxE = jxE % d;
                fputc(tbL->chr[labCode], wr);
              }

            /* Write tuple: */
            fputc(' ', wr);
            d = ipow(tbB->n, k);
            int64_t jxT = ixT;
            for (r = 0; r < k; r++)
              { d /= tbB->n;  int basCode = jxT / d; jxT = jxT % d;
                fputc(tbB->chr[basCode], wr);
              }

            /* Write raw count: */
            fprintf(wr, " %10lld", ctP[ixP]);

            /* Write probability: */
            double prob = ((double)ctP[ixP]+1)/((double)ctE[ixE]+nT); 
            fprintf(wr, " %10.7f", prob);

            fputc('\n', wr);
          }
      }

    fclose(wr);
  }

Options_t *get_options(int argc, char **argv)
  {
    /* Get command-line parameters: */
    Options_t *o = (Options_t *)notnull(malloc(sizeof(Options_t)), "no mem");
    
    if (argc != 5)
      { fprintf(stderr, "wrong number of arguments\n");
        fprintf(stderr, "usage: %s\n",PROG_HELP);
        exit(1);
      }

    o->windowSize = atoi(argv[1]);
    o->labSet = argv[2];
    o->labEqv = argv[3];
    
    if (strlen(o->labSet) != strlen(o->labEqv))
      { fprintf(stderr, "label lists have different lengths\n");
        exit(1);
      }
    
    o->bankName = argv[4];
    
    return o;
  }