#!/usr/bin/python
#author: Bryan Bishop <kanzure@gmail.com>
#date: 2010-08-24
#license: gpl2+
import math
import Numeric
bryan_message = "bryan hasn't got this far yet"
NURBSError = "NURBSError"

#somewhat useful links
#url: http://www.rhino3d.com/nurbs.htm
#url: http://www.mactech.com/articles/develop/issue_25/schneider.html
#url: http://runten.tripod.com/NURBS/
#url: http://runten.tripod.com/NURBS/Nurbs-0.1.zip
#see: runten-pynurbs/Nurbs-0.1/Nurbs/Crv.py
#see: runten-pynurbs/Nurbs-0.1/Nurbs/Srf.py
#see: opennurbs/opennurbs/opennurbs_nurbscurve.cpp
#see: pyopengl/PyOpenGL-Demo-3.0.1b1/PyOpenGL-Demo/proesch/nurbs/nurbs.py
#see: pyopengl/PyOpenGL-Demo-3.0.1b1/PyOpenGL-Demo/proesch/nurbsCurve/nurbsCircle.py
#see: pyopengl/PyOpenGL-Demo-3.0.1b1/PyOpenGL-Demo/proesch/nurbsCurve/nurbsCurve.py
#see: csg-boole/boole-1.1/surface/perf_csg.c #perform_CSG
#see: breplibrary/cpp-simple/CSG/main.cpp
#
#######################################
#  *** format for control_points ***
#
#   control_points[dimension, len(u), len(v)] = magnitude or length
#
#   where dimension means "axis" such that the 0th axis is x
#       0 = x
#       1 = y
#       2 = z
#       3 = w or weight
#   
#   where u is the parametric u direction
#         v is the parametric v direction
#
#   at control_point["x", N, M] you have the x coordinate or value of the N-by-Mth control point
#######################################

def uniformknots(cntrlpts, degree):
    knots = Numeric.zeros(degree + cntrlpts + 1, Numeric.Float)
    knots[cntrlpts:] = 1.
    knots[degree+1:cntrlpts] = Numeric.arange(1., cntrlpts-degree)* (1./(cntrlpts - degree))
    return knots

def translate(txyz):
    ret = Numeric.identity(4).astype(Numeric.Float)
    ret[0:len(txyz), 3] = txyz
    return ret

def scale(sxyz):
    ret = Numeric.identity(4).astype(Numeric.Float)
    s = Numeric.ones(3, Numeric.Float)
    s[0:len(sxyz)] = sxyz
    ret[0,0] = s[0]
    ret[1,1] = s[1]
    ret[2,2] = s[2]
    return ret

def deg2rad(angle):
    return math.pi * angle / 180
4
def rad2deg(angle):
    return angle * 180/math.pi

def rotx(angle):
    ret = Numeric.identity(4).astype(Numeric.Float)
    ret[1,1] = ret[2,2] = Numeric.cos(angle)
    sn = Numeric.sin(angle)
    ret[1,2] = -sn
    ret[2,1] = sn
    return ret

def roty(angle):
    ret = Numeric.identity(4).astype(Numeric.Float)
    ret[0,0] = ret[2,2] = Numeric.cos(angle)
    sn = Numeric.sin(angle)
    ret[0,2] = sn
    ret[2,0] = -sn
    return ret

def rotz(angle):
    ret = Numeric.identity(4).astype(Numeric.Float)
    ret[0,0] = ret[1,1] = Numeric.cos(angle)
    sn = Numeric.sin(angle)
    ret[0,1] = -sn
    ret[1,0] = sn
    return ret

def rotxyz(angles):
    ret = Numeric.identity(4).astype(Numeric.Float)

    ret[1,1] = ret[2,2] = Numeric.cos(angles[0])
    sn = Numeric.sin(angles[0])
    ret[1,2] = -sn
    ret[2,1] = sn

    cs = Numeric.cos(angles[1])
    ret[0,0] = cs
    ret[2,2] += cs
    sn = Numeric.sin(angles[1])
    ret[0,2] = sn
    ret[2,0] = -sn

    cs = Numeric.cos(angles[2])
    ret[0,0] += cs
    ret[1,1] += cs
    sn = Numeric.sin(angles[2])
    ret[0,1] = -sn
    ret[1,0] = sn
    return ret

class ControlPoint:
    def __init__(self, x=0, y=0, z=0, weight=1):
        self.x, self.y, self.z, self.weight = x, y, z, weight
    def to_list(self, weight=False):
        if weight: return [self.x, self.y, self.z, self.weight]
        else: return [self.x, self.y, self.z]

def to_array(control_points):
    """converts from a list of control points to a Numeric array

    input: [A, B, C]
    output: [[A.x, B.x, C.x], [A.y, B.y, C.y], [A.z, B.z, C.z]]

    input [[A, B, C], [D, E, F]] such that at u=0,v=0 the point is A and at u=2,v=2 the point is F

    .. except as a Numeric array."""
    assert isinstance(control_points, list), "must be a list"
    return Numeric.transpose(control_points)        

def from_array(control_points):
    return Numeric.transpose(control_points).tolist()

class NurbsCurve:
    def __init__(self, control_points, knots, degree=3, display=False):
        """
        control_points is a list of (x,y,z,weight) points in the u parametric direction

        degree 1 (linear): lines and polylines
        degree 2 (quadratic): circles
        degree 3 (cubic): free-form curves
        degree 5 (quintic): free-form curves

        spline degree
            order = degree + 1
            degree = order - 1"""
        #FIXME: how do you calculate the degree of the curve?
        #one way to do it is: degree = len(knots) - len(control_points) + 1
        #another way to do it is: 
        #   nku = uknots.shape[0]
        #   nu = cntrl.shape[1]
        #   degree = nku - nu - 1
        if isinstance(control_points, list): control_points = to_array(control_points)

        #does the array specify enough places for xyz and a weight?
        shape = control_points.shape
        dimension = shape[0]
        if dimension < 2 or dimension > 4: raise NURBSError, "control points must be specified with at least (x,y) coordinates"
        if dimension < 4: #fill in the empty
            temp_points = Numeric.zeros((4, control_points.shape[1]))
            temp_points[0:dimension,:] = control_points
            temp_points[-1,:] = Numeric.ones((control_points.shape[1],))
            control_points = temp_points

        ctrlptslst = control_points.tolist()

        #are there enough control points?
        if len(ctrlptslst) < degree: raise ValueError, "there aren't enough control points for this degree"
    
        #knots must meet minimum sequence length requirements
        if len(knots) < (degree + len(ctrlptslst) - 1): raise ValueError, "knot sequence doesn't have enough values"
        
        #knot sequence must be monotically increasing
        if not knots == sorted(knots): raise ValueError, "knot sequence must be monotonically increasing"

        #limit the number of duplicate values to no more than the degree
        dupe_count = max(knots.count(item) for item in set(item2 for item2 in knots))
        if dupe_count > degree: raise ValueError, "knot sequence has too many duplicate values"

        #a common misconception is that each knot is paired with a control point
        #this is true only for degree 1 NURBS (polylines)
        if degree == 1 and not len(knots) == len(ctrlptslst): raise ValueError, "polylines must have the same number of knots as control points"
          
        self.degree = degree #FIXME: verify the degree of the curve somehow?
        self.control_points = control_points
        self.knots = knots
        self.display = display

    def trans(self, matrix):
        """apply the 4D transform matrix to the NURB control points"""
        self.control_points = Numeric.dot(matrix, self.control_points)

class Line(NurbsCurve):
    """straight line segment"""
    def __init__(self, p1 = [0, 0, 0], p2 = [1, 0, 0]):
        NurbsCurve.__init__(self, to_array([p1,p2]), [0,0,1,1])

class PolyLine(NurbsCurve):
    """c = Polyline([[0,0],[5,2],[10,8]])"""
    def __init__(self, points):
        points = to_array(points) #Numeric.transpose(Numeric.asarray(points))
        
        num_points = points.shape[1]
        if num_points < 3: raise NURBSError, "there must be at least three points in a polyline" #or just load it off to Line

        control_points = Numeric.zeros((points.shape[0], 2 * num_points - 2))
        control_points[:,0] = points[:,0]
        control_points[:,-1] = points[:,-1]
        control_points[:,1:-2:2] = points[:,1:-1]
        control_points[:,2:-1:2] = points[:,1:-1]
        
        uknots = Numeric.zeros(num_points * 2)
        uknots[0::2] = Numeric.arange(num_points)
        uknots[1::2] = Numeric.arange(num_points)

        NurbsCurve.__init__(self, control_points, uknots)

class UnitCircle(NurbsCurve):
    """NURBS representation of a unit circle in the xy plane"""
    def __init__(self):
        r22 = Numeric.sqrt(2)/2
        uknots = [0, 0, 0, 0.25, 0.25, 0.5, 0.5, 0.75, 0.75, 1., 1, 1]
        control_points = [[0,  r22,  1, r22, 0, -r22, -1, -r22, 0],
                          [-1, -r22, 0, r22, 1, r22,  0,  -r22, -1],
                          [0,  0,    0, 0,   0, 0,    0,  0,    0],
                          [1,  r22,  1, r22, 1, r22,  1,  r22,  1]]
        control_points = [
                          [0, -1, 0, 1],
                          [r22, -r22, 0, r22],
                          [1, 0, 0, 1],
                          [r22, r22, 0, r22],
                          [0, 1, 0, 1],
                          [-r22, r22, 0, r22],
                          [-1, 0, 0, 1],
                          [-r22, -r22, 0, r22],
                          [0, -1, 0, 1]
                         ]
        NurbsCurve.__init__(self, control_points, uknots)
        self.uknots = uknots

class Circle(UnitCircle):
    """NURBS representation of a circle in the xy plane
    with a given radius (default = 1) and optional center"""
    def __init__(self, radius=1, center=None):
        UnitCircle.__init__(self)
        if radius != 1: self.trans(scale([radius, radius]))
        if center: self.trans(translate(center))

class Arc(NurbsCurve):
    """NURBS representation of an arc in the xy plane"""
    def __init__(self, radius=1, center=None, start_angle=0, end_angle=2*math.pi):
        sweep = end_angle - start_angle #sweep angle of arc

        if sweep < 0: 
            sweep = 2*math.pi + sweep
        if abs(sweep) <= math.pi/2.:
            arc_segments = 1 #number of arc segments
            knots = [0, 0, 0, 1, 1, 1] 
        elif abs(sweep) <= math.pi:
            arc_segments = 2 
            knots = [0, 0, 0, 0.5, 0.5, 1, 1, 1] 
        elif abs(sweep) <= 3*math.pi/2:
            arc_segments = 3 
            knots = [0, 0, 0, 1/3, 1/3, 2/3, 2/3, 1, 1, 1] 
        else:
            arc_segments = 4 
            knots = [0, 0, 0, 0.25, 0.25, 0.5, 0.5, 0.75, 0.75, 1, 1, 1] 
        
        dsweep = sweep / (2 * arc_segments) #arc segment sweep angle/2
        
        #determine middle control point and weight
        wm = math.cos(dsweep)
        x = radius * math.cos(dsweep)
        y = radius * math.sin(dsweep)
        xm = x + y * math.tan(dsweep)

 
 #arc segment control points (what is the format?)
        ctrlpt = Numeric.array([[x, wm*xm, x], [-y, 0, y], [0, 0, 0], [1, wm, 1]], Numeric.Float)

        #build up complete arc from rotated segments
        coefs = Numeric.zeros((4, 2 * arc_segments + 1), Numeric.Float) #nurb control points of arc

        #rotate to start angle
        coefs[:,0:3] = Numeric.dot(rotz(sang + dsweep), ctrlpt)

        xx = rotz(2*dsweep)
        for ms in range(2, 2*narcs,2):
            coefs[:,ms:ms+3] = Numeric.dot(xx, coefs[:,ms-2:ms+1])
        if center:
            xx = translate(center)
            coefs = Numeric.dot(xx, coefs)
        
        NurbsCurve.__init__(self, coefs, knots)

class NurbsSurface:
    def __init__(self, control_points, knots, vknots, display=False):
        """
        control_points is an array with the following shape:
            haha

        uknots is a list of knots along the parametric u direction
        vknots is a list of knots along the parametric v direction
        knots should be a list [uknots, vknots]

        not implemented:
            p is the degree in the u direction
            q is the degree in the v direction
        """

        shape = Numeric.asarray(control_points).shape
        #if not shape[2]  == 3: raise ValueError, "control points must be specified with 3D coordinates"

        #FIXME 2011-06-22 wow this is fucked up
        #just get it over with and add uknots/vknots to the params
        if type(knots[0]) == list:
            uknots = knots[0]
            vknots = knots[1]
        elif vknots:
            uknots = knots
        else:
            uknots, vknots = knots, knots

        #knot sequence must be monotically increasing
        if not uknots == sorted(uknots): raise ValueError, "u-direction knot sequence must be monotonically increasing"
        if not vknots == sorted(vknots): raise ValueError, "v-direction knot sequence must be monotonically increasing"

        #FIXME: figure out how to calculate the degree?
        #limit the number of duplicate values to no more than the degree
        #dupe_count = max(uknots.count(item) for item in set(item2 for item2 in uknots))
        #if dupe_count > degree: raise ValueError, "u-direction knot sequence has too many duplicate values"
        #dupe_count = max(vknots.count(item) for item in set(item2 for item2 in vknots))
        #if dupe_count > degree: raise ValueError, "v-direction knot sequence has too many duplicate values"
        
        self.control_points = control_points
        self.cntrl = control_points
        self.knots = knots
        self.uknots, self.vknots = knots, vknots
        self.display = display
    
    def trans(self, matrix):
        """apply the 4D transformation matrix to the NURB control points"""
        for v in range(self.control_points.shape[2]):
            self.control_points[:,:,v] = Numeric.dot(matrix, self.control_points[:,:,v])

def _extrude(curve, vector):
    """extrudes a nurbs curve along an extrusion vector
    curve = NurbsCurve()
    vector = [x, y, z]"""
    if not isinstance(curve, NurbsCurve): raise ValueError, "curve must be a NurbsCurve"

    #copy all of the control points in the u direction into the new surface
    coefs = Numeric.zeros((4,curve.control_points.shape[1],2), Numeric.Float)
    #if you want weights for each point, set the weights to 1

    #fill the v parametric direction with control points
    coefs[:,:,0] = curve.control_points
    #calculate these points by the dot product between the given vector and the entire curve
    #i.e.: Numeric.dot(translate(vector), curve.control_points)
    coefs[:,:,1] = Numeric.dot(translate(vector), curve.control_points)

    surface = NurbsSurface(coefs, curve.uknots, [0, 0, 1, 1])
    return surface
NurbsCurve.extrude = _extrude #retroactive definitions ftw

def _revolve(curve, point = [0, 0, 0], vector = [1, 0, 0], theta = 2*math.pi):
    """construct a surface by revolving the profile curve
    around an axis defined by a point and a unit vector.

    curve = NurbsCurve()
    point = coordinate of point to revolve around
    vector = rotation axis
    theta = angle to revolve curve"""
    if not isinstance(curve, NurbsCurve): raise NURBSError, "curve must be a NurbsCurve"

    #translate and rotate the curve into alignment with the z-axis
    T = translate(-Numeric.asarray(point, Numeric.Float))
    
    #normalize vector
    vector = Numeric.asarray(vector, Numeric.Float)
    length = Numeric.sqrt(Numeric.add.reduce(vector*vector))
    if length == 0: raise ZeroDivisionError, "can't normalize a zero-length vector"
    vector = vector / length
    
    if vector[0] == 0: angle_x = 0
    else: angle_x = math.atan2(vector[0], vector[2])
    RY = roty(-angle_x)
    temp_vector = Numeric.ones((4,), Numeric.Float)
    temp_vector[0:3] = vector
    temp_vector = Numeric.dot(RY, temp_vector)

    if temp_vector[1] == 0: angle_y = 0
    else: angle_y = math.atan(vector[1], vector[2])
    RX = rotx(angle_y)
    curve.trans(Numeric.dot(RX, Numeric.dot(RY, T)))
    
    arc = Arc(1, [0, 0, 0], 0, theta)
    
    narc = arc.control_points.shape[1]
    ncrv = curve.control_points.shape[1]
    coefs = Numeric.zeros((4, narc, ncrv), Numeric.Float)
    angle = Numeric.arctan2(curve.control_points[1,:], curve.control_points[0,:])
    temp_vec = curve.control_points[0:2,:]
    radius = Numeric.sqrt(Numeric.add.reduce(temp_vec * temp_vec))

    for i in xrange(0, ncrv):
        coefs[:,:,i] = Numeric.dot(rotz(angle[i]),
                                   Numeric.dot(translate((0, 0, curve.control_points[2,i])),
                                               Numeric.dot(scale((radius[i], radius[i])), arc.control_points)
                                              )
                                  )
        coefs[3,:,i] = coefs[3,:,i] * curve.control_points[3,i]

    #FIXME 2011-06-22 parameters to NurbsSurface are wrong (vknots)
    surface = NurbsSurface(coefs, [arc.uknots, curve.uknots], display=curve.display)

    T = translate(point)
    RX = rotx(-angle_y)
    RY = roty(angle_x)
    surface.trans(Numeric.dot(T, Numeric.dot(RY, RX)))

    return surface
NurbsCurve.revolve = _revolve