/****************************************************************************
 *                                                                          *
 *     Loki - Programs for genetic analysis of complex traits using MCMC    *
 *                                                                          *
 *             Simon Heath - University of Washington                       *
 *                                                                          *
 *                       July 1997                                          *
 *                                                                          *
 * alloc_loki.c:                                                            *
 *                                                                          *
 * 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"

int n_genetic_groups=1;
/* n_tloci refers to how much storage for trait loci we currently have, not
 * how many are actually in the model, nor is it the maximum possible */
int n_tloci=8;
struct TraitLocus *tlocus=0;

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "AllocLokiStruct"
void AllocLokiStruct(void)
{
	int i,j,k,*tp1,*tp2,n_all,**tipp,*ff;
#ifdef DEBUG
	int k2;
#endif
	
	lk_ulong **tlpp,*tlp;
	double **tpp,*tp;
	signed char *tcp,**tcpp;
	struct Locus *loc;
	
	if(n_markers) {
		for(i=0;i<n_markers;i++) {
			ff=founder_flag[i];
			for(j=0;j<ped_size;j++) {
				if(id_array[j].pruned_flag[i]) ff[j]=2;
				else {
					k=id_array[j].sire;
					if(k) {
						if(id_array[k-1].pruned_flag[i]) ff[j]=1;
						else ff[j]=0;
					} else ff[j]=1;
#ifdef DEBUG
					k=id_array[j].dam;
					if(k) {
						if(id_array[k-1].pruned_flag[i]) k2=1;
						else k2=0;
					} else k2=1;
					if(k2!=ff[j]) ABT_FUNC("Bad pruning - half pruned family\n");
#endif
				}
			}
		}
		if(!(tp1=malloc(sizeof(int)*n_markers*3*ped_size))) ABT_FUNC(MMsg);
		for(i=0;i<n_markers*3*ped_size;i++) tp1[i]=-1;
		for(i=0;i<n_markers;i++) {
			loc=&marker[i].locus;
			loc->seg[0]=tp1;
			loc->seg[1]=tp1+ped_size;
			marker[i].m_flag=tp1+2*ped_size;
			ff=founder_flag[i];
			for(j=0;j<ped_size;j++) {
				switch(ff[j]) {
				 case 0:
					loc->seg[0][j]=loc->seg[1][j]= -2;
					break;
				 case 1:
					loc->seg[0][j]=loc->seg[1][j]= -1;
					break;
				 default:
					loc->seg[0][j]=loc->seg[1][j]= -1;
					break;
				}
			}
			tp1+=3*ped_size;
		}
		for(j=i=0;i<n_markers;i++) j+=marker[i].locus.n_alleles;
		if(j) {
			if(!(tpp=malloc(sizeof(void *)*n_genetic_groups*n_markers))) ABT_FUNC(MMsg);
			if(!(tcpp=malloc(sizeof(void *)*n_genetic_groups*n_markers))) ABT_FUNC(MMsg);
			if(!(tipp=malloc(sizeof(void *)*n_markers*ped_size*2))) ABT_FUNC(MMsg);
			if(!(tlpp=malloc(sizeof(void *)*n_markers*ped_size))) ABT_FUNC(MMsg);
			if(!(tcp=malloc((size_t)j*n_genetic_groups*n_markers))) ABT_FUNC(MMsg);
			if(!(tp=malloc(sizeof(double)*j*n_genetic_groups*n_markers))) ABT_FUNC(MMsg);
			if(!(tp1=malloc(sizeof(int)*n_genetic_groups*n_markers))) ABT_FUNC(MMsg);
			if(!(tp2=malloc(sizeof(int)*5*ped_size*n_markers))) ABT_FUNC(MMsg);
			if(!(tlp=malloc(sizeof(lk_ulong)*2*ped_size*n_markers))) ABT_FUNC(MMsg);
			RemBlock=AddRemem(tpp,RemBlock);
			RemBlock=AddRemem(tipp,RemBlock);
			RemBlock=AddRemem(tlpp,RemBlock);
			RemBlock=AddRemem(tcpp,RemBlock);
			RemBlock=AddRemem(tcp,RemBlock);
			RemBlock=AddRemem(tlp,RemBlock);
			RemBlock=AddRemem(tp,RemBlock);
			RemBlock=AddRemem(tp1,RemBlock);
			RemBlock=AddRemem(tp2,RemBlock);
			for(i=0;i<n_markers;i++) {
				loc=&marker[i].locus;
				n_all=loc->n_alleles;
				loc->aff_freq=loc->diff_freq=0;
				if(!n_all) {
					loc->freq=0;
					marker[i].freq_set=0;
					marker[i].count_flag=0;
					continue;
				}
				loc->freq=tpp;
				marker[i].freq_set=tcpp;
				marker[i].count_flag=tp1;
				marker[i].nhaps=tipp;
				marker[i].lump=tipp+ped_size;
				marker[i].temp=tlpp;
				tipp+=2*ped_size;
				tlpp+=ped_size;
				marker[i].nhaps[0]=tp2;
				marker[i].lump[0]=tp2+2*ped_size;
				marker[i].ngens=tp2+4*ped_size;
				marker[i].temp[0]=tlp;
				tp2+=5*ped_size;
				tlp+=2*ped_size;
				for(j=1;j<ped_size;j++) {
					marker[i].nhaps[j]=marker[i].nhaps[j-1]+2;
					marker[i].lump[j]=marker[i].lump[j-1]+2;
					marker[i].temp[j]=marker[i].temp[j-1]+2;
				}
				tcpp+=n_genetic_groups;
				tpp+=n_genetic_groups;
				tp1+=n_genetic_groups;
				for(j=0;j<n_genetic_groups;j++) {
					marker[i].count_flag[j]=0;
					loc->freq[j]=tp;
					marker[i].freq_set[j]=tcp;
					tp+=n_all;
					tcp+=n_all;
					for(k=0;k<n_all;k++) {
						loc->freq[j][k]=0.0;
						marker[i].freq_set[j][k]=0;
					}
				}
			}
		} else for(i=0;i<n_markers;i++) {
			marker[i].locus.freq=0;
			marker[i].freq_set=0;
		}
		for(i=0;i<n_markers;i++) {
			marker[i].pos_set=0;
			marker[i].locus.pos[0]=marker[i].locus.pos[1]=0.0;
		}
	}
}

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "get_new_traitlocus"
int get_new_traitlocus(const int n_all)
{
	int i,j,k,*ff;
	double *tp,**tpp;
	struct Locus *loc;
	
	if(!tlocus) return -1;
	if(n_all<2) return -1;
	for(i=0;i<n_tloci;i++) if(!tlocus[i].locus.flag) break;
	if(i==n_tloci)	{
		if(i<max_tloci) {
			n_tloci*=1.5;
			if(n_tloci>max_tloci) n_tloci=max_tloci;
			if(!(tlocus=realloc(tlocus,sizeof(struct TraitLocus)*n_tloci))) ABT_FUNC(MMsg);
			k=n_tloci-i;
			for(j=i;j<n_tloci;j++) {
				loc=&tlocus[j].locus;
				loc->freq=0;
				loc->flag=0;
				loc->seg[0]=0;
				loc->lk_store=0;
				loc->variance=0;
				loc->gt=0;
				tlocus[j].eff=0;
				tlocus[j].model_flag=0;
			}
		} else return -1;
	}
	j=n_all*(n_all+1)/2-1;
	k=n_models*(n_models+1)/2;
	if(!(tpp=malloc(sizeof(void *)*(n_models+n_genetic_groups)))) ABT_FUNC(MMsg);
	if(!(tp=malloc(sizeof(double)*(j*n_models+k+n_genetic_groups*n_all)))) ABT_FUNC(MMsg);
	loc=&tlocus[i].locus;
	loc->freq=tpp;
	tlocus[i].eff=loc->freq+n_genetic_groups;
	loc->variance=tp;
	tp+=k;
	for(k=0;k<n_models;k++) {
		tlocus[i].eff[k]=tp;
		tp+=j;
	}
	for(k=0;k<n_genetic_groups;k++) loc->freq[k]=tp+n_all*k;
	loc->n_alleles=n_all;
	if(!(loc->seg[0])) {
		if(!(loc->seg[0]=calloc((size_t)4*ped_size,sizeof(int)))) ABT_FUNC(MMsg);
		loc->seg[1]=loc->seg[0]+ped_size;
		loc->genes[0]=loc->seg[1]+ped_size;
		loc->genes[1]=loc->genes[0]+ped_size;
	}
	ff=founder_flag[n_markers];
	for(j=0;j<ped_size;j++) {
		switch(ff[j]) {
		 case 0:
			loc->seg[0][j]=loc->seg[1][j]= -2;
			break;
		 case 1:
			loc->seg[0][j]=loc->seg[1][j]= -1;
			break;
		 default:
			loc->seg[0][j]=loc->seg[1][j]= -1;
			break;
		}
	}
	if(!loc->gt) {
		if(!(loc->gt=calloc((size_t)ped_size,sizeof(int)))) ABT_FUNC(MMsg);
	}
	if(!loc->lk_store) {
		if(!(loc->lk_store=malloc(sizeof(double)*n_comp))) ABT_FUNC(MMsg);
	}
	return i;
}

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "delete_traitlocus"
void delete_traitlocus(const int tl)
{
	struct Locus *loc;
	if(tl<0 || tl>=n_tloci) ABT_FUNC("Internal error - illegal trait locus\n");
	loc=&tlocus[tl].locus;
	loc->flag=0;
 	if(loc->variance) free(loc->variance);
	if(loc->freq) free(loc->freq);
	loc->freq=0;
	tlocus[tl].eff=0;
	loc->variance=0;
}

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "AllocEffects"
void AllocEffects(void)
{
	int i,j,k,k1,type,*tip,mod,n_all;
	struct Variable *var;
	double *tp;
	struct id_data *data;
	
	if(n_markers) {
		if(!(tip=calloc((size_t)n_markers*ped_size,sizeof(int)))) ABT_FUNC(MMsg);
		RemBlock=AddRemem(tip,RemBlock);
		for(i=0;i<n_markers;i++) {
			marker[i].locus.gt=tip;
			tip+=ped_size;
		}
	}
	/* Count amount of storage required for residuals, BV etc. */
	for(mod=0;mod<n_models;mod++) {
		j=0;
		type=models[mod].var.type;
		k=models[mod].var.var_index;
		if(models[mod].polygenic_flag) j+=ped_size; /* Storage for BV's */
		for(i=0;i<ped_size;i++) {
			if(type&ST_CONSTANT)	{
				if(id_array[i].data) data=id_array[i].data+k;
				else data=0;
				if(data && data->flag) {
					j+=1+use_student_t; /* Storage for residual + nu */
					if((type&ST_CENSORED)&&(data->flag&2)) j++; /* Storage for censored value */
				}
			} else if(id_array[i].n_rec) {
				j+=id_array[i].n_rec*(1+use_student_t); /* Storage for residual + nu */
				if(type&ST_CENSORED) j+=id_array[i].n_rec; /* Storage for censored value */
			}
		}
		if(j)	{
			if(!(tp=malloc(sizeof(double)*j))) ABT_FUNC(MMsg);
			RemBlock=AddRemem(tp,RemBlock);
			for(i=0;i<j;i++) tp[i]=0.0;
			for(i=0;i<ped_size;i++)	{	
				if(type&ST_CONSTANT)	{
					if(id_array[i].data) data=id_array[i].data+k;
					else data=0;
					if(data && data->flag) {
						id_array[i].res[mod]=tp++;
						if(use_student_t) id_array[i].vv[mod]=tp++;
						if((type&ST_CENSORED)&&(data->flag&2)) id_array[i].cens[mod]=tp++;
					}
				} else {
					if(id_array[i].n_rec) {
						id_array[i].res[mod]=tp;
						tp+=id_array[i].n_rec;
						if(use_student_t) {
							id_array[i].vv[mod]=tp;
							tp+=id_array[i].n_rec;
						}
						if(type&ST_CENSORED)	{
							id_array[i].cens[mod]=tp;
							tp+=id_array[i].n_rec;
						}
					}
				}
			}
		}
		for(k=i=0;i<models[mod].n_terms;i++) {
			models[mod].term[i].df=1;
			if(models[mod].term[i].n_vars>1) ABT_FUNC("Sorry - interaction terms not currently supported\n");
			type=models[mod].term[i].vars[0].type;
			j=models[mod].term[i].vars[0].var_index;
			if(type&ST_TRAITLOCUS) models[mod].term[i].df=2;
			else if(type&ST_ID) models[mod].term[i].df=ped_size;
			else if(type&ST_MARKER) {
				k1=marker[j].locus.n_alleles;
				models[mod].term[i].df=k1*(k1+1)/2-1;
				marker[j].mterm[mod]=models[mod].term+i;
				if(!(marker[j].locus.variance)) if(!(marker[j].locus.variance=malloc(sizeof(double)*n_models*(n_models+1)/2))) ABT_FUNC(MMsg);
			} else {
				if(type&ST_CONSTANT) var=id_variable+j;
				else var=nonid_variable+j;
				if(type&ST_FACTOR) {
					if(type&ST_RANDOM) models[mod].term[i].df=var->n_levels;
					else models[mod].term[i].df=var->n_levels-1;
				}
			}
			if(!(type&ST_TRAITLOCUS)) k+=models[mod].term[i].df;
		}
		if(!k) continue;
		if(!(tp=malloc(sizeof(double)*k))) ABT_FUNC(MMsg);
		RemBlock=AddRemem(tp,RemBlock);
		for(i=0;i<k;i++) tp[i]=0.0;
		for(i=0;i<models[mod].n_terms;i++) if(!(models[mod].term[i].vars[0].type&ST_TRAITLOCUS)) {
			models[mod].term[i].eff=tp;
			tp+=models[mod].term[i].df;
		}
	}
	if(est_aff_freq) {
		for(j=i=0;i<n_markers;i++) j+=marker[i].locus.n_alleles;
		if(j) {
			if(!(tp=malloc(sizeof(double)*j*2))) ABT_FUNC(MMsg);
			RemBlock=AddRemem(tp,RemBlock);
			for(i=0;i<j*2;i++) tp[i]=0.0;
			for(i=0;i<n_markers;i++) {
				n_all=marker[i].locus.n_alleles;
				marker[i].locus.aff_freq=tp;
				marker[i].locus.diff_freq=tp+n_all;
				tp+=2*n_all;
			}
		}
	}
}
