#ifndef	_MODEL_H
#define	_MODEL_H

#include "constants.h"
#include "dataset.h"

using namespace std;

// LDA model
class model {
public:
    // fixed options
    string wordmapfile;		// file that contains word map [string -> integer id]
	string authormapfile;	// file that contains author map [string -> integer id]
    string txassign_suffix;	// suffix for topic-author assignment file
    string theta_suffix;	// suffix for theta file
    string phi_suffix;		// suffix for phi file
	string psi_suffix;		// suffix for psi file
    string others_suffix;	// suffix for file containing other parameters
    string twords_suffix;	// suffix for file containing words-per-topics
	string tauthors_suffix; // suffix for file containing authors-per-topics

    string dir;				// model directory
    string dfile;			// data file    
    string model_name;		// model name

    dataset *ptrndata;	// pointer to training dataset object

    mapid2str id2word; // word map [int => string]
	mapid2str id2author; // author map [int => string]
    
    // --- model parameters and variables ---    
    int M; // dataset size (i.e., number of docs)
    int V; // vocabulary size (i.e., number of unique words)
    int K; // number of topics
	int A; // number of unique authors
    double alpha, beta; // AToT hyperparameters 
	int seed;	// random number seed
    int niters; // number of Gibbs sampling iterations
    int liter; // the iteration at which the model was saved
    int savestep; // saving period
    int twords; // print out top words per each topic
	int tauthors; // print our top authors per each topic

    double *pt; // temp variable for sampling topic
	double *pa; // temp variable for sampling author
	double **pbeta; // temp varialbe for Beta PDF
    int **z; // topic assignments for words, size M x doc.size()
	int **x; // author assignments for authors, size M x doc.size()
    int **nw; // nw[i][j]: number of instances of word/term i assigned to topic j, size V x K
    int **na; // na[i][j]: number of instances of topic j for author i, size A x K
	int **nd; // nd[i][j]: number of words in document i assigned to topic j, size M x K
    int *nwsum; // nwsum[j]: total number of words assigned to topic j, size K
    int *nasum; // nasum[i]: total number of topics for author i, size A
    double **theta; // theta: author-topic distributions, size A x K
    double **phi; // phi: topic-word distributions, size K x V
	double **psi; // psi: Beta distributions of time, size K x 2
    
	// constructor
    model() {
		set_default_values(); 
    }
    
	// destructor
    ~model(); 
    
    // initialize the model
    int init(void);

    // estimate AT model using Gibbs sampling
    void estimate(void); 
    
private: 
	// set default values for variables
    void set_default_values(void); 

    // init for estimation
    int init_est(void); 

	// normalize the timestamp
	int timestamp_norm(void); 

	// update the pbeta
	void update_beta(void); 

	// save AToT model to files
	// model_name.tassign: topic assignments for words in docs
	// model_name.xassign: author assignments for words in docs
    // model_name.theta: author-topic distributions
    // model_name.phi: topic-word distributions
    // model_name.others: containing other parameters of the model (alpha, beta, M, V, K, A)
	// model_name.twords: top nwords topic words for each topic
	// model_name.tauthors: top nauthors authors for each topic
	int save_model(string model_name); 
	int save_model_txassign(string filename); 
    int save_model_theta(string filename); 
    int save_model_phi(string filename); 
	int save_model_psi(string filename); 
    int save_model_others(string filename); 
    int save_model_twords(string filename); 
	int save_model_tauthors(string filename); 

	// estimate AT model using Gibbs sampling
	void sampling(int m, int n, int &topic, int &author); 
    void compute_theta(void); 
    void compute_phi(void); 
	void compute_psi(void); 
};

#endif

