/*
 * xyz2ps - convert an XYZ file to a binocular 3D Postscript file
 * (c) 1996  Will Ware
 *
 * These Postscript files can be viewed with a cheap stereoscope sold by
 * Edmund Scientific.
 *
 * GNU General Public License, blah blah, legalese, blah blah blah,
 * you have the right to an attorney, blah blah blah, Stallman is God,
 * blah blah, wear Birkenstocks and tie-dye, listen to the Dead, and eat
 * Ben and Jerry's, blah blah blah blah, Amen.
 */

#include "stdio.h"
#include "math.h"

typedef unsigned long ulong;

#define LINELENGTH 300

ulong num_atoms;
struct atom
{
  char name[8];
  double x, y, z;
} *atoms;
char line[LINELENGTH], compound_name[LINELENGTH];
double xmax, ymax, zmax;

/*
 * import_xyz is a no-brainer: find out how many atoms there are, allocate
 * an array for them, and fill the array.
 */

void import_pdb(char *filename)
{
  FILE *inf;
  ulong i;
  double xmin, ymin, zmin, xcenter, ycenter, zcenter;
  struct atom *p;

  /* one pass to get the number of atoms */
  inf = fopen(filename, "r");
  if (inf == NULL)
    {
      fprintf(stderr, "pdb2ps: can't open PDB file %s\n", filename);
      exit(1);
    }

  num_atoms = 0;
  compound_name[0] = '\0';
  while (fgets(line, LINELENGTH, inf) != NULL)
    {
      line[6] = '\0';
      if (strcmp(line, "ATOM  ") == 0 ||
	  strcmp(line, "HETATM") == 0)
	num_atoms++;
      if (strcmp(line, "COMPND") == 0)
	{
	  strcpy(compound_name, line + 10);
	  compound_name[62] = '\0';
	  if (compound_name[strlen(compound_name) - 1] == '\n')
	    compound_name[strlen(compound_name) - 1] = '\0';
	}
    }
  fclose(inf);

  atoms = (struct atom *) malloc(num_atoms * sizeof(struct atom));
  if (atoms == NULL)
    {
      fprintf(stderr, "Insufficient memory\n");
      exit(1);
    }

  /* another pass to get the data for the atoms */
  inf = fopen(filename, "r");
  i = 0;
  while (fgets(line, LINELENGTH, inf) != NULL)
    {
      line[6] = '\0';
      if (strcmp(line, "ATOM  ") == 0 ||
	  strcmp(line, "HETATM") == 0)
	{
	  p = &atoms[i];
	  line[16] = '\0';
	  {
	    char *s = line + 11, *t = p->name;
	    while (*s == ' ') s++;
	    while (*s != ' ' && *s != '\0') *t++ = *s++;
	    *t = '\0';
	  }
	  sscanf(line + 32, "%lf %lf %lf",
		 &(p->x), &(p->y), &(p->z));
	  if (i == 0)
	    {
	      xmin = xmax = p->x;
	      ymin = ymax = p->y;
	      zmin = zmax = p->z;
	    }
	  else
	    {
	      if (p->x < xmin) xmin = p->x;
	      if (p->x > xmax) xmax = p->x;
	      if (p->y < ymin) ymin = p->y;
	      if (p->y > ymax) ymax = p->y;
	      if (p->z < zmin) zmin = p->z;
	      if (p->z > zmax) zmax = p->z;
	    }
	  i++;
	}
    }
  fclose(inf);

  /* center the structure around the origin */
  xcenter = 0.5 * (xmin + xmax);
  ycenter = 0.5 * (ymin + ymax);
  zcenter = 0.5 * (zmin + zmax);

  for (i = 0; i < num_atoms; i++)
    {
      struct atom *p = &atoms[i];
      p->x -= xcenter;
      p->y -= ycenter;
      p->z -= zcenter;
    }

  /* update xmax, ymax, zmax */
  xmax -= xcenter;
  ymax -= ycenter;
  zmax -= zcenter;
}

/*
 * Postscript renders using a dumb painter's algorithm, so order atoms
 * by Z coordinate, furthest ones first (positive z), nearest ones last
 * (negative z).
 */

void reorder_atoms(void)
{
  ulong i, imin, j;
  double zmin;
  struct atom temp;

  for (j = num_atoms - 1; j > 0; j--)
    {
      imin = -1;
      for (i = 0; i <= j; i++)
	if (imin == -1 || atoms[i].z < zmin)
	  {
	    imin = i;
	    zmin = atoms[i].z;
	  }
      if (imin != j)
	{
	  memcpy(&temp, &atoms[j], sizeof(struct atom));
	  memcpy(&atoms[j], &atoms[imin], sizeof(struct atom));
	  memcpy(&atoms[imin], &temp, sizeof(struct atom));
	}
    }
}

const char
*carbon_names[] =
{
  "CA", "CB", "CD", "CD1", "CD2", "CE", "CE1", "CE2",
  "CG", "CG1", "CG2", "CH", "CH1", "CH2", "CZ"
},
*oxygen_names[] =
{
  "OA", "OB", "OD", "OD1", "OD2", "OE", "OE1", "OE2",
  "OG", "OG1", "OG2", "OH", "OH1", "OH2", "OXT", "OZ"
},
*nitrogen_names[] =
{
  "NA", "NB", "ND", "ND1", "ND2", "NE", "NE1", "NE2",
  "NG", "NG1", "NG2"  "NH", "NH1", "NH2", "NZ"
},
*hydrogen_names[] =
{
  ""
};

void cleanup_atom_names(void)
{
  ulong i;
  int j;

  for (i = 0; i < num_atoms; i++)
    {
      for (j = 0; j < sizeof(hydrogen_names) / sizeof(char*); j++)
	if (strcmp(atoms[i].name, hydrogen_names[j]) == 0)
	  {
	    strcpy(atoms[i].name, "H");
	    goto next_atom;
	  }

      for (j = 0; j < sizeof(carbon_names) / sizeof(char*); j++)
	if (strcmp(atoms[i].name, carbon_names[j]) == 0)
	  {
	    strcpy(atoms[i].name, "C");
	    goto next_atom;
	  }

      for (j = 0; j < sizeof(oxygen_names) / sizeof(char*); j++)
	if (strcmp(atoms[i].name, oxygen_names[j]) == 0)
	  {
	    strcpy(atoms[i].name, "O");
	    goto next_atom;
	  }

      for (j = 0; j < sizeof(nitrogen_names) / sizeof(char*); j++)
	if (strcmp(atoms[i].name, nitrogen_names[j]) == 0)
	  {
	    strcpy(atoms[i].name, "N");
	    goto next_atom;
	  }

    next_atom: ;
    }
}




const char *postscript_preamble =
"%%!PS-Adobe\n"
"/Courier findfont 12 scalefont setfont\n";

/* Postscript measures everything in 1/72's of an inch */
#define PS_RESOLUTION  72.0
#define MIN_RADIUS 0.5  /* Postscript coordinates, don't disappear */
#define NOMINAL_RADIUS 0.5  /* in Angstroms */

void postscript_draw_atom(struct atom *p, double scale, FILE *outf)
{
  double pf, xl, xr, xv, y, zv, r;
  char *colorstring;

  /* Perspective */
  xv = 0.5 * xmax;  /* x coordinate of vanishing point */
  zv = 30.0 * zmax;  /* z coordinate of vanishing point */
  pf = zv / (zv + p->z);  /* perspective factor */
  xl = pf * p->x - (1.0 - pf) * xv;
  xr = pf * p->x + (1.0 - pf) * xv;
  y = pf * p->y;
  r = pf * NOMINAL_RADIUS;

  xl = 2.75 * PS_RESOLUTION + scale * xl;
  xr = 5.75 * PS_RESOLUTION + scale * xr;
  y = 8.0 * PS_RESOLUTION - scale * y;  /* Postscript flips vertically */
  r *= scale;
  if (r < MIN_RADIUS) r = MIN_RADIUS;

  /* different colors for different elements */
  if      (strcmp(p->name, "C") == 0) colorstring = "0.4 setgray ";
  else if (strcmp(p->name, "H") == 0) colorstring = "1 setgray ";
  else if (strcmp(p->name, "N") == 0) colorstring = "0 0 1 setrgbcolor ";
  else if (strcmp(p->name, "O") == 0) colorstring = "1 0 0 setrgbcolor ";
  else                                colorstring = "0 1 0 setrgbcolor ";

  /* left eye view */
  /* interior color of atom */
  fprintf(outf, "%s %lf %lf %lf 0 360 arc closepath fill\n",
	  colorstring, xl, y, r);
  /* black circle for the edge */
  fprintf(outf, "0 setgray %lf %lf %lf 0 360 arc stroke\n", xl, y, r);

  /* right eye view */
  /* interior color of atom */
  fprintf(outf, "%s %lf %lf %lf 0 360 arc closepath fill\n",
	  colorstring, xr, y, r);
  /* black circle for the edge */
  fprintf(outf, "0 setgray %lf %lf %lf 0 360 arc stroke\n", xr, y, r);
}

void export_postscript(FILE *outf)
{
  double scale1, scale2, scale;
  ulong i;

  /*
   * p->x,y,z are in angstroms, we need Postscript coordinates
   * compute scale factor so both X and Y fit comfortably on page
   */
  scale1 = 1.25 * PS_RESOLUTION / (xmax + NOMINAL_RADIUS);
  scale2 = 1.75 * PS_RESOLUTION / (ymax + NOMINAL_RADIUS);
  if (scale1 < scale2) scale = scale1;
  else scale = scale2;

  fprintf(outf, postscript_preamble);
  for (i = 0; i < num_atoms; i++)
    postscript_draw_atom(&atoms[i], scale, outf);

  fprintf(outf, "%lf %lf moveto (%s) show\n",
	  1.50 * PS_RESOLUTION, 5.75 * PS_RESOLUTION, compound_name);

  fprintf(outf, "showpage\n");
}

const char *usage_string =
"Usage: pdb2ps stuff.pdb > stuff.ps\n"
"stuff.pdb is a Brookhaven Protein Databank file\n"
"stuff.ps is a Postscript file, suitable for stereoscopic viewing when\n"
"printed on a color printer that understands Postscript\n";

void usage(void)
{
  printf(usage_string);
  exit(1);
}

void main(int argc, char *argv[])
{
  if (argc < 2) usage();
  import_pdb(argv[1]);
  reorder_atoms();
  cleanup_atom_names();
  export_postscript(stdout);
  exit(0);
}