/* Philip T.L.C. Clausen Jan 2017 s123580@student.dtu.dk */

/*
 Copyright (c) 2017, Philip Clausen, Technical University of Denmark
 All rights reserved.

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <ctype.h>
#include <errno.h>


int strpos_last(const char* str1, const char* str2) {
	char* strp;
	int i, len1, len2;
	
	len1 = strlen(str1);
	len2 = strlen(str2);
	if(len1 == 0 || len2 == 0 || len1 - len2 < 0) {
		return -1;
	}
	
	strp = (char*)(str1 + len1 - len2);
	for(i = len1 - len2; i >= 0; i--) {
		if(*strp == *str2) {
			if(strncmp(strp,str2,len2)==0)
				return i;
		}
		strp--;
	}
	return -1;
}

long unsigned fget_to(char *line, long unsigned line_size, FILE *file, char stopChar) {
	int c, stop;
	long unsigned i;
	
	stop = 0;
	i = 0;
	while(!stop) {
		while(i < line_size) {
			if((c = fgetc(file)) != EOF) {
				line[i] = c;
			} else {
				line[i] = '\0';
				return line_size;
			}
			if(c != stopChar) {
				i++;
			} else {
				break;
			}
		}
		if(line[i] != stopChar) {
			line_size *= 2;
			line = realloc(line, line_size);
		} else {
			stop = 1;
		}
	}
	line[i] = '\0';
	return line_size;
}

void helpMessage(int exeStatus) {
	FILE *helpOut;
	if(exeStatus == 0) {
		helpOut = stdout;
	} else {
		helpOut = stderr;
	}
	fprintf(helpOut, "# KmerResistance-2.0 correlates genes to species from WGS data, by mapping reads with KMA\n");
	fprintf(helpOut, "# Options are:\t\tDesc:\t\t\t\tDefault:\t\tRequirements:\n");
	fprintf(helpOut, "#\n");
	fprintf(helpOut, "#\t-i\t\tInput/query file name\t\t\t\t\tREQUIRED\n");
	fprintf(helpOut, "#\t-o\t\tOutput file\t\t\t\t\t\tREQUIRED\n");
	fprintf(helpOut, "#\t-t_db\t\tTemplate DB\t\t\t\t\t\tREQUIRED\n");
	fprintf(helpOut, "#\t-s_db\t\tSpecies DB\t\t\t\t\t\tREQUIRED\n");
	fprintf(helpOut, "#\t-id\t\tID threshhold\t\t\t70.0\n");
	fprintf(helpOut, "#\t-dct\t\tDepth correction threshhold\t0.1\n");
	fprintf(helpOut, "#\t-kma\t\talternative KMA\t\tkma\n");
	fprintf(helpOut, "#\n# KMA options:\t\tDesc:\t\t\t\tDefault:\t\tRequirements:\n");
	fprintf(helpOut, "#\t-i\t\tInput/query file name\t\tNone\t\tREQUIRED\n");
	fprintf(helpOut, "#\t-o\t\tOutput file\t\t\tNone\t\tREQUIRED\n");
	fprintf(helpOut, "#\t-t_db\t\tTemplate DB\t\t\tNone\t\tREQUIRED\n");
	fprintf(helpOut, "#\t-k\t\tKmersize\t\t\t16\n");
	fprintf(helpOut, "#\t-e\t\tevalue\t\t\t\t0.05\n");
	fprintf(helpOut, "#\t-delta\t\tAlign in pieces of delta\t511\n");
	fprintf(helpOut, "#\t-mem_mode\tUse kmers to choose best\n#\t\t\ttemplate, and save memory\tFalse\n");
	fprintf(helpOut, "#\t-ex_mode\tSearh kmers exhaustively\tFalse\n");
	fprintf(helpOut, "#\t-deCon\t\tRemove contamination\t\tFalse\n");
	fprintf(helpOut, "#\t-dense\t\tDo not allow insertions\n#\t\t\tin assembly\t\t\tFalse\n");
	fprintf(helpOut, "#\t-ref_fsa\tConsensus sequnce will\n#\t\t\thave \"n\" instead of gaps\tFalse\n");
	fprintf(helpOut, "#\t-matrix\t\tPrint assembly matrix\t\tFalse\n");
	fprintf(helpOut, "#\t-mp\t\tMinimum phred score\t\t30\n");
	fprintf(helpOut, "#\t-5p\t\tCut a constant number of\n#\t\t\tnucleotides from the 5 prime.\t0\n");
	fprintf(helpOut, "#\t-Sparse\t\tRun KmerFinder\t\t\tFalse\n");
	fprintf(helpOut, "#\t-ID\t\tMinimum ID\t\t\t1.0%%\n");
	fprintf(helpOut, "#\t-ss\t\tSparse sorting (q,c,d)\t\tq\n");
	fprintf(helpOut, "#\t-shm\t\tUse shared DB made by kma_shm\t0 (lvl)\n");
	fprintf(helpOut, "#\t-swap\t\tSwap DB to disk\t\t\t0 (lvl)\n");
	fprintf(helpOut, "#\t-1t1\t\tSkip HMM\t\t\tFalse\n");
	fprintf(helpOut, "#\t-boot\t\tBootstrap sequence\t\tFalse\n");
	fprintf(helpOut, "#\t-mrs\t\tMinimum alignment score,\n#\t\t\tnormalized to alignment length\t0.5\n");
	fprintf(helpOut, "#\t-reward\t\tScore for match\t\t\t1\n");
	fprintf(helpOut, "#\t-penalty\tPenalty for mismatch\t\t-2\n");
	fprintf(helpOut, "#\t-gapopen\tPenalty for gap opening\t\t-3\n");
	fprintf(helpOut, "#\t-gapextend\tPenalty for gap extension\t-1\n");
	fprintf(helpOut, "#\t-h\t\tShows this help message\n");
	fprintf(helpOut, "#\n");
	exit(exeStatus);
}

int main(int argc, char *argv[]) {
	
	int i, args, KMA_call_len, pos, line_size, out_len, HMM, stop;
	char *outputfilename, *t_db, *s_db, *KMA_dest, *KMA_call, *line, *q_id;
	char *template_name, *score, *expected, *t_len, *q, *p, *t_cov, *q_cov;
	double depth_corr_t, ID_t, ID, depth, depth_corr, expDepth;
	FILE *output, *input;
	/* set defaults */
	depth_corr_t = 0.1;
	ID_t = 70.0;
	outputfilename = 0;
	t_db = 0;
	s_db = 0;
	KMA_dest = strdup("kma");
	if(!KMA_dest) {
		fprintf(stderr, "OOM\n");
		exit(-1);
	}
	KMA_call_len = 11;
	
	/* parse cmd-line */
	args = 1;
	while(args < argc) {
		if(strcmp(argv[args], "-o") == 0) {
			args++;
			if(args < argc) {
				KMA_call_len += (strlen(argv[args]) + 4);
				out_len = strlen(argv[args]);
				outputfilename = (char*) calloc(out_len + strlen(".KmerRes") + 1, sizeof(char));
				strcpy(outputfilename, argv[args]);
				strcat(outputfilename, ".KmerRes");
			}
		} else if(strcmp(argv[args], "-t_db") == 0) {
			args++;
			if(args < argc) {
				t_db = (char*) calloc(strlen(argv[args]) + 1, sizeof(char));
				strcpy(t_db, argv[args]);
				KMA_call_len += (strlen(argv[args]) + 1);
			}
		} else if(strcmp(argv[args], "-s_db") == 0) {
			args++;
			if(args < argc) {
				s_db = (char*) calloc(strlen(argv[args]) + 1, sizeof(char));
				strcpy(s_db, argv[args]);
				KMA_call_len += (strlen(argv[args]) + 1);
			}
		} else if(strcmp(argv[args], "-id") == 0) {
			args++;
			if(args < argc) {
				ID_t = atof(argv[args]);
			}
			if(ID_t < 0) {
				fprintf(stderr, "Invalid threshold\n");
				exit(1);
			}
		} else if(strcmp(argv[args], "-dct") == 0) {
			args++;
			if(args < argc) {
				depth_corr_t = atof(argv[args]);
			}
			if(depth_corr_t < 0) {
				fprintf(stderr, "Invalid threshold\n");
				exit(1);
			}
		} else if(strcmp(argv[args], "-kma") == 0) {
			args++;
			if(args < argc) {
				free(KMA_dest);
				KMA_dest = strdup(argv[args]);
				if(!KMA_dest) {
					fprintf(stderr, "OOM\n");
					exit(-1);
				}
				KMA_call_len += strlen(argv[args]);
			}
		} else if(strcmp(argv[args], "-h") == 0) {
			helpMessage(0);
		} else {
			KMA_call_len += strlen(argv[args]) + 1;
			if(HMM == 0 && strcmp(argv[args], "-i") == 0) {
				stop = 0;
				args++;
				while(stop == 0 && args < argc) {
					if(strncmp(argv[args], "-", 1) != 0) {
						KMA_call_len += strlen(argv[args]) + 1;
						if(strpos_last(argv[args], "fsa") != -1 || strpos_last(argv[args], "fna") != -1 || strpos_last(argv[args], "fasta") != -1) {
							HMM = 1;
						}
						args++;
					} else {
						stop = 1;
					}
				}
				args--;
			}
		}
		args++;
	}
	if(outputfilename == 0 || t_db == 0 || s_db == 0) {
		fprintf(stderr, "# Too few arguments handed\n");
		fprintf(stderr, "# Printing help message:\n");
		helpMessage(1);
	}
	
	line_size = 4096 * sizeof(char);
	line = malloc(line_size);
	
	/* prepare KMA alignment */
	KMA_call_len += 5;
	pos = strpos_last(argv[0], "/") + 1;
	KMA_call = calloc(KMA_call_len + pos + 2, sizeof(char));
	
	/* add options to KMA */
	*KMA_call = 0;
	strcat(KMA_call, KMA_dest);
	strcat(KMA_call, " ");
	args = 1;
	while(args < argc) {
		if(strcmp(argv[args], "-t_db") == 0) {
			args++;
		} else if(strcmp(argv[args], "-s_db") == 0) {
			args++;
		} else if(strcmp(argv[args], "-id") == 0) {
			args++;
		} else if(strcmp(argv[args], "-dct") == 0) {
			args++;
		} else if(strcmp(argv[args], "-kma") == 0) {
			args++;
		} else if(strcmp(argv[args], "-h") != 0) {
			strcat(KMA_call, argv[args]);
			strcat(KMA_call, " ");
		}
		args++;
	}
	strcat(KMA_call, "-t_db ");
	
	/* mapp genes */
	fprintf(stderr, "# Mapping gene data.\n");
	KMA_call_len = strlen(KMA_call);
	sprintf(KMA_call, "%s%s", KMA_call, t_db);
	//fprintf(stderr, "%s\n", KMA_call);
	input = popen(KMA_call, "r");
	if(!input) {
		fprintf(stderr, "Invalid KMA specified\n");
		exit(-1);
	}
	if(pclose(input) != 0) {
		fprintf(stderr, "# Gene mapping failed\n");
		exit(1);
	}
	KMA_call[KMA_call_len] = '\0';
	
	/* mapp species */
	fprintf(stderr, "# Finding species\n");
	sprintf(KMA_call, "%s%s -Sparse\n", KMA_call, s_db);
	KMA_call[strpos_last(KMA_call, "\n")] = '\0';
	input = popen(KMA_call, "r");
	if(!input) {
		fprintf(stderr, "Invalid KMA specified\n");
		exit(-1);
	}
	if(pclose(input) != 0) {
		fprintf(stderr, "# Species mapping failed\n");
		exit(1);
	}
	KMA_call[KMA_call_len] = '\0';
	
	/* correlate species */
	fprintf(stderr, "# Mapping done\n");
	fprintf(stderr, "# Correlating results\n");
	template_name = malloc(256 * sizeof(char));
	score = malloc(256 * sizeof(char));
	expected = malloc(256 * sizeof(char));
	t_len = malloc(256 * sizeof(char));
	q = malloc(256 * sizeof(char));
	p = malloc(256 * sizeof(char));
	t_cov = malloc(256 * sizeof(char));
	q_id = malloc(256 * sizeof(char));
	q_cov = malloc(256 * sizeof(char));
	
	output = fopen(outputfilename, "w");
	if(!output) {
		fprintf(stderr, "Error: %d (%s)\n", errno, strerror(errno));
		exit(-1);
	}
	
	fprintf(output, "#Template\tScore\tExpected\ttemplate length\tq_value\tp_value\ttemplate_id\ttemplate_coverage\tquery_id\tquery_coverage\tdepth\tdepth_corr\n");
	strncpy((outputfilename + out_len), ".spa", 4);
	outputfilename[out_len + 4] = '\0';
	input = fopen(outputfilename, "r");
	if(!output) {
		fprintf(stderr, "Error: %d (%s)\n", errno, strerror(errno));
		exit(-1);
	}
	
	/* retrieve species information */
	fgets(line, line_size, input);
	fget_to(template_name, 256, input, '\t');
	line_size = fget_to(line, line_size, input, '\t');
	fget_to(score, 256, input, '\t');
	fget_to(expected, 256, input, '\t');
	fget_to(t_len, 256, input, '\t');
	
	fget_to(q_cov, 256, input, '\t');
	
	fget_to(t_cov, 256, input, '\t');
	ID = atof(t_cov);
	
	line_size = fget_to(line, line_size, input, '\t');
	expDepth = atof(line);
	line_size = fget_to(line, line_size, input, '\t');
	line_size = fget_to(line, line_size, input, '\t');
	line_size = fget_to(line, line_size, input, '\t');
	fget_to(q, 256, input, '\t');
	fget_to(p, 256, input, '\n');
	ID_t *= (1 - pow(2.718281828, (-0.5) * expDepth));
	fprintf(output, "%s\t%s\t%s\t%s\t%s\t%s\t%3.2f\t%s\t%s\t%s\t%5.2f\t%1.4f\n", template_name, score, expected, t_len, q, p, ID, t_cov, q_cov, q_cov, expDepth, (1 - pow(2.718281828, -1)));
	fclose(input);
	
	fprintf(stderr, "# Adjusted ID threshold: %f\n", ID_t);
	
	/* retrieve template information */
	strncpy((outputfilename + out_len), ".res", 4);
	outputfilename[out_len + 4] = '\0';
	input = fopen(outputfilename, "r");
	if(!input) {
		fprintf(stderr, "Error: %d (%s)\n", errno, strerror(errno));
		exit(-1);
	}
	
	fgets(line, line_size, input);
	while(! feof(input)) {
		fget_to(template_name, 256, input, '\t');
		fget_to(score, 256, input, '\t');
		fget_to(expected, 256, input, '\t');
		fget_to(t_len, 256, input, '\t');
		line_size = fget_to(line, line_size, input, '\t');
		ID = atof(line);
		
		fget_to(t_cov, 256, input, '\t');
		fget_to(q_id, 256, input, '\t');
		fget_to(q_cov, 256, input, '\t');
		
		
		
		line_size = fget_to(line, line_size, input, '\t');
		depth = atof(line);
		fget_to(q, 256, input, '\t');
		fget_to(p, 256, input, '\n');
		//fgets(line, line_size, input);
		depth_corr = (1 - pow(2.718281828, (-1) * depth / expDepth));
		if(depth_corr > depth_corr_t && ID > ID_t) {
			fprintf(output, "%s\t%s\t%s\t%s\t%s\t%s\t%3.2f\t%s\t%s\t%s\t%5.2f\t%1.4f\n", template_name, score, expected, t_len, q, p, ID, t_cov, q_id, q_cov, depth, depth_corr);
		}
	}
	fclose(input);
	
	fprintf(stdout, "DONE");
	fclose(output);
	exit(0);
}

