"""This source manages a VTK dataset given to it.  When this source is
pickled or persisted, it saves the data given to it in the form of a
gzipped string.
"""
# Author: Prabhu Ramachandran <prabhu_r@users.sf.net>
# Copyright (c) 2005-2015, Enthought, Inc.
# License: BSD Style.

import sys
import os
import tempfile

# Enthought library imports.
from traits.api import Instance, List, Str, Bool, Int
from traitsui.api import View, Group, Item
from apptools.persistence.state_pickler \
     import gzip_string, gunzip_string, set_state
from tvtk.api import tvtk
from tvtk import messenger

# Local imports.
from tvtk.common import is_old_pipeline, configure_input_data
from mayavi.core.source import Source
from mayavi.core.common import handle_children_state
from mayavi.core.trait_defs import DEnum
from mayavi.core.pipeline_info import (PipelineInfo,
        get_tvtk_dataset_name)
from .vtk_xml_file_reader import get_all_attributes


######################################################################
# Utility functions.
######################################################################
def write_dataset_to_string(data):
    """Given a dataset, convert the dataset to an ASCII string that can
    be stored for persistence.
    """
    w = tvtk.DataSetWriter(write_to_output_string=1)
    warn = w.global_warning_display
    configure_input_data(w, data)
    w.global_warning_display = 0
    w.update()
    if w.output_string_length == 0:
        # Some VTK versions (5.2) have a bug when writing structured
        # grid datasets and produce empty output.  We work around this
        # by writing to a file and then reading that output.
        w.write_to_output_string = 0
        fh, fname = tempfile.mkstemp('.vtk')
        os.close(fh); os.remove(fname)
        w.file_name = fname
        w.write()
        # Read the data and delete the file.
        sdata = open(fname).read()
        os.remove(fname)
    else:
        sdata = w.output_string
    w.global_warning_display = warn
    return sdata

def has_attributes(dataset):
    """Returns `True` when the given TVTK `dataset` has any attribute
    arrays in point and cell data and `False` otherwise.
    """
    pd = dataset.point_data
    if pd is not None and pd.number_of_arrays > 0:
        return True
    cd = dataset.cell_data
    if cd is not None and cd.number_of_arrays > 0:
        return True
    return False


######################################################################
# `VTKDataSource` class
######################################################################
class VTKDataSource(Source):

    """This source manages a VTK dataset given to it.  When this
    source is pickled or persisted, it saves the data given to it in
    the form of a gzipped string.

    Note that if the VTK dataset has changed internally and you need
    to notify the mayavi pipeline to flush the data just call the
    `modified` method of the VTK dataset and the mayavi pipeline will
    update automatically.
    """

    # The version of this class.  Used for persistence.
    __version__ = 0

    # The VTK dataset to manage.
    data = Instance(tvtk.DataSet, allow_none=False)

    # Information about what this object can produce.
    output_info = PipelineInfo(datasets=['any'],
                               attribute_types=['any'],
                               attributes=['any'])

    ########################################
    # Dynamic traits: These traits are dynamic and are updated on the
    # _update_data method.

    # The active point scalar name.
    point_scalars_name = DEnum(values_name='_point_scalars_list',
                               desc='scalar point data attribute to use')
    # The active point vector name.
    point_vectors_name = DEnum(values_name='_point_vectors_list',
                               desc='vectors point data attribute to use')
    # The active point tensor name.
    point_tensors_name = DEnum(values_name='_point_tensors_list',
                               desc='tensor point data attribute to use')

    # The active cell scalar name.
    cell_scalars_name = DEnum(values_name='_cell_scalars_list',
                               desc='scalar cell data attribute to use')
    # The active cell vector name.
    cell_vectors_name = DEnum(values_name='_cell_vectors_list',
                               desc='vectors cell data attribute to use')
    # The active cell tensor name.
    cell_tensors_name = DEnum(values_name='_cell_tensors_list',
                               desc='tensor cell data attribute to use')

    ########################################
    # Our view.

    view = View(Group(Item(name='point_scalars_name'),
                      Item(name='point_vectors_name'),
                      Item(name='point_tensors_name'),
                      Item(name='cell_scalars_name'),
                      Item(name='cell_vectors_name'),
                      Item(name='cell_tensors_name'),
                      Item(name='data'),
                      ))

    ########################################
    # Private traits.

    # These private traits store the list of available data
    # attributes.  The non-private traits use these lists internally.
    _point_scalars_list = List(Str)
    _point_vectors_list = List(Str)
    _point_tensors_list = List(Str)
    _cell_scalars_list = List(Str)
    _cell_vectors_list = List(Str)
    _cell_tensors_list = List(Str)

    # This filter allows us to change the attributes of the data
    # object and will ensure that the pipeline is properly taken care
    # of.  Directly setting the array in the VTK object will not do
    # this.
    _assign_attribute = Instance(tvtk.AssignAttribute, args=(),
                                 allow_none=False)

    # Toggles if this is the first time this object has been used.
    _first = Bool(True)

    # The ID of the observer for the data.
    _observer_id = Int(-1)

    ######################################################################
    # `object` interface
    ######################################################################
    def __get_pure_state__(self):
        d = super(VTKDataSource, self).__get_pure_state__()
        for name in ('_assign_attribute', '_first', '_observer'):
            d.pop(name, None)
        for name in ('point_scalars', 'point_vectors',
                     'point_tensors', 'cell_scalars',
                     'cell_vectors', 'cell_tensors'):
            d.pop('_' + name + '_list', None)
            d.pop('_' + name + '_name', None)
        data = self.data
        if data is not None:
            sdata = write_dataset_to_string(data)
            if sys.version_info[0] > 2:
                z = gzip_string(sdata.encode('ascii'))
            else:
                z = gzip_string(sdata)
            d['data'] = z
        return d