#tiny-tracer-no-vector-class. 2x speedup.

#original pastebin: http://pastebin.com/F8f5GHJZ
#original thread:   http://www.reddit.com/r/tinycode/comments/169ri9/ray_tracer_in_140_sloc_of_python_with_picture/

import timeit

from math import sqrt, pow, pi
import Image, sys

def vMul(a,b):  return ( a[0]*b[0], a[1]*b[1], a[2]*b[2] )
def vMulD(a,b): return ( a[0]*b,    a[1]*b,    a[2]*b    )
def vAdd(a,b):  return ( a[0]+b[0], a[1]+b[1], a[2]+b[2] )
def vAddD(a,b): return ( a[0]+b,    a[1]+b,    a[2]+b    )
def vSub(a,b):  return ( a[0]-b[0], a[1]-b[1], a[2]-b[2] )
def vSubD(a,b): return ( a[0]-b,    a[1]-b,    a[2]-b    )
def vNeg(a):    return (-a[0],     -a[1],     -a[2]      )
def vPow(a,b):  return ( a[0]**b,   a[1]**b,   a[2]**b   )

def vDot(a,b):  return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]
def vSqr(a):    return vDot(a,a)
def vLen(a):    return sqrt(vSqr(a))
def vNorm(a):   return vMulD(a, 1.0/vLen(a))

class Sphere( object ):

        def __init__(self, center, radius, color):
                self.c = center
                self.r = radius
                self.col = color

        def intersection(self, l):
                q = vDot(l.d, vSub(l.o, self.c))**2 - vSqr(vSub(l.o, self.c)) + self.r**2
                if q < 0:
                        return Intersection( (0,0,0), -1, (0,0,0), self)
                else:
                        d = -vDot(l.d, vSub(l.o, self.c))
                        d1 = d - sqrt(q)
                        d2 = d + sqrt(q)
                        if 0 < d1 and ( d1 < d2 or d2 < 0):
                                return Intersection(vAdd(l.o,vMulD(l.d,d1)), d1, self.normal(vAdd(l.o,vMulD(l.d,d1))), self)
                        elif 0 < d2 and ( d2 < d1 or d1 < 0):
                                return Intersection(vAdd(l.o,vMulD(l.d,d2)), d2, self.normal(vAdd(l.o,vMulD(l.d,d2))), self)
                        else:
                                return Intersection( (0,0,0), -1, (0,0,0), self)

        def normal(self, b):
                return vNorm(vSub(b,self.c))

class Plane( object ):

        def __init__(self, point, normal, color):
                self.n = normal
                self.p = point
                self.col = color

        def intersection(self, l):
                d = vDot(l.d, self.n)
                if d == 0:
                        return Intersection( (0,0,0), -1, vector(0,0,0), self)
                else:
                        d = vDot(vSub(self.p,l.o),self.n) / d
                        return Intersection(vAdd(l.o,vMulD(l.d,d)), d, self.n, self)

class Ray( object ):

        def __init__(self, origin, direction):
                self.o = origin
                self.d = direction

class Intersection( object ):

        def __init__(self, point, distance, normal, obj):
                self.p = point
                self.d = distance
                self.n = normal
                self.obj = obj

def testRay(ray, objects, ignore=None):
        intersect = Intersection( (0,0,0), -1, (0,0,0), None)

        for obj in objects:
                if obj is not ignore:
                        currentIntersect = obj.intersection(ray)
                        if currentIntersect.d > 0 and intersect.d < 0:
                                intersect = currentIntersect
                        elif 0 < currentIntersect.d < intersect.d:
                                intersect = currentIntersect
        return intersect

def trace(ray, objects, light, maxRecur):
        if maxRecur < 0:
                return (0,0,0)
        intersect = testRay(ray, objects)
        if intersect.d == -1:
                col = (AMBIENT,AMBIENT,AMBIENT)
        elif vDot(intersect.n, vSub(light,intersect.p)) < 0:
                col = vMulD(intersect.obj.col, AMBIENT)
        else:
                lightRay = Ray(intersect.p, vNorm(vSub(light,intersect.p)))
                if testRay(lightRay, objects, intersect.obj).d == -1:
                        lightIntensity = 1000.0/(4*pi*vLen(vSub(light,intersect.p))**2)
                        col = vMulD(intersect.obj.col, max(vDot(vNorm(intersect.n),vMulD(vNorm(vSub(light,intersect.p)),lightIntensity)), AMBIENT))
                else:
                        col = vMulD(intersect.obj.col, AMBIENT)
        return col

def gammaCorrection(color,factor):
        return (int(pow(color[0]/255.0,factor)*255),
                        int(pow(color[1]/255.0,factor)*255),
                        int(pow(color[2]/255.0,factor)*255))

AMBIENT = 0.1
GAMMA_CORRECTION = 1/2.2
MAX_RECURSION = 10

#these are left global so we don't have to pass them to apply_async each pixel
objs = ( Sphere((-2, 0,-10), 2,   (0,255,0)),
         Sphere(( 2, 0,-10), 3.5, (255,0,0)),
         Sphere(( 0,-4,-10), 3,   (0,0,255)),
         Plane((0,0,-12), (0,0,1), (255,255,255)) )

lightSource = (0,10,0)
cameraPos = (0,0,20)

def main():
        img = Image.new("RGB",(500,500))
        for x in range(500):
                results = []
                for y in range(500):
                        ray =  Ray(cameraPos, vNorm(vSub((x/50.0-5, y/50.0-5, 0), cameraPos)))
                        col = trace(ray, objs, lightSource, MAX_RECURSION)
                        img.putpixel((x,499-y),gammaCorrection(col,GAMMA_CORRECTION))
        img.save("trace.bmp","BMP")

if __name__ == '__main__':
        timer = timeit.Timer(stmt='main()',setup="from __main__ import main")
        t = timer.timeit(number=1)
        print "t = %f"%t