/**********************************************************************************************************************************
 * hydropho_calc.c : 
 *       an hydrophobicity calculator for PDB files
 * 
 * Author: Yifeng Liu (yifeng@cs.ualberta.ca)
 *         Wishart Research Group (Department of Computing Science, University of Alberta)
 *
 * Date: May 8th, 2010
 * 
 * Synopsis:
 * 
 * hydropho_calc.c takes a PDB file as input and outputs the averaged
 * hydrophobicity value for each residue based on user specified window size
 * and hydrophobicity scale.
 *
 * Algorithm:
 *
 * 1. Read the first chain of the PDB file into a char array; each char represents a single amino acid in the chain.
 * 2. Convert each residue into its corresonding hydrophobicity with the selected hydrophobicity scale
 * 3. Slide a window of odd length from the begining of the chain to the end of the chain and calculate the averaged hydrophobicity.
 * 4. Store the result into an output file.
 * 
 * Hydrophobicity scales:
 * Currently 7 hydrophobicity scales are supported:
 *
 * 1. Kyte-Doolittle (Kyte and Doolittle, 1982)
 * 2. Engelman (Engelman et al., 1986) 
 * 3. Eisenberg (Eisenberg et al., 1984)
 * 4. Hopp-Woods (Hoop and Woods, 1983)
 * 5. Cornette (Cornette et al., 1987)
 * 6. Rose (Rose et al., 1985)
 * 7. Janin (Janin, 1997)
 *
 * The hydrophobicity indices implemented in this program is taken from 
 * http://www.clcbio.com/index.php?id=43
 *
 * Detail Hydrophobicity Scale Information:
 *    aa	aa 	        Kyte-Doolittle 	Hopp-Woods 	Cornette 	Eisenberg 	Rose 	Janin 	Engelman (GES)
 *    A 	Alanine 	1.80 		-0.50 		0.20 		0.62 		0.74 	0.30 	1.60
 *    C 	Cysteine 	2.50 		-1.00 		4.10 		0.29 		0.91 	0.90 	2.00
 *    D 	Aspartic acid 	-3.50 		3.00 		-3.10 		-0.90 		0.62 	-0.60 	-9.20
 *    E 	Glutamic acid 	-3.50 		3.00 		-1.80 		-0.74 		0.62 	-0.70 	-8.20
 *    F 	Phenylalanine 	2.80 		-2.50 		4.40 		1.19 		0.88 	0.50 	3.70
 *    G 	Glycine 	-0.40 		0.00 		0.00 		0.48 		0.72 	0.30 	1.00
 *    H 	Histidine 	-3.20 		-0.50 		0.50 		-0.40 		0.78 	-0.10 	-3.00
 *    I 	Isoleucine 	4.50 		-1.80 		4.80 		1.38 		0.88 	0.70 	3.10
 *    K 	Lysine 		-3.90 		3.00 		-3.10 		-1.50 		0.52 	-1.80 	-8.80
 *    L 	Leucine 	3.80 		-1.80 		5.70 		1.06 		0.85 	0.50 	2.80
 *    M 	Methionine 	1.90 		-1.30 		4.20 		0.64 		0.85 	0.40 	3.40
 *    N 	Asparagine 	-3.50 		0.20 		-0.50 		-0.78 		0.63 	-0.50 	-4.80
 *    P 	Proline 	-1.60 		0.00 		-2.20 		0.12 		0.64 	-0.30 	-0.20
 *    Q 	Glutamine 	-3.50 		0.20 		-2.80 		-0.85 		0.62 	-0.70 	-4.10
 *    R 	Arginine 	-4.50 		3.00 		1.40 		-2.53 		0.64 	-1.40 	-12.3
 *    S 	Serine 		-0.80 		0.30 		-0.50 		-0.18 		0.66 	-0.10 	0.60
 *    T 	Threonine 	-0.70 		-0.40 		-1.90 		-0.05 		0.70 	-0.20 	1.20
 *    V 	Valine 		4.20 		-1.50 		4.70 		1.08 		0.86 	0.60 	2.60
 *    W 	Tryptophan 	-0.90 		-3.40 		1.00 		0.81 		0.85 	0.30 	1.90
 *    Y 	Tyrosine 	-1.30 		-2.30 		3.20 		0.26 		0.76 	-0.40 	-0.70
 *
 **********************************************************************************************************************************/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
// #include </usr/include/string.h> // for mac os x
#include <limits.h>

#include "hydropho_cal.h"

Hydroscale kyte_doolittle = {1.80,2.50,-3.50,-3.50,2.80,-0.40,-3.20,4.50,-3.90,3.80,1.90,-3.50,-1.60,-3.50,-4.50,-0.80,-0.70,4.20,-0.90,-1.30};
Hydroscale engelman = {1.60,2.00,-9.20,-8.20,3.70,1.00,-3.00,3.10,-8.80,2.80,3.40,-4.80,-0.20,-4.10,-12.3,0.60,1.20,2.60,1.90,-0.70};
Hydroscale eisenberg = {0.62,0.29,-0.90,-0.74,1.19,0.48,-0.40,1.38,-1.50,1.06,0.64,-0.78,0.12,-0.85,-2.53,-0.18,-0.05,1.08,0.81,0.26};
Hydroscale hoop_woods = {-0.50,-1.00,3.00,3.00,-2.50,0.00,-0.50,-1.80,3.00,-1.80,-1.30,0.20,0.00,0.20,3.00,0.30,-0.40,-1.50,-3.40,-2.30};
Hydroscale cornette = {0.20,4.10,-3.10,-1.80,4.40,0.00,0.50,4.80,-3.10,5.70,4.20,-0.50,-2.20,-2.80,1.40,-0.50,-1.90,4.70,1.00,3.20};
Hydroscale rose = {0.74,0.91,0.62,0.62,0.88,0.72,0.78,0.88,0.52,0.85,0.85,0.63,0.64,0.62,0.64,0.66,0.70,0.86,0.85,0.76};
Hydroscale janin = {0.30,0.90,-0.60,-0.70,0.50,0.30,-0.10,0.70,-1.80,0.50,0.40,-0.50,-0.30,-0.70,-1.40,-0.10,-0.20,0.60,0.30,-0.40};

/* convert a residue from 3 letter code to 1 letter code */
char aa3to1(char *residue){
  char aa;
  if(strcmp(residue, "ALA")==0)     { aa = 'A'; }
  else if(strcmp(residue, "CYS")==0){ aa = 'C'; }
  else if(strcmp(residue, "ASP")==0){ aa = 'D'; }
  else if(strcmp(residue, "GLU")==0){ aa = 'E'; }
  else if(strcmp(residue, "PHE")==0){ aa = 'F'; }
  else if(strcmp(residue, "GLY")==0){ aa = 'G'; }
  else if(strcmp(residue, "HIS")==0){ aa = 'H'; }
  else if(strcmp(residue, "ILE")==0){ aa = 'I'; }
  else if(strcmp(residue, "LYS")==0){ aa = 'K'; }
  else if(strcmp(residue, "LEU")==0){ aa = 'L'; }
  else if(strcmp(residue, "MET")==0){ aa = 'M'; }
  else if(strcmp(residue, "ASN")==0){ aa = 'N'; }
  else if(strcmp(residue, "PRO")==0){ aa = 'P'; }
  else if(strcmp(residue, "GLN")==0){ aa = 'Q'; }
  else if(strcmp(residue, "ARG")==0){ aa = 'R'; }
  else if(strcmp(residue, "SER")==0){ aa = 'S'; }
  else if(strcmp(residue, "THR")==0){ aa = 'T'; }
  else if(strcmp(residue, "VAL")==0){ aa = 'V'; }
  else if(strcmp(residue, "TRP")==0){ aa = 'W'; }
  else if(strcmp(residue, "TYR")==0){ aa = 'Y'; }
  else { 
     aa='?'; 
     printf("Hydrophobicity: Non-standard residue type: %s\n", residue); 
  }

  return aa;
}

/* convert a residue from 1 letter code to 3 letter code */
char* aa1to3(char residue){
  char *aa = NULL;
  switch (residue){
  case 'A': aa = "ALA"; break;
  case 'C': aa = "CYS"; break;
  case 'D': aa = "ASP"; break;
  case 'E': aa = "GLU"; break;
  case 'F': aa = "PHE"; break;
  case 'G': aa = "GLY"; break;
  case 'H': aa = "HIS"; break;
  case 'I': aa = "ILE"; break;
  case 'K': aa = "LYS"; break;
  case 'L': aa = "LEU"; break;
  case 'M': aa = "MET"; break;
  case 'N': aa = "ASN"; break;
  case 'P': aa = "PRO"; break;
  case 'Q': aa = "GLN"; break;
  case 'R': aa = "ARG"; break;
  case 'S': aa = "SER"; break;
  case 'T': aa = "THR"; break;
  case 'V': aa = "VAL"; break;
  case 'W': aa = "TRP"; break;
  case 'Y': aa = "TYR"; break;
  }
  return aa;
}

/* store the selected hydrophobicity scale in struct scale */
void hydroscale_selector(int scaleid, Hydroscale *scale){
  switch(scaleid){
  case 1: *scale = kyte_doolittle; /* 1. Kyte-doolittle scale */
    break; 
  case 2: *scale = engelman; /* 2. engelman scale */
    break;
  case 3: *scale = eisenberg; /* 3. eisenberg scale */
    break;
  case 4: *scale = hoop_woods; /* 4. Hoop-Wood scale */
    break;
  case 5: *scale = cornette; /* 5. cornette scale */
    break;
  case 6: *scale = rose; /* 6. rose scale */
    break;
  case 7: *scale = janin; /* 7. Janin scale */
    break;
  default: *scale = kyte_doolittle;
  }
}

/* translate residue to hydroscore */
float hydroscore(Hydroscale *scale, char residue){
  float score = 0;
  switch(residue){
  case 'A': score = (*scale).resA; break;
  case 'C': score = (*scale).resC; break;
  case 'D': score = (*scale).resD; break;
  case 'E': score = (*scale).resE; break;
  case 'F': score = (*scale).resF; break;
  case 'G': score = (*scale).resG; break;
  case 'H': score = (*scale).resH; break;
  case 'I': score = (*scale).resI; break;
  case 'K': score = (*scale).resK; break;
  case 'L': score = (*scale).resL; break;
  case 'M': score = (*scale).resM; break;
  case 'N': score = (*scale).resN; break;
  case 'P': score = (*scale).resP; break;
  case 'Q': score = (*scale).resQ; break;
  case 'R': score = (*scale).resR; break;
  case 'S': score = (*scale).resS; break;
  case 'T': score = (*scale).resT; break;
  case 'V': score = (*scale).resV; break;
  case 'W': score = (*scale).resW; break;
  case 'Y': score = (*scale).resY; break;
  default : score = 0;
  }
  return score;
}

/* function for copying string */
char * strndup (const char *s, size_t n)
{
  char *result;
  size_t len = strlen (s);

  if (n < len)
    len = n;

  result = (char *) malloc (len + 1);
  if (!result)
    return 0;

  result[len] = '\0';
  return (char *) memcpy (result, s, len);
}


/* function for substring */
char* substring(const char* str, size_t begin, size_t len)
{
  if (str == 0 || strlen(str) == 0 || strlen(str) < begin || strlen(str) < (begin+len))
    return 0;
  return (char *) strndup(str + begin, len);
}

/* calculate the average value from begin to end in the given array */
float average_score(float *array, int array_length, int begin, int end){
  /* sanity check for index */
  int tmp;
  float average = 0;
  int length = 0;
  int i;

  if (begin<0){ 
    begin=0;
  }else if (begin>=array_length){
    begin=array_length-1;
  }

  if (end<0){
    end=0;
  }else if (end>=array_length){
    end = array_length-1;
  }
 
  if (begin>end){
    tmp = end;
    end = begin;
    begin = tmp;
  }

  for (i=begin; i<=end; i++){
    average = average + array[i];
    length++;
  }
  average = average / length;

  return average;
}

/* print a usage information */
void print_usage(){
  printf("Usage: ./hydropho_cal <input_pdb> <output_file> <window_size> <scale>\n");
  printf("<input_pdb>   : the input PDB file, it must contain the SEQRES section\n");
  printf("<output_file> : output result file, existing file WILL BE OVERWRITTEN\n");
  printf("<window_size> : an odd integer: size of the sliding window\n");
  printf("<scale>       : an integer from 1 to 7, correspondes to different hydrophobicity scale: \n");
  printf("		1. Kyte-Doolittle (Kyte and Doolittle, 1982)\n");
  printf("		2. Engelman (Engelman et al., 1986)\n");
  printf("		3. Eisenberg (Eisenberg et al., 1984)\n");
  printf("		4. Hopp-Woods (Hoop and Woods, 1983)\n");
  printf("		5. Cornette (Cornette et al., 1987)\n");
  printf("		6. Rose (Rose et al., 1985)\n");
  printf("		7. Janin (Janin, 1997)\n");
}

/* 
 * get the first amino acid chain and store it into a char array: one char per residue
 * and return the length of the chain
 */
int get_chain(char* pdbfile, char *chain){
  char *buffer;
  FILE *in = NULL;
  char *prefix = NULL;
  char *prefix3 = NULL;
  int pre_resnum = INT_MIN;
  int chain_length = 0;
  int i = 0;

  in = fopen(pdbfile, "r"); 
  if(in==NULL){
    printf("Error: cannot open input file\n");
    exit(-1);
  }

  buffer = malloc(1024);
  if (buffer == NULL){
    printf("Error: malloc failed\n");
    exit(-1);
  }


  /* put all residues of the first chain into a char array */
  while (fgets (buffer, 1024, in) != NULL ){
    prefix = substring(buffer, 0, 4); /* get the word ATOM */
    prefix3 = substring(buffer, 0, 3); /* get the word TER */

    if (strcmp("TER", prefix3)==0){
      break;
    }

    if (strcmp("ATOM", prefix)==0 || strcmp("TER", prefix3)==0){
      // char *residue = substring(buffer, 17, 3);
      char *resnum_str = substring(buffer, 23, 3);
      int resnum;
      char res = (char) aa3to1(substring(buffer, 17, 3));
      sscanf(resnum_str, "%d", &resnum);

      if(pre_resnum == INT_MIN || resnum != pre_resnum){
	// printf("%s\t%c\t%s\t%d\n", residue, res, resnum_str, resnum);
	chain[i] = res;
	i++;
	pre_resnum = resnum;
      }
    }
  }
  fclose(in);
  chain_length = i;
  free(buffer);
  return chain_length;
}

/* return hydrophobicity score array with matched input residues string and number */
void calc_hydrophobicity(char *res_seq, int chain_length, float *hydro_scores, int window_size, int scale_option)
{
  float hydro_chain[MAX_LENGTH]; /* float array for storing hydrophobicity score for each residue */

  int halfwindow = (int) (window_size-1)/2; /* half a sliding window size */
  int i; /* simple counter */

  Hydroscale scale; 

  /* get hydroscale */
  hydroscale_selector(scale_option, &scale);

  /* convert each residue into corresponding hydrophobicity score */
  for(i=0; i<chain_length; i++){
    hydro_chain[i] = hydroscore(&scale, res_seq[i]);
    // printf("%c\t%f\n", res_seq[i], hydro_chain[i]);
  }

  /* slide a windows of size window_size over the array and calculate average hydrophobicity score */
  for(i=0; i<chain_length; i++){
    int begin = i-halfwindow;
    int end = i+halfwindow;
    float cur_score = average_score(hydro_chain, chain_length, begin, end);
    // char cur_residue = res_seq[i];
    // printf("%d\t%c\t%.3f\n", i+1, cur_residue, cur_score);

    // i - should use input residue no, not itself
    hydro_scores[i] = cur_score;
  }

}


void do_test() {

	int chain_length = 28;
	char *chainp = "SDKIIHLTDDSFDTDVLKADGAILVDFW";
	float hydro_scores[1024];

	calc_hydrophobicity(chainp, chain_length, hydro_scores, 7, 1);
	int i;
	for(i=0; i < chain_length; i++) {
	    printf("%d\t%c\t%.3f\n", i+1, *(chainp+i), hydro_scores[i]);
	}
}


/* main program */

int hydropho_main(int argc, char *argv[])
{
	if (argc<5){
		print_usage();
		return -1;
	}

	/* parameters */
	char infile[512];  /* input file name */
	char outfile[512]; /* output file name */

	char chain[MAX_LENGTH]; /* char array for storing amino acid residues */
	float hydro_chain[MAX_LENGTH]; /* float array for storing hydrophobicity score for each residue */

	int scale_option=1; /* hydrophobicity scale option */
	int window_size=7;  /* size of window */
	int chain_length=0; /* length of the first amino acid chain */
	int halfwindow = (int) (window_size-1)/2; /* half a sliding window size */
	int i; /* simple counter */

	FILE *out=NULL; /* output file */
	Hydroscale scale; 

	/* extract input parameters */
	strcpy(infile, argv[1]); /* input file */
	strcpy(outfile, argv[2]); /* output file */
	window_size = atoi(argv[3]); /* window size */
	scale_option = atoi(argv[4]); /* option for hydrophobicity scale */

	/* get hydroscale */
	hydroscale_selector(scale_option, &scale);


	/* get the first amino acid chain */
	chain_length = get_chain(infile, chain);

	/* convert each residue into corresponding hydrophobicity score */
	for(i=0; i<chain_length; i++){
	     hydro_chain[i] = hydroscore(&scale, chain[i]);
	     // printf("%c\t%f\n", chain[i], hydro_chain[i]);
	}


	/* slide a windows of size window_size over the array and calculate average hydrophobicity score */
	out = fopen(outfile, "w");
	for(i=0; i<chain_length; i++){
	     int begin = i-halfwindow;
	     int end = i+halfwindow;
	     float cur_score = average_score(hydro_chain, chain_length, begin, end);
	     char cur_residue = chain[i];
	     fprintf(out, "%d\t%c\t%.3f\n", i+1, cur_residue, cur_score);
	}
	close(out);


        // do_test();

        return 0;
}
