/*
 * CSSE, Monash University, Clayton, Victoria 3800, Australia
 * CSE2304 Prac 4, Group A
 * Semester 1, May 2003
 * Sample Solution  - O.R.
 *
 * dna k file1 file2 ...
 *   Outputs an adjacency matrix of the 'distance' from one DNA source to
 *   another, measured as the effectiveness of coding each sequence by a
 *   Markov model, of order k, based on all other input sequences.
 */
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>
#include <ctype.h>
#include <math.h>

/* Used to indicate a 'missing value' in the adjacency matrix */
#define EMPTY -1

/* delta, used for probability estimation */
#define delta 0.5

/* True if the character x is a valid DNA base */
#define isBase(x) (x == 'a'\
                || x == 'A'\
                || x == 'c'\
                || x == 'C'\
                || x == 'g'\
                || x == 'G'\
                || x == 't'\
                || x == 'T')

/* Maps the DNA base x to a number:
 *         a == 0
 *         c == 1
 *         g == 2
 *         t == 3
 */
#define baseEquiv(x) ((x)<'d'?((x)=='a'?0:1):((x)=='g'?2:3))

// Self-explanatory
#define log2(x) (log(x)/log(2))
#define max(x,y) (x>y?x:y)
#define min(x,y) (x<y?x:y)

/* dnaEquiv
 * Maps context, an array of DNA bases, to a unique number. Used to determine
 * the appropriate array index for a context. start is used to specify the
 * index of the first base, and len provides the array bound. From start
 * reads all bases to the end of the array, then wraps around and reads those
 * remaining.
 * Returns the unique numeric representation.
 */
unsigned long dnaEquiv(char *context, int start, int len) {
        unsigned long c = 0;
        int n = 0;

        /* Sanity-check parameters */
        assert(0 <= start && start <= len && len < 16);
        while(n++ < len) {
                /* Shift current value left by two bits
                 * to make room for next */
                c <<= 2;
                c += baseEquiv(context[start]);
                /* Move through array and wrap at end */
                start = (start + 1) % len;
        }
        return c;
}

/* matrix
 * Allocates space for an arbitrary two-dimensional array.
 * Each element is initialised to 0.
 * Returns NULL on failure.
 * Memory must later be freed using freeMatrix.
 */
void **matrix(size_t size, int rows, int cols) {
        void **t;
        int i;

        /* Require positive input */
        assert(size > 0 && rows > 0 && cols > 0);

        /* Allocate first dimension (rows) */
        if(!(t = calloc(sizeof(void *), rows)))
                /* Handle failure */
                return NULL;

        /* Allocate each row */
        for(i=0; i<rows; i++) {
                if(!(t[i] = calloc(size, cols))) {
                        /* Unwind, clean up */
                        while(--i)
                                free(t[i]);
                        free(t);
                        return NULL;
                }
                /* Clear contents of row */
                memset(t[i], 0, size * cols);
        }
        return t;
}

/* clearMatrix
 * Resets all elements of a matrix t to 0.
 */
void clearMatrix(void **t, size_t size, int rows, int cols) {
        int i;
        if(!t)
                return;

        /* Clear each row */
        for(i=0; i<rows; i++)
                memset(t[i], 0, cols*size);
}

/* freeMatrix
 * Frees memory allocated for a matrix t.
 */
void freeMatrix(void **t, int rows) {
        int i;
        if(!t)
                return;

        /* Free each row */
        for(i=0; i<rows; i++)
                free(t[i]);

        /* Free final pointer */
        free(t);
}

/* getModel
 * Reads a file of DNA bases to determine an appropriate Markov model.
 * Optionally writes the bases found to the file out.
 * Modifies the contents of model to store the result.
 * Requires that
 *         in is a file open for reading
 *         out is a file open for writing or NULL
 *         context is a character array of capacity k+1
 *         model is a matrix with 5 rows and 4^k columns
 */
void getModel(const int k, FILE *in, FILE *out,
                 char *context, int **model) {
        int end = 0, prev;
        int c, i;

        /* If |context| > 0 then fill context array */
        if(k > 0) {
                while((c = getc(in)) != EOF) {
                        if(isBase(c)) {
                                /* Have read a base */
                                /* Store in context and move on */
                                c = tolower(c);
                                if(out) {
                                        /* write to bases output file */
                                        putc(c, out);
                                }
                                context[end] = c;
                                end = (end + 1) % k;
                                if(end == 0) {
                                        /* Context is full */
                                        break;
                                }
                        }
                }
        }

        /* Finish processing input file */
        while((c = getc(in)) != EOF) {
                if(isBase(c)) {
                        /* Have read a base. Remember and overwrite earliest
                         * value in context */
                        c = tolower(c);
                        if(out) {
                                /* write to bases output file */
                                putc(c, out);
                        }
                        prev = dnaEquiv(context, end, k);
                        if(k > 0) {
                                context[end] = c;
                                end = (end + 1) % k;
                        }
                        /* Add context/base sequence to our Markov model */
                        model[baseEquiv(c)][prev]++;
                }
        }

        /* Sum columns to determine totals */
        for(i=0; i<(1 << (k << 1)); i++) {
                model[4][i] = model[0][i] + model[1][i]
                            + model[2][i] + model[3][i];
        }
}

/* outputAdjacencyMatrix
 * Prints to stdout the adjacency matrix adj using reasonably well-formatted
 * output. size is the length of both dimensions of the matrix and the
 * number of labels.
 */
void outputAdjacencyMatrix(double **adj, int size, char **labels) {
        int i, j, maxlen=7;

        /* Calculate max label length */
        for(i=0; i<size; i++)
                maxlen = max(maxlen, strlen(labels[i]));

        /* Output column headings */
        printf("%*s", maxlen, "to\\from");
        for(i=0; i<size; i++)
                if(adj[i][0] != EMPTY)
                        printf(" | %-*s",
                                max(min(10, maxlen), strlen(labels[i])),
                                labels[i]);
        printf(" |\n");

        /* Output rows */
        for(i=0; i<size; i++) {
                /* Skip 'empty' */
                if(adj[i][0] == EMPTY)
                        continue;
                printf("%-*s", maxlen, labels[i]);
                for(j=0; j<size; j++) {
                        if(adj[i][j] == EMPTY)
                                continue;
                        printf(" | %*.4f",
                                max(min(10, maxlen), strlen(labels[j])),
                                adj[i][j]);
                }
                printf(" |\n");
        }
}

/* compressedSequenceLength
 * Reads a file of DNA bases and returns the theoretical length of that
 * sequence compressed under the supplied Markov model.
 * Requires that
 *         in is a file open for reading
 *         context is a character array of capacity k+1
 *         model is a matrix with 5 rows and 4^k columns
 */
double compressedSequenceLength(const int k, FILE *in, char *context,
                int **model) {
        int end = 0, prev;
        int c;
        double sum = 0;

        /* If |context| > 0 then fill context array */
        if(k > 0) {
                while((c = getc(in)) != EOF) {
                        if(isBase(c)) {
                                /* Have read a base */
                                /* Not enough context; use default
                                 * probability */
                                sum += -log2(0.25);
                                /* Store in context and move on */
                                c = tolower(c);
                                context[end] = c;
                                end = (end + 1) % k;
                                if(end == 0)
                                        /* Context is full */
                                        break;
                        }
                }
        }

        /* Finish processing input file */
        while((c = getc(in)) != EOF) {
                if(isBase(c)) {
                        /* Have read a base. Remember and overwrite earliest
                         * value in context */
                        c = tolower(c);
                        prev = dnaEquiv(context, end, k);
                        if(k > 0) {
                                context[end] = c;
                                end = (end + 1) % k;
                        }
                        /* Calculate current base's probability */
                        sum += -log2((model[baseEquiv(c)][prev] + delta)
                                / (model[4][prev] + (4 * delta)));
                }
        }

        /* Finished. */
        return sum;
}

/* getAdjacencyMatrix
 * Determines and outputs the adjacency matrix of DNA source distances by
 * determining the probability of each base in the filenames supplied, based
 * on a Markov model of the previous k bases. Compares all files against each
 * other. Skips unknown characters. Case-insensitive. File errors are
 * non-fatal.
 * Returns 0 on success, 1 on allocation failure.
 */
int getAdjacencyMatrix(double **adj, const int k, int numfiles,
                char *filenames[]) {
        int *len = NULL;
        char *context = NULL;
        int **model = NULL;
        FILE *in = NULL;
        int i, j, c;
        double id;

        /* Preconditions */
        assert(k >= 0);
        assert(numfiles >= 0);

        /* Allocate space to store context and Markov model. */

        /* Space to store file lengths */
        len = (int *)calloc(sizeof(int), numfiles);
        if(!len) {
                /* Couldn't allocate space for lengths */
                return 1; /* failure */
        }

        /* One extra space for context list for when k=0 */
        context = (char *)calloc(sizeof(char), k+1);
        if(!context) {
                /* Couldn't allocate space for context, clean up */
                free(len);
                return 1; /* failure */
        }

        /* The model matrix stores
         *         5 rows (each of the bases plus a total)
         *         4^k columns per row
         * This comes to about 20GB when k = 15!
         */
        model = (int **)matrix(sizeof(int), 5, 1 << (k << 1));
        if(!model) {
                /* Couldn't allocate space for model, clean up */
                free(context);
                free(len);
                return 1; /* failure */
        }

        /* Memory has been successfully allocated. */
        assert(adj && model && context);

        /* For each file as a model: */
        for(i=0; i<numfiles; i++) {
                if(filenames[i][0] == '\0') {
                        /* Empty filename indicates previous error */
                        for(j=0; j<numfiles; j++)
                                adj[j][i] = EMPTY;
                        continue;
                }
                /* Open current file */
                if(!(in = fopen(filenames[i], "r"))) {
                        fprintf(stderr,
                                "warning: can't open file '%s'\n",
                                filenames[i]);

                        /* Indicate problem so file is not opened later */
                        filenames[i][0] = '\0';
                        /* Clear relevant entries in matrix */
                        for(j=0; j<numfiles; j++)
                                adj[j][i] = EMPTY;
                } else {
                        /* Read the Markov model from the input file */
                        clearMatrix((void **)model, sizeof(int),
                                        5, 1 << (k << 1));
                        getModel(k, in, NULL, context, model);
                        /* Determine and store length of the current file */
                        rewind(in);
                        j=0;
                        while((c=getc(in)) != EOF)
                                if(isBase(c)) j++;
                        len[i] = j;
                        fclose(in);

                        /* For each file based on that model: */
                        for(j=0; j<numfiles; j++) {
                                if(!(in = fopen(filenames[j], "r"))) {
                                        adj[j][i] = EMPTY;
                                        continue;
                                }
                                /* Add the compressed sequence length to the
                                 * matrix */
                                adj[j][i] = compressedSequenceLength(
                                                k, in, context, model);
                                /* Add the compressed model length */
                                adj[j][i] += 3*32*pow(4, k);
                                /* Finished with this file */
                                fclose(in);
                        }
                }
        }
        /* Calculate distances */
        for(i=0; i<numfiles; i++) {
                id = adj[i][i];
                for(j=0; j<numfiles; j++) {
                        if(adj[i][j] != EMPTY) {
                                adj[i][j] = (adj[i][j] - id)/len[i];
                        }
                }
        }

        /* All done. Clean up. */
        freeMatrix((void **)model, 5);
        free(context);
        free(len);

        return 0; /* success */
}

/* main
 * Error-checks parameters and handles error messages.
 */
int main(int argc, char *argv[]) {
        int k, result;
        double **adj;

        /* Check for required parameters */
        if(argc < 3) {
                /* Not enough; output a usage message and quit. */
                fprintf(stderr, "usage: %s k file1 file2 ...\n", argv[0]);
                return 1;
        }

        --argc; ++argv;

        /* Read k as a number */
        if(sscanf(*argv, "%d", &k) < 1) {
                /* Could not translate k to an integer; quit. */
                fprintf(stderr, "error: k must be an integer\n");
                return 2;
        }

        /* Check limits of k */
        if(k < 0 || k > 15) {
                /* Too high or too low; quit. */
                fprintf(stderr,
                        "error: k must be between zero and 15 inclusive\n");
                return 3;
        }

        --argc; ++argv;

        adj = (double **)matrix(sizeof(double), argc, argc);
        if(!adj) {
                /* Couldn't allocate space for adjacency matrix */
                fprintf(stderr, "error: could not allocate memory\n");
                return 4;
        }

        /* Call the getAdjacencyMatrix function to do all the dirty work. */
        result = getAdjacencyMatrix(adj, k, argc, argv);
        if(result == 1) {
                /* Something went wrong allocating memory; quit. */
                fprintf(stderr, "error: could not allocate memory\n");
        }

        outputAdjacencyMatrix(adj, argc, argv);

        freeMatrix((void **)adj, argc);

        /* Finished. */
        return 0;
}