/*
 * Copyright (C) 1997-2003, R3vis Corporation.
 *
 * 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 2
 * 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, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA,
 * or visit http://www.gnu.org/copyleft/gpl.html.
 *
 * Contributor(s):
 *   Wes Bethel, R3vis Corporation, Marin County, California
 *
 * The OpenRM project is located at http://openrm.sourceforge.net/.
 */
/*
 * $Id: pdbwork.c,v 1.5 2003/04/13 18:13:23 wes Exp $
 * $Revision: 1.5 $
 * $Name: OpenRM-1-5-2-RC1 $
 * $Log: pdbwork.c,v $
 * Revision 1.5  2003/04/13 18:13:23  wes
 * Updated copyright dates.
 *
 * Revision 1.4  2003/01/16 22:22:45  wes
 * Updated all source files to reflect new organization of header files: all
 * headers that were formerly located in include/rmaux, include/rmv and
 * include/rmi are now located in include/rm.
 *
 * Revision 1.3  2002/06/17 00:40:42  wes
 * updated copyright line.
 *
 * Revision 1.2  2001/03/31 16:55:18  wes
 * Added procmode.h, which defines an RMpipe processing mode used in
 * most demonstration programs. The default processing mode is
 * RM_PIPE_MULTISTAGE_VIEW_PARALLEL.
 *
 * Revision 1.1.1.1  2000/02/28 21:55:30  wes
 * OpenRM 1.2 Release
 *
 * Revision 1.6  2000/02/28 17:21:55  wes
 * RM 1.2, pre-OpenRM
 *
 */

#include <rm/rm.h>
#include <ctype.h>

#define FALSE 0
#define TRUE 1

#define MAXLINE		8192
#define MAXNAME		1024

#define NUMBLOCK 1000

#define DISP		0
#define COL		1

#define C	0
#define O	1
#define N	2
#define H	3
#define S	4
#define P	5
#define Cl	6
#define I	7
#define F	8
#define Other	9

#define C_rad 	0.77 
#define O_rad 	0.6
#define N_rad 	0.7 
#define H_rad 	0.37 
#define S_rad 	1.0
#define P_rad 	1.1       
#define Cl_rad 	0.99      
#define I_rad 	1.33
#define F_rad 	0.64
#define Other_rad  1.5

#define BOND_THRESHOLD	1.85

#define ABS(x)	(((x) >= 0) ? (x) : -(x))

typedef struct {
	char *in;
	char *out;
	char *rep;
} STATE;

typedef struct {
	char *name;
	float r,g,b;
} COLOR;

typedef struct list LIST;

struct list {
	char *name;
	LIST *next;
};

typedef struct {
	char match[2];
	char *name;
	int num;
	char *chain_name;
	char *residue_name;
	int residue_num;
	float pos[3];
	float color[3];
	float radius;
} _ATOM;

typedef struct {
	int end[2];
} BOND;

typedef struct {
	char *name;
	int num;
	char *chain_name;
	char *residue_name;
	int residue_num;
	int type;
} TEMPLATE;

typedef struct {
	int numatoms;
        _ATOM *atom;
	int numbonds;
	BOND *bond;
} DATA;

float atomic_table[][4] = {
				C_rad,		0.0, 1.0, 0.0,
				O_rad,		1.0, 0.0, 0.0,
				N_rad,		0.0, 0.0, 1.0,
				H_rad,		1.0, 1.0, 1.0,
				S_rad,		1.0, 1.0, 0.0,
				P_rad,		1.0, 1.0, 0.5,
				Cl_rad,		1.0, 1.0, 0.5,
				I_rad,		0.5, 1.0, 0.5,
				F_rad,		0.0, 1.0, 1.0,
				Other_rad,	1.0, 0.0, 1.0,
		          };

LIST *names = NULL;
LIST *chain_names = NULL;
LIST *residue_names = NULL;

DATA molecule;

void atomic_spheres(RMnode *out, float scale);
void atomic_lines(RMnode *out);

void free_names(LIST **cur) {
	LIST *scan;
	LIST *temp;

	for (scan = *cur; scan != NULL; scan = temp) {
		free(scan->name);
		temp = scan->next;

		free(scan);
	}

	*cur = NULL;
}

void add_name(LIST **cur, char *name) {
	LIST *temp;

	for (temp = *cur; temp != NULL; temp = temp->next) {
		if (strcmp(temp->name,name) == 0) {
			break;
		}
	}

	if (temp == NULL) {
		temp = (LIST *)malloc(sizeof(*temp));

		temp->name = strdup(name);
		temp->next = *cur;
		*cur = temp;
	}
}

int compare_names(const void *e1, const void *e2) {
	return(strcmp((*(LIST **)e1)->name,(*(LIST **)e2)->name));
}

void sort_names(LIST **cur) {
	LIST *scan;
	int i,num;
	LIST **temp;

	num = 0;
	for (scan = *cur; scan != NULL; scan = scan->next) {
		num++;
	}

	temp = (LIST **)malloc(num*sizeof(*temp));

	num = 0;
	for (scan = *cur; scan != NULL; scan = scan->next) {
		temp[num] = scan;
		num++;
	}

	qsort(temp,num,sizeof(*temp),compare_names);
		
	*cur = temp[0];
	for (i = 0; i < num-1; i++) {
		temp[i]->next = temp[i+1];
	}
	temp[i]->next = NULL;

	free(temp);
}

void
make_bonds(void)
{
    int i,j;
    int numatoms;
    int numbonds;
    int numbonds_this_atom;
    int num_block;
    float dx,dy,dxy,dz;
    double dhypot;
    double a1,a2;

    numatoms = molecule.numatoms;
    
    numbonds = 0;
    num_block = NUMBLOCK;
    molecule.bond = NULL;

    for (i = numatoms-1; i > 0; i--)
    {
	numbonds_this_atom = 0;

	for (j = i-1; j >= 0; j--)
	{

	    /*
	     * The outer loop index 'i' is AFTER the inner loop 'j': 'i'
	     * leads 'j' in the list: since hydrogens traditionally follow
	     * the heavy atom they're bonded to, this makes it easy to quit
	     * bonding to hydrogens after one bond is made by breaking out of
	     * the 'j' loop when 'i' is a hydrogen and we make a bond to it.
	     * Working backwards like this makes it easy to find the heavy
	     * atom that came 'just before' the Hydrogen. mp 
	     */
	    if (num_block == NUMBLOCK)
	    {
		molecule.bond = (BOND *)realloc(molecule.bond,(numbonds+NUMBLOCK)*sizeof(*(molecule.bond)));
		
		if (molecule.bond == NULL)
		{
		    fprintf(stderr,"Unable to allocate memory for %d bonds\n",numbonds+NUMBLOCK);
		    exit(-1);	/* ?? */
		}

		num_block = 0;
	    }

	    /* never bond hydrogens to each other... */
	    if (molecule.atom[i].num == H && molecule.atom[j].num == H)
		continue;

	    dx = molecule.atom[i].pos[0] - molecule.atom[j].pos[0];
	    dx = ABS(dx);
	    if (dx > BOND_THRESHOLD) continue;

	    dy = molecule.atom[i].pos[1] - molecule.atom[j].pos[1];
	    dy = ABS(dy);
	    if (dy > BOND_THRESHOLD) continue;

	    a1 = dx;
	    a2 = dy;
	    dhypot = hypot(a1,a2);
	    dxy = dhypot;
	    if (dxy > BOND_THRESHOLD) continue;

	    dz = molecule.atom[i].pos[2] - molecule.atom[j].pos[2];
	    dz = ABS(dz);
	    if (dz > BOND_THRESHOLD) continue;
	    
	    a1 = dxy;
	    a2 = dz;
	    dhypot = hypot(a1,a2);
	    if (dhypot > BOND_THRESHOLD) continue;

	    molecule.bond[numbonds].end[0] = j;
	    molecule.bond[numbonds].end[1] = i;

	    numbonds++;
	    num_block++;
	    numbonds_this_atom++;

	    if (molecule.atom[i].num == H) break; /* only one bond per H */
	}
    }
    molecule.numbonds = numbonds;
}

void freePDB(void) {
	int i,num;

	num = molecule.numatoms;

	for (i = 0; i < num; i++) {
	    free(molecule.atom[i].name);
	    free(molecule.atom[i].chain_name);
	    free(molecule.atom[i].residue_name);
	}

	free(molecule.atom);
	free(molecule.bond);

	free_names(&names);
	free_names(&chain_names);
	free_names(&residue_names);
}

void readPDB(FILE *in)
{
    int num,num_block;
    char line[MAXLINE];
    char atom[MAXNAME];
    int atom_num;
    char atom_name[MAXNAME];
    char residue_name[MAXNAME];
    char chain_id[MAXNAME];
    int residue_num;
    float x,y,z;

    num = 0;
    num_block = NUMBLOCK;
    molecule.atom = NULL;

    while (fgets(line,sizeof(line),in) != NULL)
    {
	atom[0] = '\0';
	atom_num = 0;
	atom_name[0] = '\0';
	residue_name[0] = '\0';
	residue_num = 0;
	chain_id[0] = '\0';

	sscanf(&line[ 0],"%6s",atom);
	sscanf(&line[ 6],"%d",&atom_num);
	sscanf(&line[12],"%4s",atom_name);
	sscanf(&line[17],"%3s",residue_name);
	sscanf(&line[21],"%1s",chain_id);
	sscanf(&line[22],"%d",&residue_num);
	sscanf(&line[30],"%g",&x);
	sscanf(&line[38],"%g",&y);
	sscanf(&line[46],"%g",&z);

	if (strcmp(atom,"ATOM") == 0 ||
	    strcmp(atom,"atom") == 0 ||
	    strcmp(atom,"HETATM") == 0 ||
	    strcmp(atom,"hetatm") == 0)
	{
	    char a;

	    if (num_block == NUMBLOCK)
	    {
		molecule.atom = (_ATOM *)realloc(molecule.atom,(num+NUMBLOCK)*sizeof(*(molecule.atom)));
		
		if (molecule.atom == NULL)
		{
		    fprintf(stderr,"Unable to allocate memory for %d atoms\n",num+NUMBLOCK);
		    exit(-1);	/* ?? */
		}

		num_block = 0;
	    }

	    molecule.atom[num].match[DISP] = TRUE;
	    molecule.atom[num].match[COL] = TRUE;

	    if (atom_name[1] == 'L' ||
		atom_name[1] == 'l')
		atom_name[2] = '\0';
	    else
		atom_name[1] = '\0';

	    molecule.atom[num].name = strdup(atom_name);
	    molecule.atom[num].num = atom_num;
	    add_name(&names,atom_name);

	    molecule.atom[num].chain_name = strdup(chain_id);
	    add_name(&chain_names,chain_id);
	    
	    molecule.atom[num].residue_name = strdup(residue_name);
	    molecule.atom[num].residue_num = residue_num;
	    add_name(&residue_names,residue_name);

	    molecule.atom[num].pos[0] = x;
	    molecule.atom[num].pos[1] = y;
	    molecule.atom[num].pos[2] = z;

	    a = molecule.atom[num].name[0];
	    if (islower(a)) a = toupper(a);

	    switch (a)
	    {
	    case 'C':
		a = molecule.atom[num].name[1];
		if (islower(a)) a = toupper(a);
		if (a == 'L')
		    atom_num = Cl;
		else
		    atom_num = C;
		
		break;

	    case 'O':
		atom_num = O;
		break;

	    case 'N':
		atom_num = N;
		break;

	    case 'H':
		atom_num = H;
		break;

	    case 'S':
		atom_num = S;
		break;

	    case 'P':
		atom_num = P;
		break;

	    case 'I':
		atom_num = I;
		break;

	    case 'F':
		atom_num = F;
		break;
		
	    default:
		atom_num = Other;
		break;
	    }

	    molecule.atom[num].num = atom_num;
	    molecule.atom[num].radius = atomic_table[atom_num][0];
	    molecule.atom[num].color[0] = atomic_table[atom_num][1];
	    molecule.atom[num].color[1] = atomic_table[atom_num][2];
	    molecule.atom[num].color[2] = atomic_table[atom_num][3];
	    
	    num++;
	    num_block++;
	}
    }

    molecule.numatoms = num;
    
    make_bonds();
    
    sort_names(&names);
    sort_names(&chain_names);
    sort_names(&residue_names);
}


int matchPDB(TEMPLATE *templ)
{
    char *name;
    int num;
    char *chain_name;
    char *residue_name;
    int residue_num;
    int type;
    int i,n;
    int count;

    n = molecule.numatoms;

    name = templ->name;
    num = templ->num;
    chain_name = templ->chain_name;
    residue_name = templ->residue_name;
    residue_num = templ->residue_num;

    type = templ->type;

    count = 0;
    for (i = 0; i < n; i++)
    {
	if ((name == NULL  || strcmp(name,molecule.atom[i].name) == 0 ) &&
	    (num == -1     || num == molecule.atom[i].num ) &&
	    (chain_name == NULL  || strcmp(chain_name,molecule.atom[i].chain_name) == 0   ) &&
	    (residue_name == NULL || strcmp(residue_name,molecule.atom[i].residue_name) == 0) &&
	    (residue_num == -1    || residue_num == molecule.atom[i].residue_num ))
	{
	    molecule.atom[i].match[type] = TRUE;
	    count++;
	}
	else 
	    molecule.atom[i].match[type] = FALSE;
    }
    return(count);
}

void colorPDB(float r, float g, float b)
{
    int i,n;

    n = molecule.numatoms;

    for (i = 0; i < n; i++)
    {
	if (molecule.atom[i].match[COL] == TRUE)
	{
	    molecule.atom[i].color[0] = r;
	    molecule.atom[i].color[1] = g;
	    molecule.atom[i].color[2] = b;
	}
    }
}

void radiusPDB(float radius)
{
    int i,n;

    n = molecule.numatoms;

    for (i = 0; i < n; i++)
    {
	if (molecule.atom[i].match[COL] == TRUE) 
	    molecule.atom[i].radius = radius;
    }
}

void
buildPDBModel(RMnode *out,
	      char *display)
{
    if (strcmp(display,"Sticks") == 0) 
	atomic_lines(out);
    else if (strcmp(display,"Balls") == 0)
	atomic_spheres(out,0.6);
    else if (strcmp(display,"Sticks & Balls") == 0)
    {
	atomic_lines(out);
	atomic_spheres(out,0.75);
    }
    else
	fprintf(stderr,"Unknown display type: %s\n",display);
}


void
pack_write_lines(char *match, float *points, float *colors,
		 int num, RMnode *out)
{
    int i,i3,cur,cur3;
    RMprimitive *p;

    cur = 0;
    cur3 = 0;
    for (i = 0, i3 = 0; i < num; i++, i3 += 3)
    {
	if (match[i] == TRUE) {
	    colors[cur3+0] = colors[i3+0];
	    colors[cur3+1] = colors[i3+1];
	    colors[cur3+2] = colors[i3+2];
	    
	    points[cur3+0] = points[i3+0];
	    points[cur3+1] = points[i3+1];
	    points[cur3+2] = points[i3+2];

	    cur++;
	    cur3 += 3;
	}
    }

    if (cur != 0)
    {
	p = rmPrimitiveNew(RM_LINES);

	/*
	 * the following works only on machines that don't pad to
	 * 8-byte boundaries - some 64 bit machines do.
	 */
	rmPrimitiveSetVertex3D(p, cur, (RMvertex3D *)(points),
			       RM_COPY_DATA,NULL);
			       
	rmPrimitiveSetColor3D(p,cur,(RMcolor3D *)colors,
			      RM_COPY_DATA, NULL);

	rmNodeAddPrimitive(out, p);
    }
}

void atomic_lines(RMnode *out)
{
    int i,num;
    int a1,a2;
    int cur,cur3;
    char *match;
    float *points;
    float *colors;

    num = molecule.numbonds;

    match = (char *)malloc(2*num*sizeof(*match));
    if (match == NULL)
    {
	fprintf(stderr,"Unable to allocate enough memory to draw lines\n");
	exit(-1);
    }

    points = (float *)malloc(2*num*3*sizeof(*points));
    if (points == NULL)
    {
	fprintf(stderr,"Unable to allocate enough memory to draw lines\n");
	exit(-1);
    }

    colors = (float *)malloc(2*num*3*sizeof(*colors));
    if (colors == NULL)
    {
	fprintf(stderr,"Unable to allocate enough memory to draw lines\n");
	exit(-1);
    }

    cur = 0;
    cur3 = 0;
    for (i = 0; i < num; i++)
    {
	a1 = molecule.bond[i].end[0];
	a2 = molecule.bond[i].end[1];

	match[cur] = molecule.atom[a1].match[DISP];

	colors[cur3+0] = molecule.atom[a1].color[0];
	colors[cur3+1] = molecule.atom[a1].color[1];
	colors[cur3+2] = molecule.atom[a1].color[2];

#if 0
	/* hack to make all bonds black for r/b stereo print*/
	colors[cur3+0] = colors[cur3+1] = colors[cur3+2] = 0.0;
#endif
	points[cur3+0] = molecule.atom[a1].pos[0];
	points[cur3+1] = molecule.atom[a1].pos[1];
	points[cur3+2] = molecule.atom[a1].pos[2];
	cur++;
	cur3 += 3;

	match[cur] = molecule.atom[a1].match[DISP];

	colors[cur3+0] = molecule.atom[a1].color[0];
	colors[cur3+1] = molecule.atom[a1].color[1];
	colors[cur3+2] = molecule.atom[a1].color[2];

#if 0
	/* hack to make all bonds black for r/b stereo print */
	colors[cur3+0] = colors[cur3+1] = colors[cur3+2] = 0.0;
#endif

	points[cur3+0] = (molecule.atom[a1].pos[0] + molecule.atom[a2].pos[0]) * 0.5;
	points[cur3+1] = (molecule.atom[a1].pos[1] + molecule.atom[a2].pos[1]) * 0.5;
	points[cur3+2] = (molecule.atom[a1].pos[2] + molecule.atom[a2].pos[2]) * 0.5;
	cur++;
	cur3 += 3;
    }

    pack_write_lines(match,points,colors,cur,out);

    cur = 0;
    cur3 = 0;
    for (i = 0; i < num; i++)
    {
	a1 = molecule.bond[i].end[0];
	a2 = molecule.bond[i].end[1];

	match[cur] = molecule.atom[a2].match[DISP];

	colors[cur3+0] = molecule.atom[a2].color[0];
	colors[cur3+1] = molecule.atom[a2].color[1];
	colors[cur3+2] = molecule.atom[a2].color[2];

#if 0
	colors[cur3+0] = colors[cur3+1] = colors[cur3+2] = 0.0;
#endif

	points[cur3+0] = molecule.atom[a2].pos[0];
	points[cur3+1] = molecule.atom[a2].pos[1];
	points[cur3+2] = molecule.atom[a2].pos[2];
	cur++;
	cur3 += 3;

	match[cur] = molecule.atom[a2].match[DISP];

	colors[cur3+0] = molecule.atom[a2].color[0];
	colors[cur3+1] = molecule.atom[a2].color[1];
	colors[cur3+2] = molecule.atom[a2].color[2];

#if 0
	colors[cur3+0] = colors[cur3+1] = colors[cur3+2] = 0.0;
#endif

	points[cur3+0] = (molecule.atom[a1].pos[0] + molecule.atom[a2].pos[0]) * 0.5;
	points[cur3+1] = (molecule.atom[a1].pos[1] + molecule.atom[a2].pos[1]) * 0.5;
	points[cur3+2] = (molecule.atom[a1].pos[2] + molecule.atom[a2].pos[2]) * 0.5;
	cur++;
	cur3 += 3;
    }

    pack_write_lines(match,points,colors,cur,out);

    free((void *)match);
    free((void *)points);
    free((void *)colors);
}

void pack_write_spheres(char *match, float *points, float *colors,
		        float *radii, int num, RMnode *out)
{
    int i,i3;
    int cur,cur3;
    
    cur = 0;
    cur3 = 0;
    for (i = 0, i3 = 0; i < num; i++, i3 += 3)
    {
	if (match[i] == TRUE)
	{
	    colors[cur3+0] = colors[i3+0];
	    colors[cur3+1] = colors[i3+1];
	    colors[cur3+2] = colors[i3+2];

#if 0
	    /* hack - make 'em all black */
	    colors[cur3+0] = colors[cur3+1] = colors[cur3+2] = 0.;
#endif

	    points[cur3+0] = points[i3+0];
	    points[cur3+1] = points[i3+1];
	    points[cur3+2] = points[i3+2];

	    radii[cur] = radii[i];
	    cur++;
	    cur3 += 3;
	}
    }

    if (cur != 0)
    {
	RMprimitive *p;

	p = rmPrimitiveNew(RM_SPHERES);

	/*
	 * the following works only on machines that don't pad RMvertex3D
	 * structs to 8-byte boundaries (because we're ramming in a flat
	 * array of floats that will later be processed as if it was an
	 * array of RMvertex3D structures). some 64bit machines do such
	 * padding.
	 */
	rmPrimitiveSetVertex3D(p, cur, (RMvertex3D *)points,
			       RM_COPY_DATA, NULL);

	rmPrimitiveSetColor3D(p, cur, (RMcolor3D *)colors,
			      RM_COPY_DATA, NULL);
	
	rmPrimitiveSetRadii(p, cur, radii,
			    RM_COPY_DATA, NULL);

	rmNodeAddPrimitive(out,p);
    }
}

void atomic_spheres(RMnode *out, float scale)
{
    int i,num;
    int cur,cur3;
    char *match;
    float *points;
    float *colors;
    float *radii;

    num = molecule.numatoms;

    match = (char *)malloc(num*sizeof(*match));
    if (match == NULL)
    {
	fprintf(stderr,"Unable to allocate enough memory to draw balls\n");
	exit(-1);		/* ?? */
    }

    points = (float *)malloc(num*3*sizeof(*points));
    if (points == NULL)
    {
	fprintf(stderr,"Unable to allocate enough memory to draw balls\n");
	exit(-1);
    }

    colors = (float *)malloc(num*3*sizeof(*colors));
    if (colors == NULL)
    {
	fprintf(stderr,"Unable to allocate enough memory to draw balls\n");
	exit(-1);
    }

    radii = (float *)malloc(num*sizeof(*radii));
    if (radii == NULL)
    {
	fprintf(stderr,"Unable to allocate enough memory to draw balls\n");
	exit(-1);
    }

    cur = 0;
    cur3 = 0;
    for (i = 0; i < num; i++)
    {
	match[cur] = molecule.atom[i].match[DISP];
	
	colors[cur3+0] = molecule.atom[i].color[0];
	colors[cur3+1] = molecule.atom[i].color[1];
	colors[cur3+2] = molecule.atom[i].color[2];

	points[cur3+0] = molecule.atom[i].pos[0];
	points[cur3+1] = molecule.atom[i].pos[1];
	points[cur3+2] = molecule.atom[i].pos[2];

	radii[cur] = scale * molecule.atom[i].radius;
	cur++;
	cur3 += 3;
    }

    pack_write_spheres(match,points,colors,radii,cur,out);

    free((void *)match);
    free((void *)colors);
    free((void *)points);
    free((void *)radii);
}