output individual kmers if requested

1 files changed, 35 insertions, 6 deletions

diff --git a/kmer_total_count.c b/kmer_total_count.c
index ec9a3f8..cc9e625 100644
--- a/kmer_total_count.c
+++ b/kmer_total_count.c
@@ -9,13 +9,16 @@
 int main(int argc, char **argv) {
 
   unsigned long long i = 0;
+  long long j = 0;
 
-  if(argc !=  3) {
-    fprintf(stderr, "Please supply a filename and a kmer\n");
+  if(argc <  3) {
+    fprintf(stderr, "Please supply a filename and a kmer, then as many kmers as you want (in ACGT format)\n");
     exit(EXIT_FAILURE);
   }
 
 	unsigned long kmer = atoi(argv[2]);
+	unsigned long long width = pow_four(kmer);
+
 	if(kmer == 0) { 
 		fprintf(stderr, "Error: invalid kmer.\n");
 		exit(EXIT_FAILURE);
@@ -25,8 +28,34 @@ int main(int argc, char **argv) {
 
 	// print out our counts arrray
 	// manually unrolled 4 loops to reduce fprintf calls
-  for(i = 0; i < pow_four(atoi(argv[2])); i=i+4)
-    fprintf(stdout, "%llu\n%llu\n%llu\n%llu\n", counts[i], counts[i+1], counts[i+2], counts[i+3]);
-
-  return EXIT_SUCCESS;
+	
+	if(argc == 3) {
+		for(i = 0; i < width; i=i+4)
+			fprintf(stdout, "%llu\n%llu\n%llu\n%llu\n", counts[i], counts[i+1], counts[i+2], counts[i+3]);
+	} 
+	else {
+		for(i = 3; i < argc; i++) { 
+			unsigned long k = width;	
+			size_t len = 0;
+			fprintf(stdout, "%s\t", argv[i]);
+			
+			len = strlen(argv[i]);
+			if(len != kmer) { 
+				fprintf(stdout, "ERR\n");
+				continue;
+			}
+			for(j = 0; j < len; j++)
+				argv[i][j] = alpha[argv[i][j]];
+
+			k = num_to_index(argv[i], kmer, width);
+
+			if(k == width)  {
+				fprintf(stdout, "ERR\n");
+			}
+			else {
+				fprintf(stdout, "%llu\n", counts[k]);
+			}
+		}
+	}
+	return EXIT_SUCCESS;
 }