#pragma warning (disable: 4786)

#include <stdio.h>
#include <stdlib.h>
#include "constants.h"
#include "strtokenizer.h"
#include "dataset.h"
#include "utils.h" 

using namespace std;

int dataset::write_wordmap(string wordmapfile, mapstr2id *pword2id) {
    FILE *fout; 

    if ((fout = fopen(wordmapfile.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to write!\n", wordmapfile.c_str()); 
		return EXIT_FAILURE; 
    }    
    
    mapstr2id::iterator it; 
    fprintf(fout, "%d\n", pword2id->size()); 
    for (it = pword2id->begin(); it != pword2id->end(); it++) {
		fprintf(fout, "%s %d\n", (it->first).c_str(), it->second); 
    }
    
    fclose(fout); 
    
    return EXIT_SUCCESS; 
}

int dataset::write_authormap(string authormapfile, mapstr2id *pauthor2id) {
    FILE *fout; 

    if ((fout = fopen(authormapfile.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to write!\n", authormapfile.c_str()); 
		return EXIT_FAILURE; 
    }    
    
    mapstr2id::iterator it; 
    fprintf(fout, "%d\n", pauthor2id->size()); 
    for (it = pauthor2id->begin(); it != pauthor2id->end(); it++) {
		fprintf(fout, "%s %d\n", (it->first).c_str(), it->second); 
    }
    
    fclose(fout); 
    
    return EXIT_SUCCESS; 
}

int dataset::read_wordmap(string wordmapfile, mapstr2id *pword2id) {
    pword2id->clear(); 
    
    FILE *fin; 
    if ((fin = fopen(wordmapfile.c_str(), "r")) == NULL) {
		utils::print_error("Cannot open file %s to read!\n", wordmapfile.c_str()); 
		return EXIT_FAILURE; 
    }    
    
    char buff[BUFF_SIZE_SHORT]; 
     
    fgets(buff, BUFF_SIZE_SHORT - 1, fin); 
    int nwords = atoi(buff); 
    
    for (int i = 0; i < nwords; i++) {
		fgets(buff, BUFF_SIZE_SHORT - 1, fin); 
		string line = buff; 
	
		strtokenizer strtok(line, " \t\r\n"); 
		if (strtok.count_tokens() != 2) {
			continue; 
		}
	
		pword2id->insert(pair<string, int>(strtok.token(0), atoi(strtok.token(1).c_str()))); 
    }
    
    fclose(fin); 
    
    return EXIT_SUCCESS; 
}

int dataset::read_wordmap(string wordmapfile, mapid2str *pid2word) {
    pid2word->clear(); 
    
    FILE *fin; 
    if ((fin = fopen(wordmapfile.c_str(), "r")) == NULL) {
		utils::print_error("Cannot open file %s to read!\n", wordmapfile.c_str()); 
		return EXIT_FAILURE; 
    }    
    
    char buff[BUFF_SIZE_SHORT]; 
    
    fgets(buff, BUFF_SIZE_SHORT - 1, fin); 
    int nwords = atoi(buff); 
    
    for (int i = 0; i < nwords; i++) {
		fgets(buff, BUFF_SIZE_SHORT - 1, fin); 
		string line = buff; 
	
		strtokenizer strtok(line, " \t\r\n"); 
		if (strtok.count_tokens() != 2) {
			continue; 
		}
	
		pid2word->insert(pair<int, string>(atoi(strtok.token(1).c_str()), strtok.token(0))); 
    }
    
    fclose(fin); 
    
    return EXIT_SUCCESS; 
}

int dataset::read_authormap(string authormapfile, mapstr2id *pauthor2id) {
    pauthor2id->clear(); 
    
    FILE *fin; 
    if ((fin = fopen(authormapfile.c_str(), "r")) == NULL) {
		utils::print_error("Cannot open file %s to read!\n", authormapfile.c_str()); 
		return EXIT_FAILURE; 
    }    
    
    char buff[BUFF_SIZE_SHORT]; 
     
    fgets(buff, BUFF_SIZE_SHORT - 1, fin); 
    int nauthors = atoi(buff); 
    
    for (int i = 0; i < nauthors; i++) {
		fgets(buff, BUFF_SIZE_SHORT - 1, fin); 
		string line = buff; 
	
		strtokenizer strtok(line, ",\t\r\n"); 
		if (strtok.count_tokens() != 2) {
			continue; 
		}
	
		pauthor2id->insert(pair<string, int>(strtok.token(0), atoi(strtok.token(1).c_str()))); 
    }
    
    fclose(fin); 
    
    return EXIT_SUCCESS; 
}

int dataset::read_authormap(string authormapfile, mapid2str *pid2author) {
    pid2author->clear(); 
    
    FILE *fin; 
    if ((fin = fopen(authormapfile.c_str(), "r")) == NULL) {
		utils::print_error("Cannot open file %s to read!\n", authormapfile.c_str()); 
		return EXIT_FAILURE; 
    }    
    
    char buff[BUFF_SIZE_SHORT]; 
    
    fgets(buff, BUFF_SIZE_SHORT - 1, fin); 
    int nauthors = atoi(buff); 
    
    for (int i = 0; i < nauthors; i++) {
		fgets(buff, BUFF_SIZE_SHORT - 1, fin); 
		string line = buff; 
	
		strtokenizer strtok(line, " \t\r\n"); 
		if (strtok.count_tokens() != 2) {
			continue; 
		}
	
		pid2author->insert(pair<int, string>(atoi(strtok.token(1).c_str()), strtok.token(0))); 
    }
    
    fclose(fin); 
    
    return EXIT_SUCCESS; 
}

int dataset::read_trndata(string dfile, string wordmapfile, string authormapfile) {
    mapstr2id word2id; 
	mapstr2id author2id; 
   
	FILE *fin; 
    if ((fin = fopen(dfile.c_str(), "r")) == NULL) {
		utils::print_error("Cannot open file %s to read!\n", dfile.c_str()); 
		return EXIT_FAILURE; 
    }   
    
    mapstr2id::iterator it; 
    char buff[BUFF_SIZE_LONG]; 
    
    // get the number of documents
    fgets(buff, BUFF_SIZE_LONG - 1, fin); 
    if ((M = atoi(buff)) <= 0) {
		utils::print_error("No document available!\n"); 
		return EXIT_FAILURE; 
    }
    
    // allocate memory for corpus
    if (docs != NULL) {
		deallocate(); 
    } else {
		docs = new document * [M]; 
    }
    
    // set number of words and authors to zero
    V = 0; 
	A = 0; 
    
    for (int m = 0; m < M; m++) {
		vector<int> words, authors; 
		double timestamp; 

		while (fgets(buff, BUFF_SIZE_LONG - 1, fin)) {
			string strline = buff + 2; 

			if ((buff[0] == '#') && (buff[1] == '*')) {	// text body
				strtokenizer wordstok(strline, " \t\r\n"); 
				int nwords = wordstok.count_tokens(); 

				if (nwords <= 0) {
					utils::print_error("(1) Invalid (empty) document!\n"); 
					deallocate(); 
					M = V = A = 0; 

					return EXIT_FAILURE; 
				}

				for (int i = 0; i < nwords; i++) {
					it = word2id.find(wordstok.token(i)); 
					if (it == word2id.end()) {	// word not found, i.e., new word
						words.push_back(word2id.size()); 
						word2id.insert(pair<string, int>(wordstok.token(i), word2id.size())); 
					} else {
						words.push_back(it->second); 
					}
				}
			} else if ((buff[0] == '#') && (buff[1] == '@')) {	// author list
				strtokenizer authorstok(strline, ",\t\r\n"); 
				int nauthors = authorstok.count_tokens(); 

				if (nauthors <= 0) {
					utils::print_error("(2) Invalid (empty) document!\n"); 
					deallocate(); 
					M = V = A = 0; 

					return EXIT_FAILURE; 
				}

				for (int i = 0; i < nauthors; i++) {
					it = author2id.find(authorstok.token(i)); 
					if (it == author2id.end()) {	// author not found, i.e., new author
						authors.push_back(author2id.size()); 
						author2id.insert(pair<string, int>(authorstok.token(i), author2id.size())); 
					} else {
						authors.push_back(it->second); 
					}
				}
			} else if ((buff[0] == '#') && (buff[1] == 't')) { //timestamp for a document
				if ((timestamp = atof(buff + 2)) < 0.0 ) {
					timestamp = 0.0; 
				}
			} else if ((buff[0] == '#') && (buff[1] == '!')) {	// end of a document
				break; 
			} else {
				; // do nothing
			}
		}

		// allocate new document
		document *pdoc = new document(words, authors, timestamp); 

		// add new doc to the corpus
		add_doc(pdoc, m); 
    }
    
    fclose(fin); 
    
    // write word map to file
    if ((write_wordmap(wordmapfile, &word2id) == EXIT_FAILURE) ||
		(write_authormap(authormapfile, &author2id) == EXIT_FAILURE)) {
		return EXIT_FAILURE; 
    }
    
    // update number of words and authors
    V = word2id.size(); 
	A = author2id.size(); 
    
    return EXIT_SUCCESS; 
}