#! /usr/bin/python -t 
# _*_ coding: iso-8859-1 _*_
# Last edited on 2009-07-07 22:03:05 by stolfi

PROG_NAME = "make-chemform-figure"
PROG_DESC = "Generates an SVG 'roadkill' molecular diagram for Wikipedia articles on chem substances"
PROG_VERS = "1.0"

import sys
import re
import os
import copy
import math
from math import sqrt,sin,cos
sys.path[1:0] = [ sys.path[0] + '/../lib', os.path.expandvars('${STOLFIHOME}/lib'), '.' ] 
import argparser; from argparser import ArgParser
import mformula

from datetime import date

PROG_COPYRIGHT = "Copyright © 2009-05-02 by the State University of Campinas (UNICAMP)"

PROG_HELP = \
  PROG_NAME+ " \\\n" \
  "    -back {BOOL} \\\n" \
  +argparser.help_info_HELP+ " \\\n" \
  "    < {FIGURE}.txt \\\n" \
  "    > {FIGURE}.svg"
  
PROG_INFO = \
  "NAME\n" \
  "  " +PROG_NAME+ " - " +PROG_DESC+  ".\n" \
  "\n" \
  "SYNOPSIS\n" \
  "  " +PROG_HELP+ "\n" \
  "\n" \
  "DESCRIPTION\n" \
  "  " +PROG_DESC+ ".\n" \
  "\n" \
  "OPTIONS\n" \
  "  -back {BOOL} \n" \
  "    If {BOOL} is 1, paints the background with a nonwhite color. If {BOOL} is 0," \
  " leaves it transparent. This option is meant to debug the image size.\n" \
  "\n" \
  "INPUT FILE FORMAT\n" \
  "  The input file describes the molecular diagra, as follows. First, some general parameters:\n" \
  "\n" \
  "    fontheight = {FONT_HEIGHT_PX}\n" \
  "    atomradius = {ATOM_RADIUS_PX}\n" \
  "    bondlength = {BOND_LENGTH_PX}\n" \
  "\n" \
  "  Then, the number of chemical species (element) that will be used:\n" \
  "\n" \
  "    nelems = {NUM_CHEMICAL_SPECIES}\n" \
  "\n" \
  "  Then, one line for each chemical species, in the format:\n" \
  "\n" \
  "      {IORD} {ELEM_SYMBOL} {RADIUS} {COLOR_R} {COLOR_G} {COLOR_B} {SHOW_SYMBOL}\n" \
  "\n" \
  "  Then, the number of atoms in the formula:\n" \
  "\n" \
  "    natoms = {NUM_ATOMS}\n" \
  "\n" \
  "  Then, one line for each atom:\n" \
  "\n" \
  "      {IORD} {ELEM_SYMBOL} {CTR_X} {CTR_Y}\n" \
  "\n" \
  "  These lines are followed by the number of chemical bonds in the formula:\n" \
  "\n" \
  "    nbonds = {NUM_BONDS}\n" \
  "\n" \
  "  Then, one line for each chemical bond:\n" \
  "\n" \
  "      {IORD} {ATOM_1_INDEX} {ATOM_2_INDEX} {VALENCY}\n" \
  "\n" \
  "  The {RADIUS} will be multiplied by {ATOM_RADIUS_PX} to get" \
  " the actual atom radius. The {SHOW_SYMBOL} is a flag (0 or 1).\n" \
  "\n" \
  "  The coordintes {CTR_X,CTR_Y} will be multiplied by the {BOND_LENGTH_PX}" \
  " parameter.  The origin is assumed to be at the LOWER left, and the Y" \
  " axis points up.\n" \
  "\n" \
  "  The {VALENCY} is 1, 2, or 3 for covalent bonds (solid lines), and 0.5, 1.5, or" \
  " 2.5 for weak or aromatic bonds (the '.5' stands for a dashed line).\n" \
  "\n" \
  "DOCUMENTATION OPTIONS\n" \
  +argparser.help_info_INFO+ "\n" \
  "\n" \
  "SEE ALSO\n" \
  "  cat(1).\n" \
  "\n" \
  "AUTHOR\n" \
  "  Created 2009-07-06 by Jorge Stolfi, IC-UNICAMP.\n" \
  "\n" \
  "MODIFICATION HISTORY\n" \
  "  2009-07-06 by J. Stolfi, IC-UNICAMP: created.\n" \
  "\n" \
  "WARRANTY\n" \
  "  " +argparser.help_info_NO_WARRANTY+ "\n" \
  "\n" \
  "RIGHTS\n" \
  "  " +PROG_COPYRIGHT+ ".\n" \
  "\n" \
  "  " +argparser.help_info_STANDARD_RIGHTS

# COMMAND ARGUMENT PARSING
pp = ArgParser(sys.argv, sys.stderr, PROG_HELP, PROG_INFO)

class Options :
  back = None;
  err = None;
  #----------------------------------------------------------------------
 
def arg_error(msg):
  "Prints the error message {msg} about the command line arguments, and aborts."
  sys.stderr.write("%s\n" % msg);
  sys.stderr.write("usage: %s\n" % PROG_HELP);
  sys.exit(1)
  #----------------------------------------------------------------------

def parse_args(pp) :
  "Parses command line arguments.\n" \
  "\n" \
  "  Expects an {ArgParser} instance {pp} containing the arguments," \
  " still unparsed.  Returns an {Options} instance {op}, where" \
  " {op.err} is the parsing error messages (a string or {None})."
  
  op = Options();
 
  pp.get_keyword("-back");
  op.back = pp.get_next_int(0, 1);
  
  return op
  #----------------------------------------------------------------------

class Dimensions :
  "Plot dimensions." 
  
  def __init__(dim, op,fm) :

    # Font size:
    dim.font_wx = 0.75*fm.fontheight;
    dim.font_wy = 1.00*fm.fontheight;
    
    # Displacement from edge:
    dim.mole_dx = 4;
    dim.mole_dy = 4;

    # Figure dimensions:
    min_x = 9999999.0;
    max_x = -9999999.0;
    min_y = 9999999.0;
    max_y = -9999999.0;
    for i in range(fm.natoms) :
      xi = fm.atom_center_x[i]*fm.bondlength;
      yi = fm.atom_center_y[i]*fm.bondlength;
      k = fm.elem_symbol.index(fm.atom_symbol[i]);
      ri = fm.elem_radius[k]*fm.atomradius;
      if xi - ri < min_x : min_x = xi - ri;
      if xi + ri > max_x : max_x = xi + ri;
      if yi - ri < min_y : min_y = yi - ri;
      if yi + ri > max_y : max_y = yi + ri;
    dim.mole_min_x = min_x;
    dim.mole_max_x = max_x;
    dim.mole_min_y = min_y;
    dim.mole_max_y = max_y;
    
    # !!! Must be computed from {fm} !!!
    dim.fig_wx = 2*dim.mole_dx + (dim.mole_max_x - dim.mole_min_x);  # Total figure width.
    dim.fig_wy = 2*dim.mole_dy + (dim.mole_max_y - dim.mole_min_y);  # Total figure height.
    
    dim.scale = 2.0;   # Final scale factor (can be chaged without changing any other dim).
    #----------------------------------------------------------------------

  #----------------------------------------------------------------------
  
def output_figure(op,fm) :
  "Writes the figure to {stdout}." \
  "\n" \
  "  Expects an {Options} instance {op} and a {GeomData} instance {fm}."

  # Computes the sizes of things and the perspective map:
  dim = Dimensions(op,fm)

  output_figure_preamble(op,fm,dim)
  sys.stdout.write('\n')
  output_figure_obj_defs(op,fm,dim)
  sys.stdout.write('\n')
  output_figure_body(op,fm,dim)
  sys.stdout.write('\n')
  output_figure_postamble(op,fm,dim)
  sys.stderr.write("done.\n")
  #----------------------------------------------------------------------
  
def output_figure_preamble(op,fm,dim) :
  "Writes the SVG preamble to {stdout}."
  
  sys.stdout.write( \
    '<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n' \
    '<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN" "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">\n' \
    '<!-- Created on ' + date.isoformat(date.today()) + ' by Jorge Stolfi with the script ' + PROG_NAME + ' -->\n' \
    '<!-- This file is declared PUBLIC DOMAIN by its creator -->\n' \
    '\n' \
    '<svg\n' \
    '  id="fig"\n' \
    '  xmlns="http://www.w3.org/2000/svg"\n' \
    '  xmlns:xlink="http://www.w3.org/1999/xlink"\n' \
    '\n' \
    '  fill="none"\n' \
    '  fill-opacity="1"\n' \
    '  fill-rule="evenodd"\n' \
    '\n' \
    '  stroke-linecap="round"\n' \
    '  stroke-linejoin="round"\n' \
    '  stroke-dasharray="none"\n' \
    '  stroke-opacity="1"\n' \
    '\n' \
    '  font-style="normal"\n' \
    '  font-weight="bold"\n' \
    '  font-size="' + ifm(dim.font_wy) + 'px"\n' \
    '  font-family="Bitstream Vera"\n' \
    '\n' \
    '  width="' + ifm(dim.scale*dim.fig_wx) + '"\n' \
    '  height="' + ifm(dim.scale*dim.fig_wy) + '"\n' \
    '>\n' \
  )
  sys.stdout.write('\n')
  if (op.back) :
    sys.stdout.write( \
      '    <!-- Rectangle to test the figure size -->\n' \
      '    <rect x="000" y="000" width="' + ifm(dim.scale*dim.fig_wx) + '" height="' + ifm(dim.scale*dim.fig_wy) + '"' \
             ' stroke="#000000" stroke-width="2px" fill="#ffcc55"' \
             ' />\n' \
    )
  sys.stdout.write('\n')

  # Scale everything and position the diagram:
  sys.stdout.write( \
    '  <g\n' \
    '    transform="scale(' + ifm(dim.scale) + ')"\n' \
    '  >\n' \
  )
  #----------------------------------------------------------------------

def output_figure_obj_defs(op,fm,dim) :
  "Writes the main object definitions to {stdout}."

  sys.stdout.write( \
    '  <defs>\n' \
    '  </defs>\n' \
  )
  #----------------------------------------------------------------------

def output_figure_body(op,fm,dim) :
  "Writes the body of the figure to {stdout}."

  # Plot the bonds:
  for i in range(fm.nbonds) :
    output_bond(op,fm,dim,i);
    
  # Paint, draw, and label the atoms:
  for i in range(fm.natoms) :
    output_atom(op,fm,dim,i);
  #----------------------------------------------------------------------

def output_figure_postamble(op,fm,dim) :
  "Writes the SVG postamble to {stdout}."

  sys.stdout.write( \
      '  </g>\n' \
      '</svg>\n' \
  )

def output_bond(op,fm,dim,i):
  "Draws bond number {i}."
  
  k1 = fm.bond_atom_1_index[i];
  k2 = fm.bond_atom_2_index[i];
  vc = fm.bond_valency[i];
  
  x1 = fm.atom_center_x[k1]*fm.bondlength;
  y1 = fm.atom_center_y[k1]*fm.bondlength;
  
  x2 = fm.atom_center_x[k2]*fm.bondlength;
  y2 = fm.atom_center_y[k2]*fm.bondlength;
  
  draw_bond_lines(op,fm,dim,x1,y1,x2,y2,vc);
  #----------------------------------------------------------------------

def output_atom(op,fm,dim,i):
  "Draws atom number {i}, including its label."
  
  # Get the element index {k}:
  sym = fm.atom_symbol[i];
  k = fm.elem_symbol.index(sym);
  
  # Get the atom's radius {r}:
  rad = fm.elem_radius[k]*fm.atomradius;
  
  # Get the element's color {R,G,B}:
  R = fm.elem_color_R[k];
  G = fm.elem_color_G[k];
  B = fm.elem_color_B[k];
  
  # Get the atom's center {(x,y)}:
  cx = fm.atom_center_x[i]*fm.bondlength;
  cy = fm.atom_center_y[i]*fm.bondlength;
  
  if fm.elem_show_symbol[k] :
    label = sym;
  else :
    label = '';
  draw_atom(op,fm,dim,cx,cy,rad,R,G,B,label);
  #----------------------------------------------------------------------
  
def draw_bond_lines(op,fm,dim,x1,y1,x2,y2,vc):
  "Draws a set of bond lines from {(x1,y1)} to {(x2,y2)} (px, user axes).\n" \
  "\n" \
  "  The integer part of {vc} specifies the number of solid lines.  If" \
  " the fractional part of {vc} is not zero, adds a dashed line."
  
  # Number of lines, total and solid:
  nlines = int(vc + 0.999);
  nsolid = int(vc);
  
  # Lateral displacement between lines:
  tad = 1.5*fm.atomradius/(nlines + 1);
  dx = y2 - y1;
  dy = x1 - x2;
  d = math.hypot(dx, dy);
  dx = tad * (dx / d);
  dy = tad * (dy / d);
  hn = 0.5*(nlines - 1.0);
  
  for j in range(nlines) :
    # Draw line {j}:
    dashed = (j >= nsolid);
    fr = (j - hn);
    ex = dx * fr;
    ey = dy * fr;
    draw_segment(op,fm,dim,x1+ex,y1+ey,x2+ex,y2+ey,dashed);
  #----------------------------------------------------------------------
  
def draw_atom(op,fm,dim,cx,cy,rad,R,G,B,label):
  "Fills and draws a circle with center {(cx,cy)} and radius {rad} (in px, user axes).\n" \
  "\n" \
  "  The circle will have a stroked outline of color {R,G,B} and" \
  " will be filled with a washed-out version of the same color.  Also" \
  " draws the label {label} in contrasting color."
  
  # Compute foreground color:
  Rfg = 0.3*0.000 + 0.7*R;
  Gfg = 0.3*0.000 + 0.7*G;
  Bfg = 0.3*0.000 + 0.7*B;
  
  # Compute background color:
  Rbg = R;
  Gbg = G;
  Bbg = B;
  
  # Convert to image axes:
  cxp = (cx - dim.mole_min_x) + dim.mole_dx;
  cyp = (dim.mole_max_y - cy) + dim.mole_dy;
  radp = rad;
  
  sys.stdout.write( \
    '    <g stroke-width="' + wdfm(fm.ringwidth) + 'px"' \
           ' stroke="rgb(' + rgbfm(Rfg,Gfg,Bfg) + ')"' \
           ' fill="rgb('+ rgbfm(Rbg,Gbg,Bbg) + ')">\n' \
    '      <circle ' + xyfm(cxp,cyp,'c') + ' r="' + radfm(radp) + '"/>\n' \
    '    </g>\n' \
  );

  if label != '' :
    # Print label string with proper font formatting:
    str = make_roman_style(label);
    white = ((0.300*Rbg + 0.600*Gbg + 0.100*Bbg) < 0.290);
    draw_label(op,fm,dim,cx,cy,0,0,str,white);
  #----------------------------------------------------------------------

def draw_label(op,fm,dim,cx,cy,dx,dy,str,white) :
  "Ouputs a label {str} at {(cx,cy)} (in px, user axes).\n" \
  "\n" \
  "  The label will be displaced by {dx,dy} (in font cell" \
  " units, user axes) from the reference" \
  " point {cx,cy}.  The reference point is at start, middle, or end\n" \
  " of text depending on the sign of {dx}. If {white} is true, uses white text."

  if (str != '') :
    # Convert to image axes:
    cxp = (cx - dim.mole_min_x) + dim.mole_dx;
    cyp = (dim.mole_max_y - cy) + dim.mole_dy;
    dxp = dx*dim.font_wx*len(str);
    dyp = dy*dim.font_wy;
    # Compute anchor {anch}:
    if (dxp < 0) :
      anch = 'end'
    elif (dxp > 0) :
      anch = 'start'
    else :
      anch = 'middle';
    if (dyp > 0) :
      # Adjust {dpy} for font height:
      dyp += 0.6*dim.font_wy; 
    elif (dy == 0) :
      dyp += 0.3*dim.font_wy; 
    cxp = cxp + dxp;
    cyp = cyp + dyp;
    # text fill and outline color:
    fill_color = "rgb(" + rgbfm(white,white,white) + ")";
    draw_color = "none";
    draw_width = 0.0;
    sys.stdout.write( \
      '    <text ' + xyfm(cxp,cyp,'') + ' stroke="' + draw_color + '"' \
             ' stroke-width="' + wdfm(draw_width) + 'px"' \
             ' fill="' + fill_color + '" text-anchor="' + anch + '">' + str + '</text>\n' \
    );
  #----------------------------------------------------------------------

def draw_segment(op,fm,dim,x1,y1,x2,y2,dashed) :
  "Draws a straight line from {(x1,y1)} to {(x2,y2)} (in px, user axes).\n" \
  "\n" \
  "  If {dashed} is 1, the line is dashed, else it is solid."

  # Convert to pixel coordinates:
  x1p = (x1 - dim.mole_min_x) + dim.mole_dx;
  y1p = (dim.mole_max_y - y1) + dim.mole_dy;
    
  x2p = (x2 - dim.mole_min_x) + dim.mole_dx;
  y2p = (dim.mole_max_y - y2) + dim.mole_dy;
  
  # Parameters:
  if (dashed) :
    dash_style = 'stroke-dasharray="4" stroke-dashoffset="2"';
  else :
    dash_style = '';

  sys.stdout.write( \
    '    <path d="M ' + ptfm(x1p,y1p) + ' L ' + ptfm(x2p,y2p) + '" ' \
           'stroke-width="' + wdfm(fm.bondwidth) + 'px" ' \
           'stroke="rgb(0,0,0)" ' \
           '' + dash_style + '/> -->\n' \
  );
  #----------------------------------------------------------------------

def rgbfm(R,G,B) :
  "Formats {R,G,B} as a color triplet."
  return "%d,%d,%d" % (int(255*R+0.5),int(255*G+0.5),int(255*B+0.5));

def ptfm(x,y) :
  "Formats {x,y} as two coordinates separated by a comma."
  return "%+06.1f,%+06.1f" % (x,y);

def radfm(r) :
  "Formats {r} as a radius."
  return "%06.1f" % r;

def xyfm(x,y,tag) :
  "Formats {x,y} as '{tag}x=\"{x}\" {tag}y=\"{y}\"'."
  return "%sx=\"%+06.1f\" %sy=\"%+06.1f\"" % (tag,x,tag,y);

def wdfm(wd) :
  "Formats {wd} as a line width."
  return "%.2f" % wd;

def ifm(x) :
  "Converts the decimal number {x} to string."
  return ("%r" % x)
  #----------------------------------------------------------------------
  
def make_italic_style(str) :
  "Packages a string with italic style markup."
  if (str != '') :
    return '<tspan font-style="italic">' + str + '</tspan>'
  else :
    return ''

def make_roman_style(str) :
  "Packages a string with normal style (non-italic) markup."
  if (str != '') :
    return '<tspan>' + str + '</tspan>'
  else :
    return ''
  
#----------------------------------------------------------------------

# Main program:
op = parse_args(pp);
fm = parse_input();
output_figure(op,fm);