"""This module provides classes that allow numpy style access
to VTK datasets. See examples at bottom.
"""

from vtkmodules.vtkCommonCore import vtkPoints
from vtkmodules.vtkCommonDataModel import (
    vtkCellArray,
    vtkDataObject,
    vtkFieldData,
    vtkDataSetAttributes,
    vtkPointData,
    vtkCellData,
    vtkDataObject,
    vtkImageData,
    vtkPolyData,
    vtkUnstructuredGrid,
    vtkPartitionedDataSet
)

from vtkmodules.numpy_interface import dataset_adapter as dsa
import numpy
import weakref

class FieldDataBase(object):
    def __init__(self):
        self.association = None
        self.dataset = None

    def __getitem__(self, idx):
        """Implements the [] operator. Accepts an array name or index."""
        return self.get_array(idx)

    def __setitem__(self, name, value):
        """Implements the [] operator. Accepts an array name or index."""
        return self.set_array(name, value)

    def get_array(self, idx):
        "Given an index or name, returns a VTKArray."
        if isinstance(idx, int) and idx >= self.GetNumberOfArrays():
            raise IndexError("array index out of range")
        vtkarray = super().GetArray(idx)
        if not vtkarray:
            vtkarray = self.GetAbstractArray(idx)
            if vtkarray:
                return vtkarray
            return dsa.NoneArray
        array = dsa.vtkDataArrayToVTKArray(vtkarray, self.dataset)
        array.Association = self.association
        return array

    def keys(self):
        """Returns the names of the arrays as a list."""
        kys = []
        narrays = self.GetNumberOfArrays()
        for i in range(narrays):
            name = self.GetAbstractArray(i).GetName()
            if name:
                kys.append(name)
        return kys

    def values(self):
        """Returns the arrays as a list."""
        vals = []
        narrays = self.GetNumberOfArrays()
        for i in range(narrays):
            a = self.get_array(i)
            if a.GetName():
                vals.append(a)
        return vals

    def set_array(self, name, narray):
        """Appends a new array to the dataset attributes."""
        if narray is dsa.NoneArray:
            # if NoneArray, nothing to do.
            return

        if self.association == vtkDataObject.POINT:
            arrLength = self.dataset.GetNumberOfPoints()
        elif self.association == vtkDataObject.CELL:
            arrLength = self.dataset.GetNumberOfCells()
        elif self.association == vtkDataObject.ROW \
          and self.dataset.GetNumberOfColumns() > 0:
            arrLength = self.dataset.GetNumberOfRows()
        else:
            if not isinstance(narray, numpy.ndarray):
                arrLength = 1
            else:
                arrLength = narray.shape[0]

        # Fixup input array length:
        if not isinstance(narray, numpy.ndarray) or numpy.ndim(narray) == 0: # Scalar input
            dtype = narray.dtype if isinstance(narray, numpy.ndarray) else type(narray)
            tmparray = numpy.empty(arrLength, dtype=dtype)
            tmparray.fill(narray)
            narray = tmparray
        elif narray.shape[0] != arrLength: # Vector input
            components = 1
            for l in narray.shape:
                components *= l
            tmparray = numpy.empty((arrLength, components), dtype=narray.dtype)
            tmparray[:] = narray.flatten()
            narray = tmparray

        shape = narray.shape

        if len(shape) == 3:
            # Array of matrices. We need to make sure the order  in memory is right.
            # If column order (c order), transpose. VTK wants row order (fortran
            # order). The deep copy later will make sure that the array is contiguous.
            # If row order but not contiguous, transpose so that the deep copy below
            # does not happen.
            size = narray.dtype.itemsize
            if (narray.strides[1]/size == 3 and narray.strides[2]/size == 1) or \
                (narray.strides[1]/size == 1 and narray.strides[2]/size == 3 and \
                 not narray.flags.contiguous):
                narray  = narray.transpose(0, 2, 1)

        # If array is not contiguous, make a deep copy that is contiguous
        if not narray.flags.contiguous:
            narray = numpy.ascontiguousarray(narray)

        # Flatten array of matrices to array of vectors
        if len(shape) == 3:
            narray = narray.reshape(shape[0], shape[1]*shape[2])

        # this handle the case when an input array is directly appended on the
        # output. We want to make sure that the array added to the output is not
        # referring to the input dataset.
        copy = dsa.VTKArray(narray)
        try:
            copy.VTKObject = narray.VTKObject
        except AttributeError: pass
        arr = dsa.numpyTovtkDataArray(copy, name)
        self.AddArray(arr)

@vtkFieldData.override
class vtkFieldData(FieldDataBase, vtkFieldData):
    pass

class DataSetAttributesBase(FieldDataBase):
    pass

@vtkDataSetAttributes.override
class DataSetAttributes(DataSetAttributesBase, vtkDataSetAttributes):
    pass

@vtkPointData.override
class PointData(DataSetAttributesBase, vtkPointData):
    pass

@vtkCellData.override
class CellData(DataSetAttributesBase, vtkCellData):
    pass

class CompositeDataSetAttributesIterator(object):
    def __init__(self, cdsa):
        self._cdsa = cdsa
        if cdsa:
            self._itr = iter(cdsa.keys())
        else:
            self._itr = None

    def __iter__(self):
        return self

    def __next__(self):
        if not self._cdsa:
            raise StopIteration

        name = next(self._itr)
        return self._cdsa[name]

    def next(self):
        return self.__next__()

class CompositeDataSetAttributes(object):
    """This is a python friendly wrapper for vtkDataSetAttributes for composite
    datasets. Since composite datasets themselves don't have attribute data, but
    the attribute data is associated with the leaf nodes in the composite
    dataset, this class simulates a DataSetAttributes interface by taking a
    union of DataSetAttributes associated with all leaf nodes."""

    def __init__(self, dataset, association):
        self.DataSet = dataset
        self.Association = association
        self.ArrayNames = []
        self.Arrays = {}

        # build the set of arrays available in the composite dataset. Since
        # composite datasets can have partial arrays, we need to iterate over
        # all non-null blocks in the dataset.
        self.__determine_arraynames()

    def __determine_arraynames(self):
        array_set = set()
        array_list = []
        for dataset in self.DataSet:
            dsa = dataset.GetAttributes(self.Association)
            for array_name in dsa.keys():
                if array_name not in array_set:
                    array_set.add(array_name)
                    array_list.append(array_name)
        self.ArrayNames = array_list

    def modified(self):
        """Rescans the contained dataset to update the
        internal list of arrays."""
        self.__determine_arraynames()

    def keys(self):
        """Returns the names of the arrays as a list."""
        return self.ArrayNames

    def __getitem__(self, idx):
        """Implements the [] operator. Accepts an array name."""
        return self.get_array(idx)

    def __setitem__(self, name, narray):
        """Implements the [] operator. Accepts an array name."""
        return self.set_array(name, narray)

    def set_array(self, name, narray):
        """Appends a new array to the composite dataset attributes."""
        if narray is dsa.NoneArray:
            # if NoneArray, nothing to do.
            return

        added = False
        if not isinstance(narray, dsa.VTKCompositeDataArray): # Scalar input
            for ds in self.DataSet:
                ds.GetAttributes(self.Association).set_array(name, narray)
                added = True
            if added:
                self.ArrayNames.append(name)
                # don't add the narray since it's a scalar. GetArray() will create a
                # VTKCompositeArray on-demand.
        else:
            for ds, array in zip(self.DataSet, narray.Arrays):
                if array is not None:
                    ds.GetAttributes(self.Association).set_array(name, array)
                    added = True
            if added:
                self.ArrayNames.append(name)
                self.Arrays[name] = weakref.ref(narray)

    def get_array(self, idx):
        """Given a name, returns a VTKCompositeArray."""
        arrayname = idx
        if arrayname not in self.ArrayNames:
            return dsa.NoneArray
        if arrayname not in self.Arrays or self.Arrays[arrayname]() is None:
            array = dsa.VTKCompositeDataArray(
                dataset = self.DataSet, name = arrayname, association = self.Association)
            self.Arrays[arrayname] = weakref.ref(array)
        else:
            array = self.Arrays[arrayname]()
        return array

    def __iter__(self):
        "Creates an iterator for the contained arrays."
        return CompositeDataSetAttributesIterator(self)

#class DataSet(DataObjectBase):
class DataSet(object):
    @property
    def point_data(self):
        pd = super().GetPointData()
        pd.dataset = self
        # TODO: temporary hack
        pd.dataset.VTKObject = pd.dataset
        pd.association = self.POINT
        return pd

    @property
    def cell_data(self):
        cd = super().GetCellData()
        cd.dataset = self
        cd.association = self.CELL
        return cd

    def convert_to_unstructured_grid(self):
        from vtkmodules.vtkFiltersCore import vtkExtractCells

        ecells = vtkExtractCells()
        ecells.SetInputData(self)
        ecells.ExtractAllCellsOn()
        ecells.Update()
        return ecells.GetOutput()

class PointSet(DataSet):
    @property
    def points(self):
        pts = self.GetPoints()
        if not pts or not pts.GetData():
            return None
        return dsa.vtkDataArrayToVTKArray(pts.GetData())

    @points.setter
    def points(self, points):
        pts = dsa.numpyTovtkDataArray(points, "points")
        vtkpts = vtkPoints()
        vtkpts.SetData(pts)
        self.SetPoints(vtkpts)

@vtkUnstructuredGrid.override
class vtkUnstructuredGrid(PointSet, vtkUnstructuredGrid):
    @property
    def cells(self):
        ca = self.GetCells()
        conn_vtk = ca.GetConnectivityArray()
        conn = dsa.vtkDataArrayToVTKArray(conn_vtk)
        offsets_vtk = ca.GetOffsetsArray()
        offsets = dsa.vtkDataArrayToVTKArray(offsets_vtk)
        ct_vtk = self.GetCellTypesArray()
        ct = dsa.vtkDataArrayToVTKArray(ct_vtk)
        return { 'connectivity' : conn, 'offsets' : offsets , 'cell_types' : ct}

    @cells.setter
    def cells(self, cells):
        ca = vtkCellArray()
        conn_vtk = dsa.numpyTovtkDataArray(cells['connectivity'])
        offsets_vtk = dsa.numpyTovtkDataArray(cells['offsets'])
        cell_types_vtk = dsa.numpyTovtkDataArray(cells['cell_types'])
        print(cells['cell_types'][1])
        ca.SetData(offsets_vtk, conn_vtk)
        self.SetCells(cell_types_vtk, ca)

@vtkImageData.override
class vtkImageData(DataSet, vtkImageData):
    pass

@vtkPolyData.override
class vtkPolyData(PointSet, vtkPolyData):
    @property
    def polygons(self):
        ca = self.GetPolys()
        conn_vtk = ca.GetConnectivityArray()
        conn = dsa.vtkDataArrayToVTKArray(conn_vtk)
        offsets_vtk = ca.GetOffsetsArray()
        offsets = dsa.vtkDataArrayToVTKArray(offsets_vtk)
        return { 'connectivity' : conn, 'offsets' : offsets }

class CompositeDataIterator(object):
    """Wrapper for a vtkCompositeDataIterator class to satisfy
       the python iterator protocol. This iterator iterates
       over non-empty leaf nodes. To iterate over empty or
       non-leaf nodes, use the vtkCompositeDataIterator directly.
       """

    def __init__(self, cds):
        self.Iterator = cds.NewIterator()
        if self.Iterator:
            self.Iterator.UnRegister(None)
            self.Iterator.GoToFirstItem()

    def __iter__(self):
        return self

    def __next__(self):
        if not self.Iterator:
            raise StopIteration

        if self.Iterator.IsDoneWithTraversal():
            raise StopIteration
        retVal = self.Iterator.GetCurrentDataObject()
        self.Iterator.GoToNextItem()
        return retVal

    def next(self):
        return self.__next__()

    def __getattr__(self, name):
        """Returns attributes from the vtkCompositeDataIterator."""
        return getattr(self.Iterator, name)

class CompositeDataSetBase(object):
    """A wrapper for vtkCompositeData and subclasses that makes it easier
    to access Point/Cell/Field data as VTKCompositeDataArrays. It also
    provides a Python type iterator."""

    def __init__(self, **kwargs):
        self._PointData = None
        self._CellData = None
        self._FieldData = None
        self._Points = None
        super().__init__(**kwargs)

    def __iter__(self):
        "Creates an iterator for the contained datasets."
        return CompositeDataIterator(self)

    def get_attributes(self, type):
        """Returns the attributes specified by the type as a
        CompositeDataSetAttributes instance."""
        return CompositeDataSetAttributes(self, type)

    @property
    def point_data(self):
        "Returns the point data as a DataSetAttributes instance."
        if self._PointData is None or self._PointData() is None:
            pdata = self.get_attributes(vtkDataObject.POINT)
            self._PointData = weakref.ref(pdata)
        return self._PointData()

    @property
    def cell_data(self):
        "Returns the cell data as a DataSetAttributes instance."
        if self._CellData is None or self._CellData() is None:
            cdata = self.get_attributes(DataObject.CELL)
            self._CellData = weakref.ref(cdata)
        return self._CellData()

    @property
    def field_data(self):
        "Returns the field data as a DataSetAttributes instance."
        if self._FieldData is None or self._FieldData() is None:
            fdata = self.get_attributes(DataObject.FIELD)
            self._FieldData = weakref.ref(fdata)
        return self._FieldData()

    @property
    def points(self):
        "Returns the points as a VTKCompositeDataArray instance."
        if self._Points is None or self._Points() is None:
            pts = []
            for ds in self:
                try:
                    _pts = ds.Points
                except AttributeError:
                    _pts = None

                if _pts is None:
                    pts.append(NoneArray)
                else:
                    pts.append(_pts)
            if len(pts) == 0 or all([a is NoneArray for a in pts]):
                cpts = NoneArray
            else:
                cpts = dsa.VTKCompositeDataArray(pts, dataset=self)
            self._Points = weakref.ref(cpts)
        return self._Points()

@vtkPartitionedDataSet.override
class vtkPartitionedDataSet(CompositeDataSetBase, vtkPartitionedDataSet):
    def append(self, dataset):
        self.SetPartition(self.GetNumberOfPartitions(), dataset)