<commit_before>#include "kseq.h"

#include <zlib.h>

#include <cinttypes>
#include <cstdio>
#include <iostream>
#include <limits>
#include <vector>
#include <algorithm>
#include <set>
#include <cassert>
#include <sstream>
#include <unordered_set>
#include <getopt.h>

//typedef __uint128_t nkmer_t;
typedef uint64_t nkmer_t;
typedef std::set<nkmer_t> set_t;

const int32_t fasta_line_length=60;
const int32_t max_contig_length=10000000;
const int32_t max_allowed_kmer_length=sizeof(nkmer_t)*4;
const int32_t default_k=31;

static const uint8_t nt4_nt256[] = "ACGTN";

static const uint8_t nt256_nt4[] = {
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 0, 4, 1,  4, 4, 4, 2,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  3, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 0, 4, 1,  4, 4, 4, 2,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  3, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4
	};


KSEQ_INIT(gzFile, gzread)


void print_help(){
	std::cerr <<
		"\n" <<
		"Program:  prophyle-assembler (greedy assembler for ProPhyle)\n" <<
		"Contact:  <NAME> <<EMAIL>>\n" <<
		"\n" <<
		"Usage:    prophyle-assembler [options]\n" <<
		"\n" <<
		"Examples: prophyle-assembler -i f1.fa -i f2.fa -x fx.fa\n" <<
		"             - compute intersection of f1 and f2\n" <<
		"          prophyle-assembler -i f1.fa -i f2.fa -x fx.fa -o g1.fa -o g2.fa\n" <<
		"             - compute intersection of f1 and f2, and subtract it from them\n" <<
		"          prophyle-assembler -i f1.fa -o g1.fa\n" <<
		"             - re-assemble f1 to g1\n" <<
		"\n" <<
		"Command-line parameters:\n" <<
		" -i FILE  Input FASTA file (can be used multiple times).\n" <<
		" -o FILE  Output FASTA file (if used, must be used as many times as -i).\n" <<
		" -x FILE  Compute intersection, subtract it, save it.\n" <<
		" -k INT   K-mer size. [" << default_k << "]\n" <<
		"\n" <<
		"Note that '-' can be used for standard input/output. \n" <<
		std::endl;
}


template<typename _nkmer_T>
int32_t encode_forward(const char *kmers, const int32_t k, _nkmer_T &nkmer){
	nkmer=0;
	for (int32_t i=0;i<k;i++){
		uint8_t nt4 = nt256_nt4[static_cast<int32_t>(kmers[i])];
		if (nt4==4){
			return -1;
		}

		nkmer <<= 2;
		nkmer |= nt4;
	}
	return 0;
}

template<typename _nkmer_T>
int32_t encode_reverse(const char *kmers, const int32_t k, _nkmer_T &nkmer){
	nkmer=0;
	for (int32_t i=0;i<k;i++){
		uint8_t nt4 = nt256_nt4[static_cast<int32_t>(kmers[k-i-1])];
		if (nt4==4){
			return -1;
		}

		//complement
		nt4 = 3-nt4;

		nkmer <<= 2;
		nkmer |= nt4;
	}
	return 0;
}

template<typename _nkmer_T>
int32_t encode_canonical(const char *kmers, const int32_t k, _nkmer_T &nkmer){
	_nkmer_T nkmer_f;
	_nkmer_T nkmer_r;

	int32_t error_code;

	error_code=encode_forward(kmers, k, nkmer_f);
	if(error_code!=0){
		return error_code;
	}

	error_code=encode_reverse(kmers, k, nkmer_r);
	if(error_code!=0){
		return error_code;
	}

	nkmer=std::min(nkmer_f,nkmer_r);

	return 0;
}

template<typename _nkmer_T>
int32_t decode_kmer(_nkmer_T nkmer, int32_t k, std::string &kmer){
	kmer.resize(k);
	for(int32_t i=0;i<k;i++){
		uint8_t nt4 = nkmer & 0x3;
		nkmer >>=2;
		kmer[k-i-1]=nt4_nt256[nt4];
	}

	return 0;
}

void reverse_complement_in_place(std::string &kmer){
	//std::cerr << "before reverse complementing " << kmer << std::endl;
    std::reverse(kmer.begin(), kmer.end());
	for (int32_t i=0;i<static_cast<int32_t>(kmer.size());i++){
		char nt4=nt256_nt4[static_cast<int32_t>(kmer[i])];
		if (nt4<4){
			nt4=3-nt4;
		}
		kmer[i]=nt4_nt256[static_cast<int32_t>(nt4)];
	}
	//std::cerr << "after reverse complementing " << kmer << std::endl;
}

template<typename _set_T>
void debug_print_kmer_set(_set_T &set, int k){
	std::string kmer;
	for(auto x: set){
		decode_kmer(x, k, kmer);
		std::cerr << x << " " << kmer << ";  ";
	}
	std::cerr<<std::endl;
}


struct contig_t{
	int32_t k;

	/* contig buffer */
	char *seq_buffer;

	/* the first position of the contig */
	char *l_ext;

	/* the last position of the contig +1 (semiopen ) */
	char *r_ext;

	/* min possible value of l_ext */
	char *l_ext_border;

	/* max possible value of l_ext */
	char *r_ext_border;

	contig_t(uint32_t _k){
		this->k=_k;
		seq_buffer=new char[k+2*max_contig_length+1]();
		l_ext_border=seq_buffer;
		r_ext_border=seq_buffer+2*max_contig_length;
	}

	int32_t new_contig(const char *base_kmer){
		assert(static_cast<int32_t>(strlen(base_kmer))==k);

		l_ext = r_ext = &seq_buffer[max_contig_length];
		*r_ext='\0';

		for(int32_t i=0;i<k;i++){
			r_extend(base_kmer[i]);
		}
		return 0;
	}

	int32_t r_extend(char c){
		uint8_t nt4 = nt256_nt4[static_cast<int32_t>(c)];

		if (nt4==4){
			return -1;
		}

		*r_ext=nt4_nt256[nt4];
		++r_ext;
		*r_ext='\0';
		return 0;
	}

	int32_t l_extend(char c){
		uint8_t nt4 = nt256_nt4[static_cast<int32_t>(c)];

		if (nt4==4){
			return -1;
		}

		--l_ext;
		*l_ext=nt4_nt256[3-nt4];
		return 0;
	}

	~contig_t(){
		delete[] seq_buffer;
	}

	bool is_full(){
		return (r_ext>=r_ext_border) || (l_ext<=l_ext_border);
	}

	int32_t print_to_fasta(FILE* fasta_file, const char* contig_name, const char *comment=nullptr) const {
		if (comment==nullptr){
			fprintf(fasta_file,">%s\n",contig_name);
		} else {
			fprintf(fasta_file,">%s %s\n",contig_name,comment);
		}

		char print_buffer[fasta_line_length+1]={0};

		for (char *p=l_ext;p<r_ext;p+=fasta_line_length){
			strncpy(print_buffer,p,fasta_line_length);
			fprintf(fasta_file, "%s\n", print_buffer);
		}       

		return 0;
	}
};


/*
	TODO: test if kmer is correct
*/

//template<typename _nkmer_T, typename _set_T>
template<typename _set_T>
int kmers_from_fasta(const std::string &fasta_fn, _set_T &set, int32_t k){

	std::cerr << "Loading " << fasta_fn << std::endl;

	set.clear();

	kseq_t *seq;
	int64_t l;

	FILE *instream = NULL;
	if(fasta_fn=="-"){
		instream = stdin;
	}
	else {
		instream = fopen(fasta_fn.c_str(), "r");
	}
	gzFile fp = gzdopen(fileno(instream), "r");
	seq = kseq_init(fp);

	typename _set_T::value_type nkmer;

	for(int32_t seqid=0;(l = kseq_read(seq)) >= 0;seqid++) {
		//std::cerr << "kmers from fasta" << std::endl;

		//std::cerr << "starting iterator" << std::endl;
		for(char *kmer=seq->seq.s; kmer < (seq->seq.s) + (seq->seq.l) -k +1 ;kmer++){
			int c = encode_canonical(kmer, k, nkmer);
			if (c==0){
				set.insert(nkmer);
			}
			else{
				//std::cerr << "problem" <<std::endl;
			}
		}
	}

	std::cerr << "  " << fasta_fn << " loaded - #kmers: " << set.size()<<std::endl;
	//std::cerr << "iterator finished" << std::endl;

	kseq_destroy(seq);
	gzclose(fp);

	return 0;
}

template<typename _set_T>
int32_t find_intersection(const std::vector<_set_T> &sets, _set_T &intersection){
	assert(sets.size()>0);

	/* 
		1) Find the smallest set from sets.
	*/

	int32_t min=std::numeric_limits<int32_t>::max();
	int32_t i_min=-1;

	//std::cerr << "searching min" << std::endl;

	for(int32_t i=0;i<static_cast<int32_t>(sets.size());i++){
		if (static_cast<int32_t>(sets[i].size())<min){
			min=sets[i].size();
			i_min=i;
			//std::cerr << "new min" << i << std::endl;
		}
	}

	assert(i_min != std::numeric_limits<int32_t>::max() && i_min!=-1);

	/*
		2) Take it as the intersection.
	*/

	//std::cerr << "2" << std::endl;

	intersection.clear();
	//std::cerr << "2.1" << std::endl;
	std::copy(
		sets[i_min].cbegin(), sets[i_min].cend(),
		std::inserter(intersection,intersection.end())
	);

	/*
		3) Remove elements from intersection present in other sets.
	*/

	for(const _set_T &current_set : sets) {

		for(auto it = intersection.begin(); it !=intersection.end();){

			if(current_set.find(*it) == current_set.cend()){
				it=intersection.erase(it);
			}
			else{
				++it;
			}
		}
	}

	return 0;
}


template<typename _set_T, typename _subset_T>
int32_t remove_subset(std::vector<_set_T> &sets, const _subset_T &subset){

	for(int32_t i=0;i<static_cast<int32_t>(sets.size());i++){

		_set_T &current_set = sets[i];

		for(const auto &nkmer : subset){
			current_set.erase(nkmer);
		}
	}

	return 0;
}


template<typename _set_T>
int assemble(const std::string &fasta_fn, _set_T &set, int32_t k){
	std::cerr << "Assembling " << fasta_fn << " from " << set.size() << " kmers" << std::endl;

	FILE *file;
	if(fasta_fn=="-"){
		file=stdout;
	}
	else{
		file=fopen(fasta_fn.c_str(),"w+");
	}
	char kmer_str[max_allowed_kmer_length+1];
	contig_t contig(k);
	const std::vector<char> nucls = {'A','C','G','T'};

	//int32_t i=0;
	int32_t contig_id=1;
	while(set.size()>0){

		const auto central_nkmer=*(set.begin());
		set.erase(central_nkmer);

		std::string central_kmer_string;
		decode_kmer(central_nkmer,k,central_kmer_string);
		contig.new_contig(central_kmer_string.c_str());

		typename _set_T::value_type nkmer;


		for (int direction=0;direction<2;direction++){

			//std::cerr << "direction " << direction << std::endl;

			if (direction==0){
				// forward
			}
			else{
				// reverse
				reverse_complement_in_place(central_kmer_string);
			}

			strncpy(kmer_str,central_kmer_string.c_str(),k);
			kmer_str[k]='\0';

			bool extending = true;

			//std::cerr << "central k-mer: " << central_kmer_string << std::endl;


			while (extending){
				for(int32_t i=0;i<k;i++){
					kmer_str[i]=kmer_str[i+1];
				}
				kmer_str[k]='\0';

				extending=false;
				for(const char &c : nucls){
					kmer_str[k-1]=c;

					encode_canonical(kmer_str, k, nkmer);

					if(set.count( nkmer )){
						//std::cerr << "extending " << c << std::endl;
						//debug_print_kmer_set(set,k);
						//std::cerr << std::string(contig.l_ext) << c << std::endl;
						//std::cerr << std::endl;
						if(direction==0){
							contig.r_extend(c);
						}
						else{
							contig.l_extend(c);
						}
						set.erase(nkmer);

						if(!contig.is_full()){
							extending=true;
						}
						break;
					}
				}
			}
		}

		//std::cerr << "====================" << std::endl;

		std::stringstream ss;
		ss<<"contig_"<<contig_id;
		const std::string contig_name(ss.str());
		contig.print_to_fasta(file,contig_name.c_str());
		contig_id++;
	}

	fclose(file);

	std::cerr << "   assembly finished (" << contig_id << " contigs)" << std::endl;

	return 0;

}


int main (int argc, char* argv[])
{
	int32_t k=default_k;

	std::string intersection_fn;
	std::vector<std::string> in_fns;
	std::vector<std::string> out_fns;

	if (argc<2){
		print_help();
		exit(1);
	}

	bool compute_intersection=false;
	bool compute_output=false;
	int32_t no_sets=0;

	int c;
	while ((c = getopt(argc, (char *const *)argv, "hi:o:x:k:")) >= 0) {
		switch (c) {
			case 'h': {
				print_help();
				exit(0);
				break;
			}
			case 'i': {
				in_fns.push_back(std::string(optarg));
				no_sets+=1;
				break;
			}
			case 'o': {
				out_fns.push_back(std::string(optarg));
				compute_output=true;
				break;
			}
			case 'x': {
				intersection_fn=std::string(optarg);
				compute_intersection=true;
				break;
			}
			case 'k': {
				k = atoi(optarg);
				break;
			}
			case '?': {
				std::cerr<<"Unknown error"<<std::endl;
				exit(1);
				break;
			}
		}
	}

	if (k <= 1 || max_allowed_kmer_length<k){
		std::cerr << "K-mer size must satisfy 1 < k <= " << max_allowed_kmer_length << "." << std::endl;
		return EXIT_FAILURE;
	}

	if (compute_output && (out_fns.size()!=no_sets)){
		std::cerr << "If -o is used, it must be used as many times as -i (" << no_sets << "!=" << out_fns.size() << ")." << std::endl;
		return EXIT_FAILURE;
	}

	std::vector< std::unordered_set<nkmer_t> > full_sets(no_sets);

	std::cerr << "=====================" << std::endl;
	std::cerr << "1) Loading references" << std::endl;
	std::cerr << "=====================" << std::endl;


	std::vector<int32_t> in_sizes;
	std::vector<int32_t> out_sizes;

	for(int32_t i=0;i<no_sets;i++){
		kmers_from_fasta(in_fns[i],full_sets[i],k);
		//debug_print_kmer_set(full_sets[i],k);
		in_sizes.insert(in_sizes.end(),full_sets[i].size());
	}

	std::cerr << "===============" << std::endl;
	std::cerr << "2) Intersecting" << std::endl;
	std::cerr << "===============" << std::endl;


	std::unordered_set<nkmer_t> intersection;

	int32_t intersection_size = 0;

	if(compute_intersection){
		std::cerr << "2.1) Computing intersection" << std::endl;

		find_intersection(full_sets, intersection);
		intersection_size  = intersection.size();
		std::cerr << "   intersection size: " <<  intersection_size << std::endl;
		if(compute_output){
			std::cerr << "2.2) Removing this intersection from all kmer sets" << std::endl;
			remove_subset(full_sets, intersection);
		}
	}

	if(compute_output){
		for (int32_t i=0;i<no_sets;i++){
			out_sizes.insert(out_sizes.end(),full_sets[i].size());
			assert(in_sizes[i]==out_sizes[i]+intersection_size);
			std::cerr << in_sizes[i] << " " << out_sizes[i] << " ...inter:" << intersection_size << std::endl;
		}
	}

	std::cerr << "=============" << std::endl;
	std::cerr << "3) Assembling" << std::endl;
	std::cerr << "=============" << std::endl;

	if(compute_output){
		for(int32_t i=0;i<static_cast<int32_t>(in_fns.size());i++){
			assemble(out_fns[i],full_sets[i],k);
		}
	}
	if(compute_intersection){
		assemble(intersection_fn,intersection,k);
	}

	return 0;
}
<commit_msg>Parsable output (k-mer counts)<commit_after>#include "kseq.h"

#include <zlib.h>

#include <cinttypes>
#include <cstdio>
#include <iostream>
#include <limits>
#include <vector>
#include <algorithm>
#include <set>
#include <cassert>
#include <sstream>
#include <unordered_set>
#include <getopt.h>

//typedef __uint128_t nkmer_t;
typedef uint64_t nkmer_t;
typedef std::set<nkmer_t> set_t;

const int32_t fasta_line_length=60;
const int32_t max_contig_length=10000000;
const int32_t max_allowed_kmer_length=sizeof(nkmer_t)*4;
const int32_t default_k=31;

static const uint8_t nt4_nt256[] = "ACGTN";

static const uint8_t nt256_nt4[] = {
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 0, 4, 1,  4, 4, 4, 2,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  3, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 0, 4, 1,  4, 4, 4, 2,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  3, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,
		4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4
	};


KSEQ_INIT(gzFile, gzread)


void print_help(){
	std::cerr <<
		"\n" <<
		"Program:  prophyle-assembler (greedy assembler for ProPhyle)\n" <<
		"Contact:  <NAME> <<EMAIL>>\n" <<
		"\n" <<
		"Usage:    prophyle-assembler [options]\n" <<
		"\n" <<
		"Examples: prophyle-assembler -i f1.fa -i f2.fa -x fx.fa\n" <<
		"             - compute intersection of f1 and f2\n" <<
		"          prophyle-assembler -i f1.fa -i f2.fa -x fx.fa -o g1.fa -o g2.fa\n" <<
		"             - compute intersection of f1 and f2, and subtract it from them\n" <<
		"          prophyle-assembler -i f1.fa -o g1.fa\n" <<
		"             - re-assemble f1 to g1\n" <<
		"\n" <<
		"Command-line parameters:\n" <<
		" -i FILE  Input FASTA file (can be used multiple times).\n" <<
		" -o FILE  Output FASTA file (if used, must be used as many times as -i).\n" <<
		" -x FILE  Compute intersection, subtract it, save it.\n" <<
		" -k INT   K-mer size. [" << default_k << "]\n" <<
		"\n" <<
		"Note that '-' can be used for standard input/output. \n" <<
		std::endl;
}


template<typename _nkmer_T>
int32_t encode_forward(const char *kmers, const int32_t k, _nkmer_T &nkmer){
	nkmer=0;
	for (int32_t i=0;i<k;i++){
		uint8_t nt4 = nt256_nt4[static_cast<int32_t>(kmers[i])];
		if (nt4==4){
			return -1;
		}

		nkmer <<= 2;
		nkmer |= nt4;
	}
	return 0;
}

template<typename _nkmer_T>
int32_t encode_reverse(const char *kmers, const int32_t k, _nkmer_T &nkmer){
	nkmer=0;
	for (int32_t i=0;i<k;i++){
		uint8_t nt4 = nt256_nt4[static_cast<int32_t>(kmers[k-i-1])];
		if (nt4==4){
			return -1;
		}

		//complement
		nt4 = 3-nt4;

		nkmer <<= 2;
		nkmer |= nt4;
	}
	return 0;
}

template<typename _nkmer_T>
int32_t encode_canonical(const char *kmers, const int32_t k, _nkmer_T &nkmer){
	_nkmer_T nkmer_f;
	_nkmer_T nkmer_r;

	int32_t error_code;

	error_code=encode_forward(kmers, k, nkmer_f);
	if(error_code!=0){
		return error_code;
	}

	error_code=encode_reverse(kmers, k, nkmer_r);
	if(error_code!=0){
		return error_code;
	}

	nkmer=std::min(nkmer_f,nkmer_r);

	return 0;
}

template<typename _nkmer_T>
int32_t decode_kmer(_nkmer_T nkmer, int32_t k, std::string &kmer){
	kmer.resize(k);
	for(int32_t i=0;i<k;i++){
		uint8_t nt4 = nkmer & 0x3;
		nkmer >>=2;
		kmer[k-i-1]=nt4_nt256[nt4];
	}

	return 0;
}

void reverse_complement_in_place(std::string &kmer){
	//std::cerr << "before reverse complementing " << kmer << std::endl;
    std::reverse(kmer.begin(), kmer.end());
	for (int32_t i=0;i<static_cast<int32_t>(kmer.size());i++){
		char nt4=nt256_nt4[static_cast<int32_t>(kmer[i])];
		if (nt4<4){
			nt4=3-nt4;
		}
		kmer[i]=nt4_nt256[static_cast<int32_t>(nt4)];
	}
	//std::cerr << "after reverse complementing " << kmer << std::endl;
}

template<typename _set_T>
void debug_print_kmer_set(_set_T &set, int k){
	std::string kmer;
	for(auto x: set){
		decode_kmer(x, k, kmer);
		std::cerr << x << " " << kmer << ";  ";
	}
	std::cerr<<std::endl;
}


struct contig_t{
	int32_t k;

	/* contig buffer */
	char *seq_buffer;

	/* the first position of the contig */
	char *l_ext;

	/* the last position of the contig +1 (semiopen ) */
	char *r_ext;

	/* min possible value of l_ext */
	char *l_ext_border;

	/* max possible value of l_ext */
	char *r_ext_border;

	contig_t(uint32_t _k){
		this->k=_k;
		seq_buffer=new char[k+2*max_contig_length+1]();
		l_ext_border=seq_buffer;
		r_ext_border=seq_buffer+2*max_contig_length;
	}

	int32_t new_contig(const char *base_kmer){
		assert(static_cast<int32_t>(strlen(base_kmer))==k);

		l_ext = r_ext = &seq_buffer[max_contig_length];
		*r_ext='\0';

		for(int32_t i=0;i<k;i++){
			r_extend(base_kmer[i]);
		}
		return 0;
	}

	int32_t r_extend(char c){
		uint8_t nt4 = nt256_nt4[static_cast<int32_t>(c)];

		if (nt4==4){
			return -1;
		}

		*r_ext=nt4_nt256[nt4];
		++r_ext;
		*r_ext='\0';
		return 0;
	}

	int32_t l_extend(char c){
		uint8_t nt4 = nt256_nt4[static_cast<int32_t>(c)];

		if (nt4==4){
			return -1;
		}

		--l_ext;
		*l_ext=nt4_nt256[3-nt4];
		return 0;
	}

	~contig_t(){
		delete[] seq_buffer;
	}

	bool is_full(){
		return (r_ext>=r_ext_border) || (l_ext<=l_ext_border);
	}

	int32_t print_to_fasta(FILE* fasta_file, const char* contig_name, const char *comment=nullptr) const {
		if (comment==nullptr){
			fprintf(fasta_file,">%s\n",contig_name);
		} else {
			fprintf(fasta_file,">%s %s\n",contig_name,comment);
		}

		char print_buffer[fasta_line_length+1]={0};

		for (char *p=l_ext;p<r_ext;p+=fasta_line_length){
			strncpy(print_buffer,p,fasta_line_length);
			fprintf(fasta_file, "%s\n", print_buffer);
		}       

		return 0;
	}
};


/*
	TODO: test if kmer is correct
*/

//template<typename _nkmer_T, typename _set_T>
template<typename _set_T>
int kmers_from_fasta(const std::string &fasta_fn, _set_T &set, int32_t k){

	std::cerr << "Loading " << fasta_fn << std::endl;

	set.clear();

	kseq_t *seq;
	int64_t l;

	FILE *instream = NULL;
	if(fasta_fn=="-"){
		instream = stdin;
	}
	else {
		instream = fopen(fasta_fn.c_str(), "r");
	}
	gzFile fp = gzdopen(fileno(instream), "r");
	seq = kseq_init(fp);

	typename _set_T::value_type nkmer;

	for(int32_t seqid=0;(l = kseq_read(seq)) >= 0;seqid++) {
		//std::cerr << "kmers from fasta" << std::endl;

		//std::cerr << "starting iterator" << std::endl;
		for(char *kmer=seq->seq.s; kmer < (seq->seq.s) + (seq->seq.l) -k +1 ;kmer++){
			int c = encode_canonical(kmer, k, nkmer);
			if (c==0){
				set.insert(nkmer);
			}
			else{
				//std::cerr << "problem" <<std::endl;
			}
		}
	}

	std::cout << fasta_fn << "\t" << set.size()<<std::endl;
	//std::cerr << "iterator finished" << std::endl;

	kseq_destroy(seq);
	gzclose(fp);

	return 0;
}

template<typename _set_T>
int32_t find_intersection(const std::vector<_set_T> &sets, _set_T &intersection){
	assert(sets.size()>0);

	/* 
		1) Find the smallest set from sets.
	*/

	int32_t min=std::numeric_limits<int32_t>::max();
	int32_t i_min=-1;

	//std::cerr << "searching min" << std::endl;

	for(int32_t i=0;i<static_cast<int32_t>(sets.size());i++){
		if (static_cast<int32_t>(sets[i].size())<min){
			min=sets[i].size();
			i_min=i;
			//std::cerr << "new min" << i << std::endl;
		}
	}

	assert(i_min != std::numeric_limits<int32_t>::max() && i_min!=-1);

	/*
		2) Take it as the intersection.
	*/

	//std::cerr << "2" << std::endl;

	intersection.clear();
	//std::cerr << "2.1" << std::endl;
	std::copy(
		sets[i_min].cbegin(), sets[i_min].cend(),
		std::inserter(intersection,intersection.end())
	);

	/*
		3) Remove elements from intersection present in other sets.
	*/

	for(const _set_T &current_set : sets) {

		for(auto it = intersection.begin(); it !=intersection.end();){

			if(current_set.find(*it) == current_set.cend()){
				it=intersection.erase(it);
			}
			else{
				++it;
			}
		}
	}

	return 0;
}


template<typename _set_T, typename _subset_T>
int32_t remove_subset(std::vector<_set_T> &sets, const _subset_T &subset){

	for(int32_t i=0;i<static_cast<int32_t>(sets.size());i++){

		_set_T &current_set = sets[i];

		for(const auto &nkmer : subset){
			current_set.erase(nkmer);
		}
	}

	return 0;
}


template<typename _set_T>
int assemble(const std::string &fasta_fn, _set_T &set, int32_t k){
	std::cout << fasta_fn << "\t" << set.size() << std::endl;

	FILE *file;
	if(fasta_fn=="-"){
		file=stdout;
	}
	else{
		file=fopen(fasta_fn.c_str(),"w+");
	}
	char kmer_str[max_allowed_kmer_length+1];
	contig_t contig(k);
	const std::vector<char> nucls = {'A','C','G','T'};

	//int32_t i=0;
	int32_t contig_id=1;
	while(set.size()>0){

		const auto central_nkmer=*(set.begin());
		set.erase(central_nkmer);

		std::string central_kmer_string;
		decode_kmer(central_nkmer,k,central_kmer_string);
		contig.new_contig(central_kmer_string.c_str());

		typename _set_T::value_type nkmer;


		for (int direction=0;direction<2;direction++){

			//std::cerr << "direction " << direction << std::endl;

			if (direction==0){
				// forward
			}
			else{
				// reverse
				reverse_complement_in_place(central_kmer_string);
			}

			strncpy(kmer_str,central_kmer_string.c_str(),k);
			kmer_str[k]='\0';

			bool extending = true;

			//std::cerr << "central k-mer: " << central_kmer_string << std::endl;


			while (extending){
				for(int32_t i=0;i<k;i++){
					kmer_str[i]=kmer_str[i+1];
				}
				kmer_str[k]='\0';

				extending=false;
				for(const char &c : nucls){
					kmer_str[k-1]=c;

					encode_canonical(kmer_str, k, nkmer);

					if(set.count( nkmer )){
						//std::cerr << "extending " << c << std::endl;
						//debug_print_kmer_set(set,k);
						//std::cerr << std::string(contig.l_ext) << c << std::endl;
						//std::cerr << std::endl;
						if(direction==0){
							contig.r_extend(c);
						}
						else{
							contig.l_extend(c);
						}
						set.erase(nkmer);

						if(!contig.is_full()){
							extending=true;
						}
						break;
					}
				}
			}
		}

		//std::cerr << "====================" << std::endl;

		std::stringstream ss;
		ss<<"contig_"<<contig_id;
		const std::string contig_name(ss.str());
		contig.print_to_fasta(file,contig_name.c_str());
		contig_id++;
	}

	fclose(file);

	std::cerr << "   assembly finished (" << contig_id << " contigs)" << std::endl;

	return 0;

}


int main (int argc, char* argv[])
{
	int32_t k=default_k;

	std::string intersection_fn;
	std::vector<std::string> in_fns;
	std::vector<std::string> out_fns;

	if (argc<2){
		print_help();
		exit(1);
	}

	bool compute_intersection=false;
	bool compute_output=false;
	int32_t no_sets=0;

	int c;
	while ((c = getopt(argc, (char *const *)argv, "hi:o:x:k:")) >= 0) {
		switch (c) {
			case 'h': {
				print_help();
				exit(0);
				break;
			}
			case 'i': {
				in_fns.push_back(std::string(optarg));
				no_sets+=1;
				break;
			}
			case 'o': {
				out_fns.push_back(std::string(optarg));
				compute_output=true;
				break;
			}
			case 'x': {
				intersection_fn=std::string(optarg);
				compute_intersection=true;
				break;
			}
			case 'k': {
				k = atoi(optarg);
				break;
			}
			case '?': {
				std::cerr<<"Unknown error"<<std::endl;
				exit(1);
				break;
			}
		}
	}

	if (k <= 1 || max_allowed_kmer_length<k){
		std::cerr << "K-mer size must satisfy 1 < k <= " << max_allowed_kmer_length << "." << std::endl;
		return EXIT_FAILURE;
	}

	if (compute_output && (out_fns.size()!=no_sets)){
		std::cerr << "If -o is used, it must be used as many times as -i (" << no_sets << "!=" << out_fns.size() << ")." << std::endl;
		return EXIT_FAILURE;
	}

	std::vector< std::unordered_set<nkmer_t> > full_sets(no_sets);

	std::cerr << "=====================" << std::endl;
	std::cerr << "1) Loading references" << std::endl;
	std::cerr << "=====================" << std::endl;


	std::vector<int32_t> in_sizes;
	std::vector<int32_t> out_sizes;

	for(int32_t i=0;i<no_sets;i++){
		kmers_from_fasta(in_fns[i],full_sets[i],k);
		//debug_print_kmer_set(full_sets[i],k);
		in_sizes.insert(in_sizes.end(),full_sets[i].size());
	}

	std::cerr << "===============" << std::endl;
	std::cerr << "2) Intersecting" << std::endl;
	std::cerr << "===============" << std::endl;


	std::unordered_set<nkmer_t> intersection;

	int32_t intersection_size = 0;

	if(compute_intersection){
		std::cerr << "2.1) Computing intersection" << std::endl;

		find_intersection(full_sets, intersection);
		intersection_size  = intersection.size();
		std::cerr << "   intersection size: " <<  intersection_size << std::endl;
		if(compute_output){
			std::cerr << "2.2) Removing this intersection from all kmer sets" << std::endl;
			remove_subset(full_sets, intersection);
		}
	}

	if(compute_output){
		for (int32_t i=0;i<no_sets;i++){
			out_sizes.insert(out_sizes.end(),full_sets[i].size());
			assert(in_sizes[i]==out_sizes[i]+intersection_size);
			std::cerr << in_sizes[i] << " " << out_sizes[i] << " ...inter:" << intersection_size << std::endl;
		}
	}

	std::cerr << "=============" << std::endl;
	std::cerr << "3) Assembling" << std::endl;
	std::cerr << "=============" << std::endl;

	if(compute_output){
		for(int32_t i=0;i<static_cast<int32_t>(in_fns.size());i++){
			assemble(out_fns[i],full_sets[i],k);
		}
	}
	if(compute_intersection){
		assemble(intersection_fn,intersection,k);
	}

	return 0;
}

{'fraction_non_alphanumeric': 0.19994134467336314, 'fraction_numerical': 0.04215851602023609, 'mean_word_length': 3.6871134020618554, 'pattern_counts': {'":': 0, '<': 665, '<?xml version=': 0, '>': 129, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3217006', 'n_tokens_mistral': 13923, 'n_tokens_neox': 12044, 'n_words': 2452}