# This module deals with input and output of configurations.
#
# Written by Konrad Hinsen
#

"""
I/O of molecular configurations

Input: Z-Matrix and Cartesian
Output: VRML
PDB files are handled in :class:`~MMTK.PDB`.
"""

__docformat__ = 'restructuredtext'

from MMTK import PDB, Units, Utility
from Scientific.Geometry.Objects3D import Sphere, Cone, Plane, Line, \
                                          rotatePoint
from Scientific.Geometry import Vector
from Scientific.Visualization import VRML
from Scientific import N
import os

#
# This class represents a Z-Matrix. Z-Matrix data consists of a list
# with one element for each atom being defined. Each entry is a
# list containing the data defining the atom.
#
class ZMatrix(object):

    """
    Z-Matrix specification of a molecule conformation

    ZMatrix objects can be used in chemical database entries
    to specify molecule conformations by internal coordinates.
    With the exception of the three first atoms, each atom is
    defined relative to three previously atoms by a distance,
    an angle, and a dihedral angle.
    """

    def __init__(self, data):
        """
        :param data: a list of atom definitions. Each atom definition (except
                     for the first three ones) is a list containing seven
                     elements:

                       - the atom to be defined
                       - a previously defined atom and the distance to it
                       - another previously defined atom and the angle to it
                       - a third previously defined atom and the dihedral
                         angle to it

                     The definition of the first atom contains only the first
                     element, the second atom needs the first three elements,
                     and the third atom is defined by the first five elements.
        """
        self.data = data
        self.coordinates = {}

    _substitute = True

    def findPositions(self):
        # First atom at origin
        self.coordinates[self.data[0][0]] = Vector(0,0,0)
        # Second atom along x-axis
        self.coordinates[self.data[1][0]] = Vector(self.data[1][2],0,0)
        # Third atom in xy-plane
        try:
            pos1 = self.coordinates[self.data[2][1]]
        except KeyError:
            raise ValueError("atom %d has no defined position"
                              % self.data[2][1].number)
        try:
            pos2 = self.coordinates[self.data[2][3]]
        except KeyError:
            raise ValueError("atom %d has no defined position"
                              % self.data[2][3].number)
        sphere = Sphere(pos1, self.data[2][2])
        cone = Cone(pos1, pos2-pos1, self.data[2][4])
        plane = Plane(Vector(0,0,0), Vector(0,0,1))
        points = sphere.intersectWith(cone).intersectWith(plane)
        self.coordinates[self.data[2][0]] = points[0]
        # All following atoms defined by distance + angle + dihedral
        for entry in self.data[3:]:
            try:
                pos1 = self.coordinates[entry[1]]
            except KeyError:
                raise ValueError("atom %d has no defined position"
                                  % entry[1].number)
            try:
                pos2 = self.coordinates[entry[3]]
            except KeyError:
                raise ValueError("atom %d has no defined position"
                                  % entry[3].number)
            try:
                pos3 = self.coordinates[entry[5]]
            except KeyError:
                raise ValueError("atom %d has no defined position"
                                  % entry[5].number)
            distance = entry[2]
            angle = entry[4]
            dihedral = entry[6]
            sphere = Sphere(pos1, distance)
            cone = Cone(pos1, pos2-pos1, angle)
            plane123 = Plane(pos3, pos2, pos1)
            points = sphere.intersectWith(cone).intersectWith(plane123)
            for p in points:
                if Plane(pos2, pos1, p).normal * plane123.normal > 0:
                    break
            p = rotatePoint(p, Line(pos1, pos2-pos1), dihedral)
            self.coordinates[entry[0]] = p

    def applyTo(self, object):
        """
        Define the positions of the atoms in a chemical object by the
        internal coordinates of the Z-Matrix.

        :param object: the object to which the Z-Matrix is applied
        """
        if not len(self.coordinates):
            self.findPositions()
        for entry in self.coordinates.items():
            object.setPosition(entry[0], entry[1])
        object.normalizePosition()

#
# This class represents a dictionary of Cartesian positions
#
class Cartesian(object):

    """
    Cartesian specification of a molecule conformation

    Cartesian objects can be used in chemical database entries
    to specify molecule conformations by Cartesian coordinates.
    """
    
    def __init__(self, data):
        """
        :param data: a dictionary mapping atoms to tuples of length three
                     that define its Cartesian coordinates
        """
        self.dict = data

    _substitute = True

    def applyTo(self, object):
        """
        Define the positions of the atoms in a chemical object by the
        stored coordinates.

        :param object: the object to which the coordinates are applied
        """
        for a, r in self.dict.items():
            object.setPosition(a, Vector(r[0], r[1], r[2]))

#
# VRML output
#
class VRMLWireframeFile(VRML.VRMLFile):

    def __init__(self, filename, color_values = None):
        VRML.VRMLFile.__init__(self, filename, 'w')
        self.warning = 0
        self.color_values = color_values
        if self.color_values is not None:
            lower = N.minimum.reduce(color_values.array)
            upper = N.maximum.reduce(color_values.array)
            self.color_scale = VRML.ColorScale((lower, upper))

    def write(self, object, configuration = None, distance = None):
        from MMTK.ChemicalObjects import isChemicalObject
        from MMTK.Universe import InfiniteUniverse
        if distance is None:
            try:
                distance = object.universe().distanceVector
            except AttributeError:
                distance = InfiniteUniverse().distanceVector
        if not isChemicalObject(object):
            for o in object:
                self.write(o, configuration, distance)
        else:
            for bu in object.bondedUnits():
                for a in bu.atomList():
                    self.writeAtom(a, configuration)
                if hasattr(bu, 'bonds'):
                    for b in bu.bonds:
                        self.writeBond(b, configuration, distance)

    def close(self):
        VRML.VRMLFile.close(self)
        if self.warning:
            Utility.warning('Some atoms are missing in the output file ' + \
                            'because their positions are undefined.')
            self.warning = 0

    def atomColor(self, atom):
        if self.color_values is None:
            return atom.color
        else:
            return self.color_scale(self.color_values[atom])

    def writeAtom(self, atom, configuration):
        pass

    def writeBond(self, bond, configuration, distance):
        p1 = bond.a1.position(configuration)
        p2 = bond.a2.position(configuration)
        if p1 is not None and p2 is not None:
            bond_vector = 0.5*distance(bond.a1, bond.a2, configuration)
            cut = bond_vector != 0.5*(p2-p1)
            color1 = self.atomColor(bond.a1)
            color2 = self.atomColor(bond.a2)
            material1 = VRML.EmissiveMaterial(color1)
            material2 = VRML.EmissiveMaterial(color2)
            if color1 == color2 and not cut:
                c = VRML.Line(p1, p2, material = material1)
                c.writeToFile(self)
            else:
                c = VRML.Line(p1, p1+bond_vector, material = material1)
                c.writeToFile(self)
                c = VRML.Line(p2, p2-bond_vector, material = material2)
                c.writeToFile(self)


class VRMLHighlight(VRMLWireframeFile):

    def writeAtom(self, atom, configuration):
        try:
            highlight = atom.highlight
        except AttributeError:
            highlight = 0
        if highlight:
            p = atom.position(configuration)
            if p is None:
                self.warning = 1
            else:
                s = VRML.Sphere(p, 0.1*Units.Ang,
                                material = VRML.DiffuseMaterial(atom.color),
                                reuse = 1)
                s.writeToFile(self)


class VRMLBallAndStickFile(VRMLWireframeFile):

    def writeAtom(self, atom, configuration):
        p = atom.position(configuration)
        if p is None:
            self.warning = 1
        else:
            color = self.atomColor(atom)
            s = VRML.Sphere(p, 0.1*Units.Ang,
                            material = VRML.DiffuseMaterial(color),
                            reuse = 1)
            s.writeToFile(self)

    def writeBond(self, bond, configuration, distance):
        p1 = bond.a1.position(configuration)
        p2 = bond.a2.position(configuration)
        if p1 is not None and p2 is not None:
            bond_vector = 0.5*distance(bond.a1, bond.a2, configuration)
            cut = bond_vector != 0.5*(p2-p1)
            color1 = self.atomColor(bond.a1)
            color2 = self.atomColor(bond.a2)
            material1 = VRML.EmissiveMaterial(color1)
            material2 = VRML.EmissiveMaterial(color2)
            if color1 == color2 and not cut:
                c = VRML.Cylinder(p1, p2, 0.03*Units.Ang,
                                  material = material1)
                c.writeToFile(self)
            else:
                c = VRML.Cylinder(p1, p1+bond_vector, 0.03*Units.Ang,
                                  material = material1)
                c.writeToFile(self)
                c = VRML.Cylinder(p2, p2-bond_vector, 0.03*Units.Ang,
                                  material = material2)
                c.writeToFile(self)


class VRMLChargeFile(VRMLWireframeFile):

    color_scale = VRML.SymmetricColorScale(1.)

    def writeAtom(self, atom, configuration):
        p = atom.position(configuration)
        c = atom.charge()
        c = max(min(c, 1.), -1.)
        if p is None:
            self.warning = 1
        else:
            s = VRML.Sphere(p, 0.1*Units.Ang,
                            material = VRML.Material(diffuse_color =
                                                     self.color_scale(c)))
            s.writeToFile(self)

    bond_material = VRML.DiffuseMaterial('black')

    def writeBond(self, bond, configuration, distance):
        p1 = bond.a1.position(configuration)
        p2 = bond.a2.position(configuration)
        if p1 is not None and p2 is not None:
            bond_vector = 0.5*distance(bond.a1, bond.a2, configuration)
            cut = bond_vector != 0.5*(p2-p1)
            if not cut:
                c = VRML.Line(p1, p2, material = self.bond_material)
                c.writeToFile(self)
            else:
                c = VRML.Line(p1, p1+bond_vector, material = self.bond_material)
                c.writeToFile(self)
                c = VRML.Line(p2, p2-bond_vector, material = self.bond_material)
                c.writeToFile(self)

VRMLFile = VRMLWireframeFile

#
# Recognize some standard file types by their extensions
#
def fileFormatFromExtension(filename):
    filename, ext = os.path.splitext(filename)
    if ext in _file_compressions:
        filename, ext = os.path.splitext(filename)
    try:
        return _file_formats[ext]
    except KeyError:
        raise IOError('Unknown file format')

_file_formats = {'.pdb': 'pdb', '.wrl': 'vrml'}
_file_compressions = ['.gz', '.Z']

#
# Output file for a specified format
#
def OutputFile(filename, format = None):
    if format is None:
        format = fileFormatFromExtension(filename)
    format = tuple(format.split('.'))
    try:
        return _file_types[format](filename)
    except KeyError:
        if len(format) == 1:
            return _file_types[format[0]](filename)
        else:
            _file_types[format[0]](filename, format[1])

_file_types = {'pdb': PDB.PDBOutputFile,
               ('vrml',): VRMLFile,
               ('vrml', 'wireframe'): VRMLWireframeFile,
               ('vrml', 'highlight'): VRMLHighlight,
               ('vrml', 'ball_and_stick'): VRMLBallAndStickFile,
               ('vrml', 'charge'): VRMLChargeFile}