/****************************************************************  
  Represents a PDB File -- automates parsing and data access.

	Author:  Aaron Davidson
	Created: Dec 2nd, 2002
	
  Copyright (c) BioTools 2002

  To Do:
    * check all SAFEMALLOC results and report errors.
    * handle multiple chains
 ****************************************************************/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <ctype.h>

#include "pdbfile.h"
#include "util.h"



/* PDB uses three letter codes for residue names. This table
   can be used to lookup the code based on the residue letter */

char PDB_RES_NAMES[26][3] = {
/*  A      B      C	 D	E   */
  "ALA", "???", "CYS", "ASP", "GLU",
/*  F	   G	  H	 I	J   */
  "PHE", "GLY", "HIS", "ILE", "???",
/*  K      L      M      N      O   */
  "LYS", "LEU", "MET", "ASN", "???",
/*  P      Q      R      S      T   */
  "PRO", "GLN", "ARG", "SER", "THR",
/*  U      V      W      X      Y   */
  "???", "VAL", "TRP", "???", "TYR",
/*  Z                               */
  "???"  };

/* Use this macro to convert a residue letter to it's 3 letter code: */
#define GET_RES_NAME(r)    (PDB_RES_NAMES[toupper(r) - 'A'])

/* Funciton prototypes. */
void static initPDBFile(PDBFile *pdb);
void static ensureCapacity(PDBFile *pdb);
void static parseAtom(Atom *a);
int static addAtom(PDBFile *pdb, char *line);
static void repairResName(char *tag);
static int is_selenomethionine(char *tag);
static void convert_selenomethionine_line_to_methionine(char *line);



/*************************************************************/
/* Get the next atom after index n, of the given type.       */
/* Returns -1 if none found.                                 */
/*************************************************************/
int getNextAtomOfType(PDBFile *pdb, int n, int type) {
   int resNum = pdb->atoms[n].resNum;
   while (pdb->atoms[n].resNum == resNum) {
      if (pdb->atoms[n].type == type) return n;
      if (++n >= pdb->numAtoms) return -1;
   }
   return -1;
}


/*
 * Get the next atom of the given name after or including the nth atom in
 * the file. 'mode' must be one of GET_ATOM_ANY_RES, GET_ATOM_SAME_RES, or
 * GET_ATOM_NEXT_RES.
 *
 * GET_ATOM_ANY_RES: Get atom from residue containing nth atom, or later.
 * GET_ATOM_SAME_RES: Get atom from the same residue as the nth atom.
 * GET_ATOM_NEXT_RES: Get atom from the residue following the nth atom.
 */
int getNextAtomIndex(PDBFile *pdb, int n, char *atom_name, int mode) {
	int i;
	int starting_res_id; // Residue ID of nth atom.
	int in_starting_residue = TRUE;
	int next_res_id = -999999; // Residue ID of the residue following the one with the nth atom.
	
	if (strlen(atom_name) != 4) {
		fprintf(stderr, "getNextAtomIndex(): atom_name parameter does not have length 4 (pdbfile.c)\n");
		exit(1);
	}
	
	starting_res_id = pdb->atoms[n].resNum;
	
	for (i=n; i<pdb->numAtoms; i++)  {
		
		if (mode == GET_ATOM_SAME_RES) {
			if (pdb->atoms[i].resNum != starting_res_id) {
				return -1;
			}
		}
		else if (mode == GET_ATOM_NEXT_RES) {
			if (in_starting_residue) {
				if (pdb->atoms[i].resNum != starting_res_id) {
					in_starting_residue = FALSE;
					next_res_id = pdb->atoms[i].resNum;
				}
				else {					
					continue;
				}
			}
			
			if (!in_starting_residue && (pdb->atoms[i].resNum != next_res_id)) {
				// We're in the residue 2 residues after the one with the nth atom.
				return -1;
			}
		}
		else if (mode != GET_ATOM_ANY_RES) {
			fprintf(stderr, "getNextAtomIndex(): Illegal mode parameter.\n"); exit(1);
		}
		
		if (strncmp(pdb->atoms[i].name, atom_name, 4) == 0) {
			return i;
		}
	}
	return -1;
}

/**
 * Free memory associated with this PDBFile
 */
void freePDBFile(PDBFile *pdb) {
  int i;
  for (i=0; i<pdb->numAtoms; i++) {
    SAFEFREE(pdb->atoms[i].line);
  }
  SAFEFREE(pdb->atoms);
}

/**
 * Initialize memory for a PDBFile.
 * This must be called on your PDBFile before using it.
 */
void static initPDBFile(PDBFile *pdb) {
  pdb->numAtoms = 0;
  pdb->numRes = 0;
  pdb->capacity = 1000;
  pdb->atoms = (Atom *)SAFEMALLOC(sizeof(Atom) * pdb->capacity);
}

/**
 * Ensure that the capacity of the atoms array
 * can hold another atom. If not, grow the array.
 */
void static ensureCapacity(PDBFile *pdb) {
  if ((pdb->numAtoms + 1) >= pdb->capacity) {
    int nc = pdb->capacity * 2;
    Atom *atomPointer = pdb->atoms;
    pdb->atoms = (Atom *)SAFEMALLOC(sizeof(Atom) * nc);
    memcpy(pdb->atoms, atomPointer, pdb->capacity * sizeof(Atom));
    SAFEFREE(atomPointer);
    pdb->capacity = nc;
  }
}


/**
 * returns true if this is a CA Atom (Alpha Carbon)
 */
 /*
int atomIsCA(Atom *a) {
  return (a->type == CA_TYPE);
}
*/

/**
 * Get the 3D coordinates of an atom,
 * Places the position into the array.
 */
 /*
void getAtomPosition(Atom *a, double pos[3]) {
  pos[0] = a->point[0];
  pos[1] = a->point[1];
  pos[2] = a->point[2];
}
*/


/*  
     Example PDB ATOM Line, with ruler:
0         10        20       30        40        50        60
ATOM    388  CA  GLY G  89      37.970   4.298  27.320  1.00  0.00
*/
/**
 * Takes an ATOM tag line and creates a parsed Atom struct.
 *
 * Modified June 11, 2007 by David Arndt to use 4-character atom name abbreviations (this is
 * how atom names should be treated in PDB files) and to change any unusually-formatted atom
 * names to a more standard format.
 * See http://www.biochem.ucl.ac.uk/~roman/procheck/manual/manappb.html for information
 * on which character columns in a PDB file are reserved for which data.
 */
void static parseAtom(Atom *a) {
	char tmp[9];
	
	// Look at a->line for parsing.
	
	
	// Parse atom name, and change any unusually-formatted atom names into a more standard format.
	// Also set the atom type.
	
	if ( a->line[12] == 'H' && isupper(a->line[13]) && isdigit(a->line[14]) && isdigit(a->line[15]) ) {
		// Atom name such as "HB12", which should be changed to "2HB1".
		a->name[0] = a->line[15];
		a->name[1] = a->line[12];
		a->name[2] = a->line[13];
		a->name[3] = a->line[14];
		a->type = H_TYPE;
	}
	else if ( strncmp(&(a->line[12]), " HT", 3) == 0 && isdigit(a->line[15]) ) {
		// Atom name such as " HT1" (unusual designation for N-terminal hydrogens),
		// which should be changed to "1H  ".
		a->name[0] = a->line[15];
		a->name[1] = 'H';
		a->name[2] = ' ';
		a->name[3] = ' ';
		a->type = HN_TYPE;
	}
	else if ( a->line[12] == ' ' && a->line[13] == 'H' && isdigit(a->line[14]) && a->line[15] == ' ') {
		// Atom name such as " H1 " (seen as a designation for N-terminal hydrogens),
		// which should be changed to "1H  ".
		a->name[0] = a->line[14];
		a->name[1] = 'H';
		a->name[2] = ' ';
		a->name[3] = ' ';
		a->type = HN_TYPE;
	}
	//else if ( strncmp(&(a->line[12]), " H", 2) == 0 && isupper(a->line[14]) && isdigit(a->line[15])) {
	//	// Atom name such as " HB2", which should be changed to "2HB ".
	//	a->name[0] = a->line[15];
	//	a->name[1] = a->line[13];
	//	a->name[2] = a->line[14];
	//	a->name[3] = ' ';
	//	a->type = HN_TYPE;
	//}
	else if ( strncmp(&(a->line[12]), " HN ", 4) == 0 ) {
		// Atom name " HN ", which should be changed to " H  ".
		strcpy(a->name, " H  ");
		a->type = HN_TYPE;
	}
	else if ( strncmp(&(a->line[12]), " OT1", 4) == 0 ) {
		// First C-terminal O.
		strcpy(a->name, " O  ");
		a->type = O_TYPE;
	}
	else if ( strncmp(&(a->line[12]), " OT2", 4) == 0 ) {
		// Second C-terminal O.
		strcpy(a->name, " OXT");
		a->type = OXT_TYPE;
	}
	else if ( strncmp(&(a->line[12]), " O1 ", 4) == 0 ) {
		// First C-terminal O.
		strcpy(a->name, " O  ");
		a->type = O_TYPE;
	}
	else if ( strncmp(&(a->line[12]), " O2 ", 4) == 0 ) {
		// Second C-terminal O.
		strcpy(a->name, " OXT");
		a->type = OXT_TYPE;
	}
	else {
		// Assume any other atom names can be used as-is.
		
		strncpy(a->name, &(a->line[12]), 4);
		
		if ( strncmp(&(a->line[12]), " CA ", 4) == 0 ) {
			a->type = CA_TYPE;
		}
		else if ( strncmp(&(a->line[12]), " N  ", 4) == 0 ) {
			a->type = N_TYPE;
		}
		else if ( strncmp(&(a->line[12]), " C  ", 4) == 0 ) {
			a->type = C_TYPE;
		}
		else if ( strncmp(&(a->line[12]), " H  ", 4) == 0 ) {
			a->type = HN_TYPE;
		}
		else if ( strncmp(&(a->line[13]), "H  ", 3) == 0 && isdigit(a->line[12])) {
			// Atom name e.g. "1H  ".
			a->type = HN_TYPE;
		}
		else if ( strncmp(&(a->line[12]), " HA ", 4) == 0 || strncmp(&(a->line[12]), "1HA ", 4) == 0 ) {
			a->type = HA_TYPE;
		}
		else if ( strncmp(&(a->line[12]), "2HA ", 4) == 0 ) {
			a->type = HA2_TYPE;
		}
		else if ( a->line[13] == 'H' ) {
			// Atom name matching regex ".H.."
			strncpy(a->name, &(a->line[12]), 4);
			a->type = H_TYPE;
		}
		else if ( strncmp(&(a->line[12]), " CB ", 4) == 0 ) {
			a->type = CB_TYPE;
		}
		else if ( strncmp(&(a->line[13]), "CG", 2) == 0 || strncmp(&(a->line[13]), "OG", 2) == 0 ||
				strncmp(&(a->line[13]), "SG", 2) == 0 ) {
			a->type = CG_TYPE;
		}
		else if ( strncmp(&(a->line[12]), " O  ", 4) == 0 ) {
			a->type = O_TYPE;
		}
		else if ( strncmp(&(a->line[12]), " OXT", 4) == 0 ) {
			a->type = OXT_TYPE;
		}
		else {
			a->type = UNKNOWN_TYPE;
		}
	}
	
	// Add NULL terminator
	a->name[4] = (char) NULL;
	
	
	// Parse residue name, and repair residue name if necessary.
	strncpy(a->resName, &(a->line[17]), 3);
	a->resName[3] = (char) NULL;
	repairResName(a->resName);
	
	// Parse residue number.
	sscanf(&a->line[22], "%d", &a->resNum);

    // Add below tow lines to keep original residue number (2011.08.24)
    a->origResNum = a->resNum;
	sscanf(&a->line[26], "%c", &a->insCode);

	// Parse distances.
	tmp[8] = (char) NULL;
	strncpy(tmp, &a->line[30], 8);
	sscanf(tmp,"%lf", &(a->point[0]));
	strncpy(tmp, &a->line[38], 8);
	sscanf(tmp,"%lf", &(a->point[1]));
	strncpy(tmp, &a->line[46], 8);
	sscanf(tmp,"%lf", &(a->point[2]));
	
	// debug
	//fprintf(stderr, "%3s %d: %4s %d\n", a->resName, a->resNum, a->name, a->type);
}


/**
 * Adds the ATOM tag (line) into the PDB file as a new Atom struct.
 * Returns TRUE if an atom was added, FALSE otherwise.
 */
int static addAtom(PDBFile *pdb, char *line) {
	Atom *a;
	char alternative_location_code;
	
	// Check for alternative location code. Added by David Arndt, Feb 12, 2007.
	alternative_location_code = line[16];
	if ( !(alternative_location_code == ' ' || alternative_location_code == 'A') ) {
		// Alternative location code is not blank and is not 'A', so we will return FALSE so
		// that this line of the PDB file will be skipped. (The alternative location code is
		// sometimes used when there are multiple possible positions of a residue's side
		// chain that are specified in the PDB file, and thus multiple possible locations for
		// the same atom).
		return FALSE;
	}

	ensureCapacity(pdb);
	a = &pdb->atoms[pdb->numAtoms];
	a->line = (char *)SAFEMALLOC(strlen(line)+1);
	strcpy(a->line, line);
	parseAtom(a);
	pdb->numAtoms++;
	
	return TRUE;
}


/**
 * Parse a PDB file in to a PDBFile record.
 * Currently reads only the first chain.
 * Pass chain of '*' for all chains and '?' 
 * for the first chain only.
 */
int parsePDBFile(char *name, PDBFile *pdb, char chain) {
	int rc = TRUE;
	FILE *fp;
	char *status;
	char line[MAX_LINE_LEN];
	int lastRes = -100000;
	int startedTargetChain = 0;
	int is_aa_line = FALSE;
	
	if ((fp = fopen(name,"r")) != NULL) {
		initPDBFile(pdb);
		status = fgets(line, MAX_LINE_LEN, fp);
		while (status != 0) {
			
			is_aa_line = FALSE;
			if (strncmp("ATOM ", line, 4) == 0) {
				if (is_selenomethionine(&line[17])) {
					convert_selenomethionine_line_to_methionine(line);
				}
				is_aa_line = TRUE;
			}
			else if (strncmp("HETATM ", line, 6) == 0) {
				if (strlen(line) >=	35 && is_selenomethionine(&line[17])) {
					convert_selenomethionine_line_to_methionine(line);
					is_aa_line = TRUE;
				}
			}
			
			if (is_aa_line) {
				if ((res_str_to_char(&line[17]) != '?') &&
						(line[21] == chain || chain == '*' || chain == '?')) {
					// Residue 3-letter abbreviation must be from one of the standard 20 amino acids.
					// Chain must be the same as given, or all or any chain.
					
					startedTargetChain = 1;

					if (addAtom(pdb, line)) {
                        /////////////////////////////////////////////////////////////////////////////
                        // Here may need to modify for supporting original residue info
                        /////////////////////////////////////////////////////////////////////////////
						if (pdb->atoms[pdb->numAtoms-1].resNum != lastRes) {
							lastRes = pdb->atoms[pdb->numAtoms-1].resNum;
							pdb->numRes++;
						}
					}
					
					if (chain != '*') {
						chain = line[21];
					}
				}
			}
			else if (startedTargetChain && strncmp("TER ", line, 3) == 0) {
				if (chain != '*') {
					
					// debug
					//printf("Stopped parsing PDB file at line: %s\n", line);
					
					// '?' means only read the first chain. We are done reading the first chain,
					// so break out of the while loop.
					break;
				}
			}
			status = fgets(line, MAX_LINE_LEN, fp);
		}
		
		fclose(fp);
	} else {
		doError("Could not open PDB file", __FILE__, __LINE__);
		rc = ERROR_FLAG;
	}
	return rc;
}


/*
 * Returns the single-letter amino acid abbreviation corresponding to the given 3-letter
 * amino acid abbreviation. If the given 3-letter abbreviation is not from one of the
 * 20 standard amino acids, '?' is returned.
 *
 * Some non-standard abbreviations are interpreted as standard abbreviations: e.g. "HSD" and
 * "HID" for "HIS", and "CYX" for "CYS". These unusual abbreviations are used by some
 * common programs, e.g. "HSD" is output by NAMD.
 */
char res_str_to_char(char *three_letter_abbrev) {
	if (strncmp("ALA", three_letter_abbrev, 3) == 0) return 'A';
	if (strncmp("CYS", three_letter_abbrev, 3) == 0 || strncmp("CYX", three_letter_abbrev, 3) == 0) return 'C';
	if (strncmp("ASP", three_letter_abbrev, 3) == 0) return 'D';
	if (strncmp("GLU", three_letter_abbrev, 3) == 0) return 'E';
	if (strncmp("PHE", three_letter_abbrev, 3) == 0) return 'F';
	if (strncmp("GLY", three_letter_abbrev, 3) == 0) return 'G';
	if (strncmp("HIS", three_letter_abbrev, 3) == 0 || strncmp("HSD", three_letter_abbrev, 3) == 0 ||
		strncmp("HID", three_letter_abbrev, 3) == 0 || strncmp("HIE", three_letter_abbrev, 3) == 0) return 'H';
	if (strncmp("ILE", three_letter_abbrev, 3) == 0) return 'I';
	if (strncmp("LYS", three_letter_abbrev, 3) == 0) return 'K';
	if (strncmp("LEU", three_letter_abbrev, 3) == 0) return 'L';
	if (strncmp("MET", three_letter_abbrev, 3) == 0) return 'M';
	if (strncmp("ASN", three_letter_abbrev, 3) == 0) return 'N';
	if (strncmp("PRO", three_letter_abbrev, 3) == 0) return 'P';
	if (strncmp("GLN", three_letter_abbrev, 3) == 0) return 'Q';
	if (strncmp("ARG", three_letter_abbrev, 3) == 0) return 'R';
	if (strncmp("SER", three_letter_abbrev, 3) == 0) return 'S';
	if (strncmp("THR", three_letter_abbrev, 3) == 0) return 'T';
	if (strncmp("VAL", three_letter_abbrev, 3) == 0) return 'V';
	if (strncmp("TRP", three_letter_abbrev, 3) == 0) return 'W';
	if (strncmp("TYR", three_letter_abbrev, 3) == 0) return 'Y';
	
	// debug
	//fprintf(stderr, "Unknown residue \'%.3s\'.\n", three_letter_abbrev);
	
	return '?';
}


/***********************************************************
 * Converts the residue's single character representation
 * to its 3 character representation.
 ***********************************************************/
char *char_to_res_str(char single_letter_abbrev) {
  
  single_letter_abbrev = (char) toupper(single_letter_abbrev);
  
  switch(single_letter_abbrev){
  case 'A': return "ALA"; break;
  case 'R': return "ARG"; break;
  case 'N': return "ASN"; break;
  case 'D': return "ASP"; break;
  case 'C': return "CYS"; break;
  case 'Q': return "GLN"; break;
  case 'E': return "GLU"; break;
  case 'G': return "GLY"; break;
  case 'H': return "HIS"; break;
  case 'I': return "ILE"; break;
  case 'L': return "LEU"; break;
  case 'K': return "LYS"; break;
  case 'M': return "MET"; break;
  case 'F': return "PHE"; break;
  case 'P': return "PRO"; break;
  case 'S': return "SER"; break;
  case 'T': return "THR"; break;
  case 'W': return "TRP"; break;
  case 'Y': return "TYR"; break;
  case 'V': return "VAL"; break;
  default : return "FIL"; 
  }/* END _switch_ */

}/* END char_to_res_str () */


/*
 * Given a 3-letter amino acid abbreviation, change it from a non-standard abbreviation
 * to a standard abbreviation, if necessary.
 */
static void repairResName(char *tag) {
	if (strncmp("CYX", tag, 3) == 0) {
		strncpy(tag, "CYS", 3);
	}
	else if (strncmp("HSD", tag, 3) == 0 || strncmp("HID", tag, 3) == 0 ||
			strncmp("HIE", tag, 3) == 0) {
		strncpy(tag, "HIS", 3);
	}
}


/*
 * Given a 3-letter amino acid abbreviation, returns TRUE if it refers to a
 * selenomethionine residue, FALSE otherwise. Selenomethionine is a non-standard
 * amino acid that is commonly used in X-ray structure determination.
 */
static int is_selenomethionine(char *tag) {
	if (strncmp("MSE", tag, 3) == 0 || strncmp("SEM", tag, 3) == 0 || strncmp("SME", tag, 3) == 0) {
		return TRUE;
	}
	else {
		return FALSE;
	}
}


/*
 * Given a line from a PDB file for a selenomethionine residue, transform the line
 * into a line for a methionine residue.
 */
static void convert_selenomethionine_line_to_methionine(char *line) {
	
	// Change residue name.
	strncpy(&line[17], "MET", 3);
	
	// Change selenium atom to sulfur atom (delta position).
	if (strncmp(&line[12], "SE ", 3) == 0 || strncmp(&line[12], " SE", 3) == 0) {
		strncpy(&line[12], " SD", 3);
	}
}


int getFirstAtomOfResidue(PDBFile *pdb, int resNum) {
	int i=0, n=0;
	int last = pdb->atoms[0].resNum;
	for (i=0; i<pdb->numAtoms; i++) {
		if (pdb->atoms[i].resNum != last) {
			n++;
			last = pdb->atoms[i].resNum;
		}
		if (n == resNum) return i;
	}
	return -1;
}


/*
 * Given a PDB file, this function writes a trimmed PDB file with only the first chain or
 * model of the original file. If remove_H_atoms is TRUE, no hydrogen atoms will be output
 * to the trimmed PDB file.
 *
 * input_PDB_filename: The name of the initial PDB file.
 * trimmed_PDB_filename: The desired name of the trimmed PDB file.
 * remove_H_atoms: Boolean. Whether to remove hydrogen atoms from the structure.
 */
void trim_PDB_file(char *input_PDB_filename, char *trimmed_PDB_filename, int remove_H_atoms) {
	FILE *pdb_fp;
	FILE *trimmed_fp;
	char line[MAX_LINE_LEN];
	
	pdb_fp = fopen(input_PDB_filename, "r");
	trimmed_fp = fopen(trimmed_PDB_filename, "w");
	
	while (fgets(line, MAX_LINE_LEN, pdb_fp) != NULL) {
		if (strncmp(line, "ATOM", 4) == 0) {
			
			// Filter to not pass lines with hydrogens.
			if (remove_H_atoms && (line[12] == 'H' || line[13] == 'H')) {
				continue;
			}
			
			fprintf(trimmed_fp, "%s", line);
		}
		else if (strncmp(line, "TER", 3) == 0 || strncmp(line, "END", 3) == 0) {
			fprintf(trimmed_fp, "%s", line);
			break;
		}
	}
	
	fclose(pdb_fp);
	fclose(trimmed_fp);
}

