import os

import numpy as nm

try:
    from enthought.tvtk.api import tvtk
    from enthought.mayavi.sources.vtk_data_source import VTKDataSource
    from enthought.pyface.timer.api import Timer

except:
    from tvtk.api import tvtk
    from mayavi.sources.vtk_data_source import VTKDataSource
    from pyface.timer.api import Timer

from dataset_manager import DatasetManager

from sfepy.base.base import Struct, basestr
from sfepy.postprocess.utils import mlab
from sfepy.discrete.fem import Mesh
from sfepy.discrete.fem.meshio import MeshIO, vtk_cell_types, supported_formats

def create_file_source(filename, watch=False, offscreen=True):
    """Factory function to create a file source corresponding to the
    given file format."""
    kwargs = {'watch' : watch, 'offscreen' : offscreen}

    if isinstance(filename, basestr):
        fmt = os.path.splitext(filename)[1]
        is_sequence = False

    else: # A sequence.
        fmt = os.path.splitext(filename[0])[1]
        is_sequence = True

    fmt = fmt.lower()

    if fmt == '.vtk':
        # VTK is supported directly by Mayavi, no need to use MeshIO.
        if is_sequence:
            return VTKSequenceFileSource(filename, **kwargs)
        else:
            return VTKFileSource(filename, **kwargs)

    elif fmt in supported_formats.keys():
        if is_sequence:
            if fmt == '.h5':
                raise ValueError('format .h5 does not support file sequences!')
            else:
                return GenericSequenceFileSource(filename, **kwargs)
        else:
            return GenericFileSource(filename, **kwargs)

    else:
        raise ValueError('unknown file format! (%s)' % fmt)

class FileSource(Struct):
    """General file source."""

    def __init__(self, filename, watch=False, offscreen=True):
        """Create a file source using the given file name."""
        mlab.options.offscreen = offscreen
        self.watch = watch
        self.filename = filename
        self.reset()

    def __call__(self, step=0):
        """Get the file source."""
        if self.source is None:
            self.source = self.create_source()
            if self.watch:
                self.timer = Timer(1000, self.poll_file)

        return self.source

    def reset(self):
        """Reset."""
        self.mat_id_name = None
        self.source = None
        self.notify_obj = None
        self.steps = []
        self.times = []
        self.step = 0
        self.time = 0.0
        if self.watch:
            self.last_stat = os.stat(self.filename)

    def setup_mat_id(self, mat_id_name='mat_id', single_color=False):
        self.mat_id_name = mat_id_name
        self.single_color = single_color

    def get_step_time(self, step=None, time=None):
        """
        Set current step and time to the values closest greater or equal to
        either step or time. Return the found values.
        """
        if (step is not None) and len(self.steps):
            step = step if step >= 0 else self.steps[-1] + step + 1
            ii = nm.searchsorted(self.steps, step)
            ii = nm.clip(ii, 0, len(self.steps) - 1)

            self.step = self.steps[ii]
            if len(self.times):
                self.time = self.times[ii]

        elif (time is not None) and len(self.times):
            ii = nm.searchsorted(self.times, time)
            ii = nm.clip(ii, 0, len(self.steps) - 1)

            self.step = self.steps[ii]
            self.time = self.times[ii]

        return self.step, self.time

    def get_ts_info(self):
        return self.steps, self.times

    def get_mat_id(self, mat_id_name='mat_id'):
        """
        Get material ID numbers of the underlying mesh elements.
        """
        if self.source is not None:
            dm = DatasetManager(dataset=self.source.outputs[0])

            mat_id = dm.cell_scalars[mat_id_name]
            return mat_id

    def file_changed(self):
        pass

    def setup_notification(self, obj, attr):
        """The attribute 'attr' of the object 'obj' will be set to True
        when the source file is watched and changes."""
        self.notify_obj = obj
        self.notify_attr = attr

    def poll_file(self):
        """Check the source file's time stamp and notify the
        self.notify_obj in case it changed. Subclasses should implement
        the file_changed() method."""
        if not self.notify_obj:
            return

        s = os.stat(self.filename)
        if s[-2] == self.last_stat[-2]:
            setattr(self.notify_obj, self.notify_attr, False)
        else:
            self.file_changed()
            setattr(self.notify_obj, self.notify_attr, True)
            self.last_stat = s

class VTKFileSource(FileSource):
    """A thin wrapper around mlab.pipeline.open()."""

    def create_source(self):
        """Create a VTK file source """
        return mlab.pipeline.open(self.filename)

    def get_bounding_box(self):
        bbox = nm.array(self.source.reader.unstructured_grid_output.bounds)
        return bbox.reshape((3,2)).T

    def set_filename(self, filename, vis_source):
        self.filename = filename
        vis_source.base_file_name = filename

        # Force re-read.
        vis_source.reader.modified()
        vis_source.update()
        # Propagate changes in the pipeline.
        vis_source.data_changed = True

class VTKSequenceFileSource(VTKFileSource):
    """A thin wrapper around mlab.pipeline.open() for VTK file sequences."""

    def __init__(self, *args, **kwargs):
        FileSource.__init__(self, *args, **kwargs)

        self.steps = nm.arange(len(self.filename), dtype=nm.int32)

    def create_source(self):
        """Create a VTK file source """
        return 
mlab.pipeline.open(self.filename[0])