"""
This file contains the implementation of several functions that can be used to
read/write VolumeData objects to various file-types.
"""
import os
import numpy as np
import schrodinger.infra.mminit as mminit
import schrodinger.infra.mmsurf as mmsurf
from . import vdexception
from . import volumedata as volumedata
_X = 0
_Y = 1
_Z = 2
_VX = np.array([1.0, 0.0, 0.0])
_VY = np.array([0.0, 1.0, 0.0])
_VZ = np.array([0.0, 0.0, 1.0])
_MUST_SPECIFY_FILE = "Either stream or filename must be specified."
_NON_ORTHOGONAL = "Non-orthogonal axes are currently not permitted."
_DEFAULT_NAME = "VolumeData"
_PYTHON = "PYTHON"
_mmInitialiser = mminit.Initializer([mmsurf.mmvisio_initialize],
                                    [mmsurf.mmvisio_terminate])
class _MMVol(object):
    """
    The _MMVol class acts as a wrapper around low-level MMVol objects. The
    class implements Python's context-manager interface to ensure that we
    don't 'leak' MMVol objects in the event of an error.
    """
    def __init__(self, mmvol=None):
        """
        This function creates a new _MMVol object. It either explicitly
        instantiates a MMVol object, or assumes control of the object passed
        via mmvol.
        :param mmvol: A valid MMVol handle, if specified
        :type mmvol: `MMVol handle`
        """
        if mmvol is None:
            self._mmvol = mmsurf.mmvol_new()
        else:
            self._mmvol = mmvol
    def __enter__(self):
        """
        Called by the context-manager as part of a 'with' statement.
        """
        return self._mmvol
    def __exit__(self, type, value, traceback):
        """
        Called by the context manager as part of a 'with' statement.
        """
        mmsurf.mmvol_delete(self._mmvol)
        return False
def _MMVolToVolumeData(mmvol, handleNonOrthogonalAxes=False):
    """
    This function converts the specified mmvol object to a VolumeData instance.
    :param mmvol: The MMVol object to be loaded into a VolumeData instance
    :param mmvol: `MMVol handle`
    :param handleNonOrthogonalAxes: True if the function should handle
        non-orthogonal axes. False otherwise. [Currently this parameter is
        ignored and the function will throw an exception if it encounters non-
        orthogonal axes.
    :type handleNonOrthogonalAxes: `bool`
    :return: The newly constructed VolumeData instance
    :rtype: ``VolumeData``
    """
    # Query the object for its various dimensions and other data.
    name = mmsurf.mmvol_get_name(mmvol)
    data = mmsurf.mmvol_get_grid_data(mmvol)
    N = np.array(mmsurf.mmvol_get_dimensions_array(mmvol))
    resolution = np.array(mmsurf.mmvol_get_axis_scales_array(mmvol))
    origin = np.array(mmsurf.mmvol_get_origin_array(mmvol))
    # The orientations of the axes is important, currently we only handle
    # orthogonal axes, but the hooks are in place for handling the non-
    # orthogonal case.
    axisOrientationX = mmsurf.mmvol_get_orientation_array(mmvol, _X)
    axisOrientationY = mmsurf.mmvol_get_orientation_array(mmvol, _Y)
    axisOrientationZ = mmsurf.mmvol_get_orientation_array(mmvol, _Z)
    if not np.all(np.equal(axisOrientationX, _VX)) or \
       not np.all(np.equal(axisOrientationY, _VY)) or \
       not np.all(np.equal(axisOrientationZ, _VZ)):
        if handleNonOrthogonalAxes:
            # DDR: Should it ever be required; code to handle non-orthogonal
            # axes can be included here. I'd suggest using scipy's
            # interpolation algorithms to map the non-orthogonal data to
            # an orthogonal system.
            pass
        else:
            raise vdexception.VDException(_NON_ORTHOGONAL)
        # Create the VolumeData object, set the data. We also set the name here,
        # although that isn't technically part of the VolumeData specification.
    ret = volumedata.VolumeData(N, resolution, origin)
    ret.Name = name
    ret.setData(data)
    return ret
[docs]def LoadVisFile(filename):
    """
    This function retrieves a VolumeData object from a Schrodinger .vis file.
    :param filename: A Schrodinger .vis file containing a valid mmvol object
    :type filename: `string`
    :return: The VolumeData extracted from the contents of the .vis file
    :rtype: ``VolumeData``
    """
    # Load the mmvol object from the file.
    with _MMVol(mmvol=mmsurf.mmvisio_read_volume_from_file(filename)) as mmvol:
        return _MMVolToVolumeData(mmvol) 
def _VolumeDataToMMVol(volumeData, mmvol):
    """
    This function configures the specified mmvol object to be equivalent to
    the specified VolumeData.
    :param volumeData: The VolumeData object to be converted
    :type volumeData: ``VolumeData``
    :param mmvol: A pre-allocated mmvol handle which awaits configuration
    :type mmvol: `MMVol handle`
    """
    # BUGBUG: I'm not sure how many of these settings are required, or if they
    # are complete. These settings were arrived at through trial and error.
    # Name and trivial properties.
    if hasattr(volumeData, "Name"):
        mmsurf.mmvol_set_name(mmvol, volumeData.Name)
    else:
        mmsurf.mmvol_set_name(mmvol, _DEFAULT_NAME)
    mmsurf.mmvol_set_comment(mmvol, _DEFAULT_NAME)
    mmsurf.mmvol_set_property_name(mmvol, _DEFAULT_NAME)
    mmsurf.mmvol_set_author_name(mmvol, _PYTHON)
    # Data type and overall dimensionality.
    mmsurf.mmvol_set_grid_type(mmvol, mmsurf.MMVOL_REGULAR)
    mmsurf.mmvol_set_data_type(mmvol, mmsurf.MMVOL_FLOAT)
    mmsurf.mmvol_set_num_dimensions(mmvol, 3)
    mmsurf.mmvol_set_num_components(mmvol, 1)
    # Axes directions.
    mmsurf.mmvol_set_orientation(mmvol, _X, _VX)
    mmsurf.mmvol_set_orientation(mmvol, _Y, _VY)
    mmsurf.mmvol_set_orientation(mmvol, _Z, _VZ)
    # Number of data points and axis scaling.
    mmsurf.mmvol_set_dimensions(mmvol, volumeData.CoordinateFrame.N)
    mmsurf.mmvol_set_axis_scales(mmvol, volumeData.CoordinateFrame.resolution)
    mmsurf.mmvol_set_origin(mmvol, volumeData.CoordinateFrame.origin)
    # Actual data.
    mmsurf.mmvol_set_grid_data(mmvol,
                               np.ascontiguousarray(volumeData.getData()))
    # Visibility.
    mmsurf.mmvol_set_default_surf_visibility(mmvol, True)
    mmsurf.mmvol_set_visible(mmvol, True)
    # Data range. Currently contoured at the mean of the data.
    # BUGBUG: Possibly better? Single contour for all +ve or all -ve data
    # at the mean. Two contours at min / 2, max / 2 for +/- data.
    min = np.min(volumeData.getData())
    max = np.max(volumeData.getData())
    sigma = np.std(volumeData.getData())
    hints = np.array([(max + min) * 0.5])
    colours = np.array([mmsurf.MMVOL_COLOR_CYAN])
    #   # Handle +/- data.
    #    if np.sign( min ) == np.sign( max ):
    #        hints = np.array( [ ( max + min ) * 0.5 ] )
    #        colours = np.array( [ mm.MMVOL_COLOR_CYAN ] )
    #    else:
    #        hints = np.array( [ min * 0.5, max * 0.5 ] )
    #        colours = np.array( [ mm.MMVOL_COLOR_RED, \
    #                              mm.MMVOL_COLOR_CYAN ] )
    mmsurf.mmvol_set_min_sample_value(mmvol, min)
    mmsurf.mmvol_set_max_sample_value(mmvol, max)
    mmsurf.mmvol_set_has_sample_range(mmvol, True)
    mmsurf.mmvol_set_sigma(mmvol, sigma)
    mmsurf.mmvol_set_contour_hints(mmvol, len(hints), hints, colours)
[docs]def SaveVisFile(volumeData, filename):
    """
    This function saves a VolumeData object to a Schrodinger .vis file.
    :param volumeData: The VolumeData to be serialised
    :type volumeData: ``volumeData``
    :param filename: Where to store the .vis file
    :type filename: `string`
    """
    with _MMVol() as mmvol:
        _VolumeDataToMMVol(volumeData, mmvol)
        # Attempt the save, before we can actually save we have to remove
        # any file that has the specified name. This appears to be due to
        # a complexity of the mmvisio library.
        if os.path.isfile(filename):
            os.remove(filename)
        mmsurf.mmvisio_write_volume_to_file(filename, mmvol) 
[docs]def SaveCCP4File(volumeData, filename):
    """
    This function saves a VolumeData object to a CCP4 file.
    :param volumeData: The VolumeData to be serialised
    :type volumeData: ``VolumeData``
    :param filename: Where to store the .CCP4 file
    :type filename: `string`
    """
    with _MMVol() as mmvol:
        _VolumeDataToMMVol(volumeData, mmvol)
        mmsurf.mmvisio_write_ccp4(mmvol, filename,
                                  mmsurf.MMSURF_CNS_MAP_TYPE_CCP4,
                                  mmsurf.MMSURF_CNS_MAP_TYPE_COMMENT_CCP4) 
[docs]def LoadCNSFile(filename):
    """
    This function loads a VolumeData object from a CNS file.
    :param filename: The .cns formatted file
    :type filename: `string`
    :return: The VolumeData object extracted from the .cns file.
    :rtype: ``VolumeData``
    """
    # Attempt to load and parse the .cns file.
    mmvol, centre, type, comments = mmsurf.mmvisio_read_cns(filename)
    with _MMVol(mmvol=mmvol) as mmvol:
        return _MMVolToVolumeData(mmvol)