/* ARIADNE V.1.3

Copyright 

Richard Mott 2000

Wellcome Trust Centre For Human Genetics
Univeristy of Oxford
Roosevelt Drive
Oxford OX3 7AD
UK

The software package ARIADNE is distributed in the hope that it will be
useful, but in order that the University as a charitable foundation
protects its assets for the benefit of its educational and research
purposes, the University makes clear that no condition is made or to
be implied, nor is any warranty given or to be implied, as to the
accuracy of ARIADNE, or that it will be suitable for
any particular purpose or for use under any specific conditions, or that
the content or use of ARIADNE will not constitute
or result in infringement of third-party rights.  Furthermore, the
University disclaims all responsibility for the use which is made of 
ARIADNE.



*/

#include<limits.h>
#include<math.h>
#include<stdio.h>
#include<stdlib.h>
#include"seq_util.h"
#include"profile.h"
#include"cl.h"
#include"gapstat.h"
#include"ariadne.h"

typedef enum {
  SEQUENCES_TO_SELF, SEQUENCE_TO_SEQUENCES, SEQUENCES_TO_SEQUENCES, SEQUENCE_TO_PROFILES, PROFILE_TO_SEQUENCES, PROFILES_TO_SEQUENCES, SEQUENCES_TO_PROFILES, SEARCH_MODE_TYPES 
} SEARCH_MODES;


char *search_mode_names[SEARCH_MODE_TYPES] = {"ss2self", "s2ss", "ss2ss", "s2pp", "p2ss", "pp2ss", "ss2pp" };

#define TRAINING_SIZE 200

int main( int argc, char **argv ) {

  FILE *seqfp, *seq2fp, *profp;
  char seqfilename[256];
  char profilename[256];
  char matfilename[256];
  char seqfilename2[256];
  int **matrix=NULL;
  int A = 11;
  int B = 1;
  float pthresh1;
  float pthresh2;
  int mode=-1;
  char modestr[256];
  int count = 0;
  int total = 0;
  int align = 1;
  float dbsize = 1.0;
  float ethresh = 0.1;
  int errors = 0;
  int all=0;
  PPARAMS params;

  modestr[0] = 0;
  GetRestrictedArg("-mode", modestr, argc, argv, SEARCH_MODE_TYPES, search_mode_names, &mode );
  seqfp = ArgFile("-seq", "r", argc, argv, seqfilename );
  seq2fp = ArgFile("-seq2", "r", argc, argv, seqfilename2 );
  profilename[0] = 0;
  profp = ArgFile("-pro", "r", argc, argv, profilename );
  matfilename[0] = 0;
  GetArg("-matrix", matfilename, argc, argv );
  GetInt("-A", &A, argc, argv );
  GetInt("-B", &B, argc, argv );
  GetFloat("-ethresh", &ethresh, argc, argv);
  GetFloat("-dbsize", &dbsize, argc, argv);
  GetBoolean("-align", &align, argc, argv );
  GetBoolean("-all", &all, argc, argv );
  AddHelp("ariadne -mode [p2ss|s2pp|pp2ss|ss2pp] -pro [file-of-profiles] -seq [fasta-file] -A [gap open] -B [gap extend] -[no]align");
  GetHelp(argc,argv);


  if ( dbsize < 1 ) dbsize = 1;
  pthresh2 = ethresh/dbsize; /* the pvalue threshold corresponding to the evalue threshold. If the database size is 1 then this is the pvalue thresold */

  pthresh1 = pthresh2*10; /* the initial threshold for determining whether to compute more accurate stats */

  if ( mode < 0 ) {
    int n;
    fprintf( stderr, "ERROR: must set mode to one of " );
    for(n=0;n<SEARCH_MODE_TYPES;n++) 
      fprintf( stderr, " %s", search_mode_names[n] );
    fprintf(stderr, "\n");
    errors++;
  }

  if ( ! seqfp  ) {
    fprintf( stderr, "ERROR: cannot open sequence file %s\n", seqfilename );
    errors++;
  }

  params.A = A;
  params.B = B;
  params.matrix = NULL;
  params.pthresh1 = pthresh1;
  params.pthresh2 = pthresh2;
  params.dbsize = dbsize;
  params.align = align;
  params.hits = 0;
  params.total = 0;
  params.self = ( mode == SEQUENCES_TO_SELF );
  params.all = all;

  if ( mode > SEQUENCES_TO_SEQUENCES && ! profp  ) {
    fprintf( stderr, "ERROR: cannot open profile file %s\n", profilename );
    errors++;
  }
  else if ( mode <= SEQUENCES_TO_SEQUENCES ) {
    if ( ! (matrix = ReadBlastMatrix(matfilename) ) ) {
      fprintf(stderr, "ERROR - no matrix %s\n", matfilename );
      errors++;
    }
    params.matrix = matrix;
    if ( mode == SEQUENCES_TO_SELF ) 
      seq2fp = fopen(seqfilename,"r");
    
    if ( ! seq2fp ) {
      fprintf( stderr, "ERROR: cannot open sequence2 file %s\n", seqfilename2 );
      errors++;
    }
  }

  if ( errors ) 
    PrintUsage( argc, argv, 1 );

  printf( "***** ARIADNE v%s  %s *****\n\n", ARIADNE_VERSION, Now() );
  printf( "Copyright Richard Mott, 2000, Wellcome Trust Centre for Human Genetics, University of Oxford\n\n");
  printf( "%s", ARIADNE_PREAMBLE );

  printf("using mode %s\n", modestr );
  printf("using sequences in %s\n", seqfilename );
  if ( profilename[0] ) 
    printf("using profiles in %s\n", profilename );
  printf("using gap penalty %d+%dk\n", A, B );
  if ( matfilename[0] ) printf("using matrix %s\n", matfilename );
  printf("using database size %.0f sequences for assessing significance\n", dbsize );
  printf("using evalue threshold %e [ pvalue threshold %e initial threshold %e]\n\n", ethresh, pthresh2, pthresh1 );


  if ( mode == SEQUENCE_TO_SEQUENCES ) {
    CompareSequenceToSequences( seqfp, seq2fp, &params );
  }
  else if ( mode == SEQUENCES_TO_SEQUENCES ) {
    while( CompareSequenceToSequences( seqfp, seq2fp, &params ) );
  }
  else if ( mode == SEQUENCES_TO_SELF ) {
    while( CompareSequenceToSequences( seqfp, seq2fp, &params ) );
  }
  else if ( mode == SEQUENCE_TO_PROFILES ) {
    CompareSequenceToProfiles( seqfp, profp, &params );
  }
  else if ( mode == SEQUENCES_TO_PROFILES ) {
    while( CompareSequenceToProfiles( seqfp, profp, &params ) );
  }
  else if ( mode == PROFILE_TO_SEQUENCES ) {
    CompareProfileToSequences( seqfp, profp, &params );
  }
  else if ( mode == PROFILES_TO_SEQUENCES ) {
    while( CompareProfileToSequences( seqfp, profp, &params ) );
  }
  else 
    return 1;
  
  printf( "\n\n Done. %d Comparisons; %d Hits with evalue < %g \n", params.total, params.hits, ethresh );

  return 0;
}

int CompareSequenceToProfiles( FILE *seqfp, FILE *profp, PPARAMS *params  ) {

  SEQUENCE *seq = GetFastaSeq( seqfp );
  
  if ( seq ) {
    int sw_score;
    double *freq = PseudoResidueFrequencies( seq, 100 );
    double lambda0, K0, Kplus, H, r, s, alpha;
    double theta, logkappa;
    double K0mean=0.0, lambda0mean=0.0, Hmean=0.0, alpha_mean=0.0;
    int count=0;
    double pval1, pval2;
    PROFILE *pro;

    while( (pro=ReadProfile( profp ) ) ) {
      pval1 = pval2 = 1.0;
      sw_score = SWScoreProteinToProfile( seq, pro, params->A, params->B );
      params->total++;
      if ( ++count > TRAINING_SIZE ) {
	pval2 = pval1 = EmpiricalGEM( lambda0mean, K0mean, Hmean, alpha_mean, seq->len, pro->length, &theta, &logkappa, sw_score );
	if ( pval1 < params->pthresh1 ) {
	  if ( KarlinAltschulProfileStatistics( pro, freq, &lambda0, &K0, &Kplus, &H, &r, &s ) ) {
	    alpha = 2*s*exp(-lambda0*(params->A+params->B))/(1-exp(-lambda0*params->B));
	    pval2 = EmpiricalGEM( lambda0, K0, H, alpha, seq->len, pro->length, &theta, &logkappa, sw_score ); 
	  }
	}
      }
      else {
	if ( KarlinAltschulProfileStatistics( pro, freq, &lambda0, &K0, &Kplus, &H, &r, &s ) ) {
	  alpha = 2*s*exp(-lambda0*(params->A+params->B))/(1-exp(-lambda0*params->B));
	  pval2 = EmpiricalGEM( lambda0, K0, H, alpha, seq->len, pro->length, &theta, &logkappa, sw_score ); 
	  lambda0mean += lambda0;
	  K0mean += K0;
	  Hmean += H;
	  alpha_mean += alpha;
	  
	  if ( count == TRAINING_SIZE ) {
	    lambda0mean /= TRAINING_SIZE;
	    K0mean /= TRAINING_SIZE;
	    Hmean /= TRAINING_SIZE;
	    alpha_mean /= TRAINING_SIZE;
	  }
	}
	else {
	  printf( "ERROR sequence %s has positive expected match score with profile %s - ignored\n", seq->name, pro->name );
	  pval2 = 1.0;
	}
      }
      if ( pval2 < params->pthresh2 ) {
	int length;
	double percent;

	if ( params->all ) {
	  params->hits += TopAlignProteinToProfile( seq, pro, params->A, params->B, 0, params->pthresh2, params->self, lambda0, K0, H, alpha );
	}
	else {
	  PROFILE_ALIGNMENT *align = AlignProteinToProfile( seq, pro, params->A, params->B, &sw_score, &length, &percent ); /* redo the alignment */
	  PrintProteinProfileAlignment( stdout, params->total, pro, seq, align, length, sw_score,  params->dbsize*pval2, percent, K0*exp(logkappa), lambda0*theta, H, alpha, params->align );
	  fflush( stdout );
	  params->hits++;
	  free(align);
	}
      }
      FreeProfile( pro );
    }
    rewind(profp);
    FreeSeq( seq );
    return count;
  }
  else
    return 0;
}

int CompareProfileToSequences( FILE *seqfp, FILE *profp, PPARAMS *params  ) {

  PROFILE *pro=ReadProfile( profp );
  int count=0;

  if ( pro ) {
    rewind(seqfp);
    count = CompareProfileToSequences2( seqfp, pro, params  );
    FreeProfile( pro );
  }
  return count;
}


int CompareProfileToSequences2( FILE *seqfp, PROFILE *pro, PPARAMS *params  ) {
  
  int count=0;

  if ( pro ) {
    int sw_score = 0;
    double *freq_mean = RobinsonResidueFrequencies ();
    double lambda0, K0, Kplus, H, r, s, alpha;
    double theta, logkappa;
    double K0mean=0.0, lambda0mean=0.0, Hmean=0.0, alpha_mean=0.0;
    double pval1=1.0, pval2=1.0;
    SEQUENCE *seq;
    
    /*    printf( "profile %s %d\n", pro->name, pro->length ); */

    /* initial estimates for K, L etc */

    if ( KarlinAltschulProfileStatistics( pro, freq_mean, &lambda0mean, &K0mean, &Kplus, &Hmean, &r, &s ) ) {
      alpha_mean = 2*s*exp(-lambda0mean*(params->A+params->B))/(1-exp(-lambda0mean*params->B));
    }
    free(freq_mean);

    while( (seq = GetFastaSeq( seqfp ) ) ) {
      int A=params->A, B=params->B;
      if ( sw_score = SWScoreProteinToProfile( seq, pro, params->A, params->B ) ) {
	params->total++;
	pval1 = pval2 = EmpiricalGEM( lambda0mean, K0mean, Hmean, alpha_mean, seq->len, pro->length, &theta, &logkappa, sw_score );
	count++;
	if ( pval1 < params->pthresh1 ) {
	  double *freq = PseudoResidueFrequencies( seq, 100 );
	  if ( KarlinAltschulProfileStatistics( pro, freq, &lambda0, &K0, &Kplus, &H, &r, &s ) ) {
	    alpha = 2*s*exp(-lambda0*(params->A+params->B))/(1-exp(-lambda0*params->B));
	    if ( alpha > 0.25 ) {
	      IterativelyIncreaseGapPenalty( seq, pro, &A, &B, lambda0, K0, H, s, &alpha, &sw_score );
	    }
	    pval2 = EmpiricalGEM( lambda0, K0, H, alpha, seq->len, pro->length, &theta, &logkappa, sw_score );
	  }
	  else {
	    printf( "ERROR sequence %s has positive expected match score with profile %s - ignored\n", seq->name, pro->name );
	    pval2 = 1.0;
	  }
	  free(freq);
	}
	else {
	  pval2 = pval1;
	}
	if ( pval2 < params->pthresh2 ) {
	  if ( params->all ) {
	    params->hits += TopAlignProteinToProfile( seq, pro, params->A, params->B, 0, params->pthresh2, params->self, lambda0, K0, H, alpha );
	  }
	  else {
	    int length;
	    double percent;
	    PROFILE_ALIGNMENT *align = AlignProteinToProfile( seq, pro, A, B, &sw_score, &length, &percent );
	    PrintProteinProfileAlignment( stdout, params->total, pro, seq, align, length, sw_score,  params->dbsize*pval2, percent, K0*exp(logkappa), lambda0*theta, H, alpha, params->align );
	    fflush( stdout );
	    params->hits++;
	    free(align);
	  }
	}
	FreeSeq( seq );
      }
    }
  }
  rewind(seqfp);
  return count;
}

int CompareSequenceToSequences( FILE *seq2fp, FILE *seqfp, PPARAMS *params ) {
  
  SEQUENCE *seq2 = GetFastaSeq( seq2fp );  /* read in the query seq */
  int count=0, k;
  long file_pos=ftell(seq2fp);

  if ( seq2 ) {
    double *freq2 = PseudoResidueFrequencies( seq2, 100 );
    double *freq_mean = RobinsonResidueFrequencies ();
    double lambda0, K0, Kplus, H, r, s, alpha;
    double theta, logkappa;
    double K0mean=0.0, lambda0mean=0.0, Hmean=0.0, alpha_mean=0.0;
    double pval1=1.0, pval2=1.0;
    int sw_score;
    SEQUENCE *seq;
    PROFILE *pro = SeqToProfile(seq2, params->matrix); /* turn the query into a profile */

    /* initial estimates for K, L etc */

    if ( KarlinAltschulStatistics( params->matrix, freq2, freq_mean, &lambda0mean, &K0mean, &Kplus, &Hmean, &r, &s ) ) {
      alpha_mean = 2*s*exp(-lambda0mean*(params->A+params->B))/(1-exp(-lambda0mean*params->B));
    }
    free(freq_mean);

    if ( params->self )
      fseek(seqfp,file_pos,SEEK_SET);

    while( (seq = GetFastaSeq( seqfp ) ) ) {
      int A=params->A, B=params->B;
      if ( sw_score = SWScoreProteinToProfile( seq, pro, params->A, params->B ) ) {
	params->total++;
	pval1 = pval2 = EmpiricalGEM( lambda0mean, K0mean, Hmean, alpha_mean, seq->len, pro->length, &theta, &logkappa, sw_score );
	count++;
	if ( pval1 < params->pthresh1 ) {
	  double *freq = PseudoResidueFrequencies( seq, 100 );
	  if ( KarlinAltschulStatistics( params->matrix, freq2, freq, &lambda0, &K0, &Kplus, &H, &r, &s ) ) {
	    alpha = 2*s*exp(-lambda0*(params->A+params->B))/(1-exp(-lambda0*params->B));
	    if ( alpha > 0.25 ) 
	      IterativelyIncreaseGapPenalty( seq, pro, &A, &B, lambda0, K0, H, s, &alpha, &sw_score );
	    pval2 = EmpiricalGEM( lambda0, K0, H, alpha, seq->len, pro->length, &theta, &logkappa, sw_score );
	  }
	  else {
	    printf( "ERROR sequence %s has positive expected match score with profile %s - ignored\n", seq->name, pro->name );
	    pval2 = 1.0;
	  }
	  free(freq);
	}
	else {
	  pval2 = pval1;
	}
	if ( pval2 < params->pthresh2 ) {

	  if ( params->all ) {
	    params->hits += TopAlignProteinToProfile( seq, pro, params->A, params->B, 0, params->pthresh2, params->self, lambda0, K0, H, alpha );
	  }
	  else {
	    int length;
	    double percent;
	    
	    PROFILE_ALIGNMENT *align = AlignProteinToProfile( seq, pro, A, B, &sw_score, &length, &percent );
	    PrintProteinProfileAlignment( stdout, params->total, pro, seq, align, length, sw_score,  params->dbsize*pval2, percent, K0*exp(logkappa), lambda0*theta, H, alpha, params->align );
	    fflush( stdout );
	    params->hits++;
	    free(align);
	  }
	}
      }
      FreeSeq( seq );
    }
    FreeProfile(pro);
    free(freq2);
    FreeSeq(seq2);
  }
  rewind(seqfp);
  return count;
}

