#! /usr/bin/python3 -t 

MODULE_NAME = "mformula"
MODULE_DESC = "Style-independent geometric description of 'roadkill' molecular diagrams"
MODULE_VERS = "1.0"

MODULE_COPYRIGHT = "Copyright © 2009-07-06 by the State University of Campinas (UNICAMP)"

import sys
import os
import math
import copy
from math import sqrt,sin,cos,pi
# sys.path[1:0] = [ sys.path[0] + '/../lib', os.path.expandvars('${STOLFIHOME}/lib'), '.' ] 
# sys.stderr.write("%s.py: path = %r\n" % (MODULE_NAME, sys.path))
import rn

class obj :
  """Formula parameters and geometry.
  
  The origin is arbitrary, the X axis points to the right, and the Y axis points UP.
  The {Z} coordinate is usually 0 for true "roadkill", but if nonzero it will affect the 
  size of atoms and the style of bonds ("coming out" or "dipping in")."""
  
  def __init__(fm) :
    # Atom data:
    fm.natoms = 0              # Number of atoms in molecule.
    fm.atom_symbol = []        # Symbol of element of each atom (string).
    fm.atom_ctr = []           # Center coordinates (list of three floats).
    fm.atom_chnum = []         # Numerator of charge (0 if neutral).
    fm.atom_chden = []         # Denominator of charge (1 if integer).

    # Bond data:
    fm.nbonds = 0              # Number of chemical bonds.
    fm.bond_atom_1_index = []  # Index of source atom (int).
    fm.bond_atom_2_index = []  # Index of destination atom (int or None).
    fm.bond_angle = []         # Angle of exit, for open bonds (radians or None).
    fm.bond_valency = []       # Multiplicity of bond (float).
    
    # Labels:
    fm.nlabels = 0             # Number of labels.
    fm.label_text = []         # Text of each label.
    fm.label_ctr = []          # Center coords of each label (two floats, {X} and {Y} only).
    fm.label_ftsize = []       # Relative font height.
    fm.label_bold = []         # True for bold font.
    fm.label_color = []        # Label color.
    # !!! Add background oval? !!!
    
    fm.err = None              # Error message, if needed.
    #----------------------------------------------------------------------
    
  def add_atom(fm, sym,ctr,chnum,chden) :
    """Adds a new atom to the formula {fm}.
  
    The atom is an instance of the element with symbol {sym}, and its
    center is {ctr} (a list of three floats {[x,y,z]}, or just two foats
    {with z=0} implied). If {sym} is '.' the atom will be invisible, but
    bonds can still connect to it to make dangling bonds or generic
    lines. If {chnum} is not 0, the atom will have a charge of
    {chnum/chden}. Returns the index of the atom in the atom list."""
    
    i = fm.natoms
    sys.stderr.write("! adding atom %d:" % i)
    sys.stderr.write(" symbol = %r  center = %r  ch = %r/%r\n" % (sym,ctr,chnum,chden))
    
    fm.atom_symbol.append(sym)
    assert len(ctr) == 2 or len(ctr) == 3, "invalid center coords"
    if len(ctr) == 2: ctr = list(ctr) + [ 0 ]
    fm.atom_ctr.append(copy.copy(ctr))
    if (chden == 0) : chden = 1
    fm.atom_chnum.append(chnum)
    fm.atom_chden.append(chden)
    fm.natoms = i+1
    return i
    #----------------------------------------------------------------------

  def add_bond(fm, aix1,aix2,val) :
    """Adds a new chemical bond to the formula {fm}.
  
    The bond connects the two atoms with indices {aix1,aix2}. The
    parameter {val} is a float that specifies how many parallel lines to
    draw for the bond. For a dangling bond, add an atom with symbol '.'
    and connect to it. The integer part of {val} is the number of solid
    lines; if the fractional part is nonzero, there will be an
    additional dashed line. Returns the index of the bond in the bond
    list."""
    
    j = fm.nbonds
    sys.stderr.write("! adding bond %d:  org = %r  dst = %r  valency = %r\n" % (j,aix1,aix2,val))
    fm.bond_atom_1_index.append(aix1)
    fm.bond_atom_2_index.append(aix2)
    fm.bond_angle.append(None)
    fm.bond_valency.append(val)
    fm.nbonds = j+1
    return j
    #----------------------------------------------------------------------
        
  def add_free_bond(fm, aix,ang,val) :
    """Adds one or more dots representing unfilled orbitals to the formula {fm}.
  
    The dots are placed next to the atom with index {aix} at {ang}
    radians counterclockwise from the X-axis. The parameter {val} is an
    integer that specifies how many dots to draw. Returns the index of
    the bond in the bond list."""
    
    j = fm.nbonds
    sys.stderr.write("! adding free bond %d:  org = %r  ang = %r  valency = %r\n" % (j,aix,ang,val))
    fm.bond_atom_1_index.append(aix)
    fm.bond_atom_2_index.append(None)
    fm.bond_angle.append(ang)
    fm.bond_valency.append(val)
    fm.nbonds = j+1
    return j
    #----------------------------------------------------------------------
        
  def add_label(fm, txt,ctr,fh,bold,fRGB) :
    """Adds a new label to the formula {fm}.
  
    The label will be the text {txt}, rendered in the standard font with
    nominal font height {fh} in color {fRGB}, and cenetered at {ctr} (a
    list of two floats {[x,y]}). The font height is relative to the
    height of the font used for atom names. Returns the index of the
    label in the label list."""
    
    i = fm.nlabels
    sys.stderr.write("! adding label %d:" % i)
    sys.stderr.write(" text = '%s' center = %r  fh = %r fRGB = %r\n" % (txt,ctr,fh,fRGB))
    fm.label_text.append(txt)
    assert len(ctr) == 2, "invalid label center coords"
    fm.label_ctr.append(copy.copy(ctr))
    fm.label_ftsize.append(fh)
    fm.label_bold.append(bold)
    fm.label_color.append(fRGB)
    fm.nlabels = i+1
    return i
    #----------------------------------------------------------------------

  def add_subformula(fm, gm,xflip,rot,dsp) :
    """Appends to {fm} a copy of formula {gm}, modified by {xflip,rot,dsp}.
    Returns the index of the first atom of the subformula in formula
    {fm}.
  
    The global parameters and elements of {gm} are ignored. The atoms,
    bonds, and labels of {gm} are appended to the respective lists of
    {fm}, so that the old atoms, bonds, and labels of {fm} keep their
    original numbers. First, if {xflip} is true, the X coordinate of each
    atom or label is negated. In any case, the coordinates of each atom
    or label are then rotated by {rot} degrees couterclockwise around
    the origin, and finally displaced by {dsp} (a list of 2 or 3 floats).
  
    Returns the index of the first atom of the copy of {gm} in the formula {fm}."""
    
    rot_rads = 3.1415926/180*rot # Rotation in radians.
    
    assert len(dsp) == 2 or len(dsp) == 3, "invalid displacement vector"
    if len(dsp) == 2: dsp = list(dsp) + [ 0 ]
    
    # Append the atoms.
    first_atom_index = fm.natoms       # Index of first copied atom.
    for i_g in range(gm.natoms) :
      sym = gm.atom_symbol[i_g]
      ctr = copy.copy(gm.atom_ctr[i_g])
      chN = gm.atom_chnum[i_g]
      chD = gm.atom_chden[i_g]
      if xflip != 0 : ctr[0] = - ctr[0]
      if rot_rads != 0: ctr = rn.rotate2(ctr, rot_rads)
      assert len(ctr) == 3
      ctr = rn.add(ctr, dsp)
      i_f = fm.add_atom(sym, ctr, chN,chD)
      
    # Append the labels.
    for i_g in range(gm.nlabels) :
      txt = gm.label_text[i_g]
      ctr = copy.copy(gm.label_ctr[i_g])
      fsz = gm.label_ftsize[i_g]
      fbo = gm.label_bold[i_g]
      fco = gm.label_color[i_g]
      if xflip != 0 : ctr[0] = - ctr[0]
      ctr = rn.add(rn.rotate2(ctr, rot_rads), dsp[0:2])
      assert len(ctr) == 2
      i_f = fm.add_label(txt, ctr, fsz,bold,fco)
      
    # Append the bonds.
    for j_g in range(gm.nbonds) :
      aix1 = gm.bond_atom_1_index[j_g]
      aix2 = gm.bond_atom_2_index[j_g]
      ang = gm.bond_angle[j_g]
      val = gm.bond_valency[j_g]
      if (aix2 != None) :
        aix1 = aix1 + first_atom_index
        aix2 = aix2  + first_atom_index
        # If the formula is flipped, reverse direction so that dashes are flipped too:
        if xflip : 
          fm.add_bond(aix2, aix1, val)
        else :
          fm.add_bond(aix1, aix2, val)
      else :
        aix1 = aix1 + first_atom_index
        fm.add_free_bond(aix1, ang+(rot*pi/180), val)
      
    return first_atom_index
    #----------------------------------------------------------------------
   
  def compute_atom_valences_and_dirs(fm):
    """Given a formula {fm} with {n} atoms, returns a list 
    of pairs {(v, d)}, one for each atom, where {v} is the total valence of the 
    bonds to that atom, and {d} is the sum of the direction vectors from it to 
    the bonded atoms.  (In this sum, each bond is counted once with unit length, 
    ignoring its valence and length.)  Note that {d} is a 3D vector."""
    
    n = fm.natoms
    m = fm.nbonds
    vals = [0]*n
    dirs = [(0,0,0)]*n
    for kb in range(m):
      bv = fm.bond_valency[kb]
      
      ka1 = fm.bond_atom_1_index[kb]
      assert ka1 != None, "invalid atom 1 index"
      assert type(ka1) is int
      vals[ka1] += bv
      pa1 = fm.atom_ctr[ka1]
      
      ka2 = fm.bond_atom_2_index[kb]
      if ka2 == None:
        ang2 = fm.bond_angle[ka2]*pi/180
        d12 = (cos(ang2), sin(ang2), 0)
      else:
        assert type(ka2) is int
        vals[ka2] += bv
        pa2 = fm.atom_ctr[ka2]
        d12 = rn.sub(pa2, pa1)
        d12, dlen = rn.dir(d12)
        dirs[ka2] = rn.sub(dirs[ka2], d12)
      dirs[ka1] = rn.add(dirs[ka1], d12)
    return list(zip(vals, dirs))
    #----------------------------------------------------------------------
        
def greek_nums():
  """Returns a table that maps a number {n} in {1..99} to the 
  corresponding greek numeric prefix.  Except for entries {1..3},
  they end in a consonant so may require an epenthetic vowel."""
  grnum = [ None ]*101
  grnum[ 2] = "mono"
  grnum[ 2] = "di"
  grnum[ 3] = "tri"
  grnum[ 4] = "tetr"
  grnum[ 5] = "pent"
  grnum[ 6] = "hex"
  grnum[ 7] = "hept"
  grnum[ 8] = "oct"
  grnum[ 9] = "non"
  grnum[10] = "dec"
  grnum[11] = "hendec"
  grnum[12] = "dodec"
  grnum[13] = "tridec"
  for k in range(4,10):
    grnum[10+k] = grnum[k] + "adec"
  grnum[20] = "icos"
  grnum[21] = "hencos"
  grnum[22] = "docos"
  grnum[23] = "tricos"
  for k in range(4,10):
    grnum[20+k] = grnum[k] + "acos"
  for t in range(3,10):
    grt = grnum[t] + "acont"
    for k in range(0,10):
      if k == 0:
        grk = ""
      elif k == 1:
        grk = "hen"
      elif k == 2:
        grk = "di"
      else:
        grk = grnum[k] + "a"
      grnum[10*t+k] = grk + grt
  return grnum
  #----------------------------------------------------------------------
    
def alkane_nums():
  """Returns a table that maps the number {n} of carbons in a chain to the
  corresponding IUPAC alkane number prefixes.  These are the same as the
  greek numeric prefixes, except for {n} in {1..4}.
  They all end in a consonant so may require an epenthetic vowel."""
  aknum = [ None ]*101 
  aknum[1] = "meth"
  aknum[2] = "eth"
  aknum[3] = "prop"
  aknum[4] = "but"
  grnum = greek_nums()
  for k in range(5,100):
    aknum[k] = grnum[k]
  return aknum
  #----------------------------------------------------------------------