import core.modules.module_registry
from core.modules.vistrails_module import Module, ModuleError
import vtk, math

class SwitchCoord(Module):
    """SwitchCoord let you transform between different coord systems"""

    def compute(self):
        minp = self.getInputFromPort("min_density")
        omega = self.getInputFromPort("omega")
        dataset=self.getInputFromPort("dataset")

        output = self.create_instance_of_type('edu.utah.sci.vistrails.vtk', 'vtkStructuredGrid')
        output.vtkInstance = vtk.vtkStructuredGrid()
        mydata=output.vtkInstance
        mydata.DeepCopy(dataset.vtkInstance)

        self.op(mydata, minp, omega)
        self.setResult("changed_dataset", output)

    def op(self, mydata, minp, omega):
        extent=mydata.GetExtent()
        print extent

        pcoords = mydata.GetPoints().GetData()
        density = mydata.GetPointData().GetScalars("Density")
        momentum = mydata.GetPointData().GetVectors("Momentum")


        #manipulate data test

        maxnorm = 0.0
        for i in range(0, mydata.GetNumberOfPoints()):
            [x, y, z] = pcoords.GetTuple3(i)  
            [_v1, _v2, _v3] = momentum.GetTuple3(i)
            p = density.GetValue(i)

            r = math.sqrt(x*x + y*y)         
            phi = math.atan2(y, x)
            #if phi < 0: phi = 2*math.pi + phi
            #print "i, density, radius, phi", i, p, r, phi

            if p < minp: 
                vx=vy=vz=0
            else:
                vr = _v1 / p
                vphi = _v2 / (p) + r * omega
                vz = _v3 / p

                vx = vr * math.cos(phi) - vphi * math.sin(phi)
                vy = vr * math.sin(phi) + vphi * math.cos(phi)

                norm = math.sqrt(vx*vx + vy*vy + vz*vz)
                if norm > maxnorm:
                    maxnorm = norm
        #        vx = vx/norm
        #        vy = vy/norm
        #        vz = vz/norm

            momentum.SetTuple3(i, vx, vy, vz)

        print maxnorm
        for i in range(0, mydata.GetNumberOfPoints()):
            [vx, vy, vz] = momentum.GetTuple3(i)
            vx = vx/maxnorm
            vy = vy/maxnorm
            vz = vz/maxnorm
            momentum.SetTuple3(i, vx, vy, vz)


def initialize(*args, **keywords):
    reg=core.modules.module_registry.registry

    reg.add_module(SwitchCoord)

#     reg.add_input_port(PythonCalc, "value1",
#                        (core.modules.basic_modules.Float, 'the first argument'))
#     reg.add_input_port(PythonCalc, "value2",
#                        (core.modules.basic_modules.Float, 'the second argument'))
#     reg.add_input_port(PythonCalc, "op", 
#                        (core.modules.basic_modules.String, 'the operation'))
#     reg.add_output_port(PythonCalc, "value", 
#                         (core.modules.basic_modules.Float, 'the result'))

    reg.add_input_port(SwitchCoord, "scalar_range", 
                       [core.modules.basic_modules.Float, 
                        core.modules.basic_modules.Float]) 
    reg.add_input_port(SwitchCoord, "min_density", core.modules.basic_modules.Float)
    reg.add_input_port(SwitchCoord, "omega", core.modules.basic_modules.Float)
    reg.add_input_port(SwitchCoord, "dataset", 
                       (reg.get_descriptor_by_name(
                'edu.utah.sci.vistrails.vtk', 'vtkStructuredGrid').module) )
    reg.add_output_port(SwitchCoord, "changed_dataset", 
                        (reg.get_descriptor_by_name(
                'edu.utah.sci.vistrails.vtk', 'vtkStructuredGrid').module) )