/*  trinculo, Flexible tools for multinomial association
    Copyright (C) 2014  Luke Jostins

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.*/

#include <iostream>
#include <string.h>
#include <stdlib.h>

using namespace std;

class inputs {
 // Class for parsing and storing user input from the command line 
 public:
  int normalizecovar;
  int modelSelect;
  int outputErrs;
  int PhenoPs;
  int defaultPrior;
  int empiricalCovar;
  int noAssoc;
  int ses;
  int outputFullMd;
  int threads;
  int nooutput;
  int precision;
  int fititer;
  double fitabstol;
  double fitreltol;
  double covarsigma;
  double priorsigma;
  double priorcor;
  string bedprefix;
  string outprefix;
  string covarfile;
  string phenofile;
  string missingpheno;
  string basepheno;
  string phenoname;
  string usage;
  string help;
  string priorfile;
  string dosagefile;
  string conditionfile;
  inputs();
  int parseinput(int argc, char *argv[]); 
};

inputs::inputs(){
  // initialize the object with all the defaults
  bedprefix = "";
  nooutput = 0;
  outprefix = "trinculo";
  covarfile = "";
  phenofile = "";
  missingpheno = "NA";
  basepheno = "";
  phenoname = "";
  priorfile="";
  covarsigma = 1;
  defaultPrior=0;
  precision=3;
  
  fititer=20;
  fitreltol = 1e-6;
  fitabstol = 1e-8;
  
  empiricalCovar=0;
  priorsigma=0.04;
  priorcor=0.0;
  noAssoc = 0;
  outputErrs = 0;
  PhenoPs = 0;
  ses = 0;
  outputFullMd = 0;
  threads=1;

  usage = "Usage: trinculo MODE [--bfile GENOTYPES_PREFIX | --dosage GENOTYPES.DOSE ] --pheno PHENOTYPES.TXT --phenoname PHENOTYPE [ --basepheno REF --missingpheno MISSING ] [--covar COVARIATES.TXT --normalize] [--priors PRIORS.TXT | --defaultprior --priorsigma PRIORVAR --priorcor PRIORCOR --covarsigma PRIORVAR | --empiricalprior ] [ --select --phenops --errormat --fullmodel --waldstats ] [ --noassoc ] [--threads n] [ --precision DIGITS ] [ --fitIter NITER --fitAbsTol ABSTOL --fitRelTol RELTOL ][--help]";
  help = "\nUsage:\n trinculo MODE  [--bfile GENOTYPES_PREFIX | --dosage GENOTYPES.DOSE ] [--out OUTPUT_PREFIX] --pheno PHENOTYPES.TXT --phenoname PHENOTYPE [--basepheno REF --missingpheno MISSING ] [--covar COVARIATES.TXT --normalize] [--condition SNPS.TXT] [--priors PRIORS.TXT | --defaultprior --priorsigma PRIORVAR --priorcor PRIORCOR --covarsigma PRIORVAR | --empiricalprior  ] [ --select --phenops --errormat --fullmodel --waldstats ] [ --noassoc ] [--threads n] [ --precision DIGITS ] [ --fitIter NITER --fitAbsTol ABSTOL --fitRelTol RELTOL ] [--help]\n\nOptions:\n\nMODE\t- Either multinom or ordinal for multinomial or ordinal logistic regression\n\n--bfile FILE_PREFIX\t- Genotype input in binary plink format (FILE_PREFIX.bed, FILE_PREFIX.bim and FILE_PREFIX.fam)\n--dosage FILE\t- Genotype input in dosage format\n--out PREFIX\t- The prefix used for output files [default is \"trinculo\"]\n\n--pheno FILE\t- Phenotype input (same format as plink)--phenoname STRING\t-The phenotype in the phenotype file to analyse\n--basepheno STRING\t- The reference catagory to use for the multinomial analysis [defaults to first catagory encountered]\n--missingpheno STRING\t- The symbol to be use to indicate a missing phenotype in the phenotype file\n\n--condition FILE\t- A file containing a list of SNPs (one per line) to conditon on (i.e. include as covariates)\n\n--covar FILE\t- A file containing covariates to condition on (same format as plink)\n--normalize\t- Normalizes covariates to have a variance of 1 [off by default]\n\n--priors FILE\t- A prior covariance matrix on effect sizes\n--defaultpriors\t- Construct a prior matrix based on default values\n--priorsigma VALUE\t- Set the prior variance on effect sizes within a disease[default is 0.04]\n--priorcor VALUE\t- Set the prior correlation on effect sizes across diseases [default is 0]\n--covarsigma VALUE\t- Prior on the covariate effect size [Default is 1]\n--empiricalprior\t - Calculate an empirical prior directly from the data\n--noassoc\t-do not do any association testing, just infer the prior\n\n--select\t-Do Bayesian model selection (i.e. calculate marginal likelihoods for each sharing model)\n--phenops\t-Do a likelihood ratio test on each phenotype individually\n--waldstats\t-Calculate standard errors, Z scores and Wald p-values\n--errormat\t-Output the variance-covariance matrix for the log odds ratio\n--fullmodel\t-Output all parameter estimates and the full error matrix\n\n--threads n\tUse up to n cores\n\n--precision DIGITS\tPrint output to DIGITS significant figures\n\n--fitIter NITER\tTry up to NITER iterations when fitting models\n--fitAbsTol ABSTOL\tTerminate model fitting if results change by less than ABSTOL\n--fitRelTol RELTOL\tTerminate model fitting if results change by less than (RELTOL x value)\n\n--help\t- Prints this message\n\n";
  normalizecovar = 0;
  modelSelect = 0;
  dosagefile = "";
  conditionfile = "";
  
} 

int inputs::parseinput(int argc, char * argv[]){
  // parse the command line input and store the user commands in the object
  argc--;
  argv++;
   
  if (!argc){
    cout << "Error: no options entered" << endl;
    cout << usage << endl;
    throw 20;
  }

  while (argc){
    if (!strcmp(*argv,"--bfile")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --bfile" << endl;
	throw 20;  
      }
      if (dosagefile != ""){
	cout << "Error: both dosage and plink files provided. Use one or the other!" << endl;
	throw 20;
      }
      if (bedprefix != ""){
	cout << "Error: Multiple plink file prefixes provided!" << endl;
	throw 20;
      }
      
      bedprefix=*argv; 
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--select")){
	argv++; argc--;
	modelSelect=1;
	continue;
    }
    if (!strcmp(*argv,"--nooutput")){
	argv++; argc--;
	nooutput=1;
	continue;
    }
    if (!strcmp(*argv,"--waldstats")){
	argv++; argc--;
	ses=1;
	continue;
    }
    if (!strcmp(*argv,"--fullmodel")){
	argv++; argc--;
	outputFullMd=1;
	continue;
    }
    if (!strcmp(*argv,"--phenops")){
	argv++; argc--;
	PhenoPs=1;
	continue;
    }
    if (!strcmp(*argv,"--errormat")){
	argv++; argc--;
	outputErrs=1;
	continue;
    }
    if (!strcmp(*argv,"--noassoc")){
	argv++; argc--;
	noAssoc=1;
	continue;
    }
    if (!strcmp(*argv,"--defaultprior")){
	argv++; argc--;
	defaultPrior=1;
	continue;
    }
    if (!strcmp(*argv,"--empiricalprior")){
	argv++; argc--;
	empiricalCovar=1;
	continue;
    }
	
    if (!strcmp(*argv,"--dosage")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --dosage" << endl;
	throw 20;  
      }
      if (bedprefix != ""){
	cout << "Error: both dosage and plink files provided. Use one or the other!" << endl;
	throw 20;
      }
      if (dosagefile != ""){
	cout << "Error: Multiple dosage files provided!" << endl;
	throw 20;
      }
      
      dosagefile=*argv; 
      argv++; argc --; 
      continue;
    }

    if (!strcmp(*argv,"--condition")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --condition" << endl;
	throw 20;  
      } 
      conditionfile=*argv; 
      argv++; argc --; 
      continue;
    }
    
    if (!strcmp(*argv,"--priors")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --priors" << endl;
	throw 20;  
      }
      priorfile=*argv; 
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--covarsigma")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --covarsigma" << endl;
	throw 20;  
      }
      covarsigma=atof(*argv); 
	  if (covarsigma <= 0) {
		  cout << "Error: invalid value provided to --covarsigma (" << *argv << ")" << endl;
		  throw 20;
	  }
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--precision")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --precision" << endl;
		throw 20;  
      }
      precision=atoi(*argv); 
	  if (precision < 1) {
		  cout << "Error: invalid value for precision (" << *argv << ")" << endl;
		  throw 20;
	  }
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--fitRelTol")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --fitRelTol" << endl;
		throw 20;  
      }
      fitreltol=atof(*argv); 
	  if (!(fitreltol < 1 && fitreltol > 0)) {
		  cout << "Error: invalid value for relative tolerance (" << *argv << ")" << endl;
		  throw 20;
	  }
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--fitAbsTol")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --fitAbsTol" << endl;
		throw 20;  
      }
      fitabstol=atof(*argv); 
	  if (!(fitabstol < 1 && fitabstol > 0)) {
		  cout << "Error: invalid value for absolute tolerance (" << *argv << ")" << endl;
		  throw 20;
	  }
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--fitIter")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --fitIter" << endl;
		throw 20;  
      }
      fititer=atoi(*argv);
	  if (fititer < 1) {
		  cout << "Error: invalid number of model fit iterations (" << *argv << ")" << endl;
		  throw 20;
	  }
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--priorcor")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --priorcor" << endl;
	throw 20;  
      }
      priorcor=atof(*argv);
	  if (priorcor <= 0) {
		  cout << "Error: invalid value for --priorcor (" << *argv << ")" << endl;
		  throw 20;
	  }
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--file")){ cout << "Error: --file (reading in .ped/.map files) not currently supported. Use binary plink files with --bfile.\n"; throw 20;
}
    if (!strcmp(*argv,"--help") || !strcmp(*argv,"-h")) {cout << help << endl; throw 1;}
    if (!strcmp(*argv,"--out")) {
      argv++; argc--;
      if (!argc){
        cout << "Error: No argument provided to --out" << endl;
	throw 20;
      }
      
      outprefix=*argv; 
      argv++; argc--; 
      continue;
    }
    if (!strcmp(*argv,"--covar")) {
      argv++; argc--;
      if (!argc){
        cout << "Error: No argument provided to --covar" << endl;
	throw 20;
      }
      covarfile=*argv; 
      argv++; argc--;
      continue;
    }
    if (!strcmp(*argv,"--pheno")) {
      argv++; argc--;
      if (!argc){
        cout << "Error: No argument provided to --pheno" << endl;
	throw 20;
      }
      phenofile=*argv; 
      argv++; argc--;
      continue;
    }
    if (!strcmp(*argv,"--missingpheno")) {
      argv++; argc--;
      if (!argc){
        cout << "Error: No argument provided to --missingpheno" << endl;
	throw 20;
      }
      missingpheno=*argv; 
      argv++; argc--;
      continue;
    }
    if (!strcmp(*argv,"--phenoname")) {
      argv++; argc--;
      if (!argc){
        cout << "Error: No argument provided to --phenoname" << endl;
	throw 20;
      }
      phenoname=*argv; 
      argv++; argc--;
      continue;
    }
    if (!strcmp(*argv,"--basepheno")) {
      argv++; argc--;
      if (!argc){
        cout << "Error: No argument provided to --basepheno" << endl;
        throw 20;
      }
      basepheno=*argv; 
      argv++; argc--;
      continue;
    }
    if (!strcmp(*argv,"--threads")){ 
      argv++; argc--;
      if (!argc){
        cout << "Error: No arguent provided to --threads" << endl;
	throw 20;  
      }
      threads=atoi(*argv);
	  if (threads < 1) {
		  cout << "Error: invalid number of cores (" << *argv << ")" << endl;
		  throw 20;
	  }
      argv++; argc --; 
      continue;
    }
    if (!strcmp(*argv,"--normalize")) {argv++; argc--; normalizecovar = 1; continue;}
    cout << "Error: Option " << *argv << " not found" << endl;
    throw 20;
  }

  if (defaultPrior && priorfile != ""){
    cout << "Error: Default priors selected but prior file provided. Please select one or the other." << endl;
    throw 20;
  }
  
  if (empiricalCovar && priorfile != ""){
    cout << "Error: Emperical prior selected but prior file provided. Please select one or the other." << endl;
    throw 20;
  }
  
  if (empiricalCovar && defaultPrior){
    cout << "Error: Both empirical and default priors selected. Please select one or the other." << endl;
    throw 20;
  }
  
  return(0);
  
}
