# Demonstration script
# written by Jean M. Favre, Swiss National Supercomputing Centre
# tested Wed  8 Oct 08:13:10 CEST 2025

import vtk
import numpy as np
from vtk.numpy_interface import dataset_adapter as dsa
#from vtk.util.vtkAlgorithm import VTKPythonAlgorithmBase
from paraview.util.vtkAlgorithm import *
from vtk.numpy_interface import algorithms as algs

@smproxy.source(name='StructuredGrid_Source',
label='StructuredGrid_Source')
class StructuredGridSource(VTKPythonAlgorithmBase):
    def __init__(self):
        VTKPythonAlgorithmBase.__init__(self,
            nInputPorts=0, nOutputPorts=1, outputType='vtkStructuredGrid')

    def RequestInformation(self, request, inInfo, outInfo):
        dims = [13, 27, 5]
        info = outInfo.GetInformationObject(0)
        info.Set(vtk.vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT(),
            (0, dims[0]-1, 0, dims[1]-1, 0, dims[2]-1), 6)
        return 1

    def RequestData(self, request, inInfo, outInfo):
        output = dsa.WrapDataObject(vtk.vtkStructuredGrid.GetData(outInfo))
        info = outInfo.GetInformationObject(0)
        exts = info.Get(vtk.vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT())
        dims = [exts[1]-exts[0]+1, exts[3]-exts[2]+1, exts[5]-exts[4]+1]
        output.SetExtent(exts)
        Raxis = np.linspace(1., 2., dims[0])
        Thetaaxis = np.linspace(0.,np.pi*0.5, dims[1])
        Zaxis = np.linspace(0., 2.0, dims[2])
        Z, t, r = np.meshgrid(Zaxis, Thetaaxis, Raxis, indexing="ij")
        X = r * np.cos(t)
        Y = r * np.sin(t)
        coordinates = algs.make_vector(X.ravel(),Y.ravel(),Z.ravel())
        output.Points = coordinates
        output.PointData.append(r.ravel(), "radius")
        output.PointData.append(t.ravel(), "angle")
        output.PointData.append(Z.ravel(), "z-elevation")
        return 1