/****************************************************************************
 *                                                                          *
 *     Loki - Programs for genetic analysis of complex traits using MCMC    *
 *                                                                          *
 *             Simon Heath - University of Washington                       *
 *                                                                          *
 *                       July 1997                                          *
 *                                                                          *
 * 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 <config.h>
#include <stdlib.h>
#include <string.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/types.h>
#ifdef USE_MONITOR
#include <sys/ipc.h>
#include <sys/shm.h>
#endif
#include <sys/time.h>
#include <signal.h>

#include "version.h"
#include "ranlib.h"
#include "utils.h"
#include "libhdr.h"
#include "loki.h"
#include "loki_peel.h"
#include "seg_pen.h"
#include "loki_ibd.h"
#include "loki_output.h"
#include "sample_rand.h"
#include "loki_tlmoves.h"
#include "sample_effects.h"
#include "sample_nu.h"
#include "calc_var_locus.h"
#include "meiosis_scan.h"
#include "loki_monitor.h"
#include "output_recomb.h"
#ifdef USE_MONITOR
#include "count_dbr.h"
#endif

unsigned int RunID;
int nrm_flag,catch_sigs,sig_caught,multiple_rec;
int map_function,singleton_flag,est_aff_freq,*ibd_mode;
char *Filter,*Seedfile;
struct remember *RemBlock;
static struct remember *FirstRemBlock;
static loki_time lt;
#ifdef TRACE_PEEL
int *peel_trace;
#endif
static struct peel_mem peel_workspace;
static int ranseed_set;
static char *LogFile;

/* Called when program receives signal */
static void int_handler(int i)
{
	char *sigs[]={"SIGTERM","SIGQUIT","SIGHUP","SIGINT","SIGALRM","SIGVTALRM","SIGPROF","UNKNOWN"};
	int j,signals[]={SIGTERM,SIGQUIT,SIGHUP,SIGINT,SIGALRM,SIGVTALRM,SIGPROF};
	
	for(j=0;j<7;j++) if(i==signals[j]) break;
	if(catch_sigs && !sig_caught) {
		sig_caught=i;
		if(j>=4 && j<7) (void)fprintf(stderr,"Received signal %s: %s closing down.\n",sigs[j],LOKI_NAME);
		else (void)fprintf(stderr,"Received signal %s: please wait while %s closes down.\nRepeat signal to abort immediately.\n",sigs[j],LOKI_NAME);
		return;
	}
	(void)fprintf(stderr,"Exiting on signal %s (%d)\n",sigs[j],i);
	exit(i);
}

void FreeStuff(void)
{
	int i,j;
	struct Locus *loc;

#ifdef USE_MONITOR
	if(lpar) free_dbr_count();
	if(lmon_shm_id>=0) {
		(void)shmctl(lmon_shm_id,IPC_RMID,0);
		lmon_shm_id=-1;
	}
#else
	if(lpar) free(lpar);
#endif
	if(from_abt) return;
	if((ranseed_set&3)==2) {
		if(Seedfile) (void)writeseed(Seedfile,1);
		else (void)writeseed("seedfile",1);
	}
	if(n_markers && marker[0].locus.seg[0]) free(marker[0].locus.seg[0]);
	if(Seedfile) free(Seedfile);
	if(LogFile) free(LogFile);
	if(Filter) free(Filter);
	if(Dumpfile) free(Dumpfile);
	if(Polyfile) free(Polyfile);
	if(Freqfile) free(Freqfile);
	if(Haplofile) free(Haplofile);
	if(Outputfile) free(Outputfile);
	if(OutputIBDfile) free(OutputIBDfile);
	if(OutputIBDdir) free(OutputIBDdir);
	if(OutputPosfile) free(OutputPosfile);
	if(id_array) free(id_array);
	if(founder_flag) {
 		if(founder_flag[0]) free(founder_flag[0]);
		free(founder_flag);
	}
	if(marker) {
		if(marker[0].mterm) free(marker[0].mterm);
		for(i=0;i<n_markers;i++) {
			if(marker[i].locus.variance) free(marker[i].locus.variance);
		}
		free(marker);
	}
	if(linkage) {
		for(i=0;i<n_links;i++) {
			if(linkage[i].ibd_list) {
				free(linkage[i].ibd_list->pos);
				free(linkage[i].ibd_list);
			}
		}
		free(linkage);
	}
	OutputSample(0,0);
	free_IBD();
	free_rand();
	free_sample_effects();
	if(use_student_t) free_sample_nu();
	seg_dealloc();
	peel_dealloc(&peel_workspace); 
	free_complex_mem();
	TL_Free();
	if(tlocus) {
		for(i=0;i<n_tloci;i++) {
			loc=&tlocus[i].locus;
			if(loc->seg[0]) free(loc->seg[0]);
			if(loc->gt) free(loc->gt);
			if(loc->lk_store) free(loc->lk_store);
			if(loc->freq) {
				if(loc->variance) free(loc->variance);
				free(loc->freq);
			}
		}
		free(tlocus);
	}
	if(all_set) {
		if(n_markers && req_set[0][0]) free(req_set[0][0]);
		if(n_markers && allele_trans[0])	{
			if(allele_trans[0][0]) free(allele_trans[0][0]);
			free(allele_trans[0]);
		}
		if(n_markers && all_set[0]) free(all_set[0]);
		if(r_func[0]) {
			for(i=0;i<n_markers+tlocus_flag;i++)
			  for(j=0;j<n_comp;j++)
				 if(r_func[i][j]) free(r_func[i][j]);
			free(r_func[0]);
		}
		if(peelseq_head[0]) {
			for(i=0;i<n_markers+tlocus_flag;i++)
			  for(j=0;j<n_comp;j++)
				 free_peelseq(peelseq_head[i]+j);
			free(peelseq_head[0]);
		}
		free(all_set);
	}
	if(comp_npeel) free(comp_npeel);
	calc_var_locus(n_markers);
	if(AIMatrix) {
		for(i=0;i<n_comp;i++) if(AIMatrix[i]) free(AIMatrix[i]);
		free(AIMatrix);
	}
	if(comp_size) free(comp_size);
	if(models) {
		for(i=0;i<n_models;i++) if(models[i].term) free(models[i].term);
		free(models);
	}
	if(residual_var) free(residual_var);
	if(res_var_set) free(res_var_set);
	if(id_variable) free(id_variable);
	else if(nonid_variable) free(nonid_variable);
	FreeRemem(FirstRemBlock);
	return;
}

void ignore_handler(int i)
{
	i=i;
	/* Do nothing */
}

#ifdef FUNC_NAME
#undef FUNC_NAME
#endif
#define FUNC_NAME "main"
int main(int argc,char *argv[])
{
	FILE *fptr;
	int i=0,j,c,*si_mode,read_dump_flag=0,append_output_flag=0;
	struct sigaction s_action;
	char *tfile=0,*tname;
	struct itimerval itimerval;
	double sec,frac;
	
	lt.start_time=time(0);
	while((c=getopt(argc,argv,"arvf:p:d:"))!=-1) switch(c) {
	 case 'f':
		i=strlen(optarg);
		if(tfile) free(tfile);
		if(!(tfile=malloc(i+1))) ABT_FUNC(MMsg);
		(void)strcpy(tfile,optarg);
		read_dump_flag=1;
		break;
	 case 'd':
		if((j=set_file_dir(optarg))) {
			fprintf(stderr,"Error setting default file directory: %s\n",j==UTL_BAD_STAT?strerror(errno):utl_error(j));
			exit(EXIT_FAILURE);
		}
		break;
	 case 'p':
		if((j=set_file_prefix(optarg))) {
			fprintf(stderr,"Error setting default file prefix: %s\n",utl_error(j));
			exit(EXIT_FAILURE);
		}
		break;
	 case 'r':
		if(!tfile) tfile=make_file_name(".dump");
		read_dump_flag=1;
		break;
	 case 'a':
		append_output_flag=1;
		break;
	 case 'v':
		(void)printf("%s\n",LOKI_NAME);
		return 0;
	}
	if(optind>=argc) abt(__FILE__,__LINE__,"No parameter file specified\n");
	/* If program terminates normally or via a signal I want it to go through
	 * FreeStuff() so that Memory can be cleared so I can check for unfreed blocks
	 */
	if(!(FirstRemBlock=malloc(sizeof(struct remember)))) ABT_FUNC(MMsg);
	RemBlock=FirstRemBlock;
	RemBlock->pos=0;
	RemBlock->next=0;
	s_action.sa_handler=ignore_handler;
	s_action.sa_flags=0;
	(void)sigaction(SIGPIPE,&s_action,0L);
#ifdef USE_MONITOR
	s_action.sa_handler=reaper;
	s_action.sa_flags=0;
 	(void)sigaction(SIGCHLD,&s_action,0L);
#endif
	ReadBinFiles(&LogFile,0);
	print_start_time(LOKI_NAME,"a",LogFile,&lt);
	AllocLokiStruct();
 	if((fptr=fopen(argv[optind],"r"))) i=ReadParam(fptr,argv[optind],&ranseed_set);
	else {
		(void)printf("Couldn't open '%s' for input as parameter file\nAborting...\n",argv[optind]);
		exit(EXIT_FAILURE);
	}
	(void)fclose(fptr);
	if(i) exit(EXIT_FAILURE);
#ifdef USE_MONITOR
	/* Set up shared memory segment */
	lmon_shm_id=shmget(IPC_PRIVATE,sizeof(struct lmon_param),IPC_CREAT|0600);
	if(lmon_shm_id<0) {
		perror("Failed to get shared memory segment");
		ABT_FUNC(AbMsg);
	}
	/* Attach shared memory segment */
	lpar=shmat(lmon_shm_id,0,0);
	if(lpar==(void *)-1) {
		perror("Failed to attach shared memory segment");
		ABT_FUNC(AbMsg);
	}
#else
	if(!(lpar=malloc(sizeof(struct lmon_param)))) ABT_FUNC(MMsg);
#endif
	lpar->magic=0;
	lpar->dbr_flag=0;
	lpar->dbr_shm_id=-1;
	lpar->peel_trace=0;
	lpar->debug_level=0;
	lpar->si_mode=0;
	lpar->ibd_mode=0;
#ifdef TRACE_PEEL
	peel_trace=&lpar->peel_trace;
#endif
#ifdef DEBUG
	debug_level=&lpar->debug_level;
#endif
	si_mode=&lpar->si_mode;
	ibd_mode=&lpar->ibd_mode;
#ifdef USE_MONITOR
	/* Start up monitor process */
	start_monitor();
#endif
	if(atexit(FreeStuff)) ABT_FUNC("Unable to register exit function FreeStuff()\n");
	if(atexit(print_end_time)) ABT_FUNC("Unable to register exit function print_end_time()\n");
	s_action.sa_handler=int_handler;
	s_action.sa_flags=0;
	(void)sigaction(SIGINT,&s_action,0L);
	(void)sigaction(SIGHUP,&s_action,0L);
	(void)sigaction(SIGQUIT,&s_action,0L);
	(void)sigaction(SIGTERM,&s_action,0L);
	(void)sigaction(SIGALRM,&s_action,0L);
	(void)sigaction(SIGVTALRM,&s_action,0L);
	(void)sigaction(SIGPROF,&s_action,0L);
	if(tfile) {
		if(Dumpfile) free(Dumpfile);
		Dumpfile=tfile;
	}
	if(!(ranseed_set&2))	{
		if(!Seedfile) {
			if(getseed("seedfile",0)) init_ranf(135421);
		} else if(getseed(Seedfile,0)) init_ranf(135421);
		ranseed_set|=2;
	}
	if(LogFile && (tname=add_file_dir(LogFile)))	{
		fptr=fopen(tname,"a");
		free(tname);
		if(fptr) {
			(void)fputs("\n     Starting state of RNG (seedfile):\n\n",fptr);
			(void)dumpseed(fptr,1);
			(void)fputc('\n',fptr);
			(void)fclose(fptr);
		}
	}
	LokiSetup();
	AllocEffects();
	InitValues(si_mode);
	seg_alloc();
	peel_alloc(&peel_workspace);
	if(output_recomb) print_recomb();
	if(limit_time>0.0) {
		frac=modf(limit_time,&sec);
		itimerval.it_value.tv_sec=(time_t)sec;
		itimerval.it_value.tv_usec=(time_t)(frac*1000000.0);
		/* Set it to trigger another timer event after 5 minutes in case
		 * we get stuck in the peel operation */
		itimerval.it_interval.tv_sec=300;
		itimerval.it_interval.tv_usec=0;
		setitimer(limit_timer_type,&itimerval,0);
	}
  	SampleLoop(&peel_workspace,si_mode,read_dump_flag,append_output_flag,&lt);
	return sig_caught;
}
