"""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.output]

    ######################################################################
    # 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
