import numpy as np

import vtk
from vtkmodules.vtkCommonDataModel import vtkDataSet
from vtkmodules.util.vtkAlgorithm import VTKPythonAlgorithmBase
from vtkmodules.numpy_interface import dataset_adapter as dsa

# new module for ParaView-specific decorators.
from paraview.util.vtkAlgorithm import smproxy, smproperty, smdomain, smhint

# to add a source, instead of a filter, use the `smproxy.source` decorator.
@smproxy.reader(filename_patterns="fort.*", file_description="Turbulent Combustion DNS data", label="POSTECH Turbulent Combustion DNS Reader")
class PythonDnsDataReader(VTKPythonAlgorithmBase):
    """This is dummy VTKPythonAlgorithmBase subclass that
    simply puts out a Superquadric poly data using a vtkSuperquadricSource
    internally"""
    def __init__(self):
        VTKPythonAlgorithmBase.__init__(self,
                nInputPorts=0,
                nOutputPorts=1,
                outputType='vtkImageData')
        self._filename = None

    @smproperty.stringvector(name="FileName")
    @smdomain.filelist()
    # @smhint.filechooser(filename_patterns="fort.*", file_description="Numpy CSV files")
    def SetFileName(self, name):
        """Specify filename for the file to read."""
        if self._filename != name:
            self._filename = name
            self._ndata = None
            self._timesteps = None
            self.Modified()

    def RequestInformation(self, request, inInfoVec, outInfoVec):
        executive = self.GetExecutive()
        outInfo = outInfoVec.GetInformationObject(0)

        filename = self._filename
        t, nx, ny, nz, U, V, W, Y0, Y1 = self._read_data(filename, True)

        dx = 4*np.pi/nx
        dy = 2*np.pi/ny
        dz = 2*np.pi/nz

        outInfo.Set(executive.WHOLE_EXTENT(), 0, nx-1, 0, ny-1, 0, nz-1)
        outInfo.Set(vtk.vtkDataObject.SPACING(), dx, dy, dz)

        return 1

    def RequestData(self, request, inInfo, outInfo):
        output = dsa.WrapDataObject(self.GetOutputData(outInfo, 0))

        if self._filename is None:
            # Note, exceptions are totally fine!
            raise RuntimeError("No filename specified")

        filename = self._filename
        t, nx, ny, nz, U, V, W, Y0, Y1 = self._read_data(filename)

        u = U.ravel()
        v = V.ravel()
        w = W.ravel()
        y1 = Y1.ravel()

        dims = [nx, ny, nz]
        assert y1.shape[0] == dims[0]*dims[1]*dims[2], "dimension mismatch"

        output.SetExtent(0, dims[0]-1, 0, dims[1]-1, 0, dims[2]-1)
        output.PointData.append(y1, "yr")
        output.PointData.append(np.vstack((u,v,w)).T, "U")

        output.PointData.SetActiveScalars("yr")

        return 1

    # for anything too complex or not yet supported, you can explicitly
    # provide the XML for the method.
    @smproperty.xml("""
        <DoubleVectorProperty name="Center"
            number_of_elements="3"
            default_values="0 0 0"
            command="SetCenter">
            <DoubleRangeDomain name="range" />
            <Documentation>Set center of the superquadric</Documentation>
        </DoubleVectorProperty>""")
    def SetCenter(self, x, y, z):
        self.Modified()

    # In most cases, one can simply use available decorators.
    @smproperty.doublevector(name="Scale", default_values=[1, 1, 1])
    @smdomain.doublerange()
    def SetScale(self, x, y, z):
        self.Modified()

    @smproperty.intvector(name="ThetaResolution", default_values=16)
    def SetThetaResolution(self, x):
        self.Modified()

    @smproperty.intvector(name="PhiResolution", default_values=16)
    @smdomain.intrange(min=0, max=1000)
    def SetPhiResolution(self, x):
        self.Modified()

    @smproperty.doublevector(name="Thickness", default_values=0.3333)
    @smdomain.doublerange(min=1e-24, max=1.0)
    def SetThickness(self, x):
        self.Modified()

    def _read_data (self, fname, only_tags=False):
        import struct
        import sys
        import os
        with open(fname, 'rb') as f1 :
            f1.seek(0)

            raw_info = f1.read(4+8*6+4)[4:-4]
            t  = struct.unpack('d', raw_info[ 0: 8])[0]
            nx = struct.unpack('q', raw_info[ 8:16])[0]
            ny = struct.unpack('q', raw_info[16:24])[0]
            nz = struct.unpack('q', raw_info[24:32])[0]
            count = nx*ny*nz
            bSize = count*8 # size in bytes for a variable

            dummy_len = (4+8*3+4) + (4+8*2+4) + (4+8*2+4) + (4+8*2+4) + 4
            dummy = f1.read(dummy_len)
           #dummy = f1.read(4)

           #raw_field = f1.read(4+bSize*5+4)[4:-4]
            if only_tags:
                u = None
                v = None
                w = None
                Y0 = None
                Y1 = None
            else:
                V = np.fromfile(f1, dtype=np.double, count=(3*count)).reshape((3,nz,ny,nx))
                s = np.fromfile(f1, dtype=np.double, count=(2*count)).reshape((2,nz,ny,nx))

                u = V[0]
                v = V[1]
                w = V[2]

                Y0 = s[0]
                Y1 = s[1]

        return t, nx, ny, nz, u, v, w, Y0, Y1