/***************************************************************************
 *   Copyright (C) 2007 by Faubet Pierre   *
 *   pierre.faubet@e.ujf-grenoble.fr   *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

#ifndef MCMC_H
#define MCMC_H

// #include <QReadWriteLock>

#include <fstream>
#include <string>

#include <gsl/gsl_linalg.h>
#include <gsl/gsl_rng.h>

// #include "mtrand.h"
#include "multilocus_genotypes.h"
#include "environmental_factors.h"

using namespace std;

/**
 * The MCMC class implements the mcmc scheme for the inference of migration rates
 * @author Pierre Faubet <pierre.faubet@e.ujf-grenoble.fr>
 */
class MCMC
  {

    friend class MCMCThread;

  public:
    /**
    * Constructor
     */
    MCMC();

    ~MCMC();

    static void setMCMCData();

    static void setMCMCData(const string& MGfilename,const string& EFfilename=string(),bool interaction=false);

    static void readMCMCSettings(const string& inputfilename);

    static void writeMCMCParameters(const string& outpufilename);

    /**
     * Initialization of the mcmc run
     */
    void init();

    /**
     * Mcmc updates
     */
    void update();

    /**
     * Model update
     */
    void RJupdate();

    /**
     * Make pilot runs
     */
    void pilot_run();

    /**
     * Make a mcmc run
     */
    void run();

    /**
     * Print posterior estimates in output files
     */
    void print();

    /**
     * Write raw data in outputfiles
     * @param iteration sampled MCMC iteration
     */
    void write_iter ( int iteration );

    /**
     * The number of mcmc runs
     * @return the number of mcmc that were created
     */
    static int get_nb_chain()
    {
      return nb_mcmc;
    }

    static int getI()
    {
      return genodata->get_nbpop();
    }

    static int getJ()
    {
      return genodata->get_nbloci();
    }

    static int getn()
    {
      return genodata->get_nbindiv();
    }

    static int getR()
    {
      return factors->get_nb_factors();
    }

    static int getwithinter()
    {
      return factors->get_with_interaction();
    }

    static bool is_a_valid_model(int M)
    {
      return factors->is_a_valid_model(M);
    }

    /**
     * Burnin
     * @return the length of the burnin period
     */
    static int get_burnin()
    {
      return burnin;
    }

    /**
     * Sample size
     * @return the size of the sample
     */
    static int get_samplesize()
    {
      return samplesize;
    }

    /**
     * Thining interval
     * @return the thining interval
     */
    static int get_thining()
    {
      return thining;
    }

    /**
    * Acceptances rates
    */
    double psi_acc; /**< Acceptance rate for the paramters of the Dirichlet prior for migration rates */

    double nm_acc; /**< Acceptance rate for non migrant proportions */
    double m_acc; /**< Acceptance rate for migration rate matrix */
    double Mt_acc; /**< Acceptance rate for assignments */
    double F_acc; /**< Acceptance rate for inbreeding coefficient vector */
    double p_acc; /**< Acceptance rate for subpopulation allele frequencies */

    double pglob_acc; /**< Acceptance rate for the ancestral population allele frequencies */
    double theta_acc; /**< Acceptance rate for subpopulation distances from the ancestral population */

    double RJ_acc; /**< Acceptance rate for between model moves */

  private:
    /* Static members */
    static bool useFmodel; /**< Use correlated allele frequency model */
    static bool useReg; /**< Infer environmental factors influence */
    static bool useRJMCMC; /**< Identify environmental factors */
    static bool useinter;

    /* Output options */
    static bool printLglik;
    static bool printPsi;
    static bool printMig;
    static bool printReg;
    static bool printP;
    static bool printF;
    static bool printFst;
    static bool printPglob;
    static bool printSd;
    static bool printpopq;
    static bool printindivq;
    static bool printmtild;

    /**
     * Prior parameters
     */
    static double atau; /**< shape parameter for s2*/
    static double btau; /**< scale parameter for s2*/
    static double s2_alpha; /**< Variance for alpha */

    static double Psi; /**< Shape parameters for migration rates */

    static double lambda; /**< Shape parameters for allele frequencies */
    static double omega; /**< Mean for theta */
    static double xi; /**< Variance for theta */
    static int nb_mcmc; /**< Number of mcmc */

    static string EFinputfile;
    static string MGinputfile;

    /**
    * Genetic and environmental data
    */
    static multilocus_genotypes *genodata; /**< multilocus genotype data */
    static environmental_factors *factors; /**< environmental data */

    /**
     * MCMC settings: iterations and incremental values
     */
    static int samplesize; /**< Sample size */
    static int burnin; /**< Burnin period */
    static int thining; /**< Thining interval */

    static double e_F_init; /**< Incremental value for inbreeding coefficient proposal */
    static double e_nm_init; /**< Incremental value for non migrant proportions */
    static double e_m_init; /**< Incremental value for migration rate proposal */
    static double e_p_init; /**< Incremental value for allele frequency proposal */
    static double s2_psi_init; /**< Incremental value for the parameters of the Dirichlet prior for migration rates proposal */
    static double s2_theta_init;/**< Incremetal value for theta */

    double e_F; /**< Incremental value for inbreeding coefficient proposal */
    double e_nm; /**< Incremental value for non migrant proportions */
    double e_m; /**< Incremental value for migration rate proposal */
    double e_p; /**< Incremental value for allele frequency proposal */
    double s2_psi; /**< Incremental value for the parameters of the Dirichlet prior for migration rates proposal */
    double s2_theta;/**< Incremetal value for theta */

    static int nb_pilot; /**< Number of pilot runs */
    static int pilot_length; /**< Length of each pilot run */

//     static QReadWriteLock lock ;

    /**
     * Posterior estimates
     */
    double **m_mean; /**< Posterior mean of migration rates */
    double **m_sd; /**< Standard deviation of migration rates */
    int ***post_assign; /**< Posterior probibabilities of assignments */
    double *F_mean; /**< Posterior mean of inbreeding coefficients */
    double *F_sd; /**< Standard deviation of inbreeding coefficients */
    double ***p_mean; /**< Posterior mean of allele frequencies */
    double ***p_sd; /**< Standard deviation of allele frequencies */

    double d_assign; /**< Deviance */

    int nb_model; /**< Number of models */
    int model; /**< Current model */
    double *visit; /**< Posterior probabilities of models */

    /* Members */
    int idf; /**< Mcmc run id */
    string fileprefix;

    int nb_reg_param; /**< Number of regression parameters */


    /**
     * Parameters to estimate
     */
    double *alpha; /**< Regression parameters */
    double s2; /**< Regression fit */
    double **mu; /**< Regression */
    double **psi; /**< Parameter of the Dirichlet distribution for migration rates */

    double **m; /**< Migration rate matrix */
    int **Mt; /**< Individual migrant ancestries */
    int ***ancescount; /**< Number of individuals assigned to each ancestry class */
    double *F; /**< Inbreeding coefficient vector */
    double ***p; /**< Subopulation allele frequencies */

    double *theta; /**< Subpopulation distances from the ancestral population */
    double **glob_all_frq; /**< Ancestral population allele frequencies */

    double *freqmin;

    /**
     * Log-likelihoods
     */
    double *tmp_ind_lglik;
    double *ind_lglik;
    double lglikgeno; /**< Pr(X|S;M,t,p,F) */
    double *pop_lglik;
    double lglikMt; /**< Pr(M,t|m) */

    double lgMV;


    /**
    * Output files
    */
    ofstream resultfile; /**< Results, summary and posterior estimates */
    ofstream mcmcfile; /** Migration rate and Dirichlet hyperparameters */
    ofstream regfile; /**< Regression and model */
    ofstream Fisfile; /**< F-statistics: Fis */
    ofstream fstfile; /**< F-statistics: Fst*/
    ofstream frqfile; /**< Ancestral allele frequencies */
    ofstream pfile; /**< Population allele frequencies */
    ofstream propfile; /**< Assignment proportions */
    ofstream loglikfile; /**< Loglikelihood */
    ofstream psifile; /**< Shape parameters for the Dirichlet prior */
    ofstream mtildfile; /**< Individual migration rates */

    //MTRand generator; /**< Random number generator */
    gsl_rng *rng; /**< Random number generator */

    /**
    * Create output files
     */
    void create_output_files();

    /**
    * Update ancestral population allele frequencies
    * Metropolis-Hastings
    */
    void update_glob_all_frq();

    /**
    * Update subpopulation distances from the ancestral population
    * Metropolis-Hastings
    */
    void update_theta();

    /**
    * Multivariate move for regression parameters
    * @return the posterior loglikelihood f(alpha|...)
    */
    double MVregupdate();

    /**
    * Update regression parameters for models with only one regressor
    * Gibbs sampling scheme
    * @return the posterior loglikelihood f(alpha|...)
    */
    double update_alpha();

    /**
    * Update regression
    * Gibbs sampling scheme
    */
    void update_s2();

    /**
    * Update parameters of the Dirichlet prior for migration rates
    * Metropolis-Hastings
    */
    void update_psi();

    /**
    * Update migration rate matrix when not using environmental factors
    * Gibbs
    */
    void update_migration_rate_matrix();

    /**
    * Update non migrant proportions
    */
    void update_non_migrant_proportions();

    /**
    * Update migration rate matrix
    * Gibbs sampling scheme
    */
    void update_m();

    /**
    * Update assignments
    * Metropolis-Hastings or Gibbs
    */
    void update_Mt();

    /**
    * Update inbreeding coefficient vector
    * Metropolis-Hastings
    */
    void update_F();

    /**
    * Update subpopulation allele frequencies
    * Metropolis-Hastings
    */
    void update_p();

    /**
    * Initialization
    */
    void init_glob_all_frq();
    void init_theta();

    void init_alpha();
    void init_s2();
    void init_psi();

    void init_m();
    void init_Mt();
    void init_p();
    void init_F();

    void init_lglik();

    /**
    * Create posterior estimates
    */
    void init_stat();

    /**
    * Compute individual likelihood
    * @param h indiviual
    */
    void indiv_loglik ( int h );

    /**
    * Compute P(M,t|m)
    */
    void loglikMt();

    /**
    * Compute Pr(X|S;M,t,p,F)
    */
    void loglikgeno();

    /**
    * Compute the log-likelihood of individual h's genotype at locus j without admixture
    * @param h indiviual
    * @param j locus
    * @return The log-likelihood of individual h's genotype at locus j given its ancestry class
    */
    double phi ( int h,int j );

    /**
    * Compute the log-likelihood of individual h's genotype at locus j with admixture
    * @param h indiviual
    * @param j locus
    * @return The log-likelihood of individual h's genotype at locus j given its ancestry class
    */
    double Phi ( int h,int j );

    /**
    * Compute Pr(M,t|m) for population q
    * @param q population
    */
    void pop_loglik ( int q );

    /**
    * Create output files
    */
    void create_psi_outputfile();
    void create_F_outputfile();
    void create_prop_file();
    void create_p_file();
    void create_lglk_file();
    void create_mtild_outputfile ();
    void create_mcmc_outputfile();
    void create_frq_outputfile();
    void create_fst_outputfile();
    void create_reg_outputfile();

    void write_parameters();

    /**
    * Print acceptance rates
    */
    void print_acceptance_rates ( int nbit );

    /**
    * Compute posterior estimates
    */
    void poststats();

    /**
    * multivariate normal distribution draw
    * @param omega Mean
    * @param varcovar Variance/covariance matrix
    * @param vect a random draw
    */
    void MVNormaldev ( const gsl_matrix *omega,const gsl_matrix *varcovar, gsl_matrix *vect );

    /**
     *   Migration rate update when considering only two populations
     */
    void update_m_2pop();
    void close_outputfiles();
  };

#endif
