from core.modules.vistrails_module import Module, ModuleError
from Scientific.IO import NetCDF
import numpy
import vtk
import scipy
import scipy.interpolate

from SelfeBase import SelfeData

class Utility(object):
    namespace = 'Utilities'

class StationDataDepthInterpolator(Utility, Module):
    name = 'Interpolate Station Data With Depth'

    def interp(self, depth_ar, scalar_ar, depth_values):
        value_ar = []
        for t in range(depth_ar.shape[0]):
            interp_func = scipy.interpolate.interp1d(depth_ar[t,:], scalar_ar[t,:])
            time_ar = []
            for depth in depth_values:
                time_ar.append(interp_func(depth))

            value_ar.append(numpy.array(time_ar))

        return value_ar
            

    def compute(self):
        scalars = self.getInputFromPort("Scalars").get_array()
        depth_ar = self.getInputFromPort("Depths").get_array()
        time_ar = self.getInputFromPort("Times").get_array().squeeze()
        at_depth = self.getInputFromPort("At Depth")

        out_time_ar = []
        for t in at_depth:
            out_time_ar.append(time_ar)
        
        out_list = self.interp(depth_ar, scalars, at_depth)

        out_ar = self.create_instance_of_type('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array')
        out_ar.set_array(numpy.array(out_list).transpose())

        out_t_ar = self.create_instance_of_type('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array')
        out_t_ar.set_array(numpy.array(out_time_ar).transpose())

        self.setResult("Output", out_ar)
        self.setResult("Time Array", out_t_ar)
        
    @classmethod
    def register(cls, reg, basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)
        reg.add_input_port(cls, "Scalars", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Input Scalar Array'))
        reg.add_input_port(cls, "Depths", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Input Depth Array'))
        reg.add_input_port(cls, "Times", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Input Time Array'))
        reg.add_input_port(cls, "At Depth", (basic.List, 'At Depth'))
        reg.add_output_port(cls, "Output", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Output Array'))
        reg.add_output_port(cls, "Time Array", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Output Time Array'))

class StationDataInterpolator(Utility, Module):
    name = 'Interpolate Station Data'

    def get_param(self, depth_ar, depth):
        param = 0.
        depth_ar.shape[0]
        for d_index in range(depth_ar.shape[0]):
            d = depth_ar[d_index]
            if d <= depth:
                param += 1.
                continue

            d_low = depth_ar[d_index - 1]
            d_hi = depth_ar[d_index]
            d_range = d_hi - d_low
            d_tmp = depth - d_low

            p = abs(d_tmp / d_range)
            return d_low + p
            
    def interp(self, depth, depth_ar, scalar_ar):
        num_depths = depth_ar.shape[1]
        times = depth_ar.shape[0]
        scalar_ret = []
        for t in range(times):
            d = depth_ar[t, :].squeeze()
            param = self.get_param(d, depth)
            ind_lo = int(param)
            ind_hi = ind_lo + 1
            param = param - float(ind_lo)
            scalar_lo = scalar_ar[t, ind_lo]
            scalar_hi = scalar_ar[t, ind_hi]
            scalar = scalar_lo + (scalar_hi - scalar_lo) * param
            scalar_ret.append(scalar)
        return scalar_ret
        
    def compute(self):
        scalars = self.getInputFromPort("Scalars").get_array()
        depths = self.getInputFromPort("Depths").get_array()

        at_depth = self.getInputListFromPort("At Depth")
        num_total_depths = depths.shape[1]
        num_depths = len(at_depth)
        out_ar = []
        for depth in at_depth:
            scalars_ar = numpy.array(self.interp(depth, depths, scalars))
            out_ar.append(scalars_ar)

        out = self.create_instance_of_type('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array')
        out.set_array(numpy.array(out_ar))
        self.setResult("Output", out)

    @classmethod
    def register(cls, reg, basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)
        reg.add_input_port(cls, "Scalars", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Input Scalar Array'))
        reg.add_input_port(cls, "Depths", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Input Depth Array'))
        reg.add_input_port(cls, "At Depth", (basic.Float, 'At Depth'))
        reg.add_output_port(cls, "Output", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Output Array'))

class SelfeArrayToVTKScalars(Utility, Module):
    name = 'GridPak Array To VTK Scalars'

    def compute(self):
        s_ar = self.getInputFromPort("Scalars").get_array()
        ds = self.getInputFromPort("VTK Data")
        name = self.forceGetInputFromPort("Array Name")
        if name == None:
            name = "scalars"
            
        v_ar = vtk.vtkFloatArray()
        v_ar.SetNumberOfValues(s_ar.size)
        i = 0
        for z in range(s_ar.shape[0]):
            for x in range(s_ar.shape[2]):
                for y in range(s_ar.shape[1]):
                    v = s_ar[z,y,x]
                    v_ar.SetValue(i, v)
                    i += 1
                    
        v_ar.SetName(name)
        ds.vtkInstance.GetPointData().AddArray(v_ar)
        set = self.forceGetInputFromPort("Activate Scalars")
        if set == None:
            set = True
            
        if set:
            ds.vtkInstance.GetPointData().SetScalars(v_ar)

        self.setResult("Output Mesh", ds)

    @classmethod
    def register(cls,reg,basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)
        reg.add_input_port(cls, "VTK Data", (reg.get_module_by_name('edu.utah.sci.vistrails.vtk','vtkStructuredGrid'), 'Input Mesh'))
        reg.add_input_port(cls, "Activate Scalars", (basic.Boolean, 'Set active scalars'),True)
        reg.add_input_port(cls, "Array Name", (basic.String, 'Array name'))
        reg.add_input_port(cls, "Scalars", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Output Array'))
        reg.add_output_port(cls, "Output Mesh", (reg.get_module_by_name('edu.utah.sci.vistrails.vtk','vtkStructuredGrid'), 'Output Mesh'))

class SelfeArrayToVTKVectors(Utility, Module):
    name = 'GridPak Arrays to VTK Vectors'

    def compute(self):
        v0 = self.forceGetInputFromPort("Array0")
        v1 = self.forceGetInputFromPort("Array1")
        v2 = self.forceGetInputFromPort("Array2")
        ds = self.getInputFromPort("VTK Data")
        name = self.forceGetInputFromPort("Array Name")

        vsize = 0
        if v0:
            vsize += 1
            v0 = v0.get_array()
        if v1:
            vsize += 1
            v1 = v1.get_array()
        if v2:
            vsize += 1
            v2 = v2.get_array()
            
        if name == None:
            name = "vectors"
            
        v_ar = vtk.vtkFloatArray()
        v_ar.SetNumberOfComponents(3)
        i = 0
        for z in range(v0.shape[0]):
            for y in range(v0.shape[2]):
                for x in range(v0.shape[1]):
                    u = v0[z,x,y]
                    v = v1[z,x,y]
                    w = v2[z,x,y]
                    v_ar.InsertNextTuple3(u, v, w)
                    i += 1

        v_ar.SetName(name)
        ds.vtkInstance.GetPointData().AddArray(v_ar)
        ds.vtkInstance.GetPointData().SetVectors(v_ar)
        self.setResult("Output Mesh", ds)
        print ds.vtkInstance
        
    @classmethod
    def register(cls,reg,basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)
        reg.add_input_port(cls, "VTK Data", (reg.get_module_by_name('edu.utah.sci.vistrails.vtk','vtkStructuredGrid'), 'Input Mesh'))
        reg.add_input_port(cls, "Array Name", (basic.String, 'Array name'))
        reg.add_input_port(cls, "Array0", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Scalar Array'))
        reg.add_input_port(cls, "Array1", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Scalar Array'))
        reg.add_input_port(cls, "Array2", (reg.get_module_by_name('edu.utah.sci.vistrails.numpyscipy', 'Numpy Array', 'numpy|array'), 'Scalar Array'))
        reg.add_output_port(cls, "Output Mesh", (reg.get_module_by_name('edu.utah.sci.vistrails.vtk','vtkStructuredGrid'), 'Output Mesh'))