"""A module that displays Parametric Surfaces, which are generated by
sets of equations describing the deformation of the real plane into
the surface.
"""
# Authors: KK Rai (kk.rai [at] iitb.ac.in)
#          R. Ambareesha (ambareesha [at] iitb.ac.in)
#          Prabhu Ramachandran (prabhu [at] aero.iitb.ac.in)

# Enthought library imports.
from traits.api import Instance, Enum, Dict, Str
from tvtk.api import tvtk

# Local imports
from mayavi.core.source import Source
from mayavi.core.pipeline_info import PipelineInfo


######################################################################
# `ParametricSurface` class.
######################################################################
class ParametricSurface(Source):
    # The version of this class.  Used for persistence.
    __version__ = 0

    # Flag to set the parametric function type.
    function = Enum('boy', 'conic_spiral', 'cap', 'dini',
                    'ellipsoid', 'enneper', 'figure8klein', 'klein',
                    'mobius', 'random_hills', 'roman', 'spline',
                    'super_ellipsoid', 'super_toroid', 'torus',
                    desc='which parametric function to be used')

    # Define the trait 'parametric_function' whose value must be an instance of
    # type ParametricFunction
    parametric_function = Instance(tvtk.ParametricFunction,
                                   allow_none=False, record=True)

    # The Parametric function source which generates the data.
    source = Instance(tvtk.ParametricFunctionSource, args=(),
                      kw={'scalar_mode': 'distance'},
                      allow_none=False, record=True)

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

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

    # A dictionary that maps the function names to instances of the
    # parametric surfaces
    _function_dict = Dict(Str,
                          Instance(tvtk.ParametricFunction,
                                   allow_none=False))

    ######################################################################
    # `object` interface
    ######################################################################
    def __init__(self, **traits):
        # Setup the function dict.
        fd = {'boy': tvtk.ParametricBoy(),
              'conic_spiral': tvtk.ParametricConicSpiral(),
              'cap': tvtk.ParametricCrossCap(),
              'dini': tvtk.ParametricDini(),
              'ellipsoid': tvtk.ParametricEllipsoid(),
              'enneper': tvtk.ParametricEnneper(),
              'figure8klein': tvtk.ParametricFigure8Klein(),
              'klein': tvtk.ParametricKlein(),
              'mobius': tvtk.ParametricMobius(),
              'random_hills': tvtk.ParametricRandomHills(),
              'roman': tvtk.ParametricRoman(),
              'spline': tvtk.ParametricSpline(),
              'super_ellipsoid': tvtk.ParametricSuperEllipsoid(),
              'super_toroid': tvtk.ParametricSuperToroid(),
              'torus': tvtk.ParametricTorus()}
        self._function_dict = fd

        # Call parent class' init.
        super(ParametricSurface, self).__init__(**traits)

        # Initialize the function to the default mode's instance from
        # the dictionary
        self.parametric_function = self._function_dict[self.function]

        # Call render everytime source traits change.
        self.source.on_trait_change(self.render)
        # Setup the outputs.
        self.outputs = [self.source]

    ######################################################################
    # Non-public methods.
    ######################################################################
    def _function_changed(self, value):
        """This method is invoked (automatically) when the `function`
        trait is changed.
        """
        self.parametric_function = self._function_dict[self.function]

    def _parametric_function_changed(self, old, new):
        """This method is invoked (automatically) when the
        `parametric_function` attribute is changed.
        """
        self.source.parametric_function = self.parametric_function

        # Setup the handlers so that
        if old is not None:
            old.on_trait_change(self.render, remove=True)
        new.on_trait_change(self.render)
        self.source.update()
        self.data_changed = True
