/*
Genomik und Bioinformatik I
Projekt 3 (Greedy contig assembly)
written by Andreas Loibl (enrollment no. 1524148)

compile with:
	g++ -o projekt3 -O2 projekt3.cpp

for debugging purposes the following might be useful:
	g++ -o projekt3 -g -Wall -DDEBUG projekt3.cpp	# compile with debugging symbols
	ulimit -c unlimited				# enable core dumps

usage (example):
	./projekt3 errorfreeReads.fna > contigs.txt

*/

// enable debug output
//#define DEBUG

// the contig-merge-algorithm is (currently) not stable, it has some bugs...
// if you want to use it you may enable it using the symbol "MERGE".
// (or from the commandline by appending the "-DMERGE" option to g++)
//#define MERGE

// contigs shorter than CONTIG_MINLEN will be ignored
#define CONTIG_MINLEN 32

// the contig assembly will be stopped after STOP_PERCENT % of the reads have been processed
#define STOP_PERCENT 100

#include <iostream>
#include <fstream>
#include <inttypes.h>
#include <stdint.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
using namespace std;

#define READLEN 32
#define READ uint64_t

#define DNA_A 0x00 // 00
#define DNA_G 0x01 // 01
#define DNA_C 0x02 // 10
#define DNA_T 0x03 // 11

// 24bit hash -> 256MB table
#define TABLE_SIZE 16777216

// linked list of reads, each element can have a length up to READLEN
struct DNApiece {
	READ read;
	int bits; // how many bits of the read are used -> defines the length of the DNApiece
	struct DNApiece *next;
	struct DNApiece *prev;
};

// linked list of DNApieces which contain the data to form a DNAstring without a maximum length
struct DNAstring {
	struct DNApiece *first;
	struct DNApiece *last;
};

// linked list to store various DNAstrings (-> contigs)
struct contig {
	struct DNAstring *dna;
	struct contig *next;
	struct contig *prev;
};

// element of the hashtable
struct node {
	READ read;
	struct node *next;
	struct node *prev;
	// to use the hashtable for the merging algorithm the following matching positions are stored within the hashtable:
	DNAstring *contig; // contig which contains the read
	DNApiece *piece; // DNApiece inside the contig which contains the match
	int pos; // relative position of the match inside the DNApiece
};

// hash table (will be filled with "node" elements)
static node *hash_table[TABLE_SIZE] = {};

int read_count = 0, read_count_processed = 0; // only used to show the assembly progress as percentage...

READ reverse_complement(READ dna)
{
	READ ret = 0;
	for(int i = 0; i < READLEN; i++)
	{
		ret = (ret << 2) | (dna & 0x03);
		dna >>= 2;
	}
	return(~ret);
}

#ifdef DEBUG
void print_DNAstring(DNAstring *s)
{
	printf("DNAstring %x:", s);
	DNApiece *p = s->first;
	READ dna;
	do
	{
		printf("\n\tDNApiece %x: ", p);
		dna = ~reverse_complement(p->read);
		dna >>= 2*READLEN - p->bits;
		for(int i = 0; i < p->bits/2; i++)
		{
			char c = dna & 0x03;
			if(c == DNA_A) cout << 'A';
			else if(c == DNA_T) cout << 'T';
			else if(c == DNA_G) cout << 'G';
			else if(c == DNA_C) cout << 'C';
			dna = (dna >> 2);
		}
	} while((p = p->next));
	cout << endl;
}
#endif

// DNAstring_len: returns the length of s
int DNAstring_len(DNAstring *s)
{
	DNApiece *p = s->first;
	int len = 0;
	if(p) do len+=p->bits/2; while((p = p->next));
	return(len);
}

// DNAstring_string: converts the DNAstring to a C-string (containing "A","T","C","G" letters)
char *DNAstring_string(DNAstring *s)
{
	char *pstr, *str;
	pstr = str = (char*) malloc(sizeof(char)*(DNAstring_len(s)+1));
	DNApiece *p = s->first;
	*str = 0;
	if(!p) return(str);
	READ dna;
	do
	{
		dna = ~reverse_complement(p->read);
		dna >>= 2*READLEN - p->bits;
		for(int i = 0; i < p->bits/2; i++)
		{
			char c = dna & 0x03;
			if(c == DNA_A) *pstr++ = 'A';
			else if(c == DNA_T) *pstr++ = 'T';
			else if(c == DNA_G) *pstr++ = 'G';
			else if(c == DNA_C) *pstr++ = 'C';
			dna = (dna >> 2);
		}
	} while((p = p->next));
	*pstr = 0;
	return(str);
}

// hash: returns the first 12bit of the supplied read; used to build the hashtable
uint32_t hash(READ read)
{
	return(read >> (64-24));
}

// match_count: searchs for the supplied read in the hashtable and returns the count of the matches
int match_count(READ read)
{
        uint32_t hashkey = hash(read);
	node *tmp = hash_table[hashkey];
	int count = 0;
	while(tmp)
	{
		if(tmp->read == read && !tmp->contig) count++;
		tmp = tmp->next;
	}
	return(count);
}

// new_node: allocates memory for a new node and returns a pointer to it.
node *new_node()
{
        node *n = (node *) malloc(sizeof(node));
	n->contig = NULL;
	n->prev = n->next = NULL;
        return n;
}

// insert: insert a read into the hashtable
void insert(READ read)
{
	if(match_count(read)) return;
        uint32_t hashkey = hash(read);
        node *tmp = new_node();
        tmp->read = read;
	tmp->contig = NULL;
        tmp->next = hash_table[hashkey];
	if(tmp->next) tmp->next->prev = tmp;
        hash_table[hashkey] = tmp;
	read_count++;
}

// mark: marks a read as used and stores the position of the match inside the hashtable
//       (this information is later used by the merge-contig-algorithm)
void mark(READ read, DNAstring *contig, int pos)
{
        uint32_t hashkey = hash(read);
	node *n = hash_table[hashkey];
	while(n)
	{
		if(n->read == read)
		{
			n->contig = contig;
			n->piece = contig->last;
			n->pos = pos;
			read_count_processed++;
		}
		n = n->next;
	}
}

// new_DNApiece: allocates memory for a new DNApiece and returns a pointer to it.
DNApiece *new_DNApiece()
{
        DNApiece *p = (DNApiece *) malloc(sizeof(DNApiece));
	p->read = 0;
	p->bits = 0;
	p->next = p->prev = NULL;
        return p;
}

// new_DNAstring: allocates memory for a new DNAstring and returns a pointer to it.
DNAstring *new_DNAstring()
{
        DNAstring *s = (DNAstring *) malloc(sizeof(DNAstring));
	//s->bits = 0;
	s->first = s->last = NULL;
        return s;
}

// DNAstring_reverse_complement: reverse-complements the supplied DNAstring (inplace! the original DNAstring will be changed)
DNAstring *DNAstring_reverse_complement(DNAstring *s)
{
	DNApiece *tmp = new_DNApiece();
	tmp = s->last;
	s->last = s->first;
	s->first = tmp;
	DNApiece *p = s->first;
	do
	{
		tmp = p->prev;
		p->prev = p->next;
		p->next = tmp;
		p->read = reverse_complement(p->read) >> (2*READLEN - p->bits);
	} while((p = p->next));
	free(tmp);
	return(s);
}

/*
DNAstring *DNAstring_reverse_complement_notinplace(DNAstring *s)
{
	DNAstring *ret = new_DNAstring();
	DNApiece *p = s->last;
	ret->first = ret->last = new_DNApiece();
	do
	{
		ret->last->read = reverse_complement(p->read) >> (2*READLEN - p->bits);
		ret->last->bits = p->bits;
		if(p->prev)
		{
			ret->last->next = new_DNApiece();
			ret->last->next->prev = ret->last;
			ret->last = ret->last->next;
		}
	} while((p = p->prev));
	return(ret);
}
*/

// DNAstring_subseq: returns a part of the given DNAstring
DNAstring *DNAstring_subseq(DNAstring *s, int start, int len = READLEN)
{
	DNAstring *ret = new_DNAstring();
	DNApiece *p = s->first;
	if(!p) return(ret);
	ret->first = ret->last = new_DNApiece();
	while(2*start > p->bits)
	{
		start-=p->bits/2;
		if(!p->next) return(ret);
		p = p->next;
	}
	ret->last->read = p->read;
	ret->last->bits = p->bits;
	ret->last->bits -= 2*start;
	while(len)
	{
		if(ret->last->bits < 2*len)
		{
			len-=ret->last->bits/2;
		}
		else
		{
			ret->last->read >>= ret->last->bits - 2*len;
			ret->last->bits = 2*len;
			len=0;
		}
		if(len && p->next)
		{
			ret->last->next = new_DNApiece();
			ret->last->next->prev = ret->last;
			ret->last = ret->last->next;
			p = p->next;
			ret->last->read = p->read;
			ret->last->bits = p->bits;
		} else break;
	}
	return(ret);
}


uint64_t shift(uint64_t value, int bits)
{
	return value << bits;
}

// DNAstring_read: return the last READLEN characters of the DNAstring as a read
READ DNAstring_read(DNAstring *s)
{
	DNApiece *p = s->first;
	READ dna = 0;
	int len = 2*READLEN;
	do
	{
		dna += p->read & ~(READ)0 >> (2*READLEN - p->bits);
		if(len - p->bits >= 64) dna=0;
		dna <<= len - p->bits;
		len-=p->bits;
	} while((p = p->next));
	return(dna);
}

// DNAstring_append: appends a nucleotide to a DNAstring (inplace!)
DNAstring *DNAstring_append(DNAstring *s, int nucleotid)
{
	if(!s->last) s->first = s->last = new_DNApiece();
	if(s->last->bits == 2*READLEN)
	{
		s->last->next = new_DNApiece();
		s->last->next->prev = s->last;
		s->last = s->last->next;
	}
	//else if(!s->last->bits) s->first = s->last = new_node();
	s->last->read <<= 2;
	s->last->read |= nucleotid;
	s->last->bits+=2;
	return(s);
}

// DNAstring_join: joins s1 with s2 (inplace!)
// 	NOTE: s1 and s2 both contain the same DNApieces after the join
void DNAstring_join(DNAstring *s1, DNAstring *s2)
{
	if(!s1->last)
	{
		s1->first = s2->first;
		s1->last = s2->last;
	}
	else if(!s2->last)
	{
		s2->first = s1->first;
		s2->last = s1->last;
	}
	else
	{
		s1->last->next = s2->first;
		s2->first->prev = s1->last;
		s1->last = s2->last;
		s2->first = s1->first;
	}
}

/*
DNAstring *DNAstring_join_notinplace(DNAstring *s1, DNAstring *s2)
{
	DNAstring *ret = new_DNAstring();
	DNApiece *p = s1->first;
	ret->first = ret->last = new_DNApiece();
	do
	{
		ret->last->read = p->read;
		ret->last->bits = p->bits;
		if(p->next)
		{
			ret->last->next = new_DNApiece();
			ret->last->next->prev = ret->last;
			ret->last = ret->last->next;
		}
	} while((p = p->next));
	p = s2->first;
	do
	{
		ret->last->read = p->read;
		ret->last->bits = p->bits;
		if(p->next)
		{
			ret->last->next = new_DNApiece();
			ret->last->next->prev = ret->last;
			ret->last = ret->last->next;
		}
	} while((p = p->next));
	return(ret);
}
*/

// DNAstring_get_read: returns the first read from the hashtable that is not marked as read
DNAstring *DNAstring_get_read()
{
	node **n;
	for(n = hash_table; n < hash_table+TABLE_SIZE; n++) if(*n) if(!(*n)->contig) break;
	if(!*n) return(NULL);
	if((*n)->contig) return(NULL);
	DNAstring *dna = new_DNAstring();
	dna->first = dna->last = new_DNApiece();
	dna->first->read = (*n)->read;
	dna->first->bits=64;
	mark(reverse_complement((*n)->read), dna, -1);
	mark((*n)->read, dna, 1);
	return(dna);
}

// build_contig: allocates memory for a new contig and fills it with an unused read from the hashtable.
contig *build_contig()
{
	contig *c = (contig*) malloc(sizeof(contig));
	c->dna = DNAstring_get_read();
	if(!c->dna) return(NULL);
	return(c);
}

// DNAstring_extend: extends the supplied DNAstring to the right (inplace!)
// 	and marks the read with the matching position in the hashtable
DNAstring *DNAstring_extend(DNAstring *dna, int direction = 1)
{
	int i, pos=DNAstring_len(dna)-READLEN;
	for(;;)
	{
		int sum = 0, matches[4] = {0};
		for(i = DNA_A; i <= DNA_T; i++)
		{
			READ read = DNAstring_read(DNAstring_append(DNAstring_subseq(dna,pos+1), i));
			sum += matches[i] = match_count(read) + match_count(reverse_complement(read));
		}
		if(!sum)
		{
			// no matches
			break;
		}
		for(i = DNA_A; i <= DNA_T; i++)
			if(matches[i] == sum)
			{
				// match
				DNAstring_append(dna, i);
				pos++;
				READ read = DNAstring_read(DNAstring_subseq(dna,pos));
				// store the matching position inside the hashtable (for the merging-algorithm)
				mark(reverse_complement(read), dna, -direction*dna->last->bits/2);
				mark(read, dna, direction*dna->last->bits/2);
				i=-1;
				break;
			}
		if(i>0)
		{
			// mismatch
			read_count_processed+=sum;
			break;
		}
	}
	return(dna);
}

// DNAstring_match: returns the position of the first occurrence of needle in haystack, or -1 if needle is not part of haystack.
int DNAstring_match(DNAstring *haystack, DNAstring *needle)
{
	char *str_haystack, *str_needle;
	str_haystack = DNAstring_string(haystack);
	str_needle = DNAstring_string(needle);
	char *p = strstr(str_haystack, str_needle);
	if(p) return(p-str_haystack); else return(-1);
}

/* DNAstring_cmp: compares the DNAstring s1 to the DNAstring s2

Returns an integral value indicating the relationship between the strings:
A zero value indicates that both strings are equal, a value greater than zero indicates that the first character
that does not match has a greater value in s1 than in s2; And a value less than zero indicates the opposite.
*/
int DNAstring_cmp(DNAstring *s1, DNAstring *s2)
{
	return(strcmp(DNAstring_string(s1), DNAstring_string(s2)));
}

#ifdef DEBUG
// dump: dumps the contig list to stdout (for debugging purposes)
void dump(contig *c)
{
	printf("\nContigs:\n#########\n");
	while(c)
	{
		if(c->dna)
		{
			if(DNAstring_len(c->dna)>=CONTIG_MINLEN) print_DNAstring(c->dna);
			if(!c->next) break;
		}
		c = c->next;
	}
}
#endif

int main(int argc, char **argv)
{
	string line, filename;
	if(argc < 2)
	{
		cout << "Usage: " << argv[0] << " <filename>" << endl;
		return(0);
	}
	ifstream f_reads (argv[1]);
	if(f_reads.is_open())
	{
		// read data
		cerr << "loading data..." << endl;
		while(f_reads >> line)
		{
			if(line[0] == '>') continue;
			READ dna = 0;
			for(int i = 0; i < READLEN; i++)
			{
				dna = (dna << 2);
				if(line[i] == 'A') dna|=DNA_A;
				else if(line[i] == 'T') dna|=DNA_T;
				else if(line[i] == 'G') dna|=DNA_G;
				else if(line[i] == 'C') dna|=DNA_C;
				
			}
			insert(dna);
		}
		f_reads.close();
		// assemble contigs
		contig *contigs, *c;
		contigs = c = build_contig();
		for(int i=1;(float)100*read_count_processed/read_count<STOP_PERCENT;)
		{
			fprintf(stderr, "\rassembling contig %4d - %.2f %% of reads processed...",
				i, (float)100*read_count_processed/read_count);
			c->dna = DNAstring_reverse_complement(DNAstring_extend(c->dna, 1));
			c->dna = DNAstring_reverse_complement(DNAstring_extend(c->dna, -1));
			fprintf(stderr, "\rassembling contig %4d - %.2f %% of reads processed...",
				i, (float)100*read_count_processed/read_count);
			if(DNAstring_len(c->dna)<CONTIG_MINLEN)
			{
				// contig too short -> throw away
				c->dna = DNAstring_get_read();
				if(!c->dna) break;
			}
			else
			{
				c->next = build_contig();
				if(!c->next) break;
				c->next->prev = c;
				c = c->next;
				i++;
			}
		}
		fprintf(stderr, "\n");
#ifdef MERGE
		// merge contigs
		cerr << "merging contigs..." << endl;
		c = contigs;
		for(int ic = 1; c; ++ic)
		{
			// walk through the list of contigs and check if the beginning of the contig is also found in another contig...
			if(c->dna)
			{
				DNAstring *pattern = DNAstring_subseq(c->dna, 0, 12);
#ifdef DEBUG
				printf("\nPattern:\n");
				print_DNAstring(DNAstring_reverse_complement(pattern));
				print_DNAstring(DNAstring_reverse_complement(pattern));
#endif
				// for both directions (reverse complement and normal)
				for(char direction = -1; direction < 2; direction+=2)
				{
					// lookup the matching positions using the hashtable
					// (this is much faster than comparing each contig to each of all other contigs)
					uint32_t hashkey = hash(DNAstring_read(DNAstring_reverse_complement(pattern)));
					node *tmp = hash_table[hashkey];
					while(tmp)
					{
						// match?
						if(tmp->contig && tmp->contig != c->dna && DNAstring_len(tmp->contig) >= READLEN)
						{
							if(tmp->contig->first->prev)
							{
								// this contig pointer is not up-to-date, it points into the middle of another contig...
								break;
							}
#ifdef DEBUG
							printf("contig %i matches at pos %i, dir %i-> merge DNAstring %x with DNAstring %x at DNApiece %x\n",
								ic, tmp->pos, direction, c->dna, tmp->contig, tmp->piece);

							printf("len: a: %i, b: %i\n", DNAstring_len(c->dna), DNAstring_len(tmp->contig));
#endif
							DNAstring *overlap = new_DNAstring();
							int overlap_len = 0;
							if(direction*tmp->pos>0)
							{
								DNAstring_reverse_complement(tmp->contig);
								overlap_len = -12+direction*tmp->pos;
								overlap->first = overlap->last = tmp->piece;
								// check if the overlap pattern is split into two DNApieces
								if(overlap_len < 0)
								{
									overlap_len += READLEN;
									overlap->first = tmp->piece->next;
								}
							}
							else
							{
								DNAstring_reverse_complement(tmp->contig);
								overlap_len = -12-direction*tmp->pos;
								overlap->first = overlap->last = tmp->piece;
								// check if the overlap pattern is split into two DNApieces
								if(direction*tmp->pos>-12 && tmp->piece->prev)
								{
									overlap_len += READLEN;
									overlap->first = tmp->piece->prev;
								}
							}
							// check if the pattern matches
							if(!DNAstring_cmp(DNAstring_subseq(overlap, overlap_len, 12), DNAstring_reverse_complement(pattern)))
							{
								// splice out the overlapping (matching) pattern
								if(overlap->first->prev) overlap->first->prev->next = NULL;
								if(overlap->last->next) overlap->last->next->prev = NULL;
								tmp->contig->last = overlap->first->prev;
								overlap->first->prev = NULL;
								overlap->last->next = NULL;
								// join everything together (DNAstring1 + overlap + DNAstring2)
								DNAstring_join(tmp->contig, DNAstring_subseq(overlap, 0, overlap_len));
								DNAstring_join(tmp->contig, c->dna);
								// after the join operation we have two identical contigs, so we delete one of them:
								if(tmp->prev) tmp->prev->next = tmp->next; else hash_table[hashkey] = tmp->next;
								c->dna=NULL;
								break;
							}
						}
						tmp = tmp->next;
					}
#ifdef DEBUG
					dump(contigs);
#endif
					if(!c->dna) break;
				}
			}
			if(!c) break;
			c = c->next;
		}
#endif
		// print (merged) contigs
		c = contigs;
		while(c)
		{
			if(c->dna)
			{
				// make sure the pointer is up-to-date and not in the middle of a contig
				while(c->dna->first->prev) c->dna->first = c->dna->first->prev;
				while(c->dna->last->next) c->dna->last = c->dna->last->next;
#ifdef DEBUG
				printf("DNAstring %x\n", c->dna);
#endif
				if(DNAstring_len(c->dna)>=CONTIG_MINLEN) cout << DNAstring_string(c->dna) << endl;
				if(!c->next) break;
			}
			c = c->next;
		}
	}
	else
	{
		cerr << "Unable to open file" << endl;
		return(1);
	}
	return(0);
}