/* Accompanying code for Stark, Kheradpour, et al. (2007):                  
 * Reliable prediction of regulator targets using 12 Drosophila genomes     
 * 
 * Software not to be redistributed without permission from authors
 * Support contact email:  alex.stark / pouyak / manoli / at mit.edu
 *
 * PURPOSE:
 *   1) calculates branch length score (BLS) and its confidence
 *   2) identifies control motifs
 *   3) counts the number of instances at each confidence level
 *
 * COMPILATION:
 *   gcc -Wall -O3 -m32 -lm -o branch-sn branch-sn.c 
 *
 * SAMPLE USAGE:
 * 
 * REQUIRED FILES (. in formats indicates spacer column):
 *
 * <INPUTFILE> with lines of format:
 *    motif . . . . species_bitstring
 *    where species_bitstring is, for e.g. 5 if both the first and 
 *    third informant species species match, but not the second
 *    these represent each match to each motif and potential 
 *    control motif in the desired regions
 * 
 * <MOTIFFILE> with one motif (no whitespace) per line (should not include control motifs)
 *
 * <BLFILE> with lines of format:
 *    species_bitstring . . branch_length
 *
 * <NUMSP> is the number of species includingthe target species
 *    thus, 2^(<NUMSP>-1)-1 is the greatest species_bitstring
 *
 * First, compute <CONTORLFILE> (control motifs for each motif):
 *    branch-sn -b <BLFILE> -n <NUMSP> -e <MOTIFFILE> -pc 1 <INPUTFILE> | grep -w -f <MOTIFFILE> > <CONTROLFILE>
 *    NOTE: control motifs have the same composition and are within +/-20% of the number of instances
 *    NOTE: we do not permit one of the non-control motifs to be control motifs
 *    NOTE: this only works if the motif in <INPUTFILE> is a consensus sequence
 *
 * Next, use these and compute the branch length to confidence mapping:
 *    branch-sn -f <CONTROLFILE> -b <BLFILE> -n <NUMSP> <INPUTFILE> > <OUTFILE>
 *
 * <OUTFILE> has the format:
 *    column 1:      motif 
 *    columns 2-103: for branch lengths MAXBL*(0.00, 0.01, ..., 0.99, 1.00):
 *                   confidence;#instances;branch_length
 */


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

#define asserte1(b,a1)			do {if (!(b)) {fprintf(stderr, (a1)); exit(1);}} while(0)
#define asserte2(b,a1,a2)		do {if (!(b)) {fprintf(stderr, (a1), (a2)); exit(1);}} while(0)

#define MAXMOTLEN			50
#define ALPHSIZE			26


#define MAX(a,b)	(((a)>(b)) ? (a) : (b))
#define MIN(a,b)	(((a)<(b)) ? (a) : (b))

struct motif
{
	char str[MAXMOTLEN+1];
	unsigned char base_cnts[ALPHSIZE];
	unsigned int *bin_cnts;
	unsigned int total_cnt;
	unsigned char c_ex; /* indicates excluded from being a control */
};

/* reads a branch length file where the first column is the species bistring and the
 * forth column is the total branch length */
float* read_bl_file(char* bl_file, int num_sp);

/* compares motifs on the basis of base_cnts (i.e. motifs with the same composition will
 * have return value 0) */
int motifs_cmp (const void *a, const void *b)
{
	return memcmp(((struct motif*)a)->base_cnts, ((struct motif*)b)->base_cnts, sizeof(unsigned char)*ALPHSIZE);
}

/* read in all the motifs along with their counts at each of the blcuts
 * final motifs table is sorted such that motifs with the same composition are together */
int read_motifs (char* fn, struct motif **motifs_ptr, int num_bl, int num_bins, unsigned int *spbit_to_bin);

/* same as read_motifs except uses output from motif_match directly without use of 
 * mm2branch-sn */
int read_motifs_mm (char* fn, struct motif **motifs_ptr, int num_bl, int num_bins, unsigned int *spbit_to_bin);

double count_to_confid (unsigned int n, unsigned int nc, unsigned int n_r, unsigned int nc_r);

void print_motif_line (char* mstr, unsigned int*  m_bins, unsigned int* c_bins, float* consm_bl, int sp_mask, int num_bins);

int main (int argc, char** argv)
{
	int num_sp = 0;
	int num_motifs = 0;
	int num_bins = 101;
	struct motif* motifs;
	int num_bl;
	char* bl_file = NULL;
	char* c_file = NULL;
	char* ex_file = NULL;
	float* consm_bl = NULL;
	float tol_diff = 0.2;
	int sp_mask = -1;
	int print_controls = 0;
	int motif_match_input = 1;
	unsigned int *spbit_to_bin;

	/* same_end is actually one AFTER last with same composition */
	int same_start = 0, same_end = 0;

	int i, j;

	for (i=1; i<argc-1; i++)
		if (strcmp(argv[i],"-n") == 0)
			num_sp = atoi(argv[++i]);
		else if (strcmp(argv[i],"-m") == 0)
			num_bins = atoi(argv[++i]);
		else if (strcmp(argv[i],"-k") == 0)
			sp_mask = atoi(argv[++i]);
		else if (strcmp(argv[i], "-t") == 0)
			tol_diff = atof(argv[++i]);
		else if (strcmp(argv[i], "-e") == 0)
			ex_file = argv[++i];
		else if (strcmp(argv[i], "-f") == 0)
			c_file = argv[++i];
		else if (strcmp(argv[i], "-b") == 0)
			bl_file = argv[++i];
		else if (strcmp(argv[i], "-mm") == 0)
			motif_match_input = atoi(argv[++i]);
		else if (strcmp(argv[i], "-pc") == 0)
			print_controls = atoi(argv[++i]);
		else
			asserte1(0, "Invalid command line argument.\n");

	if (num_sp == 0 || bl_file == NULL)
	{
		printf("USAGE: %s <-n int> <-b file> [-t frac] [-n num] [-f file] [-e file] [-pc 0/1] [-mm 0/1] [-z float] <input file>\n", argv[0]);
		printf("    -n    Number of species (0 <= #Bls < 2^(NS - 1)) [required]\n");
		printf("    -f    List of controls for each motif (overrides -e, -pc, -t)\n");
		printf("    -e    List of motifs to exclude from controls\n");
		printf("    -k    Species mask (-1 for all) [default: -1]\n");
		printf("    -b    Branch length file [required]\n");
		printf("    -m    Number of bins to use for output [default: 101]\n");
		printf("    -t    Tolerance permitted between motifs and matched controls [default: 0.1]\n");
		printf("    -z    Z-value to use to create confidence intervals for ratios when generating confidence number [default: 0]\n");
		printf("    -mm   Input is output directly from motif-match [default: 0]\n");
		printf("    -pc   Prints control motifs chosen for each motif [default: 0]\n");
		return 1;
	}

	num_bl = 1 << (num_sp-1);

	if (sp_mask < 0)
		sp_mask = num_bl - 1;

	consm_bl = read_bl_file(bl_file,num_sp);

	/* conversion from the species bitstring to the corresponding bin */
	spbit_to_bin = (unsigned int*) malloc(sizeof(unsigned int)*num_bl);
	for (i=0; i<num_bl; i++)
		if ((i & sp_mask) == sp_mask)
			/* all the species is the last bin; this should already happen, but may screw up
			 * because of rounding issues with floating point */
			spbit_to_bin[i] = num_bins - 1;
		else
			spbit_to_bin[i] = (unsigned int) ((consm_bl[i & sp_mask] / consm_bl[sp_mask]) * (num_bins-1));

	/* create motifs array */
	if (!motif_match_input)
		num_motifs = read_motifs(argv[argc-1], &motifs, num_bl, num_bins, spbit_to_bin);
	else
		num_motifs = read_motifs_mm(argv[argc-1], &motifs, num_bl, num_bins, spbit_to_bin);

	/* create hash table for motifs for quick lookup */
	hcreate (num_motifs*2);
	for (i=0; i<num_motifs; i++)
	{
		ENTRY e;
		e.key = motifs[i].str;
		e.data = (void*) i;
		hsearch(e,ENTER);
	}

	if (c_file != NULL)
	{
		/* read controls explicitly from c_file */
		FILE* fp = fopen(c_file, "r");
		ENTRY e, *ep;
		size_t len = 0;
		char *line = NULL;
		ssize_t read;

		/* read in a line from the file */
	    while ((read = getline(&line, &len, fp)) != -1)
		{
			/* get rid of \n */
			if (line[read-1] == '\n')
				line[read-1] = '\0';

			/* look up the location of the main motif from the hash file */
			e.key = strtok(line, "\t ");
			ep = hsearch(e, FIND);

			if (ep != NULL)
			{
				/* if found, begin */
				struct motif* m = &motifs[(int) ep->data];
				unsigned int *c_bins = (unsigned int*) calloc(num_bins,sizeof(unsigned int));

				e.key = strtok(NULL, "\t ");
				while (e.key != NULL)
				{
					/* if found, update c_bins */
					ep = hsearch(e, FIND);
					if (ep != NULL)
						for (i=0; i<num_bins; i++)
							c_bins[i] += motifs[(int) ep->data].bin_cnts[i];
					e.key = strtok(NULL, "\t ");
				}

				/* print the line corresponding to this motif */
				print_motif_line(m->str, m->bin_cnts, c_bins, consm_bl, sp_mask, num_bins);
				free(c_bins);
			}
		}

		if (line)
			free(line);
	}
	else
	{
		/* exclude motifs */
		if (ex_file != NULL)
		{
			ENTRY e, *ep;
			FILE* fp = fopen (ex_file, "r");
			char ex_motif[MAXMOTLEN+1];
			e.key = ex_motif;

			while (fscanf(fp, "%s", e.key) != EOF)
			{
				ep = hsearch(e, FIND);
				if (ep != NULL)
					motifs[(int) ep->data].c_ex = 1;
			}
			fclose(fp);
		}

		/* scan through motifs and find matching instances */
		for (i=0; i<num_motifs; i++)
		{
			/* same_start and same_end correspond to the range of motifs with the same
			 * composition as the current one */
			/* same_end is actually one AFTER last with same composition */
			if (i >= same_end)
			{
				same_start = i;

				/* correct same_end */
				for (same_end=i+1; same_end<num_motifs && motifs_cmp(&motifs[i], &motifs[same_end]) == 0; same_end++);
			}

			if (!print_controls)
			{
				int k;
				/* print confidence at each bl cut */
				unsigned int *c_bins = (unsigned int*) calloc(num_bins,sizeof(unsigned int));

				for (j=same_start; j<same_end; j++)
					if (j!=i && !motifs[j].c_ex
							 && motifs[i].total_cnt*(1-tol_diff) <= motifs[j].total_cnt 
							 && motifs[i].total_cnt*(1+tol_diff) >= motifs[j].total_cnt
					)
						for (k=0; k<num_bins; k++)
							c_bins[k] += motifs[j].bin_cnts[k];
				print_motif_line(motifs[i].str, motifs[i].bin_cnts, c_bins, consm_bl, sp_mask, num_bins);
				
				free(c_bins);
			}
			else
			{
				printf("%s", motifs[i].str);
				for (j=same_start; j<same_end; j++)
					if (j!=i && !motifs[j].c_ex
							 && motifs[i].total_cnt*(1-tol_diff) <= motifs[j].total_cnt 
							 && motifs[i].total_cnt*(1+tol_diff) >= motifs[j].total_cnt
					)
						printf("\t%s", motifs[j].str);
				printf("\n");
			}
		}
	}

	/* free up used memory */
	for (i=0; i<num_motifs; i++)
		free(motifs[i].bin_cnts);
	free(motifs);
	free(consm_bl);
	return 0;
}

unsigned int* cumul_array (unsigned int* s, int n)
{
	unsigned int* t = (unsigned int*) malloc (n * sizeof(unsigned int));

	int i;
	t[n-1] = s[n-1];
	for (i=n-1; i>0; i--)
		t[i-1] = t[i] + s[i-1];

	return t;
}

void print_motif_line (char* mstr, unsigned int*  m_bins, unsigned int* c_bins, float* consm_bl, int sp_mask, int num_bins)
{
	int j;

	unsigned int* m_bins_c = cumul_array(m_bins, num_bins);
	unsigned int* c_bins_c = cumul_array(c_bins, num_bins);

	printf("%s", mstr);
	for (j=0; j<num_bins; j++)
	{
		printf("\t");

		if (c_bins_c[0] > 0 && m_bins_c[j] > 0)
			printf("%lf", count_to_confid(m_bins_c[0], m_bins_c[j], c_bins_c[0], c_bins_c[j]));
		else
			/* use -1 confidence to indicate unknown */
			printf("-");

		printf(";%d;%lf", m_bins_c[j], ((double)j*consm_bl[sp_mask])/(num_bins-1));
	}
	printf("\n");

	free(m_bins_c);
	free(c_bins_c);
}

double count_to_confid (unsigned int n, unsigned int nc, unsigned int n_r, unsigned int nc_r)
{
	if (nc == 0 || n_r == 0)
		return 0.0;

	return 1.0 - ((double) nc_r * n) / n_r / nc;

}

int read_motifs (char* fn, struct motif **motifs_ptr, int num_bl, int num_bins, unsigned int *spbit_to_bin)
{
	int num_motifs = 0;
	int size_reserved = 1 << 5;
	struct motif* motifs = (struct motif*) malloc(size_reserved*sizeof(struct motif));

	FILE* fp = strcmp(fn, "-")==0 ? stdin : fopen(fn, "r");
	asserte1(fp != NULL, "File could not be opened!\n");

	while (fscanf(fp, "%s", motifs[num_motifs].str) != EOF)
	{
		struct motif* m;
		int i;
		
		num_motifs++;
		
		if ((size_reserved-1) < num_motifs)
		{
			size_reserved *= 2;
			motifs = (struct motif*) realloc(motifs,size_reserved*sizeof(struct motif));
		}
		
		m = &motifs[num_motifs-1];

		m->c_ex = 0;
		memset(m->base_cnts, 0, sizeof(unsigned char)*ALPHSIZE);
		m->bin_cnts = (unsigned int*) calloc(num_bins, sizeof(unsigned int));
		m->total_cnt = 0;

		for (i=0; m->str[i]; i++)
		{
			/* note; with -f all this base composition stuff is ignored allowing for motifs 
			 * to have names that include non-alpha characters */
			int c = toupper(m->str[i]) - 'A';
			if (c >= 0 && c < 26)
				m->base_cnts[c]++;
		}

		for (i=0; i<num_bl; i++)
		{
			unsigned int c;
			fscanf(fp, "%u", &c);
			m->bin_cnts[spbit_to_bin[i]] += c;
			m->total_cnt += c;
		}
	}

	/* sort motifs so that motifs with the same composition are put together */
	qsort (motifs, num_motifs, sizeof(struct motif), motifs_cmp);

	*motifs_ptr = motifs;

	fclose(fp);

	return num_motifs;
}

struct motif_idx
{
	char str[MAXMOTLEN+1];
	int idx;
};

int motif_idx_cmp (const void *a, const void *b)
{
	return strcmp(((struct motif_idx*)a)->str, ((struct motif_idx*)b)->str);
}

int read_motifs_mm (char* fn, struct motif **motifs_ptr, int num_bl, int num_bins, unsigned int *spbit_to_bin)
{
	int num_motifs = 0;
	int size_reserved = 1 << 5;
	struct motif* motifs = (struct motif*) malloc(size_reserved*sizeof(struct motif));
	struct motif_idx m, *mp, **mpp;
	void* motif_idx_tree = NULL;
	int bl;

	FILE* fp = strcmp(fn, "-")==0 ? stdin : fopen(fn, "r");
	asserte1(fp != NULL, "File could not be opened!\n");

	while (fscanf(fp, "%s %*s %*s %*s %*s %d%*[^\n]\n", m.str, &bl) != EOF)
	{
		/* lookup in tree */
		mpp = tfind((void *) (&m), &motif_idx_tree, motif_idx_cmp);

		/* if not found, setup the motif */
		if (mpp == NULL)
		{
			int i;

			/* add to array */
			num_motifs++;
			if ((size_reserved-1) < num_motifs)
			{
				size_reserved *= 2;
				motifs = (struct motif*) realloc(motifs,size_reserved*sizeof(struct motif));
			}

			strcpy(motifs[num_motifs-1].str, m.str);
			memset(motifs[num_motifs-1].base_cnts, 0, sizeof(unsigned char)*ALPHSIZE);
			motifs[num_motifs-1].c_ex = 0;
			motifs[num_motifs-1].bin_cnts = (unsigned int*) calloc(num_bins, sizeof(unsigned int));
			motifs[num_motifs-1].total_cnt = 0;

			/* figure out base_cnts */
			for (i=0; motifs[num_motifs-1].str[i]; i++)
			{
				/* note; with -f all this base composition stuff is ignored allowing for motifs 
				 * to have names that include non-alpha characters */
				int c = toupper(motifs[num_motifs-1].str[i]) - 'A';
				if (c >= 0 && c < 26)
					motifs[num_motifs-1].base_cnts[c]++;
			}

			/* add to tree */
			mp = (struct motif_idx*) malloc(sizeof(struct motif_idx));
			strcpy(mp->str, m.str);
			mp->idx = num_motifs-1; 
			mpp = tsearch((void *) mp, &motif_idx_tree, motif_idx_cmp);
		}

		motifs[(*mpp)->idx].bin_cnts[spbit_to_bin[bl & (num_bl-1)]]++;
		motifs[(*mpp)->idx].total_cnt++;
	}

	/* NOTE: the tree is not deleted even though it could be now */

	/* sort motifs so that motifs with the same composition are put together */
	qsort (motifs, num_motifs, sizeof(struct motif), motifs_cmp);

	*motifs_ptr = motifs;

	fclose(fp);

	return num_motifs;
}


float* read_bl_file(char* bl_file, int num_sp)
{
	int num_sp_comb = 1 << (num_sp-1);
	float* consm_bl = (float*) calloc(num_sp_comb, sizeof(float));

	int index;
	float bl;

	FILE* fp = fopen(bl_file, "r");
	asserte2(fp != NULL, "Cannot open: %s\n", bl_file);

	while (fscanf(fp, "%d %*s %*s %f\n", &index, &bl) != EOF)
		if (index < num_sp_comb)
			consm_bl[index]=bl;
	fclose(fp);

	return consm_bl;
}