/*
 * xyz2rgb.c - convert an XYZ file (or a concatenation of several XYZ files)
 *             to a ImageMagick-readable bitmap file (or several bitmaps),
 *             output files are line-interlaced RGB
 *
 * (c) 1997 Will Ware 
 *
 * Orientation: when longitude and latitude are both zero (the default case)
 * the camera lies on the x axis. The camera always faces the origin. The z
 * axis always points up, and for zero longitude, the y axis points to the
 * right. Increasing longitude will swing the camera around toward the y axis.
 * There are two other parameters for controlling the camera, which are the
 * width of the viewing area and the distance from the camera to the origin
 * (in same distance units as input file, usually angstroms). Longitude,
 * latitude, and width are all angles expressed in degrees. The viewing area
 * is assumed to have an aspect ratio of 4:3.
 *
 * 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 can obtain a copy of the GNU General Public License by writing
 * to the Free Software Foundation, Inc., 675 Mass Ave, Cambridge,
 * MA 02139, USA.
 */

#include 
#include 
#include 
#include 

#define VERSION_STRING  "Version 1.0  1997/06/06\n"

#ifndef DEBUG
#define DEBUG 0
#endif

#if DEBUG
#define ASSERT(c) if (!c) \
  {fprintf(stderr,"Assertion failed: %s\n%s %d\n", #c, __FILE__, __LINE__); \
   punt(1);}

void
punt (int n)
{
  exit (n);
}

#else
#define ASSERT(c)
#endif

/* function return values */
#define OK 0
#define OUCH 1
#define ENDDATA 2

char usage_string[] =
"Usage: xyz2rgb [options]\n"
"Options:\n"
"    -i filename      input filename (XYZ format)\n"
"    -longitude a     specify horizontal angle of camera\n"
"    -latitude a      specify vertical angle of camera\n"
"    -distance x      specify distance to origin of camera\n"
"    -width a         specify wide-angleness of camera\n"
"    -size HxV        specify number of horizontal, vertical pixels\n"
"    -radius r        select a common radius for all atoms\n"
"    -o filename      output filename, files will have numbers\n"
"                       and \".rgb\" appended to their names\n"
"    -walls filename  specify a file of walls\n"
"    -walls-ok        walls should not be centered or rotated\n"
"                       (perspective is still done on walls)\n"
"    -opacity n       opacity of atoms (0.0-1.0, default 1.0)"
 ;

typedef struct
  {
    double longitude, latitude, distance, width;
  }
camera;

typedef unsigned int color[3];	/* r, g, b */

typedef struct
  {
    color c;
    double radius;		/* van der Waals radius */
    char *name;
  }
species;

typedef struct s_linked_pixel
  {
    struct s_linked_pixel *next;
    color c;
    double a, x;		/* opacity (0.0 to 1.0), x coord */
  }
buffer_link;

typedef struct
  {
    double x, y, z;
  }
vertex;

typedef struct
  {
    vertex xyz;
    double opacity, radius;
    species *spc;
  }
atom;

typedef struct
  {
    vertex xyz[4];
    color c;
    double a;
  }
wall;

species known_species[] =
{
  {
    255, 255, 255, 1.5, "H"
  }
  ,
  {
    128, 128, 128, 1.94, "C"
  }
  ,
  {
    255, 0, 0, 1.82, "O"
  }
  ,
  {
    0, 0, 255, 1.74, "N"
  }
  ,
  {
    255, 255, 0, 2.15, "S"
  }
  ,
  {
    255, 0, 255, 2.22, "P"
  }
};

int num_known_species = sizeof (known_species) / sizeof (species);

/* the background is white */
color background =
{
  255, 255, 255
};

#define MAXHSIZE 2000
buffer_link *buffer[MAXHSIZE];
double common_radius = -1.0;
double atom_opacity = 1.0;
atom *atoms = NULL;
wall *walls = NULL;
int hsize, vsize, num_atoms, num_walls;
double ymax, zmax;
camera cam;
int transform_walls = 1;	/* controls centering and rotation */
int outfile_number = 0;
char line[200], xyzname[40], wfname[40], fname_prefix[40];
char fname[40], ouch_str[50];

int
ouch (char *s)
{
  fprintf (stderr, "xyz2rgb  %s", VERSION_STRING);
  fprintf (stderr, "%s\n", s);
  return OUCH;
}

void
free_chain (buffer_link * b)
{
  buffer_link *n;

  while (b != NULL)
    {
      n = b->next;
      free (b);
      b = n;
    }
}

buffer_link *
create_link (double r, double g, double b, double a, double x,
	     buffer_link * next)
{
  buffer_link *n;

  n = malloc (sizeof (buffer_link));
  if (a > 0.99)
    {
      n->next = NULL;
      free_chain (next);
    }
  else
    n->next = next;
  n->c[0] = (unsigned int) r;
  n->c[1] = (unsigned int) g;
  n->c[2] = (unsigned int) b;
  n->a = a;
  n->x = x;
  return n;
}

void
buffer_a_pixel (int h, double r, double g, double b, double a, double x)
{
  buffer_link *ba, *bb;

  a = (a < 0.0) ? 0.0 : a;
  a = (a > 1.0) ? 1.0 : a;

  if (buffer[h] == NULL || x > buffer[h]->x)
    buffer[h] = create_link (r, g, b, a, x, buffer[h]);
  else
    {
      ba = buffer[h];
      bb = ba->next;
      while (bb != NULL && bb->x > x)
	{
	  ba = bb;
	  bb = ba->next;
	}
      ba->next = create_link (r, g, b, a, x, bb);
    }
}

double
dump_pixel_helper (int color, buffer_link * link)
{
  /* if there are no more surfaces, return background color */
  if (link == NULL)
    return background[color];
  /* if the current surface is opaque, show only it */
  if (link->a >= 0.99)
    return (double) link->c[color];
  /* else handle translucency */
  return link->a * link->c[color] +
    (1.0 - link->a) * dump_pixel_helper (color, link->next);
}

static void
transform_coords (vertex * xyz, double *r, camera * cam, int rot)
{
  double angle, tmp, pfactor;

  if (rot)
    {
      /* rotate by longitude */
      angle = cam->longitude;
      tmp = xyz->x * cos (angle) + xyz->y * sin (angle);
      xyz->y = xyz->y * cos (angle) - xyz->x * sin (angle);
      xyz->x = tmp;
      /* rotate by longitude */
      angle = cam->latitude;
      tmp = xyz->x * cos (angle) + xyz->z * sin (angle);
      xyz->z = xyz->z * cos (angle) - xyz->x * sin (angle);
      xyz->x = tmp;
    }

  /* apply perspective */
  pfactor = (cam->distance - xyz->x) / cam->distance;
  pfactor = (pfactor < 0.01) ? 100.0 : (1.0 / pfactor);
  xyz->y *= pfactor;
  xyz->z *= pfactor;

  if (r != NULL)
    *r *= pfactor;
}

static void
transform_atom_coords (atom * atm, camera * cam)
{
  transform_coords (&(atm->xyz), &(atm->radius), cam, 1);
}

static void
transform_wall_coords (wall * w, camera * cam)
{
  int i;

  for (i = 0; i < 4; i++)
    transform_coords (&(w->xyz[i]), NULL, cam, transform_walls);
}

static int
hlimit (int h)
{
  if (h < 0)
    return 0;
  if (h >= hsize)
    return hsize - 1;
  return h;
}

static void
render_atom (atom * a, double z, double ymax, int half_hsize)
{
  double local_radius, rsq, zsq, y1, y2;
  int i, i1, i2;

  rsq = a->radius * a->radius;
  zsq = (z - a->xyz.z) * (z - a->xyz.z);
  local_radius = rsq - zsq;
  if (local_radius < 0.0)
    return;
  local_radius = sqrt (local_radius);
  y1 = (a->xyz.y - local_radius) / ymax;
  y2 = (a->xyz.y + local_radius) / ymax;
  i1 = hlimit ((int) (half_hsize * (y1 + 1.0)));
  i2 = hlimit ((int) (half_hsize * (y2 + 1.0)));
  for (i = i1; i <= i2; i++)
    {
      double y, ysq, xblip, color_intensity;
      y = (i - half_hsize) * (ymax / half_hsize);
      ysq = (y - a->xyz.y) * (y - a->xyz.y);
      xblip = rsq - ysq - zsq;
      if (xblip < 0.0)
	continue;
      xblip = sqrt (xblip);
      color_intensity = xblip / a->radius;
      buffer_a_pixel (i,
		      (color_intensity * a->spc->c[0]),
		      (color_intensity * a->spc->c[1]),
		      (color_intensity * a->spc->c[2]),
		      (1.0 + 3.0 * (1.0 - color_intensity)) * atom_opacity,
		      a->xyz.x + xblip);
    }
}

static void
render_triangle (wall * w, int which_triangle,
		 double z, double ymax, int half_hsize)
{
  vertex v0, v1, v2, nv;
  double m, x, x1, x2, y, y1, y2;
  int i, i1, i2;

  if (which_triangle)
    {
      if ((w->xyz[0].z > z && w->xyz[1].z > z && w->xyz[2].z > z) ||
	  (w->xyz[0].z < z && w->xyz[1].z < z && w->xyz[2].z < z))
	return;
      memcpy (&v0, &(w->xyz[0]), sizeof (vertex));
      memcpy (&v1, &(w->xyz[1]), sizeof (vertex));
      memcpy (&v2, &(w->xyz[2]), sizeof (vertex));
    }
  else
    {
      if ((w->xyz[0].z > z && w->xyz[2].z > z && w->xyz[3].z > z) ||
	  (w->xyz[0].z < z && w->xyz[2].z < z && w->xyz[3].z < z))
	return;
      memcpy (&v0, &(w->xyz[0]), sizeof (vertex));
      memcpy (&v1, &(w->xyz[2]), sizeof (vertex));
      memcpy (&v2, &(w->xyz[3]), sizeof (vertex));
    }

  /* Generate a unit normal to the triangle (normalize X product of sides) */
  nv.x = (v1.y - v0.y) * (v2.z - v0.z) - (v2.y - v0.y) * (v1.z - v0.z);
  nv.y = (v1.z - v0.z) * (v2.x - v0.x) - (v2.z - v0.z) * (v1.x - v0.x);
  nv.z = (v1.x - v0.x) * (v2.y - v0.y) - (v2.x - v0.x) * (v1.y - v0.y);
  m = nv.x * nv.x + nv.y * nv.y + nv.z * nv.z;
  /* only care about the x component, but it's gotta be positive */
  nv.x /= sqrt (m);
  if (nv.x < 0.0)
    nv.x = -nv.x;

  if ((v0.z > z && v1.z < z && v2.z > z) ||
      (v0.z < z && v1.z > z && v2.z < z))
    {
      m = (z - v1.z) / (v0.z - v1.z);
      y1 = m * v0.y + (1.0 - m) * v1.y;
      x1 = m * v0.x + (1.0 - m) * v1.x;
      m = (z - v1.z) / (v2.z - v1.z);
      y2 = m * v2.y + (1.0 - m) * v1.y;
      x2 = m * v2.x + (1.0 - m) * v1.x;
    }
  else if ((v0.z > z && v1.z > z && v2.z < z) ||
	   (v0.z < z && v1.z < z && v2.z > z))
    {
      m = (z - v2.z) / (v0.z - v2.z);
      y1 = m * v0.y + (1.0 - m) * v2.y;
      x1 = m * v0.x + (1.0 - m) * v2.x;
      m = (z - v2.z) / (v1.z - v2.z);
      y2 = m * v1.y + (1.0 - m) * v2.y;
      x2 = m * v1.x + (1.0 - m) * v2.x;
    }
  else if (v1.z != v0.z && v2.z != v0.z)
    {
      m = (z - v0.z) / (v1.z - v0.z);
      y1 = m * v1.y + (1.0 - m) * v0.y;
      x1 = m * v1.x + (1.0 - m) * v0.x;
      m = (z - v0.z) / (v2.z - v0.z);
      y2 = m * v2.y + (1.0 - m) * v0.y;
      x2 = m * v2.x + (1.0 - m) * v0.x;
    }
  else
    return;

  if (y1 == y2)
    return;
  else if (y1 > y2)
    {
      double temp;
      temp = y1;
      y1 = y2;
      y2 = temp;
      temp = x1;
      x1 = x2;
      x2 = temp;
    }

  y1 /= ymax;
  y2 /= ymax;
  i1 = hlimit ((int) (half_hsize * (y1 + 1.0)));
  i2 = hlimit ((int) (half_hsize * (y2 + 1.0)));
  for (i = i1; i < i2; i++)
    {
      y = (((double) i) / half_hsize) - 1.0;
      x = x1 + ((y - y1) * (x2 - x1)) / (y2 - y1);
      buffer_a_pixel (i, w->c[0], w->c[1], w->c[2],
		      (1.0 + 3.0 * (1.0 - nv.x)) * w->a, x);
    }
}

static void
render_wall (wall * w, double z, double ymax, int half_hsize)
{
  render_triangle (w, 0, z, ymax, half_hsize);
  render_triangle (w, 1, z, ymax, half_hsize);
}

int
render_structure (void)
{
  int i, j, h, v, color;
  FILE *outf;

  if (strlen (fname_prefix) > 0)
    {
      sprintf (fname, "%s%03d.rgb", fname_prefix, outfile_number++);
      outf = fopen (fname, "w");
    }
  else
    outf = stdout;
  if (outf == NULL)
    return ouch ("Can't open output file");

  /* center the structure */
  if (num_atoms > 0)
    {
      double xc, yc, zc;
      xc = yc = zc = 0.0;
      for (i = 0; i < num_atoms; i++)
	{
	  xc += atoms[i].xyz.x;
	  yc += atoms[i].xyz.y;
	  zc += atoms[i].xyz.z;
	}
      xc /= num_atoms;
      yc /= num_atoms;
      zc /= num_atoms;
#if DEBUG
      fprintf (stderr, "center: %f %f %f\n", xc, yc, zc);
#endif
      for (i = 0; i < num_atoms; i++)
	{
	  atoms[i].xyz.x -= xc;
	  atoms[i].xyz.y -= yc;
	  atoms[i].xyz.z -= zc;
	}
      if (transform_walls)
	for (i = 0; i < num_walls; i++)
	  for (j = 0; j < 4; j++)
	    {
	      walls[i].xyz[j].x -= xc;
	      walls[i].xyz[j].y -= yc;
	      walls[i].xyz[j].z -= zc;
	    }
    }

  if (common_radius < 0.0)
    for (i = 0; i < num_atoms; i++)
      atoms[i].radius = atoms[i].spc->radius;
  else
    for (i = 0; i < num_atoms; i++)
      atoms[i].radius = common_radius;

  for (i = 0; i < num_atoms; i++)
    transform_atom_coords (&atoms[i], &cam);

  for (i = 0; i < num_walls; i++)
    transform_wall_coords (&walls[i], &cam);

  for (v = 0; v < vsize; v++)
    {
      double z = (-2.0 * zmax / vsize) * v + zmax;

      for (h = 0; h < hsize; h++)
	buffer[h] = NULL;

      /* draw atoms */
      for (i = 0; i < num_atoms; i++)
	render_atom (&atoms[i], z, ymax, hsize / 2);

      /* draw walls */
      for (i = 0; i < num_walls; i++)
	render_wall (&walls[i], z, ymax, hsize / 2);

      /* dump all the pixels in this line */
      for (color = 0; color < 3; color++)
	for (h = 0; h < hsize; h++)
	  fputc ((unsigned char) dump_pixel_helper (color, buffer[h]),
		 outf);

      /* free all the linked pixel structures */
      for (h = 0; h < hsize; h++)
	free_chain (buffer[h]);
    }

  if (outf != stdout)
    fclose (outf);
  return OK;
}

int
read_wall_file (FILE * inf)
{
  int i;

  fgets (line, 200, inf);
  if (sscanf (line, "%d", &num_walls) < 1)
    return 0;
#if DEBUG
  fprintf (stderr, "num_walls=%d\n", num_walls);
#endif
  walls = malloc (num_walls * sizeof (wall));
  if (walls == NULL)
    return ouch ("Not enough memory for so many walls!");
  for (i = 0; i < num_walls; i++)
    {
      int j;
      /* read in a line from the file, with a wall on it */
      fgets (line, 200, inf);
      if (sscanf (line, "%u %u %u %lf",
		  &(walls[i].c[0]),
		  &(walls[i].c[1]),
		  &(walls[i].c[2]),
		  &(walls[i].a)) < 4)
	walls[i].a = 0.5;
      for (j = 0; j < 4; j++)
	{
	  fgets (line, 200, inf);
	  sscanf (line, "%lf %lf %lf",
		  &(walls[i].xyz[j].x),
		  &(walls[i].xyz[j].y),
		  &(walls[i].xyz[j].z));
	}
    }
  return 1;
}

int
read_xyz_frame (FILE * inf)
{
  int i, j, found;
  char *s;

  /* find out how many atoms to grab */
  if (fgets (line, 200, inf) == NULL)
    return ENDDATA;
  if (sscanf (line, "%d", &num_atoms) < 1)
    return ENDDATA;
#if DEBUG
  fprintf (stderr, "num_atoms=%d\n", num_atoms);
#endif
  if (atoms != NULL)
    free (atoms);
  atoms = malloc (num_atoms * sizeof (atom));
  if (atoms == NULL)
    return ouch ("Not enough memory for so many atoms!");
  fgets (line, 200, inf);
  for (i = 0; i < num_atoms; i++)
    {
      char *t, tempstr[5];

      /* read in a line from the file, with an atom on it */
      fgets (line, 200, inf);
      s = line;

      /* extract the element name for the atom, put in tempstr */
      while (*s == ' ' || *s == '\t')
	s++;
      t = tempstr;
      while (*s != ' ' && *s != '\t')
	*t++ = *s++;
      *t = '\0';

      /* look up this element in the known_species table */
      for (j = found = 0; !found && j < num_known_species; j++)
	if (strcmp (tempstr, known_species[j].name) == 0)
	  {
	    found = 1;
	    atoms[i].spc = &known_species[j];
	  }
      if (!found)
	atoms[i].spc = &known_species[0];

      /* pick up X, Y, Z coordinates (and opacity?) for this atom */
      while (*s == ' ' || *s == '\t')
	s++;
      if (sscanf (s, "%lf %lf %lf %lf",
		  &atoms[i].xyz.x,
		  &atoms[i].xyz.y,
		  &atoms[i].xyz.z,
		  &atoms[i].opacity) < 4)
	atoms[i].opacity = atom_opacity;
    }
  return OK;
}

int
main (int argc, char *argv[])
{
  int i;
  FILE *inf;

  num_atoms = 0;
  num_walls = 0;

  /* initialize all options to default values */
  cam.longitude = 0.0;
  cam.latitude = 0.0;
  cam.distance = 30.0;
  cam.width = 30.0 * (3.14159 / 180.0);
  hsize = 640;
  vsize = 480;
  xyzname[0] = '\0';
  wfname[0] = '\0';
  fname_prefix[0] = '\0';

  if (argc < 2)
    return ouch (usage_string);

  /* collect command line options */
  for (i = 1; i < argc; i++)
    {
      if (strcmp (argv[i], "-longitude") == 0)
	{
	  double tmp;
	  if (sscanf (argv[++i], "%lf", &tmp) == 1)
	    cam.longitude = tmp * (3.14159 / 180.0);
	  else
	    return ouch ("Bad longitude");
	}
      else if (strcmp (argv[i], "-latitude") == 0)
	{
	  double tmp;
	  if (sscanf (argv[++i], "%lf", &tmp) == 1)
	    cam.latitude = tmp * (3.14159 / 180.0);
	  else
	    return ouch ("Bad latitude");
	}
      else if (strcmp (argv[i], "-distance") == 0)
	{
	  double tmp;
	  if (sscanf (argv[++i], "%lf", &tmp) == 1)
	    cam.distance = tmp;
	  else
	    return ouch ("Bad distance");
	}
      else if (strcmp (argv[i], "-width") == 0)
	{
	  double tmp;
	  if (sscanf (argv[++i], "%lf", &tmp) == 1)
	    cam.width = tmp * (3.14159 / 180.0);
	  else
	    return ouch ("Bad width");
	}
      else if (strcmp (argv[i], "-radius") == 0)
	{
	  double tmp;
	  if (sscanf (argv[++i], "%lf", &tmp) == 1)
	    common_radius = tmp;
	  else
	    return ouch ("Bad radius");
	}
      else if (strcmp (argv[i], "-opacity") == 0)
	{
	  double tmp;
	  if (sscanf (argv[++i], "%lf", &tmp) == 1)
	    atom_opacity = tmp;
	  else
	    return ouch ("Bad opacity");
	}
      else if (strcmp (argv[i], "-size") == 0)
	{
	  int tmph, tmpv;
	  if (sscanf (argv[++i], "%dx%d", &tmph, &tmpv) == 2)
	    {
	      hsize = tmph;
	      vsize = tmpv;
	    }
	  else
	    return ouch ("Bad size");
	}
      else if (strcmp (argv[i], "-i") == 0)
	{
	  strcpy (xyzname, argv[++i]);
	}
      else if (strcmp (argv[i], "-o") == 0)
	{
	  strcpy (fname_prefix, argv[++i]);
	}
      else if (strcmp (argv[i], "-walls") == 0)
	{
	  strcpy (wfname, argv[++i]);
	}
      else if (strcmp (argv[i], "-walls-ok") == 0)
	{
	  transform_walls = 0;
	}
      else
	{
	  sprintf (ouch_str, "Bad command line argument: %s", argv[i]);
	  return ouch (ouch_str);
	}
    }

  if (hsize >= MAXHSIZE)
    {
      char sorry[80];
      sprintf (sorry, "Sorry, hsize must be less than %d", MAXHSIZE);
      return ouch (sorry);
    }

#if DEBUG
  fprintf (stderr, "Options info\n");
  fprintf (stderr, "Long: %f Lat: %f Dist: %f Width: %f\n",
	   cam.longitude, cam.latitude, cam.distance, cam.width);
  fprintf (stderr, "common_radius=%f\n", common_radius);
  fprintf (stderr, "hsize=%d vsize=%d\n", hsize, vsize);
  fprintf (stderr, "xyzname=%s, fname_prefix=%s, wfname=%s\n",
	   xyzname, fname_prefix, wfname);
  fprintf (stderr, "transform_walls=%d\n", transform_walls);
#endif

  ymax = cam.distance * tan (cam.width);
  zmax = (3.0 / 4.0) * ymax;

  if (strlen (wfname) > 0)
    {
      inf = fopen (wfname, "r");
      if (inf == NULL)
	return ouch ("Can't open wall file");
      if (!read_wall_file (inf))
	return ouch ("Bad wall file");
      fclose (inf);
    }

  if (strlen (xyzname) > 0)
    {
      inf = fopen (xyzname, "r");
      if (inf == NULL)
	{
	  sprintf (ouch_str, "Can't open XYZ file %s", xyzname);
	  return ouch (ouch_str);
	}
      /* find out how many atoms to grab */
      while ((i = read_xyz_frame (inf)) == OK)
	if (render_structure () == OUCH)
	  return OUCH;
      fclose (inf);
    }
  if (i == OUCH)
    return OUCH;

  if (num_atoms == 0)
    {
      if (num_walls == 0)
	fprintf (stderr, "No atoms, no walls, producing a blank image");
      else if (render_structure () == OUCH)
	return OUCH;
    }

  return OK;
}