# A ray tracer for scenes consisting of multiple reflective ellipsoids,
# triangles, and parallelograms, lit by a single white point light source and
# some amount of white ambient light.
#
# The overall program design is as follows:
#   - Scenes provide information about the position of the light source, the
#     intensity of ambient light, and the geometric objects in the scene.
#     The light is always white, at the maximum possible intensity, although
#     ambient light can be any intensity.
#   - Geometric objects (ellipsoids, triangles, and parallelograms) are
#     instances of canonical objects. They know their color, their
#     transformation from global to canonical coordinates, and how to detect
#     intersections between themselves and rays. The result of these
#     intersection calculations is three pieces of information: the global
#     points at which intersections happen, the global normals to the object at
#     those points, and the t values at which the intersections happen. Colors
#     are represented by 5 numbers, namely the coefficients of diffuse
#     reflection, a coefficient of specular reflection for all colors, and a
#     Phong shininess exponent.
#   - Pixel colors are represented in RGB, by NumPy column vectors.
#   - Points and vectors are also column vectors, in homogeneous form.
#   - Lists of points, vectors, or colors are NumPy matrices, which each point
#     etc. in a column. See comments in "RayTracingUtilities.py" for more on
#     this list representation.
#   - The main program sets up the viewer and scene, calls a "traceForColor"
#     function to build an image from the scene, and displays the image.
#   - The traceForColor function generates a list of colors from a scene and
#     a list of rays into it. Rays are represented by a list of origin points
#     and a list of direction vectors. This function calls itself recursively
#     to process reflections, and so also takes an argument that tells it how
#     many levels of reflection to process. This function calls a more general
#     "rayTrace" function that returns information about the first intersection
#     on each of multiple rays; "traceForColor" uses this information to trace
#     shadow rays (again by calling "rayTrace") and do Phong shading for each
#     pixel in the image.
#   - The "rayTrace" function takes a list of geometric objects and a set of rays
#     as arguments, and returns 4 lists describing the first intersection each
#     ray has with an object in the scene: the index of the object intersected,
#     or -1 if there is no intersection, the global point at which the
#     intersection happens, the global normal to the object at that point, and
#     the t value at which the intersection happens. The rays given to
#     "rayTrace" are equal-length lists of origin points and direction vectors.
#   - Because the shading model can now produce color intensities outside the
#     0 to 1 range allowed in images, a "scaleColor" function adjusts the color
#     values in the final image to the required range.

# History:
#
#   November 2021 -- Modified by Doug Baldwin to add reflection to a ray tracer
#     that handled only Phong shading and shadows.
#
#   October 2021 -- Ray tracer with shading and shadows created by Doug Baldwin
#     as a demonstration for Math 384.


from RayTracingUtilities import listify, delistify, makeHomogeneous, normalize, \
    listLength, listDot

import numpy as np
import matplotlib.pyplot as plot




# A class that represents scenes. This class serves only to store and provide
# access to the information about the scene, with a constructor as an easy way
# to initialize scenes. Attributes that hold the scene information are...
#   - The position of the light source, in attribute "lightPos"
#   - The intensity of ambient light, in attribute "ambient"
#   - The list of geometric objects in the scene, in attribute "elements"

class Scene :


    # The constructor that initializes a scene from its light position (a NumPy
    # column vector in homogeneous form), ambient light intensity, and
    # elements.

    def __init__( self, light, ambient, elements ) :
        self.lightPos = light
        self.ambient = ambient
        self.elements = elements




# Classes that represent geometric objects. All geometric objects have the
# following attributes and methods:
#   - Color information, consisting of...
#       o The coefficients of diffuse reflection for red, green, and blue
#         light, in attributes "red," "green," and "blue" respectively
#       o The coefficient of specular reflection for all colors of light, in
#         attribute "specular"
#       o A shininess exponent for Phong's lighting model, in attribute "shine"
#   - An "intersect" method that takes lists of ray origin points and ray
#     directions, in the global coordinate system, as its arguments, and
#     determines where those rays intersect the object. This method returns
#     three values, namely...
#       o A NumPy row vector of ray "t" parameters at which intersections
#         happen. These values are infinity for any rays that don't intersect
#         the object.
#       o A list of global points where the intersections happen. These are
#         undefined wherever the corresponding t value is infinite.
#       o A list of global normal vectors to the intersections. Undefined where
#         t values are infinite.


# A class that represents triangles.

class Triangle :


    # In addition to the attributes that all geometric objects have, triangles
    # also store a unit-length normal column vector, in homogeneous global
    # coordinates, in attribute "normal".


    # Initialize a triangle with its vertices (NumPy column vectors, in
    # homogeneous form) and color information. Vertices should be given in
    # some counterclockwise order as seen from the side of the triangle that
    # the normal points to.

    def __init__( self, ptA, ptB, ptC, red, green, blue, specular, shine ) :

        # Store the color information.
        self.red = red
        self.green = green
        self.blue = blue
        self.specular = specular
        self.shine = shine

        # Initialize the transformation matrix. Build the canonical-to-global
        # transformation matrix described in Baldwin & Weber, then have NumPy
        # invert it.
        E1 = np.ravel( ptB - ptA )
        E2 = np.ravel( ptC - ptA )
        N = np.cross( E1[0:3], E2[0:3] )
        free = np.array( [ 1.0, 0.0, 0.0, 0.0 ] ) if abs( N[0] ) > abs( N[1] ) and abs( N[0] ) > abs( N[2] ) \
               else np.array( [ 0.0, 1.0, 0.0, 0.0 ] ) if abs( N[1] ) > abs( N[0] ) and abs( N[1] ) > abs( N[2] ) \
               else np.array( [ 0.0, 0.0, 1.0, 0.0 ] )

        self.toCanonical = np.linalg.inv( np.transpose( np.array( [ E1, E2, free, np.ravel(ptA) ] ) ) )

        # Store the normal vector, based on the one computed as part of finding
        # the transformation matrix. Note that what I found earlier is a row
        # vector in Cartesian form, not a column vector in homogeneous form,
        # nor necessarily unit length.
        self.normal = normalize( np.array( [ [N[0]], [N[1]], [N[2]], [0.0] ] ) )


    # Calculate the intersections between a set of rays and this triangle. See
    # above for a description of the arguments to and results from this method.

    def intersect( self, sources, directions ) :

        nRays = listLength( sources )

        # Transform rays to Cartesian canonical coordinates.
        canonicalSources = ( self.toCanonical @ sources )[0:3,:]
        canonicalDirections = ( self.toCanonical @ directions )[0:3,:]

        # Accumulate a mask for those rays that actually intersect the triangle.
        # Simultaneously accumulate a vector of t values at which those
        # intersections happen. Initially I consider all rays to be potential
        # intersectors, but the t values to be infinite.
        hitMask = np.full( nRays, True )
        ts = np.full( nRays, np.inf )

        # Any rays parallel to the xy plane cannot be intersectors.
        hitMask[ canonicalDirections[2,:] == 0.0 ] = False

        # Any rays that intersect the xy plane behind their starting points
        # (i.e., at negative t values) aren't really intersectors either.
        ts[ hitMask ] = -canonicalSources[ 2, hitMask ] / canonicalDirections[ 2, hitMask ]
        hitMask[ ts < 0.0 ] = False

        # Finally, intersections with the xy plane but outside the canonical
        # triangle don't intersect that triangle either. An intersection is
        # outside the canonical triangle if its x coordinate is negative, its
        # y coordinate is negative, or the sum of its x and y coordinates is
        # greater than 1.
        canonicalHits = np.ones( [ 3, nRays ] )
        canonicalHits[:,hitMask] = canonicalSources[:,hitMask] + ts[hitMask] * canonicalDirections[:,hitMask]
        hitMask[ canonicalHits[0,:] < 0.0 ] = False
        hitMask[ canonicalHits[1,:] < 0.0 ] = False
        hitMask[ canonicalHits[0,:] + canonicalHits[1,:] > 1.0 ] = False

        # The t values I actually want to return are the ones where a ray
        # intersects the canonical triangle; all other ts should be infinite.
        finalTs = np.full( nRays, np.inf )
        finalTs[ hitMask ] = ts[ hitMask ]

        # Build a list of world-coordinate points at which the rays intersect
        # this triangle.
        finalPoints = np.ones( [ 4, nRays ] )
        finalPoints[ :, hitMask ] = sources[ :, hitMask ] + finalTs[hitMask] * directions[ :, hitMask ]

        # The world-coordinate normals to the intersection points are all just
        # this triangle's normal:
        finalNormals = np.full( [ 4, nRays ], self.normal )

        # All done. Return t values, intersection points, and normals.
        return finalTs, finalPoints, finalNormals


# A class that represents paralleograms. This class is a small adaptation of
# triangle class, differing only in how it determines whether a ray intersects
# the canonical parallelogram (which is a unit square with one corner on the
# origin in the canonical coordinate frame).

class Parallelogram :


    # In addition to the attributes that all geometric objects have, parallelograms
    # also store a unit-length normal column vector, in homogeneous global
    # coordinates, in attribute "normal".


    # Initialize a parallelogram with its vertices (NumPy column vectors, in
    # homogeneous form) and color information. Vertices should be given in
    # some counterclockwise order as seen from the side of the parallelogram that
    # the normal points to.

    def __init__( self, ptA, ptB, ptC, red, green, blue, specular, shine ) :

        # Store the color information.
        self.red = red
        self.green = green
        self.blue = blue
        self.specular = specular
        self.shine = shine

        # Initialize the transformation matrix. Build the canonical-to-global
        # transformation matrix described in Baldwin & Weber, then have NumPy
        # invert it.
        E1 = np.ravel( ptB - ptA )
        E2 = np.ravel( ptC - ptA )
        N = np.cross( E1[0:3], E2[0:3] )
        free = np.array( [ 1.0, 0.0, 0.0, 0.0 ] ) if abs( N[0] ) > abs( N[1] ) and abs( N[0] ) > abs( N[2] ) \
               else np.array( [ 0.0, 1.0, 0.0, 0.0 ] ) if abs( N[1] ) > abs( N[0] ) and abs( N[1] ) > abs( N[2] ) \
               else np.array( [ 0.0, 0.0, 1.0, 0.0 ] )

        self.toCanonical = np.linalg.inv( np.transpose( np.array( [ E1, E2, free, np.ravel(ptA) ] ) ) )

        # Store the normal vector, based on the one computed as part of finding
        # the transformation matrix. Note that what I found earlier is a row
        # vector in Cartesian form, not a column vector in homogeneous form,
        # nor necessarily unit length.
        self.normal = normalize( np.array( [ [N[0]], [N[1]], [N[2]], [0.0] ] ) )


    # Calculate the intersections between a set of rays and this parallelogram. See
    # above for a description of the arguments to and results from this method.

    def intersect( self, sources, directions ) :

        nRays = listLength( sources )

        # Transform rays to Cartesian canonical coordinates.
        canonicalSources = ( self.toCanonical @ sources )[0:3,:]
        canonicalDirections = ( self.toCanonical @ directions )[0:3,:]

        # Accumulate a mask for those rays that actually intersect the triangle.
        # Simultaneously accumulate a vector of t values at which those
        # intersections happen. Initially I consider all rays to be potential
        # intersectors, but the t values to be infinite.
        hitMask = np.full( nRays, True )
        ts = np.full( nRays, np.inf )

        # Any rays parallel to the xy plane cannot be intersectors.
        hitMask[ canonicalDirections[2,:] == 0.0 ] = False

        # Any rays that intersect the xy plane behind their starting points
        # (i.e., at negative t values) aren't really intersectors either.
        ts[ hitMask ] = -canonicalSources[ 2, hitMask ] / canonicalDirections[ 2, hitMask ]
        hitMask[ ts < 0.0 ] = False

        # Finally, intersections with the xy plane but outside the canonical
        # parallelogram don't intersect that parallelogram either. An intersection is
        # outside the canonical paralleogram if its x or y coordinates are
        # negative or greater than 1.
        canonicalHits = np.ones( [ 3, nRays ] )
        canonicalHits[:,hitMask] = canonicalSources[:,hitMask] + ts[hitMask] * canonicalDirections[:,hitMask]
        hitMask[ canonicalHits[0,:] < 0.0 ] = False
        hitMask[ canonicalHits[0,:] > 1.0 ] = False
        hitMask[ canonicalHits[1,:] < 0.0 ] = False
        hitMask[ canonicalHits[1,:] > 1.0 ] = False

        # The t values I actually want to return are the ones where a ray
        # intersects the canonical triangle; all other ts should be infinite.
        finalTs = np.full( nRays, np.inf )
        finalTs[ hitMask ] = ts[ hitMask ]

        # Build a list of world-coordinate points at which the rays intersect
        # this triangle.
        finalPoints = np.ones( [ 4, nRays ] )
        finalPoints[ :, hitMask ] = sources[ :, hitMask ] + finalTs[hitMask] * directions[ :, hitMask ]

        # The world-coordinate normals to the intersection points are all just
        # this triangle's normal:
        finalNormals = np.full( [ 4, nRays ], self.normal )

        # All done. Return t values, intersection points, and normals.
        return finalTs, finalPoints, finalNormals


# A class that represents ellipsoids.

class Ellipsoid :


    # Initialize an ellipsoid with the position of its center, its "radius" in
    # each dimension, and its color. Position is a 4-component NumPy column
    # vector giving the center point for the ellipsoid in homogeneous form; the
    # other arguments are real-valued scalars.

    def __init__( self, center, xRadius, yRadius, zRadius, red, green, blue, specular, shine ) :

        # Save the color information.
        self.red = red
        self.green = green
        self.blue = blue
        self.specular = specular
        self.shine = shine

        # The world-to-canonical coordinates transformation is the inverse of a
        # transformation that scales by the appropriate radii and then
        # translates to the new position -- i.e., a transformation that
        # first translates back to the canonical origin and then scales by the
        # reciprocals of the radii.
        self.toCanonical = np.array( [ [ 1.0/xRadius,  0.0,     0.0,       -center[0,0]/xRadius ],
                                       [   0.0,   1.0/yRadius,  0.0,       -center[1,0]/yRadius ],
                                       [   0.0,        0.0,  1.0/zRadius,  -center[2,0]/zRadius ],
                                       [   0.0,        0.0,     0.0,          1.0               ] ] )


    # Calculate intersections between this ellipsoid and a set of rays. See
    # above for a detailed description of the arguments to and results from
    # this method.

    def intersect( self, sources, directions ) :

        # Start by transforming all the rays to the ellipsoid's canonical
        # coordinate system. But I don't need the transformed values in
        # homogeneous form.
        canonicalSources = (self.toCanonical @ sources)[ 0:3, : ]
        canonicalDirections = (self.toCanonical @ directions)[ 0:3, : ]

        # A ray from origin <Rx,Ry,Rz> in direction (dx,dy,dz) intersects the
        # canonical sphere where
        # (dx^2 + dy^2 + dz^2) t^2 + 2(Rxdx + Rydy + Rzdz) t + (Rx^2 + Ry^2 + Rz^2 - 1) = 0.
        # Calculate the coefficients of this equation for each direction
        # vector.
        aValues = np.sum( canonicalDirections ** 2, 0 )
        bValues = 2.0 * np.sum( canonicalSources * canonicalDirections, 0 )
        cValues = np.sum( canonicalSources ** 2, 0 ) - 1.0

        # Use the quadratic formula to solve for t values where rays intersect
        # the canonical sphere. First I need to know where the rays intersect
        # the sphere at all, i.e., where the equation has real solutions.
        discriminants = bValues ** 2 - 4.0 * aValues * cValues
        hitMask = discriminants >= 0.0

        # Now, the result is infinity where there's no intersection or both
        # solutions to the quadratic are negative, the lower solution to the
        # quadratic if that solution is positive and real, and the higher
        # solution if the lower is negative but the higher is positive.
        nRays = listLength( directions )

        lowTs = np.full( nRays, np.inf )
        lowTs[hitMask] = ( -bValues[hitMask] - np.sqrt(discriminants[hitMask]) ) / ( 2.0 * aValues[hitMask] )
        lowMask = hitMask & ( lowTs > 0.0 )

        highTs = np.full( nRays, np.inf )
        highTs[hitMask] = ( -bValues[hitMask] + np.sqrt( discriminants[hitMask] ) ) / ( 2.0 * aValues[hitMask] )
        highMask = hitMask & ( lowTs <= 0.0 ) & ( highTs > 0.0 )

        finalTs = np.full( nRays, np.inf )
        finalTs[ lowMask ] = lowTs[ lowMask ]
        finalTs[ highMask ] = highTs[ highMask ]

        hitMask = lowMask | highMask                    # Visible intersections are where at least 1 t is positive

        # Calculate global intersection points for the intersections.
        finalPoints = np.ones( [ 4, nRays ] )
        finalPoints[ :, hitMask ] = sources[ :, hitMask ] + finalTs[ hitMask ] * directions[ :, hitMask ]

        # Calculate global normals at the intersections. Start with the local
        # intersection points in Cartesian form, which, thanks to the geometry
        # of the canonical sphere, are also local normals. Multiply those local
        # normals by the transpose of the global-to-canonical matrix, i.e., the
        # transpose of the inverse of the canonical-to-global matrix, in order
        # to put them into global coordinates, and finally normalize them.
        canonicalNormals = ( canonicalSources[ :, hitMask ] + finalTs[ hitMask ] * canonicalDirections[ :, hitMask ] )[ 0:3, : ]
        toGlobal = np.transpose( self.toCanonical[ 0:3, 0:3 ] )
        finalNormals = np.zeros( [ 4, nRays ] )
        finalNormals[ :, hitMask ] = makeHomogeneous( normalize( toGlobal @ canonicalNormals ), 0 )

        # All done. Return the t values, intersection points, and normals.
        return finalTs, finalPoints, finalNormals





# The function that traces arbitrary rays through a set of objects, returning
# geometric information (object intersected, intersection point, normal vector,
# and t parameter) about intersections rather than colors.

def rayTrace( objects, sources, directions ) :


    # This function accumulates object indices, intersection points, normals,
    # and corresponding t values in "lists" that it updates with intersection
    # information about each object in turn. Initialize those lists to "no
    # object" values: indices to -1, t values to infinity, and points and
    # normals to values that are homogeneous points and vectors respectively,
    # but otherwise arbitrary.

    nRays = listLength( sources )

    ts = np.full( nRays, np.inf )
    indices = np.full( nRays, -1 )
    hitPoints = np.ones( [ 4, nRays ] )
    hitNormals = np.zeros( [ 4, nRays ] )


    # Loop through the geometric objects, intersecting the rays with each and
    # updating t values, indices, intersection points, and normals wherever an
    # intersection happens at a smaller t value than previously seen.

    index = 0                               # Index of current object
    for obj in objects :

        newTs, newPoints, newNormals = obj.intersect( sources, directions )

        mask = newTs < ts
        ts[ mask ] = newTs[ mask ]
        indices[ mask ] = index
        hitPoints[ :, mask ] = newPoints[ :, mask ]
        hitNormals[ :, mask ] = newNormals[ :, mask ]

        index += 1


    # The results have now accumulated in the lists of indices, points, and
    # normals

    return indices, hitPoints, hitNormals, ts




# The function that traces a set of rays through a scene, returning a list of
# pixel colors for an image of that scene.

def traceForColor( scene, sources, directions, reflections ) :


    # Pull object colors out of the scene and into a NumPy matrix. Each object
    # has a column, with red, green, and blue coefficients of diffuse
    # reflection in the 1st, 2nd, and 3rd rows respectively. Similarly pull
    # coefficients of specular reflection and shininess exponents out of the
    # scene's objects.

    diffuseColors = np.array( [ [ e.red for e in scene.elements ],
                                [ e.green for e in scene.elements ],
                                [ e.blue for e in scene.elements ] ] )

    specularCoefficients = np.array( [ e.specular for e in scene.elements ] )
    shineExponents = np.array( [ e.shine for e in scene.elements ] )


    # Assume light intensities are always 1 for all colors, but put that
    # assumption into a variable to make it easy to change if I ever want to.

    brightness = np.array( [[1.0], [1.0], [1.0]] )


    # Figure out how many pixels there will be in the image.

    nPixels = listLength( sources )


    # Trace the primary rays, i.e., the ones from the origin through each
    # pixel. Make masks for subsequent calculations based on where these rays
    # intersect or don't intersect anything.

    objects, hits, normals, hitTs = rayTrace( scene.elements, sources, directions )
    hitMask = objects >= 0
    missMask = objects < 0


    # Now that I know where rays really hit objects, replace all the -1 object
    # IDs where rays missed everything with dummy values of 0 so that all
    # object IDs will be valid object indices from here on.

    objects[ missMask ] = 0


    # Build object colors piece by piece, starting with everything being black.

    colors = np.zeros( [ 3, nPixels ] )


    # Add in the contribution of ambient light.

    ambients = np.full( [3, nPixels], scene.ambient )
    colors += ambients * diffuseColors[ :, objects ]


    # Add in any reflected colors.

    secondarySources = hits + 1.0e-6 * normals                      # Starting points for all secondary rays
    viewVectors = normalize( -directions )                          # Vectors from primary intersections to viewer

    if reflections > 0 :

        # To make mirror reflections work, you will need to remove the line
        # that says "pass" and uncomment the line after it, then add one or
        # more other lines of code.

        pass
        # reflectionColors = traceForColor( scene, secondarySources, reflectionDirections, reflections - 1 )


    # Figure out where shadows lie. Conceptually, do this by tracing rays from
    # each intersection of the original rays with a surface to the light,
    # taking each such ray that intersects any object to come from a point in
    # shadow. This assumes that the light is outside of the scene, i.e., that
    # it's not possible for a shadow ray to intersect an object after it passes
    # through the light. I also don't actually trace shadow rays from exactly
    # the original intersection points, but rather from those points displaced
    # slightly along the normals to the surfaces, in order to be sure a ray
    # doesn't intersect its own starting point due to roundoff.

    shadowDirections = scene.lightPos - secondarySources
    shadowObjects, shadowHits, shadowNormals, shadowTs = rayTrace( scene.elements, secondarySources, shadowDirections )
    litMask = hitMask & ( shadowTs > 1.0 )


    # Add the contribution of diffuse reflection to the colors.

    lightVectors = normalize( scene.lightPos - hits )                   # Vectors from primary intersections to light
    lightCosines = np.maximum( listDot( normals, lightVectors ), 0.0)

    colors[ :, litMask ] += ( brightness * lightCosines * diffuseColors[ :, objects ] )[ :, litMask ]


    # Add the contribution of specular reflection, using the angle between the
    # halfway vector and the normal as the basis for scaling light intensity.

    halfways = normalize( lightVectors + viewVectors )
    halfwayCosines = np.maximum( listDot( halfways, normals ), 0.0 )
    colors[ :, litMask ] += ( brightness * specularCoefficients[objects] * halfwayCosines ** shineExponents[objects] )[ :, litMask ]


    # Pixels not occupied by any object receive a background color.

    bkgndColor = np.array( [[0.35], [0.65], [1.0]] )
    colors[ :, missMask ] = bkgndColor


    # All color calculations are finished, return the result.

    return colors




# A function that scales color intensities to lie between 0 and 1. For now this
# just clamps values below 0 to 0, then scales colors so that the maximum color
# intensity in each pixel is 1. This function takes a NumPy matrix representing
# a list of colors as its argument, and returns a similar list of adjusted
# colors.

def scaleColor( pixels ) :


    # Clamp any negative colors to 0.

    positivePixels = np.maximum( pixels, 0.0 )


    # Tone map colors by scaling each pixel so that its maximum color component
    # is 1.0, if the maximum was above 1. Otherwise leave the color alone.
    # Based on Kevin Suffern's "Ray Tracing from the Ground Up" text.

    maxColors = np.amax( positivePixels, 0 )
    scaleMask = maxColors > 1.0
    positivePixels[ :, scaleMask ] = positivePixels[ :, scaleMask ] / maxColors[ scaleMask ]


    # Return the color-corrected pixel values

    return positivePixels




# Some functions that generate scenes with which to test the rest of the ray
# tracer.


# A scene that exercises each of the geometric objects, with no significant
# reflection. This scene consists of a sphere floating above a square platform,
# with a triangular wall behind it.

def basicScene() :

    leftX = -3.0                                # X coordinate of left side of platform and wall.
    rightX = 3.0                                # X coordinate of right side of platform and wall.
    centerX = ( leftX + rightX ) / 2.0          # X coordinate for center of scene
    platformY = -2.5                            # Y coordinate for platform and bottom of wall.
    frontZ = -2.0                               # Z coordinate of front of platform
    backZ = -6.0                                # Z coordinate for wall and back of platform
    centerZ = ( frontZ + backZ ) / 2.0          # Z coordinate for center of scene

    specular = 0.0                              # Coefficient of specular reflection for matte surfaces
    shine = 1.0                                 # Shininess exponent for matte surfaces

    return Scene( np.array( [ [0.0], [50.0], [10.0], [1.0] ] ), 0.6,
                  [ Parallelogram( np.array( [ [leftX], [platformY], [frontZ], [1.0] ] ),
                                   np.array( [ [rightX], [platformY], [frontZ], [1.0] ] ),
                                   np.array( [ [leftX], [platformY], [backZ], [1.0] ] ),
                                   0.0, 0.97, 0.0, specular, shine ),
                    Triangle( np.array( [ [leftX], [platformY], [backZ], [1.0] ] ),
                              np.array( [ [rightX], [platformY], [backZ], [1.0] ] ),
                              np.array( [ [centerX], [platformY + 5.0], [backZ], [1.0] ] ),
                              0.0, 0.0, 0.97, specular, shine ),
                    Ellipsoid( np.array( [ [centerX], [0.0], [centerZ], [1.0] ] ),
                               0.5, 0.5, 0.5,
                               0.97, 0.0, 0.0, specular, shine ) ] )


# A scene with a simple reflection in it: based on the basic scene from above,
# without the ball and with the platform replaced by a mirror, i.e., a highly
# reflective parallelogram (apparently) enclosed in a black frame.

def reflectionScene() :

    leftX = -3.0                                # X coordinate of left side of platform and wall.
    rightX = 3.0                                # X coordinate of right side of platform and wall.
    centerX = ( leftX + rightX ) / 2.0          # X coordinate for center of scene
    platformY = -2.5                            # Y coordinate for platform.
    wallY = platformY + 0.25                    # Y coordinate for bottom of wall.
    frontZ = -2.0                               # Z coordinate of front of platform
    backZ = -12.0                               # Z coordinate for wall and back of platform
    frameMargin = 0.2                           # Margin between frame and mirror
    frameDrop = 0.01                            # Lower the frame by this much below the mirror

    return Scene( np.array( [ [0.0], [50.0], [20.0], [1.0] ] ), 0.6,
                  [ Parallelogram( np.array( [[leftX], [platformY], [frontZ], [1.0]] ),           # Mirror
                                   np.array( [[rightX], [platformY], [frontZ], [1.0]] ),
                                   np.array( [[leftX], [platformY], [backZ], [1.0]] ),
                                   0.15, 0.15, 0.15, 0.9, 50.0 ),
                    Parallelogram( np.array( [[leftX - frameMargin], [platformY - frameDrop], [frontZ + frameMargin], [1.0]] ),   # Frame
                                   np.array( [[rightX + frameMargin], [platformY - frameDrop], [frontZ + frameMargin], [1.0]] ),
                                   np.array( [[leftX - frameMargin], [platformY - frameDrop], [backZ - frameMargin], [1.0]] ),
                                   0.0, 0.0, 0.0, 0.0, 1.0 ),
                    Triangle( np.array( [ [leftX], [wallY], [backZ], [1.0] ] ),                 # Back wall
                              np.array( [ [rightX], [wallY], [backZ], [1.0] ] ),
                              np.array( [ [centerX], [wallY + 8.0], [backZ], [1.0] ] ),
                              0.0, 0.0, 0.97, 0.0, 1.0 ) ] )


# A scene consisting of a pair of reflective balls floating above a plane.

def ballScene() :

    ballY = 1.0                         # Y coordinate for balls
    ballZ = -5.0                        # Z coordinate for balls
    left = -3.0                         # X coordinate for left edge of plane
    right = 3.0                         # X coordinate for right edge of plane
    planeY = -2.0                       # Y coordinate for plane
    front = -1.0                        # Z coordinate for front of plane
    back = -8.0                         # Z coordinate for back of plane
    radius = 1.0                        # Radius of balls
    specular = 0.97                     # Coefficient of specular reflection for shiny balls
    shine = 80.0                        # shininess exponent for shiny balls

    return Scene( np.array( [[0.0], [50.0], [10.0], [1.0]] ), 0.7,
                  [ Ellipsoid( np.array( [[-1.5], [ballY], [ballZ], [1.0]] ),
                               radius, radius, radius,
                               0.6, 0.0, 0.6, specular, shine ),
                    Ellipsoid( np.array( [[1.5], [ballY], [ballZ], [1.0]] ),
                               radius, radius, radius,
                               0.0, 0.6, 0.6, specular, shine ),
                    Parallelogram( np.array( [[left], [planeY], [front], [1.0]] ),
                                   np.array( [[right], [planeY], [front], [1.0]] ),
                                   np.array( [[left], [planeY], [back], [1.0]] ),
                                   0.0, 0.7, 0.0, 0.0, 1.0 ) ] )




# The main program. This creates a scene to render, defines the viewer's
# position and rays from that point through the pixels of a virtual image grid,
# generates an image from that information, and finally displays the image.

# Set up the scene.

scene = basicScene()


# Define the focal point and image grid. For this ray tracer, the image grid
# is a square in the z = 0 plane, centered at the origin, and the focal point
# is somewhere on the positive z axis. This assumes scenes mostly lie in the
# z < 0 halfspace.

focusZ = 2.0                            # Z coordinate of focal point
focus = np.array( [ [0.0], [0.0], [focusZ], [1.0] ] )

nPixels = 512                           # Number of pixels in each dimension of the image
totalPixels = nPixels * nPixels         # Total number of pixels in the image

imageLeft = -1.0                        # X coordinate of left edge of image
imageRight = 1.0                        # Right edge of image
imageBottom = -1.0                      # Y coordinate of bottom of image
imageTop = 1.0                          # Top of image

halfPixel = ( imageRight - imageLeft ) / ( 2.0 * nPixels )      # Half the width or height of a pixel in world units


# Build a "list" of direction vectors for the rays. This is really a NumPy
# matrix with the direction vectors in its columns, as described in the
# "RayTracingUtilities" module.

imageXs = np.linspace( imageLeft + halfPixel, imageRight - halfPixel, nPixels )
imageYs = np.linspace( imageTop - halfPixel, imageBottom + halfPixel, nPixels )

gridXs, gridYs = np.meshgrid( imageXs, imageYs )
gridZs = np.full( totalPixels, -focusZ )
gridWs = np.full( totalPixels, 0.0 )

directions = listify( [ gridXs - focus[0,0], gridYs - focus[1,0], gridZs, gridWs ] )
foci = np.full( [ 4, listLength(directions) ], focus )


# Trace the rays, using 3 "bounces" for reflection. Unpack the resulting colors
# into an image.

image = delistify( scaleColor( traceForColor( scene, foci, directions, 3 ) ), [ nPixels, nPixels ] )


# Display the image.

plot.imshow( image )
plot.show()