# To add a primvar to selected polymesh objects in Maya execute theis 
# MEL command:
#    addFloatPrimVar("curve_mode", 1);
# The code for this proc can be found here:
#   http://fundza.com/rfm/customizing/html/mel/addFloatPrimVar.mel.html
# Malcolm Kesson
# August 6 2018
  
import prman
  
class Rif(prman.Rif):
    MESH = 0
    WIRE = 1
    WIRE_AND_MESH = 2
    DEFAULT_MODE = MESH
  
    def __init__(self, ri, width, color=[1.0,0.0,1.0]):    
        self.width = float(width)
        self.color = color
        prman.Rif.__init__(self, ri)
                   
    def PointsGeneralPolygons(self, nloops, nverts, verts, params):
        if not params.has_key('constant float curve_mode'):
            self.m_ri.PointsGeneralPolygons(nloops, nverts, verts, params)
            return
        mode = Rif.DEFAULT_MODE
        if params.has_key('constant float curve_mode'):
            mode = int(params['constant float curve_mode'][0])
            #print('PointsGeneralPolygons mode is %s' % mode)
        if mode == Rif.MESH:
            self.m_ri.PointsGeneralPolygons(nloops, nverts, verts, params)
            return
        xyz_poly = params[prman.Ri.P]
        xyz_curves = []
        loop_counter = 0
        vert_counter = 0
        npolys = len(nloops)
        for n in range(npolys):
            for j in range(nloops[n]):
                num_verts = nverts[loop_counter]
                loop_counter += 1
                for k in range(num_verts):
                    vert_index = verts[vert_counter] * 3
                    xyz_curves.append(xyz_poly[vert_index])
                    xyz_curves.append(xyz_poly[vert_index + 1])
                    xyz_curves.append(xyz_poly[vert_index + 2])
                    vert_counter += 1
                                
        parms = { prman.Ri.P : xyz_curves, 
                  prman.Ri.CONSTANTWIDTH : self.width,
                  'constant color cs' : self.color 
                }
        self.m_ri.Curves(prman.Ri.LINEAR, nverts, prman.Ri.PERIODIC, parms)
        if mode == Rif.WIRE_AND_MESH:
            self.m_ri.PointsGeneralPolygons(nloops, nverts, verts, params)