/****************************************************************************
 *                                                                          *
 *     Loki - Programs for genetic analysis of complex traits using MCMC    *
 *                                                                          *
 *             Simon Heath - University of Washington                       *
 *                                                                          *
 *                      June 1997                                           *
 *                                                                          *
 * write_data.c:                                                            *
 *                                                                          *
 * Write out data files for use by Loki                                     *
 *                                                                          *
 * 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>
#include <string.h>
#ifdef USE_DMALLOC
#include <dmalloc.h>
#endif
#include <stdio.h>
#include <math.h>
#include <errno.h>
#include <sys/wait.h>
#ifdef HAVE_LIMITS_H
#include <limits.h>
#endif

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

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "get_element_index"
static int get_element_index(struct var_element *elem)
{
	int i,j;
	unsigned long tm;

	tm=ST_MODEL|ST_TRAIT;
	if(elem->type&(ST_ID|ST_DAM|ST_SIRE|ST_FAMILY|ST_TRAITLOCUS)) return 0;
	if(elem->type&ST_MARKER) {
		for(i=0;i<n_markers;i++) if(markers[i].element==elem) return i;
		ABT_FUNC("Internal error - couldn't find marker element\n");
	}
	if(elem->type&ST_CONSTANT) {
		for(j=i=0;i<n_id_records;i++) if(id_elements[i]->type&tm) {
			if(id_elements[i]==elem) return j;
			j++;
		}
		ABT_FUNC("Internal error - couldn't find id element\n");
	}
	for(j=i=0;i<n_nonid_records;i++) if(nonid_elements[i]->type&tm) {
		if(nonid_elements[i]==elem) return j;
		j++;
	}
	ABT_FUNC("Internal error - couldn't find non-id element\n");
	return -1;
}

#define DataFileError(a) New_DFE(__FILE__,__LINE__,a)
void New_DFE(const char *sfile,const int line,const char *file)
{
	(void)fprintf(stderr,"[%s:%d] Error writing to file '%s'\n",sfile,line,file);
	if(errno) perror("loki");
	exit(EXIT_FAILURE);
}

int print_id_data(struct id_data *s,int type,FILE *fptr)
{
	int er=0,c=0;
	
	if(fprintf(fptr,";%x",s->flag)<0) er=1;
	if(!er) {
		if(type&ST_FACTOR) c='I';
		else if(s->flag&ST_INTTYPE) c='I';
		else c='R';
		if(fputc(c,fptr)==EOF) er=1;
	}
	if(!er) {
		if(c=='I') {
			if(fprintf(fptr,"%lx",s->data.value)<0) er=1;
		} else er=txt_print_double(s->data.rvalue,fptr);
	}
	return er;
}

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "WriteData"
void WriteData(char *LogFile)
{
	int i,i1,j,k,k1,k2,k3,n_links,ids,idd,*perm,nrec=0,ntyped=0,n_mk;
	int id_rec_count=0,nonid_rec_count=0,n_orig_fam1,fam,type;
	int v2[14],*id_trans=0,*family_recode1=0,n_mod;
	char *fname;
	struct Link *linkp;
	struct model_list *mlist;
	struct model *model;
	struct var_element *elem;
	struct scan_data *sd;
	unsigned long tm;
	FILE *fptr;

	tm=ST_MODEL|ST_TRAIT;
	if(n_genetic_groups>1) tm|=ST_GROUP;
	for(i=0;i<n_id_records;i++) if(id_elements[i]->type&tm) id_rec_count++;
	for(i=0;i<n_nonid_records;i++) if(nonid_elements[i]->type&(ST_MODEL|ST_TRAIT)) nonid_rec_count++;
	if(id_rec_count) {
		if(!(id_trans=malloc(sizeof(int)*id_rec_count))) ABT_FUNC(MMsg);
		for(i=j=0;i<n_id_records;i++)	{
			if(id_elements[i]->type&(ST_MODEL|ST_TRAIT)) id_trans[j++]=i;
			else if(id_elements[i]->type&ST_GROUP && n_genetic_groups>1) id_trans[id_rec_count-1]=i;
		}
	}
	(void)fputs("Writing out data file...\n",stdout);
	errno=0;
	fname=make_file_name(".opt");
	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.opt:%x,%s,%s\n",RunID,LogFile?LogFile:"",Filter?Filter:"")<0) DataFileError(fname);
	if(fclose(fptr)) DataFileError(fname);
	free(fname);
	fname=make_file_name(".dat");
	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(family_id) {
		if(!(family_recode1=calloc((size_t)n_orig_families,sizeof(int)))) ABT_FUNC(MMsg);
		for(i=0;i<ped_size;i++) {
			j=ped_recode1[i];
			if(!j) continue;
			k=id_array[i].fam_code;
			if(k) family_recode1[k-1]=1;
		}
		n_orig_fam1=0;
		for(i=0;i<n_orig_families;i++) if(family_recode1[i]) family_recode1[i]=++n_orig_fam1;
	}
	if(!(perm=malloc(sizeof(int)*pruned_ped_size))) ABT_FUNC(MMsg);
	for(k2=i=0;i<ped_size;i++) {
		j=ped_recode1[i];
		if(!j) continue;
		perm[j-1]=i;
		j=n_genetic_groups>1?1:0;
		if(id_array[i].data && (id_rec_count-j)) k2+=id_rec_count-j;
		if(id_array[i].nrec && id_array[i].data1 && nonid_rec_count) {
			for(k1=0;k1<id_array[i].nrec;k1++) if(id_array[i].data1[k1])
			  for(k3=0;k3<n_nonid_records;k3++) 
				 if(nonid_elements[k3]->type&(ST_MODEL|ST_TRAIT)) {
					 k2++;
					 nrec++;
				 }
		}
		if(id_array[i].haplo[0]) ntyped++;
	}
	n_links=0;
	linkp=links;
	while(linkp) {
		n_links++;
		linkp=linkp->next;
	}
	for(n_mk=i=0;i<n_markers;i++) if(markers[i].element->n_levels) n_mk++;
	i=nrm_flag?1:0;
	i|=family_id?2:0;
	n_mod=0;
	model=Models;
	while(model) {
		n_mod++;
		model=model->next;
	}
	if(fprintf(fptr,"Loki.dat:%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x,%x\n",RunID,pruned_ped_size,
				  n_comp,n_mk,n_links,id_rec_count,nonid_rec_count,i,k2,nrec,ntyped,
				  traitlocus?1:0,n_genetic_groups,n_mod)<0) DataFileError(fname);
	for(i=0;i<n_comp;i++) if(fprintf(fptr,"%x%c",comp_size[i],i==n_comp-1?'\n':',')<0) DataFileError(fname);
	for(type=k=j=0;j<pruned_ped_size;j++)	{
		i=perm[j];
		ids=id_array[i].sire;
		idd=id_array[i].dam;
		if(ids) ids=ped_recode1[ids-1];
		if(idd) idd=ped_recode1[idd-1];
		k2=0;
		if(id_array[i].data1 && nonid_rec_count)
		  for(k1=0;k1<id_array[i].nrec;k1++) if(id_array[i].data1[k1]) k2++;
		if(fprintf(fptr,"%x,%x",ids,idd)<0) DataFileError(fname);
		if(family_id) {
			fam=id_array[i].fam_code;
			if(fam) fam=family_recode1[fam-1];
			if(fprintf(fptr,",%x",fam)<0) DataFileError(fname);
		}
		if(fprintf(fptr,",%x,%x,%x,%x",k2,(id_array[i].haplo[0]?1:0)|(id_array[i].data?2:0),
				  id_array[i].sex,id_array[i].affected)<0) DataFileError(fname);
		if(n_genetic_groups>1) if(fprintf(fptr,",%x",id_array[i].group)<0) DataFileError(fname);
		k2=n_genetic_groups>1?1:0;
		if(id_array[i].data && (id_rec_count-k2))	{
			for(k1=0;k1<n_id_records;k1++) if(id_elements[k1]->type&(ST_MODEL|ST_TRAIT)) {
				if(print_id_data(id_array[i].data+k1,id_elements[k1]->type,fptr)) DataFileError(fname);
			}
		}
		if(id_array[i].data1 && nonid_rec_count) {
			for(k1=0;k1<id_array[i].nrec;k1++) if(id_array[i].data1[k1]) {
				for(k2=0;k2<n_nonid_records;k2++) if(nonid_elements[k2]->type&(ST_MODEL|ST_TRAIT)) {
					if(print_id_data(id_array[i].data1[k1]+k2,nonid_elements[k2]->type,fptr)) DataFileError(fname);
				}
			}
		}
		if(id_array[i].haplo[0]) {
			if(fputc(';',fptr)==EOF) DataFileError(fname);
			for(k1=0;k1<n_markers;k1++) if(markers[k1].element->n_levels && fprintf(fptr,"%x,%x;",id_array[i].haplo[0][k1],id_array[i].haplo[1][k1])<0) DataFileError(fname);
		}
		if(ped_recode[i] && ped_recode[i]->type==STRING) {
			k+=strlen(ped_recode[i]->data.string);
			type=STRING;
		}
		if(fputc('\n',fptr)==EOF) DataFileError(fname);
	}
	if(k) k+=pruned_ped_size;
	if(fprintf(fptr,"LKTR:%x",k)<0) DataFileError(fname);
	if(family_id) {
		for(k1=i=0;i<n_orig_families;i++) {
			if(family_recode1[i] && family_recode[i]->type==STRING) k1+=strlen(family_recode[i]->data.string)+1;
		}
		if(fprintf(fptr,",%x\n",k1)<0) DataFileError(fname);
		for(i=0;i<n_orig_families;i++) if(family_recode1[i]) {
			if(family_recode[i]->type==STRING) {
				(void)fputs(family_recode[i]->data.string,fptr);
				(void)fputc('\n',fptr);
				if(errno) DataFileError(fname);
			} else {
				k1=(int)family_recode[i]->data.value;
				if(fprintf(fptr,"%x\n",k1)<0) DataFileError(fname);
			}
		}
		free(family_recode1);
	} else (void)fputc('\n',fptr);
	for(j=0;j<pruned_ped_size;j++) {
		i=perm[j];
		if(type==STRING)	{
			if(ped_recode[i]) (void)fputs(ped_recode[i]->data.string,fptr);
			(void)putc('\n',fptr);
			if(errno) DataFileError(fname);
		} else {
			if(ped_recode[i]) {
				k1=(int)ped_recode[i]->data.value;
				if(fprintf(fptr,"%x",k1)<0) DataFileError(fname);
			}
			(void)putc('\n',fptr);
			if(errno) DataFileError(fname);
		}
	}
	if(n_links) {
		k=0;
		linkp=links;
		while(linkp) {
			if(linkp->name) k+=strlen(linkp->name)+1;
			else k++;
			linkp=linkp->next;
		}
		if(fprintf(fptr,"LKLN:%x",k)<0) DataFileError(fname);
		linkp=links;
		while(linkp) {
			if(fprintf(fptr,",%x",linkp->type)<0) DataFileError(fname);
			linkp=linkp->next;
		}
		if(fputc('\n',fptr)==EOF) DataFileError(fname);
		linkp=links;
		while(linkp) {
			if(linkp->name) {
				if(fprintf(fptr,"%s\n",linkp->name)<0) DataFileError(fname);
			} else if(fputc('\n',fptr)==EOF) DataFileError(fname);
			linkp=linkp->next;
		}
	}
	if(n_mk) {
		for(k=k1=i=0;i<n_markers;i++)	if(markers[i].element->n_levels) {
			k+=strlen(markers[i].var->name)+1;
			k1+=markers[i].element->n_levels;
		}
		if(fprintf(fptr,"LKMK:%x,%x\n",k,k1)<0) DataFileError(fname);
		for(i=0;i<n_markers;i++) if(markers[i].element->n_levels && fprintf(fptr,"%s\n",markers[i].var->name)<0) DataFileError(fname);
		for(i=0;i<n_markers;i++) if(markers[i].element->n_levels) {
			if(markers[i].var->vtype&ST_ARRAY) k=markers[i].index;
			else k=0;
			if(fprintf(fptr,"%x",k)<0) DataFileError(fname);
			if(n_links>1) {
				if(fprintf(fptr,",%x",markers[i].link-1)<0) DataFileError(fname);
			}
			k=markers[i].element->n_levels;
			if(fprintf(fptr,",%x",k)<0) DataFileError(fname);
			k1=0;
			if(factor_recode[n_factors+i][0]->type==STRING)
			  for(j=0;j<k;j++) k1+=1+strlen(factor_recode[n_factors+i][j]->data.string);
			if(fprintf(fptr,",%x\n",k1)<0) DataFileError(fname);
			if(k1) {
				for(j=0;j<k;j++)
				  if(fprintf(fptr,"%s\n",factor_recode[n_factors+i][j]->data.string)<0) DataFileError(fname);
			} else {
				for(j=0;j<k;j++) {
					k1=(int)factor_recode[n_factors+i][j]->data.value;
					if(fprintf(fptr,"%x\n",k1)<0) DataFileError(fname);
				}
			}
		}
	}
	if(id_rec_count) {
		for(k=k1=i=0;i<id_rec_count;i++) {
			i1=id_trans[i];
			sd=id_elements[i1]->arg.var->data;
			k+=strlen(sd->name)+1;
			if(id_elements[i1]->type&ST_FACTOR) k1+=id_elements[i1]->n_levels;
		}
		if(fprintf(fptr,"LKIR:%x,%x\n",k,k1)<0) DataFileError(fname);
		for(i=0;i<id_rec_count;i++) {
			i1=id_trans[i];
			sd=id_elements[i1]->arg.var->data;
			(void)fprintf(fptr,"%s\n",sd->name);
		}
		for(i1=0;i1<id_rec_count;i1++) {
			i=id_trans[i1];
			v2[0]=id_elements[i]->type;
			sd=id_elements[i]->arg.var->data;
			if(sd->vtype&ST_ARRAY) v2[1]=id_elements[i]->oindex;
			else v2[1]=0;
			if(fprintf(fptr,"%x,%x",v2[0],v2[1])<0) DataFileError(fname);
			if(id_elements[i]->type&ST_FACTOR) {	
				k=id_elements[i]->n_levels;
				if(fprintf(fptr,",%x",k)<0) DataFileError(fname);
				if(k)	{
					for(j=0;j<n_factors;j++) if(var_factors[j]==id_elements[i]) break;
					if(j==n_factors) ABT_FUNC("Internal error - no recode information for factor\n");
					k1=0;
					if(factor_recode[j][0]->type==STRING)
					  for(k2=0;k2<k;k2++) k1+=1+strlen(factor_recode[j][k2]->data.string);
					if(fprintf(fptr,",%x\n",k1)<0) DataFileError(fname);
					if(k1) {
						for(k2=0;k2<k;k2++)
						  if(fprintf(fptr,"%s\n",factor_recode[j][k2]->data.string)<0) DataFileError(fname);
					} else {
						for(k2=0;k2<k;k2++) {
							k1=(int)factor_recode[j][k2]->data.value;
							if(fprintf(fptr,"%x\n",k1)<0) DataFileError(fname);
						}
					}
				} else if(fputc('\n',fptr)==EOF) DataFileError(fname);
			} else if(fputc('\n',fptr)==EOF) DataFileError(fname);
		}
	}
	if(nonid_rec_count) {
		for(k=k1=i=0;i<n_nonid_records;i++) if(nonid_elements[i]->type&(ST_MODEL|ST_TRAIT))	{
			sd=nonid_elements[i]->arg.var->data;
			k+=strlen(sd->name)+1;
			if(nonid_elements[i]->type&ST_FACTOR) k1+=nonid_elements[i]->n_levels;
		}
		v2[0]=k;
		v2[1]=k1;
		if(fprintf(fptr,"LKNR:%x,%x\n",k,k1)<0) DataFileError(fname);
		for(i=0;i<n_nonid_records;i++) if(nonid_elements[i]->type&(ST_MODEL|ST_TRAIT)) {
			sd=nonid_elements[i]->arg.var->data;
			(void)fprintf(fptr,"%s\n",sd->name);
		}
		for(i=0;i<n_nonid_records;i++) if(nonid_elements[i]->type&(ST_MODEL|ST_TRAIT)) {
			v2[0]=nonid_elements[i]->type;
			sd=nonid_elements[i]->arg.var->data;
			if(sd->vtype&ST_ARRAY) v2[1]=nonid_elements[i]->index;
			else v2[1]=0;
			if(fprintf(fptr,"%x,%x",v2[0],v2[1])<0) DataFileError(fname);
			if(nonid_elements[i]->type&ST_FACTOR) {	
				k=nonid_elements[i]->n_levels;	
				if(fprintf(fptr,",%x",k)<0) DataFileError(fname);
				if(k) {
					for(j=0;j<n_factors;j++) if(var_factors[j]==nonid_elements[i]) break;
					if(j==n_factors) ABT_FUNC("Internal error - no recode information for factor\n");
					k1=0;
					if(factor_recode[j][0]->type==STRING)
					  for(k2=0;k2<k;k2++) k1+=1+strlen(factor_recode[j][k2]->data.string);
					if(fprintf(fptr,",%x\n",k1)<0) DataFileError(fname);
					if(k1) {
						for(k2=0;k2<k;k2++)
						  if(fprintf(fptr,"%s\n",factor_recode[j][k2]->data.string)<0) DataFileError(fname);
					} else for(k2=0;k2<k;k2++)	{
						k1=(int)factor_recode[j][k2]->data.value;
						if(fprintf(fptr,"%x\n",k1)<0) DataFileError(fname);
					}
				} else if(fputc('\n',fptr)==EOF) DataFileError(fname);
			} else if(fputc('\n',fptr)==EOF) DataFileError(fname);
		}
	}
	model=Models;
	while(model) {
		if(model->trait->vtype&ST_ARRAY) elem=model->trait->element+model->index-1;
		else elem=model->trait->element;
		mlist=model->model_list;
		j=k1=0;
		while(mlist) {
			k=mlist->nvar;
			if(!k) ABT_FUNC("Internal error - empty model_list\n");
			j++;
			k1+=k;
			mlist=mlist->next;
		}
		k=get_element_index(elem);
		if(fprintf(fptr,"LKMD:%x,%x,%x,%x\n",elem->type,k,j,k1)<0) DataFileError(fname);
		if(j) {
			mlist=model->model_list;
			while(mlist) {
				k=mlist->nvar;
				if(fprintf(fptr,"%x",k)<0) DataFileError(fname);
				for(i=0;i<k;i++) {
					elem=mlist->element[i];
					j=get_element_index(elem);
					if(fprintf(fptr,",%x,%x",elem->type,j)<0) DataFileError(fname);
				}
				mlist=mlist->next;
				if(fputc('\n',fptr)==EOF) DataFileError(fname);
			}
		} else ABT_FUNC("Internal error - empty model\n");
		model=model->next;
	}
	if(fputs("Ldat.end\n",fptr)==EOF) DataFileError(fname);
	free(perm);
	if(fclose(fptr)) DataFileError(fname);
	free(fname);
	if(id_trans) free(id_trans);
	if(Filter) do i=wait(&j); while(i>0);
}
