/****************************************************************************
 *                                                                          *
 *     Loki - Programs for genetic analysis of complex traits using MCMC    *
 *                                                                          *
 *                   Simon Heath - MSKCC                                    *
 *                                                                          *
 *                       August 2000                                        *
 *                                                                          *
 * sample_rand.c:                                                           *
 *                                                                          *
 * Sampling routine for uncorrelated random variance components             *
 *                                                                          *
 * Copyright (C) Simon C. Heath 1997, 2000, 2002                            *
 * This is free software.  You can distribute it and/or modify it           *
 * under the terms of the Modified BSD license, see the file COPYING        *
 *                                                                          *
 ****************************************************************************/

#include <config.h>
#include <stdlib.h>
#ifdef HAVE_UNISTD_H
#include <unistd.h>
#endif
#ifdef USE_DMALLOC
#include <dmalloc.h>
#endif
#include <math.h>
#include <stdio.h>
#include <float.h>
#include <errno.h>

#include "ranlib.h"
#include "utils.h"
#include "mat_utils.h"
#include "sparse.h"
#include "loki.h"
#include "sample_rand.h"

static double *work;
double **c_var;
unsigned long *rand_flag;
struct Variable **rand_list;
int n_random;

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "init_rand"
void init_rand(void)
{
	int k,i,j,type,mod;
	struct Variable *var;
	double *tmp;
	
	if(!n_models) return;
	for(n_random=i=0;i<n_id_records;i++) if(id_variable[i].type&ST_RANDOM) n_random++;
	for(i=0;i<n_nonid_records;i++) if(nonid_variable[i].type&ST_RANDOM) n_random++;
	if(!(rand_flag=malloc(sizeof(unsigned long)*n_models))) ABT_FUNC(MMsg);
	if(n_random) {
		if(!(rand_list=malloc(sizeof(struct Variable *)*n_random))) ABT_FUNC(MMsg);
		for(j=i=0;i<n_id_records;i++) if(id_variable[i].type&ST_RANDOM) rand_list[j++]=id_variable+i;
		for(i=0;i<n_nonid_records;i++) if(nonid_variable[i].type&ST_RANDOM) rand_list[j++]=nonid_variable+i;
		for(mod=0;mod<n_models;mod++) {
			rand_flag[mod]=0;
			for(k=0;k<models[mod].n_terms;k++) {
				type=models[mod].term[k].vars[0].type;
				if(type&ST_RANDOM) {
					i=models[mod].term[k].vars[0].var_index;
					var=(type&ST_CONSTANT)?id_variable+i:nonid_variable+i;
					for(i=0;i<n_random;i++) if(rand_list[i]==var) break;
					if(i==n_random) ABT_FUNC("Internal error - random variable not found\n");
					rand_flag[mod]|=1<<i;
				}
			}
		}
		k=n_models*(n_models+1)/2;
		if(!(c_var=malloc(sizeof(double *)*n_random))) ABT_FUNC(MMsg);
		if(!(tmp=malloc(sizeof(double)*k*n_random))) ABT_FUNC(MMsg);
		for(i=0;i<k*n_random;i++) tmp[i]=0.0;
		for(i=0;i<n_random;i++) {
			c_var[i]=tmp;
			for(j=0;j<n_models;j++) if(rand_flag[j]&(1<<i)) BB(tmp,j,j)=BB(residual_var,j,j);
			tmp+=k;
		}
	}
	if(polygenic_flag) if(!(work=malloc(sizeof(double)*ped_size))) ABT_FUNC(MMsg);
}

void free_rand(void)
{
	if(c_var) {
		if(c_var[0]) free(c_var[0]);
		free(c_var);
	}
	if(rand_list) free(rand_list);
	if(rand_flag) free(rand_flag);
	if(work) free(work);
	rand_flag=0;
	rand_list=0;
	c_var=0;
	work=0;
}

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "sample_rand"
void sample_rand(void)
{
	int i,k,r,type,df;
	double *eff,y,ss,n;
	
	for(r=k=0;k<models[0].n_terms && r<n_random;k++) {
		type=models[0].term[k].vars[0].type;
		if(type&ST_RANDOM) {
			if(!(type&ST_FACTOR)) ABT_FUNC("Can't fit random effect with continuous covariates\n");
			df=models[0].term[k].df;
			eff=models[0].term[k].eff;
			ss=0.0;
			for(i=0;i<df;i++) {
				y=eff[i];
				ss+=y*y;
			}
			n=(double)df+RES_PRIOR_VC0;
			ss+=RES_PRIOR_VC0*RES_PRIOR_SC0;
			c_var[r++][0]=ss/(sgamma(n*.5)*2.0);
		}
	}
}

void sample_additive_var(void)
{
	int i,id,j,k,comp;
	double y,n,z,ss;
	struct SparseMatRec *AI;
	
	ss=0.0;
	for(id=comp=0;comp<n_comp;comp++) {
		AI=AIMatrix[comp];
		for(i=0;i<comp_size[comp];i++) {
			y=id_array[i+id].bv[0];
			work[i]=y*AI[i].val;
			for(j=AI[i].x;j<AI[i+1].x;j++) {
				k=AI[j].x;
				z=AI[j].val;
				work[k]+=y*z;
				work[i]+=id_array[k+id].bv[0]*z;
			}
		}
		for(i=0;i<comp_size[comp];i++) ss+=id_array[i+id].bv[0]*work[i];
		id+=i;
	}
	n=(double)ped_size+RES_PRIOR_VA0;
	ss+=RES_PRIOR_VA0*RES_PRIOR_SA0;
	additive_var[0]=ss/(sgamma(n*.5)*2.0);
}
