# -----------------------------------------------------------------------------
# Compute inertia tensor principle axes for molecule.
#
def inertia_ellipsoid(m):

  atoms = m.atoms
  n = len(atoms)
  from _multiscale import get_atom_coordinates
  xyz = get_atom_coordinates(atoms)
  amass = [a.element.mass for a in atoms]
  from numpy import array, identity, dot, outer, argsort, linalg
  wxyz = array(amass).reshape((n,1)) * xyz
  mass = sum(amass)

  c = wxyz.sum(axis = 0) / mass    # Center of mass
  i = (xyz*wxyz).sum()*identity(3) - dot(xyz.transpose(), wxyz)
  # Adjust inertia tensor to be about center of mass
  i -= mass*dot(c,c)*identity(3) - mass*outer(c,c)
  i /= mass

  eval, evect = linalg.eigh(i)

  # Sort by eigenvalue size.
  order = argsort(eval)
  seval = eval[order]
  sevect = evect[:,order]

  # Make rows of 3 by 3 matrix the principle axes.
  return sevect.transpose(), seval, c

# -----------------------------------------------------------------------------
#
def ellipsoid_surface(center, axes, lengths, color = (.7,.7,.7,1)):

  from Icosahedron import icosahedron_triangulation
  varray, tarray = icosahedron_triangulation(subdivision_levels = 3,
                                             sphere_factor = 1.0)
  from numpy import dot, multiply
  es = dot(varray, axes.transpose())
  ee = multiply(es, lengths)
  ev = dot(ee, axes)
  ev += center

  import _surface
  sm = _surface.SurfaceModel()
  sm.addPiece(ev, tarray, color)
  return sm

# -----------------------------------------------------------------------------
#
def show_ellipsoid(axes, d2, center, model):

  from math import sqrt
  d = [sqrt(e) for e in d2]
  elen = [sqrt(d[1]*d[2]), sqrt(d[0]*d[2]), sqrt(d[0]*d[1])]
  sm = ellipsoid_surface(center, axes, elen)
  sm.name = 'inertia ellipsoid for %s' % m.name
  from chimera import openModels as om
  om.add([sm], sameAs = model)

# -----------------------------------------------------------------------------
#
def print_axes(axes, d2, m):

  from math import sqrt
  paxes = ['\tv%d = %6.3f %6.3f %6.3f   d%d = %6.3f' %
           (a+1, axes[a][0], axes[a][1], axes[a][2], a+1, sqrt(d2[a]))
           for a in range(3)]
  from chimera.replyobj import info
  info('Inertia axes for %s\n%s\n' % (m.name, '\n'.join(paxes)))
  from Accelerators.standard_accelerators import show_reply_log
  show_reply_log()

# -----------------------------------------------------------------------------
#
from chimera import openModels as om, Molecule
mlist = om.list(modelTypes = [Molecule])
for m in mlist:
  axes, d2, center = inertia_ellipsoid(m)
  print_axes(axes, d2, m)
  show_ellipsoid(axes, d2, center, m)