#pragma warning (disable: 4786)

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

using namespace std;

model::~model() {
    if (pa != NULL) {
		delete pa; 
    }

	if (pt != NULL) {
		delete pt; 
	}

	if (pbeta != NULL) {
		for (int m = 0; m < M; m++) {
			if (pbeta[m] != NULL) {
				delete pbeta[m]; 
			}
		}
	}

    if (ptrndata != NULL) {
		delete ptrndata; 
    }

    if (z != NULL) {
		for (int m = 0; m < M; m++) {
			if (z[m] != NULL) {
				delete z[m]; 
			}
		}
    }

	if (x != NULL) {
		for (int m = 0; m < M; m++) {
			if (x[m] != NULL) {
				delete x[m]; 
			}
		}
	}
    
    if (nw != NULL) {
		for (int v = 0; v < V; v++) {
			if (nw[v] != NULL) {
				delete nw[v]; 
			}
		}
    }

    if (na != NULL) {
		for (int a = 0; a < A; a++) {
			if (na[a] != NULL) {
				delete na[a]; 
			}
		}
    }

	if (nd != NULL) {
		for (int m = 0; m < M; m++) {
			if (nd[m] != NULL) {
				delete nd[m]; 
			}
		}
	}
    
    if (nwsum != NULL) {
		delete nwsum; 
    }   
    
    if (nasum != NULL) {
		delete nasum; 
    }
    
    if (theta != NULL) {
		for (int a = 0; a < A; a++) {
			if (theta[a] != NULL) {
				delete theta[a]; 
			}
		}
    }
    
    if (phi != NULL) {
		for (int k = 0; k < K; k++) {
			if (phi[k] != NULL) {
				delete phi[k]; 
			}
		}
    }

	if (psi != NULL) {
		for (int k = 0; k < K; k++) {
			if (psi[k] != NULL) {
				delete psi[k]; 
			}
		}
	}
}

void model::set_default_values(void) {
    wordmapfile = "wordmap.txt"; 
	authormapfile = "authormap.txt"; 
    txassign_suffix = ".txassign"; 
    theta_suffix = ".theta"; 
    phi_suffix = ".phi"; 
	psi_suffix = ".psi"; 
    others_suffix = ".others"; 
    twords_suffix = ".twords"; 
	tauthors_suffix = ".tauthors"; 
    
    dir = "./"; 
    dfile = "trndocs.dat"; 
    model_name = "model-final"; 
    
    ptrndata = NULL; 
    
    M = 0; 
    V = 0; 
    K = 100; 
	A = 0; 
    alpha = 50.0 / K; 
    beta = 0.1; 
    niters = 2000; 
    liter = 0; 
    savestep = 0; 
    twords = 20; 
	tauthors = 20; 
    
    pa = NULL; 
	pt = NULL; 
	pbeta = NULL; 
    z = NULL; 
	x = NULL; 
    nw = NULL; 
    na = NULL; 
	nd = NULL; 
    nwsum = NULL; 
    nasum = NULL; 
    theta = NULL; 
    phi = NULL; 
	psi = NULL; 
}

int model::timestamp_norm(void) {
	vector<double> timestamp; 

	for (int m = 0; m < ptrndata->M; m++) {
		timestamp.push_back(ptrndata->docs[m]->timestamp); 
	}

	double max_val = utils::get_max(timestamp); 
	double min_val = utils::get_min(timestamp); 

	if (min_val == max_val) {
		return EXIT_FAILURE; 
	}

	for (m = 0; m < ptrndata->M; m++) {
		ptrndata->docs[m]->timestamp = (ptrndata->docs[m]->timestamp - min_val) * (MAX_VALUE - MIN_VALUE) 
			/ (max_val - min_val) + MIN_VALUE; 
	}

	return EXIT_SUCCESS; 
}

int model::init(void) {
	return init_est(); 
}

int model::save_model(string model_name) {
	if ((save_model_txassign(dir + model_name + txassign_suffix) == EXIT_FAILURE) ||
		(save_model_others(dir + model_name + others_suffix) == EXIT_FAILURE) ||
		(save_model_theta(dir + model_name + theta_suffix) == EXIT_FAILURE) ||
		(save_model_phi(dir + model_name + phi_suffix) == EXIT_FAILURE) ||
		(save_model_psi(dir + model_name + psi_suffix) == EXIT_FAILURE)) {
		return EXIT_FAILURE; 
	}
    
    if (twords > 0) {
		if (save_model_twords(dir + model_name + twords_suffix) == EXIT_FAILURE) {
			return EXIT_FAILURE;
		}
    }

	if (tauthors > 0) {
		if (save_model_tauthors(dir + model_name + tauthors_suffix) == EXIT_FAILURE) {
			return EXIT_FAILURE; 
		}
	}
    
    return EXIT_SUCCESS;
}

int model::save_model_txassign(string filename) {
    FILE *fout;

    if ((fout = fopen(filename.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to save!\n", filename.c_str());
		return EXIT_FAILURE;
    }

    // wirte docs with topic assignments for words
    for (int m = 0; m < ptrndata->M; m++) {    
		for (int n = 0; n < ptrndata->docs[m]->nwords; n++) {
			fprintf(fout, "%d:%d:%d ", ptrndata->docs[m]->words[n], z[m][n], x[m][n]);
		}
		fprintf(fout, "\n");
    }

    fclose(fout);
    
    return EXIT_SUCCESS;
}

int model::save_model_theta(string filename) {
    FILE *fout; 

    if ((fout = fopen(filename.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to save!\n", filename.c_str()); 
		return EXIT_FAILURE; 
    }
    
    for (int a = 0; a < A; a++) {
		for (int k = 0; k < K; k++) {
			fprintf(fout, "%g ", theta[a][k]); 
		}
		fprintf(fout, "\n"); 
    }
    
    fclose(fout); 
    
    return EXIT_SUCCESS;
}

int model::save_model_psi(string filename) {
    FILE *fout; 

    if ((fout = fopen(filename.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to save!\n", filename.c_str()); 
		return EXIT_FAILURE; 
    }
    
    for (int k = 0; k < K; k++) {
		fprintf(fout, "%g %g\n", psi[k][0], psi[k][1]); 
    }
    
    fclose(fout); 
    
    return EXIT_SUCCESS;
}

int model::save_model_phi(string filename) {
    FILE *fout;

    if ((fout = fopen(filename.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to save!\n", filename.c_str());
		return EXIT_FAILURE;
    }
    
    for (int k = 0; k < K; k++) {
		for (int v = 0; v < V; v++) {
			fprintf(fout, "%g ", phi[k][v]);
		}
		fprintf(fout, "\n");
    }
    
    fclose(fout);    
    
    return EXIT_SUCCESS;
}

int model::save_model_others(string filename) {
    FILE *fout; 

    if ((fout = fopen(filename.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to save!\n", filename.c_str()); 
		return EXIT_FAILURE; 
    }

    fprintf(fout, "alpha=%g\n", alpha); 
    fprintf(fout, "beta=%g\n", beta); 
    fprintf(fout, "ntopics=%d\n", K); 
    fprintf(fout, "ndocs=%d\n", M); 
    fprintf(fout, "nwords=%d\n", V); 
	fprintf(fout, "nauthors=%d\n", A); 
    fprintf(fout, "liter=%d\n", liter); 
	fprintf(fout, "seed=%d\n", seed); 
    
    fclose(fout); 
    
    return EXIT_SUCCESS; 
}

int model::save_model_twords(string filename) {
    FILE *fout; 

    if ((fout = fopen(filename.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to save!\n", filename.c_str()); 
		return EXIT_FAILURE; 
    }
    
    if (twords > V) {
		twords = V; 
    }
    mapid2str::iterator it; 
    
    for (int k = 0; k < K; k++) {
		vector<pair<int, double> > words_probs; 
		pair<int, double> word_prob; 
		for (int v = 0; v < V; v++) {
			word_prob.first = v; 
			word_prob.second = phi[k][v]; 
			words_probs.push_back(word_prob); 
		}
    
        // quick sort to sort word-topic probability
		utils::quicksort(words_probs, 0, words_probs.size() - 1); 
	
		fprintf(fout, "Topic %dth:\n", k); 
		for (int i = 0; i < twords; i++) {
			it = id2word.find(words_probs[i].first); 
			if (it != id2word.end()) {
				fprintf(fout, "\t%s   %g\n", (it->second).c_str(), words_probs[i].second); 
			}
		}
    }
    
    fclose(fout); 
    
    return EXIT_SUCCESS; 
}

int model::save_model_tauthors(string filename) {
    FILE *fout; 

    if ((fout = fopen(filename.c_str(), "w")) == NULL) {
		utils::print_error("Cannot open file %s to save!\n", filename.c_str()); 
		return EXIT_FAILURE; 
    }
    
    if (tauthors > A) {
		tauthors = A; 
    }
    mapid2str::iterator it; 
    
    for (int k = 0; k < K; k++) {
		vector<pair<int, double> > authors_probs; 
		pair<int, double> author_prob; 
		for (int a = 0; a < A; a++) {
			author_prob.first = a; 
			author_prob.second = theta[a][k]; 
			authors_probs.push_back(author_prob); 
		}
    
        // quick sort to sort word-topic probability
		utils::quicksort(authors_probs, 0, authors_probs.size() - 1); 
	
		fprintf(fout, "Topic %dth:\n", k); 
		for (int i = 0; i < tauthors; i++) {
			it = id2author.find(authors_probs[i].first); 
			if (it != id2author.end()) {
				fprintf(fout, "\t%s   %g\n", (it->second).c_str(), authors_probs[i].second); 
			}
		}
    }
    
    fclose(fout); 
    
    return EXIT_SUCCESS; 
}

int model::init_est(void) {
	pt = new double [K]; 

    // read training data
    ptrndata = new dataset; 
    if (ptrndata->read_trndata(dir + dfile, dir + wordmapfile, dir + authormapfile) == EXIT_FAILURE) {
		utils::print_error("Fail to read training data!\n"); 
        return EXIT_FAILURE; 
    }

	// normalize the timestamps
	if (timestamp_norm() == EXIT_FAILURE) {
		utils::print_error("Fail to normalize the timestamps!\n"); 
		return EXIT_FAILURE; 
	}
		
    // allocate memory and assign values for variables
    M = ptrndata->M; 
    V = ptrndata->V; 
	A = ptrndata->A; 

	// nw[i][j]: number of instances of word/term i assigned to topic j, size V x K
    nw = new int * [V];	
    for (int v = 0; v < V; v++) {
        nw[v] = new int [K]; 
		memset(nw[v], 0, K * sizeof(int)); 
    }
	
	// na[i][j]: number of instances of topic j for author i, size A x K
    na = new int * [A]; 
    for (int a = 0; a < A; a++) {
        na[a] = new int [K]; 
		memset(na[a], 0, K * sizeof(int)); 
    }

	// nd[i][j]: number of words in document i assigned to topic j, size M x K
	nd = new int * [M]; 
	for (int m = 0; m < M; m++) {
		nd[m] = new int [K]; 
		memset(nd[m], 0, K * sizeof(int)); 
	}
	
	// nwsum[j]: total number of words assigned to topic j, size K
    nwsum = new int [K]; 
	memset(nwsum, 0, K * sizeof(int)); 
    
	// nasum[i]: total number of topics for author i, size A
    nasum = new int [A]; 
	memset(nasum, 0, A * sizeof(int)); 

	utils::set_seed(seed);		// initialize for random number generation
    z = new int * [M];
	x = new int * [M]; 
    for (m = 0; m < ptrndata->M; m++) {
		int Am = ptrndata->docs[m]->nauthors; 
		z[m] = new int [ptrndata->docs[m]->nwords]; 
		x[m] = new int [ptrndata->docs[m]->nwords]; 
	
        // initialize for z and x
        for (int n = 0; n < ptrndata->docs[m]->nwords; n++) {
   			int topic = (int) (utils::get_rand() * K); 
			z[m][n] = (topic == K? --topic: topic); 

			int author = (Am == 1? 0: (int) (utils::get_rand() * Am)); 
			x[m][n] = (author == Am? ptrndata->docs[m]->authors[--author]: ptrndata->docs[m]->authors[author]); 
    	    
    	    // number of instances of word i assigned to topic j
    	    nw[ptrndata->docs[m]->words[n]][topic]++;
    	    // number of words in document i assigned to topic j
			nd[m][topic]++; 
			// number of instances of topic j for author i
    	    na[x[m][n]][topic]++;
    	    // total number of words assigned to topic j
    	    nwsum[topic]++; 
			// total number of topics for author i
			nasum[x[m][n]]++; 
        } 
    }
    
	// allocate memory for theta
    theta = new double * [A];
    for (a = 0; a < A; a++) {
        theta[a] = new double [K];
    }
	
	// allocate memory for phi
    phi = new double * [K];
    for (int k = 0; k < K; k++) {
        phi[k] = new double [V];
    }

	// allocate memory for psi
	psi = new double * [K]; 
	for (k = 0; k < K; k++) {
		psi[k] = new double [2]; 
	}
	compute_psi(); 

	// allocate memory for pbeta
	pbeta = new double * [M]; 
	for (m = 0; m < M; m++) {
		pbeta[m] = new double [K]; 
	}
    
    return EXIT_SUCCESS;
}

void model::update_beta(void) {
	for (int m = 0; m < M; m++) {
		for (int k = 0; k < K; k++) {
			pbeta[m][k] = utils::beta_pdf(ptrndata->docs[m]->timestamp, psi[k][0], psi[k][1]); 
		}
	}
}

void model::estimate(void) {
	if (twords > 0) {
		// print out top words per topic
		dataset::read_wordmap(dir + wordmapfile, &id2word); 
    }

	if (tauthors > 0) {
		// print out top authors per topic
		dataset::read_authormap(dir + authormapfile, &id2author); 
	}

    fprintf(stdout, "Sampling %d iterations!\n", niters); 

    int last_iter = liter; 
    for (liter = last_iter + 1; liter <= niters + last_iter; liter++) {
		fprintf(stdout, "Iteration %d ...\n", liter); 

		// update the Beta PDF
		update_beta(); 
	
		// for all z_{m, n}
		for (int m = 0; m < M; m++) {
			for (int n = 0; n < ptrndata->docs[m]->nwords; n++) {
				// sample from P(z_{m, n}, x_{m, n} | z_{\neg m, n}, x_{\neg m, n}, w, a)
				sampling(m, n, z[m][n], x[m][n]); 
			}
		}

		compute_psi(); 
	
		if (savestep > 0) {
			if (liter % savestep == 0) {
				// saving the model
				fprintf(stdout, "Saving the model at iteration %d ...\n", liter); 
				compute_theta(); 
				compute_phi(); 
				save_model(utils::generate_model_name(liter)); 
			}
		}
    }
    
    fprintf(stdout, "Gibbs sampling completed!\n"); 
    fprintf(stdout, "Saving the final model!\n"); 
    compute_theta(); 
    compute_phi(); 
    liter--; 
    save_model(utils::generate_model_name(-1)); 
}

void model::sampling(int m, int n, int &topic, int &author) {
	int w = ptrndata->docs[m]->words[n]; 

	// remove z_{m, n}, x_{m, n} from the count variables
	nw[w][topic]--; 
	na[author][topic]--; 
	nd[m][topic]--; 
	nwsum[topic]--; 
	nasum[author]--; 

	double Vbeta = V * beta; 
    double Kalpha = K * alpha; 
	int Am = ptrndata->docs[m]->nauthors; 
	if (Am > 1) {	
		pa = new double [Am]; 

		// do multinomial sampling via cumulative method
		for (int a = 0; a < Am; a++) {
			pa[a] = (na[ptrndata->docs[m]->authors[a]][topic] + alpha) / (nasum[ptrndata->docs[m]->authors[a]] + Kalpha); 
		}
		// cumulate multinomial parameters
		for (a = 1; a < Am; a++) {
			pa[a] += pa[a - 1]; 
		}
		// scaled sample because of unnormalized pa[]
		double u = utils::get_rand() * pa[Am - 1]; 
  
		for (author = 0; author < Am; author++) {
			if (pa[author] > u) {
				break; 
			}
		}
		author = (author == Am? ptrndata->docs[m]->authors[--author]: ptrndata->docs[m]->authors[author]); 

		delete pa; 
		pa = NULL; 
	}

    
    // do multinomial sampling via cumulative method
    for (int k = 0; k < K; k++) {
		pt[k] = ((nw[w][k] + beta) / (nwsum[k] + Vbeta)) * ((na[author][k] + alpha) / (nasum[author] + Kalpha)) * pbeta[m][k]; 
    }
    // cumulate multinomial parameters
    for (k = 1; k < K; k++) {
		pt[k] += pt[k - 1]; 
    }
    // scaled sample because of unnormalized pt[]
	double u = utils::get_rand() * pt[K - 1]; 
    
    for (topic = 0; topic < K; topic++) {
		if (pt[topic] > u) {
			break; 
		}
    }
	if (topic == K) {
		topic--; 
	}
    
    // add newly estimated z_{m, n}, x_{m, n} to count variables
    nw[w][topic]++; 
    na[author][topic]++; 
	nd[m][topic]++; 
    nwsum[topic]++; 
    nasum[author]++; 
}

void model::compute_theta(void) {
	double Kalpha = K * alpha; 

    for (int a = 0; a < A; a++) {
		for (int k = 0; k < K; k++) {
			theta[a][k] = (na[a][k] + alpha) / (nasum[a] + Kalpha);
		}
    }
}

void model::compute_phi(void) {
	double Vbeta = V * beta; 

    for (int k = 0; k < K; k++) {
		for (int w = 0; w < V; w++) {
			phi[k][w] = (nw[w][k] + beta) / (nwsum[k] + Vbeta);
		}
    }
}

void model::compute_psi(void) {
	double Malpha = M * alpha; 

	for (int k = 0; k < K; k++) {
		double mean = 0.0; 
		for (int m = 0; m < M; m++) {
			mean += (nd[m][k] + alpha) * ptrndata->docs[m]->timestamp; 
		}
		mean /= (nwsum[k] + Malpha); 

		double variance = 0.0; 
		for (m = 0; m < M; m++) {
			variance += (nd[m][k] + alpha) * (ptrndata->docs[m]->timestamp - mean) * (ptrndata->docs[m]->timestamp - mean); 
		}
		variance /= (nwsum[k] + Malpha); 

		double tmp = mean * (1.0 - mean) / variance -1.0; 
		psi[k][0] = mean * tmp; 
		psi[k][1] = tmp - psi[k][0]; 
	}
}