/****************************************************************************
 *                                                                          *
 *     Loki - Programs for genetic analysis of complex traits using MCMC    *
 *                                                                          *
 *             Simon Heath - University of Washington                       *
 *                                                                          *
 *                      March/April 1997                                    *
 *                                                                          *
 * calc_nrm.c:                                                              *
 *                                                                          *
 * Calculate Inverse of NRM matrix using algorithm of Quaas (1976)          *
 *                                                                          *
 * 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 <errno.h>
#include <sys/wait.h>

#include "utils.h"
#include "libhdr.h"
#include "scan.h"
#include "sparse.h"

static double *u;

int cmp_inbr(const void *s1,const void *s2)
{
	double x1,x2;
	int i;
	
	i=*((const int *)s1);
	x1=u[i];
	i=*((const int *)s2);
	x2=u[i];
	if(x1<x2) return 1;
	if(x1>x2) return -1;
	return 0;
}

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "Calculate_NRM"
/* Calculate inverse of NRM (G) matrix using Quaas/Henderson method 
 * See Numerical Recipes in C 2nd edition pps. 78-79 for a description 
 * of the sparse matrix storage. Note the matrix is half stored. */
int Calculate_NRM(char *LogFile)
{
	int id,i,j,k,k1,idd,ids,idd1,ids1,nz=0,comp,max_comp=0,n_inbr;
	char *fname,*tname;
	int *sire_list,*dam_list,*temp,*spouses=0,*perm;
	FILE *flog=0,*fptr=0;
	double d,*v,xx,d2,d4,Detl,TDetl,avg_inbr;
	struct SparseMatRec *AIMatrix=0;
	struct model *model;
	struct model_list *mlist;

	fname=make_file_name(".nrm");
	if(unlink(fname)<0) if(errno!=ENOENT) {
		perror("Calculate_NRM");
		(void)fprintf(stderr,"Can't remove file %s.\n",fname);
		free(fname);
		return -1;
	}
	if(!pruned_ped_size)	{
		free(fname);
		return 0;
	}
	/* Do we need an NRM ? */
	model=Models;
	while(model) {
		mlist=model->model_list;
		while(mlist) {
			for(i=0;i<mlist->nvar;i++) if(mlist->element[i]->type&ST_ID) break;
			if(i<mlist->nvar) break;
			mlist=mlist->next;
		}
		if(mlist) break;
		model=model->next;
	}
	/* Find maximum component size */
	for(i=0;i<n_comp;i++) if(comp_size[i]>max_comp) max_comp=comp_size[i];
	if(!(sire_list=calloc((size_t)(pruned_ped_size+max_comp*2),sizeof(int)))) ABT_FUNC(MMsg);
	dam_list=sire_list+max_comp;
	perm=dam_list+max_comp;
	if(!(u=malloc(max_comp*2*sizeof(double)))) ABT_FUNC(MMsg);
	v=u+max_comp;
	for(i=0;i<ped_size;i++)	{
		j=ped_recode1[i];
		if(j) perm[j-1]=i;
	}
	TDetl=0.0;
	if(LogFile && (tname=add_file_dir(LogFile))) {
		flog=fopen(tname,"a");
		free(tname);
	}
	if(flog)	{
		if(model) (void)fputs("\n**************** Calculation of NRM inverse *************\n\n",flog);
		else (void)fputs("\n********* Calculation of inbreeding coefficients ********\n\n",flog);
		(void)fprintf(flog,"     Pedigree size = %d\n\n",pruned_ped_size);
	}
	if(model) {
		errno=0;
		if(Filter) {
			i=child_open(WRITE,fname,Filter);
			if(!(fptr=fdopen(i,"w"))) DataFileError(fname);
			if(errno && errno!=ESPIPE) DataFileError(fname);
			errno=0;
		} else if(!(fptr=fopen(fname,"w"))) abt(__FILE__,__LINE__,"%s(): File Error.  Couldn't open '%s' for writing.\n",FUNC_NAME,fname);
		if(fprintf(fptr,"Loki.nrm:%x,%x,%x\n",RunID,pruned_ped_size,n_comp)<0) DataFileError(fname);
		(void)fputs("Generating inverse of NRM matrix\n",stdout);
	}
	/* Do the calculations one component at a time */
	for(id=comp=0;comp<n_comp;comp++) {
		(void)fflush(stdout);
		/* Calculate storage requirements */
		if(!(temp=calloc((size_t)comp_size[comp]+1,sizeof(int)))) ABT_FUNC(MMsg);
		/* Count non-zero (off diagonal) elements in matrix */
		for(j=k=0;j<comp_size[comp];j++)	{
			i=perm[j+id];
			ids=id_array[i].sire;
			idd=id_array[i].dam;
			if(ids) ids=ped_recode1[ids-1]-id;
			if(idd) idd=ped_recode1[idd-1]-id;
			sire_list[j]=ids;
			dam_list[j]=idd;
			if(model && ids!=idd) {
				k++;
				if(ids>idd) temp[ids-1]++;
				else temp[idd-1]++;
			}
		}
		if(model) {
			nz=0;
			spouses=0;
			if(k) {
				if(!(spouses=calloc((size_t)k,sizeof(int)))) ABT_FUNC(MMsg);
				for(j=k=0;j<comp_size[comp];j++)	{
					k1=k;
					k+=temp[j];
					temp[j]=k1;
					ids=sire_list[j];
					idd=dam_list[j];
					if(ids && ids!=(j+1)) nz++;
					if(idd && idd!=ids && idd!=(j+1)) nz++;
					if(ids!=idd) {
						if(ids>idd)	{
							for(k1=temp[ids-1];k1<temp[ids];k1++) {
								if(spouses[k1]==idd) break;
								if(!spouses[k1]) {
									spouses[k1]=idd;
									nz++;
									break;
								}
							}
						} else {
							for(k1=temp[idd-1];k1<temp[idd];k1++) {
								if(spouses[k1]==ids) break;
								if(!spouses[k1]) {
									spouses[k1]=ids;
									nz++;
									break;
								}
							}
						}
					}
				}
			}
			if(!(AIMatrix=malloc((nz+1+comp_size[comp])*sizeof(struct SparseMatRec)))) ABT_FUNC(MMsg);
		}
		/* Zero inbreeding count */
		avg_inbr=0.0;
		n_inbr=0;
		if(model) {
			/* Set up coordinates for non-zero elements */
			nz=comp_size[comp]+1;
			for(j=0;j<comp_size[comp];j++) {
				u[j]=0.0;
				AIMatrix[j].val=0.0;
				AIMatrix[j].x=nz;
				ids=sire_list[j];
				idd=dam_list[j];
				if(ids && ids!=(j+1)) {
					AIMatrix[nz].x=ids-1;
					AIMatrix[nz++].val=0.0;
				}
				if(idd && idd!=ids && idd!=(j+1)) {
					AIMatrix[nz].x=idd-1;
					AIMatrix[nz++].val=0.0;
				}
				if(!spouses) continue;
				for(k=temp[j];k<temp[j+1];k++) {
					if(!spouses[k]) break;
					AIMatrix[nz].x=spouses[k]-1;
					AIMatrix[nz++].val=0.0;
				}
			}
			AIMatrix[comp_size[comp]].x=nz;
			if(spouses) free(spouses);
		} else for(j=0;j<comp_size[comp];j++) u[j]=0.0;
		Detl=0.0;
		for(i=0;i<comp_size[comp];i++) {
			xx=0.0;
			ids=sire_list[i]-1;
			idd=dam_list[i]-1;
			if(ids>=0) xx=u[ids];
			if(idd>=0) xx+=u[idd];
			xx=1.0-.25*xx;
			d=1.0/xx;
			u[i]+=xx;
			xx=sqrt(xx);
			v[i]=xx;
			Detl+=log(xx);
			for(k=i+1;k<comp_size[comp];k++) {
				xx=0.0;
				ids1=sire_list[k]-1;
				idd1=dam_list[k]-1;
				if(ids1>=i) xx=v[ids1];
				if(idd1>=i) xx+=v[idd1];
				if(xx>0.0) {
					xx=.5*xx;
					u[k]+=xx*xx;
				}
				v[k]=xx;
			}
			if(model) {
				AIMatrix[i].val+=d;
				d2= -.5*d;
				d4=.25*d;
				j=AIMatrix[i].x;
				if(ids>=0) {
					AIMatrix[j++].val+=d2;
					AIMatrix[ids].val+=d4;
				}
				if(idd>=0) {
					AIMatrix[j++].val+=d2;
					AIMatrix[idd].val+=d4;
					if(ids>=0) {
						if(ids==idd) AIMatrix[ids].val+=d4;
						else if(ids>idd) {
							for(j=AIMatrix[ids].x;j<AIMatrix[ids+1].x;j++) if(AIMatrix[j].x==idd) {
								AIMatrix[j].val+=d4;
								break;
							}
						} else {
							for(j=AIMatrix[idd].x;j<AIMatrix[idd+1].x;j++) if(AIMatrix[j].x==ids) {
								AIMatrix[j].val+=d4;
								break;
							}
						}
					}
				}
			}
		}
		for(j=0;j<comp_size[comp];j++) {
			if(u[j]>1.0) {
				xx=u[j]-1.0;
				avg_inbr+=xx;
				n_inbr++;
			}
		}
		Detl+=Detl;
		TDetl+=Detl;
		if(model) {
			AIMatrix[comp_size[comp]].val=Detl;
			if(fprintf(fptr,"LKNM:%x,%x\n",comp_size[comp],nz)<0) DataFileError(fname);
			for(j=0;j<nz;j++) {
				if(fprintf(fptr,"%x,",AIMatrix[j].x)<0) DataFileError(fname);
				if(txt_print_double(AIMatrix[j].val,fptr)) DataFileError(fname);
				if(fputc('\n',fptr)==EOF) DataFileError(fname);
			}
			free(AIMatrix);
		}
		if(flog)	{
			(void)fprintf(flog,"     Component %d\n",comp+1);
			(void)fprintf(flog,"       Component size = %d\n",comp_size[comp]);
			if(model) (void)fprintf(flog,"       No. non-zero off-diagonal elements of NRM = %d\n",nz-1-comp_size[comp]);
			(void)fprintf(flog,"       Log Determinant of NRM matrix = %f\n",Detl);
			(void)fprintf(flog,"       No. inbred individuals = %d\n",n_inbr);
			if(n_inbr) {
				(void)fprintf(flog,"       Average inbreeding coefficient (overall) = %g\n",avg_inbr/(double)comp_size[comp]);
				(void)fprintf(flog,"       Average inbreeding coefficient (inbred individuals only) = %g\n",avg_inbr/(double)n_inbr);
				(void)fputs("\n       Ids and inbreeding coefficients of inbred individuals:\n\n",flog);
				for(k=j=0;j<comp_size[comp];j++) if(u[j]>1.0) {
					temp[k++]=j;
				}
				gnu_qsort(temp,(size_t)k,sizeof(int),cmp_inbr); 
				for(j=0;j<n_inbr;j++) {
					k1=temp[j];
					xx=u[k1]-1.0;
					k=perm[id+k1];
					(void)fputs("          ",flog);
					print_orig_id(flog,k+1,0);
					(void)fprintf(flog," %g\n",xx);
				}
			}
		}
		free(temp);
		id+=comp_size[comp];
	}
	free(sire_list);
	free(u);
	if(model) {
		if(fputs("Lnrm.end\n",fptr)==EOF) DataFileError(fname);
		if(fclose(fptr)) DataFileError(fname);
	}
	free(fname);
	if(flog)	{
		if(n_comp>1) (void)fprintf(flog,"\n     Log Determinant of combined NRM matrix = %f\n",TDetl);
		(void)fclose(flog);
	}
	if(model && Filter) do i=wait(&j); while(i>0);
	return model?1:0;
}
