/****************************************************************************
 *                                                                          *
 *     Loki - Programs for genetic analysis of complex traits using MCMC    *
 *                                                                          *
 *                   Simon Heath - MSKCC                                    *
 *                                                                          *
 *                       August 2000                                        *
 *                                                                          *
 * calc_var_locus.c:                                                        *
 *                                                                          *
 * Calculate variance contributed by a trait locus                          *
 *                                                                          *
 * 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 <stdlib.h>
#ifdef USE_DMALLOC
#include <dmalloc.h>
#endif
#include <stdio.h>

#include "utils.h"
#include "loki.h"
#include "mat_utils.h"
#include "calc_var_locus.h"

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "calc_var_locus"
void calc_var_locus(int locus)
{
	int i,j,k,mod,mod1;
	double z,z1;
	int *gt;
	static double *x,*mu,*n;
	unsigned long mflag=0,a,b;
	
	if(locus==n_markers) {
		if(x) free(x);
		return;
	}
	k=n_models*(n_models+1)/2;
	if(!x) {
		if(!(x=malloc(sizeof(double)*(n_models+2*k)))) ABT_FUNC(MMsg);
		n=x+k;
		mu=n+k;
	}
	if(locus<0) {
		gt=tlocus[-1-locus].locus.gt;
		mflag=tlocus[-1-locus].model_flag;
	} else gt=marker[locus].locus.gt;
	for(i=0;i<k;i++) x[i]=n[i]=0.0;
	for(i=0;i<n_models;i++) mu[i]=0.0;
	if(locus<0) {
		for(i=0;i<ped_size;i++) {
			k=gt[i]-1;
			for(a=1,mod=0;mod<n_models;mod++,a<<=1) if((mflag&a)&&id_array[i].res[mod]) {
				if(k) mu[mod]+=tlocus[-1-locus].eff[mod][k-1];
				n[mod]++;
			}
		}
		for(mod=0;mod<n_models;mod++) {
			mu[mod]=n[mod]>0.0?mu[mod]/n[mod]:0.0;
			n[mod]=0.0;
		}
		for(i=0;i<ped_size;i++) {
			k=gt[i]-1;
			for(a=1,mod=0;mod<n_models;mod++,a<<=1) if((mflag&a)&&id_array[i].res[mod]) {
				z=k?tlocus[-1-locus].eff[mod][k-1]:0.0;
				z-=mu[mod];
				for(b=1,mod1=0;mod1<=mod;mod1++,b<<=1) if((mflag&b)&&id_array[i].res[mod1]) {
					z1=k?tlocus[-1-locus].eff[mod1][k-1]:0.0;
					z1-=mu[mod];
					j=IDX(mod,mod1);
					n[j]++;
					x[j]+=z*z1;
				}
			}
		}
	} else {
 		for(i=0;i<ped_size;i++) {
			k=gt[i]-1;
			for(mod=0;mod<n_models;mod++) if(marker[locus].mterm[mod]&&id_array[i].res[mod]) {
				if(k) mu[mod]+=marker[locus].mterm[mod]->eff[k-1];
				n[mod]++;
			}
		}
		for(mod=0;mod<n_models;mod++) {
			mu[mod]=n[mod]>0.0?mu[mod]/n[mod]:0.0;
			n[mod]=0.0;
		}
 		for(i=0;i<ped_size;i++) {
			k=gt[i]-1;
			for(mod=0;mod<n_models;mod++) if(marker[locus].mterm[mod]&&id_array[i].res[mod]) {
				z=k?marker[locus].mterm[mod]->eff[k-1]:0.0;
				z-=mu[mod];
				for(mod1=0;mod1<=mod;mod1++) if(marker[locus].mterm[mod1]&&id_array[i].res[mod1]) {
					z1=k?marker[locus].mterm[mod1]->eff[k-1]:0.0;
					z1-=mu[mod1];
					j=IDX(mod,mod1);
					n[j]++;
					x[j]+=z*z1;
				}
			}
		}
	}
	for(j=mod=0;mod<n_models;mod++) {
		for(mod1=0;mod1<=mod;mod1++) {
			z=n[j];
			if(z>0.0) x[j]/=z;
			if(x[j]<1.0e-16) x[j]=0.0;
			j++;
		}
	}
	k=n_models*(n_models+1)/2;
	if(locus<0) for(i=0;i<k;i++) tlocus[-1-locus].locus.variance[i]=x[i];
	else for(i=0;i<k;i++) marker[locus].locus.variance[i]=x[i];
}

