/*
    Genome-wide Efficient Mixed Model Association (GEMMA)
    Copyright (C) 2011-2017, Xiang Zhou

    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 3 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, see <http://www.gnu.org/licenses/>.
*/

#ifndef __VC_H__
#define __VC_H__

#include "gsl/gsl_vector.h"
#include "gsl/gsl_matrix.h"
#include "param.h"
#include "io.h"

using namespace std;

class VC_PARAM {

public:
	const gsl_matrix *K;
	const gsl_matrix *W;
	const gsl_vector *y;
	gsl_matrix *P;
	gsl_vector *Py;
	gsl_matrix *KPy_mat;
	gsl_matrix *PKPy_mat;
	gsl_matrix *Hessian;
	bool noconstrain;
};

class VC {

public:
  // IO-related parameters
  size_t a_mode;
	string file_cat;
	string file_beta;
	string file_cor;
	string file_mq;
	string file_ms;

	string file_out;
	string path_out;

	set<string> setSnps;

	size_t ni_total_ref, ns_total_ref, ns_pair;
	size_t ni_total, ns_total, ns_test;
	size_t n_vc;

	double pve_total, se_pve_total;
	vector<double> v_sigma2;
	vector<double> v_se_sigma2;
	vector<double> v_pve;
	vector<double> v_se_pve;
	vector<double> v_traceG;
	vector<double> v_beta;
	vector<double> v_se_beta;

	size_t crt;
	double window_cm, window_bp, window_ns;

	double time_UtX;
	double time_opt;

	// Main functions.
	void CopyFromParam (PARAM &cPar);
	void CopyToParam (PARAM &cPar);
	void WriteFile_qs (const gsl_vector *s_vec, const gsl_vector *q_vec,
			   const gsl_vector *qvar_vec, const gsl_matrix *S_mat,
			   const gsl_matrix *Svar_mat);
	void CalcVChe (const gsl_matrix *K, const gsl_matrix *W,
		       const gsl_vector *y);
	void CalcVCreml (const bool noconstrain, const gsl_matrix *K,
			 const gsl_matrix *W, const gsl_vector *y);
	void CalcVCacl (const gsl_matrix *K, const gsl_matrix *W,
			const gsl_vector *y);
};

void CalcVCss(const gsl_matrix *Vq, const gsl_matrix *S_mat,
	      const gsl_matrix *Svar_mat, const gsl_vector *q_vec,
	      const gsl_vector *s_vec, const double df, vector<double> &v_pve,
	      vector<double> &v_se_pve, double &pve_total,
	      double &se_pve_total, vector<double> &v_sigma2,
	      vector<double> &v_se_sigma2, vector<double> &v_enrich,
	      vector<double> &v_se_enrich);

bool BimbamXwz (const string &file_geno, const int display_pace,
		vector<int> &indicator_idv, vector<int> &indicator_snp,
		const vector<size_t> &vec_cat, const gsl_vector *w,
		const gsl_vector *z, size_t ns_test, gsl_matrix *XWz);
bool PlinkXwz (const string &file_bed, const int display_pace,
	       vector<int> &indicator_idv, vector<int> &indicator_snp,
	       const vector<size_t> &vec_cat, const gsl_vector *w,
	       const gsl_vector *z, size_t ns_test, gsl_matrix *XWz);
bool MFILEXwz (const size_t mfile_mode, const string &file_mfile,
	       const int display_pace, vector<int> &indicator_idv,
	       vector<vector<int> > &mindicator_snp,
	       const vector<size_t> &vec_cat, const gsl_vector *w,
	       const gsl_vector *z, gsl_matrix *XWz);

bool BimbamXtXwz (const string &file_geno, const int display_pace,
		  vector<int> &indicator_idv, vector<int> &indicator_snp,
		  const gsl_matrix *XWz, size_t ns_test, gsl_matrix *XtXWz);
bool PlinkXtXwz (const string &file_bed, const int display_pace,
		 vector<int> &indicator_idv, vector<int> &indicator_snp,
		 const gsl_matrix *XWz, size_t ns_test, gsl_matrix *XtXWz);
bool MFILEXtXwz (const size_t mfile_mode, const string &file_mfile,
		 const int display_pace, vector<int> &indicator_idv,
		 vector<vector<int> > &mindicator_snp, const gsl_matrix *XWz,
		 gsl_matrix *XtXWz);

void CalcCIss(const gsl_matrix *Xz, const gsl_matrix *XWz,
	      const gsl_matrix *XtXWz, const gsl_matrix *S_mat,
	      const gsl_matrix *Svar_mat, const gsl_vector *w,
	      const gsl_vector *z, const gsl_vector *s_vec,
	      const vector<size_t> &vec_cat, const vector<double> &v_pve,
	      vector<double> &v_se_pve, double &pve_total,
	      double &se_pve_total, vector<double> &v_sigma2,
	      vector<double> &v_se_sigma2, vector<double> &v_enrich,
	      vector<double> &v_se_enrich);

#endif