#! /usr/bin/python

# a rewrite of my pygame port of
# youtube.com/watch?v=N8elxpSu9pw
# which is by Bisqwit
#
# (multithreaded version)
# 
# theinternetftw

import pygame, math, random, sys
from multiprocessing import Pool

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 math.sqrt(vSqr(a))
def vNorm(a):   return vMulD(a, 1.0/vLen(a))

def vMirrorAround(a, axis): 
    n = vNorm(axis)
    v = vDot(a,n)
    return vSub(vMulD(n,v+v), a)

# for color vectors (rgb)
def vLuma(a): return a[0]*0.299 + a[1]*0.587 + a[2]*0.114
def vClamp(a):
    for i in range(3):
        if   a[i] < 0.0: a[i] = 0.0
        elif a[i] > 1.0: a[i] = 1.0
    return a

def vClampWithDesaturation(a):
    # if the color represented by this triplet
    # is too bright or too dim, decrease the saturation
    # as much as required, while keeping the luma unmodified
    l = vLuma(a)
    if   l > 1.0: return [1.0, 1.0, 1.0]
    elif l < 0.0: return [0.0, 0.0, 0.0]
    # if any component is over the bounds,
    # calculate how much the saturation must
    # be reduced to achieve an in-bounds value.
    # Since the luma was verified to be in 0..1,
    # a maximum reduction of saturation to 0% will
    # always produce an in-bounds value, but usually
    # such a drastic reduction is not necessary.
    # Because we're only doing relative modifications,
    # we don't need the original saturation level of the
    # pixel.
    sat = 1.0
    for c in a:
        if   c > 1.0: sat = min(sat, (l-1.0) / (l-c))
        elif c < 0.0: sat = min(sat,  l      / (l-c))
    if sat != 1.0:
        a = vAddD(vMulD(vSubD(a,l), sat),l)
        a = vClamp(a)
    return a

def vGetRotMatrix(ang): 
    cx,cy,cz = math.cos(ang[0]),math.cos(ang[1]),math.cos(ang[2])
    sx,sy,sz = math.sin(ang[0]),math.sin(ang[1]),math.sin(ang[2])
    sxsz,cxsz = sx*sz,cx*sz
    cxcz,sxcz = cx*cz,sx*cz
    return [ [cy*cz,          cy*sz,          -sy  ],
             [sxcz*sy - cxsz, sxsz*sy + cxcz, sx*cy],
             [cxcz*sy + sxsz, cxsz*sy - sxcz, cx*cy] ]

def mTransform(m, vec):
    return [ vDot(m[0], vec), vDot(m[1], vec), vDot(m[2], vec) ]

class Plane:
    def __init__(self, norm, off):
        self.normal = norm
        self.offset = off
# declare six planes, each looks
# towards the origin and is 30 units away
planes = [ Plane( [0.0,0.0,-1.0], -30),
           Plane( [0.0,1.0,0.0],  -30),
           Plane( [0.0,-1.0,0.0], -30),
           Plane( [1.0,0.0,0.0],  -30),
           Plane( [0.0,0.0,1.0],  -30),
           Plane( [-1.0,0.0,0.0], -30) ]

class Sphere:
    def __init__(self, c, r):
        self.center = c
        self.radius = r
# declare a few spheres
spheres = [ Sphere( [0.0,0.0,0.0],    7.0), 
            Sphere( [19.4, -19.4, 0], 2.1),
            Sphere( [-19.4, 19.4, 0], 2.1),
            Sphere( [13.1, 5.1, 0.0], 1.1),
            Sphere( [-5.1, -13.1, 0], 1.1),
            Sphere( [-30.0,-30.0,15.0], 11.0),
            Sphere( [15.0, -30.0,30.0], 6.0),
            Sphere( [30.0, 15.0, -30.0],6.0) ]

class LightSource:
    def __init__(self, w, c):
        self.where = w
        self.color = c
#declare lightsources, each w/ a loc and a rgb color
lights = [ LightSource( [-28.0,-14.0,  4.0],  [ .4,.51, .9] ),
           LightSource( [-29.0,-29.0,-29.0],  [.95, .1, .1] ),
           LightSource( [ 14.0, 29.0,-14.0],  [ .8, .8, .8] ),
           LightSource( [ 29.0, 29.0, 29.0],  [1.0,1.0,1.0] ),
           LightSource( [ 28.0,  0.0, 29.0],  [ .5, .6, .1] ) ]

numPlanes = len(planes)
numSpheres = len(spheres)
numLights = len(lights)
MAXTRACE = 6

def rayFindObstacle(eye, dir, hitDist):
    #try to intersect ray w/ each object, see which gives the closest hit
    hitType = -1
    hitIndex = -1
    hitLoc = [0.0,0.0,0.0]
    hitNormal = [0.0,0.0,0.0]
    for i in range(numSpheres):
        v = vSub(eye, spheres[i].center)
        r = spheres[i].radius
        dv = vDot(dir,v)
        d2 = vSqr(dir)
        sq = dv*dv - d2*(vSqr(v) - r*r)
        #does the ray coincide w/ the sphere?
        if (sq < 1e-6):
            continue
        #if so, where?
        sqt = math.sqrt(sq)
        dist = min(-dv-sqt, -dv+sqt) / d2
        if dist < 1e-6 or dist >= hitDist:
            continue
        hitType = 1 
        hitIndex = i
        hitDist = dist
        hitLoc = vAdd(eye, vMulD(dir,hitDist))
        hitNormal = vMulD(vSub(hitLoc,spheres[i].center), 1/r)
    for i in range(numPlanes):
        dv = -vDot(planes[i].normal,dir)
        if dv > -1e-6:
            continue
        d2 = vDot(planes[i].normal,eye)
        dist = (d2 + planes[i].offset) / dv
        if dist < 1e-6 or dist >= hitDist:
            continue
        hitType = 0
        hitIndex = i
        hitDist = dist
        hitLoc = vAdd(eye, vMulD(dir,hitDist))
        hitNormal = vNeg(planes[i].normal)
    return hitType, hitIndex, hitDist, hitLoc, hitNormal

random.seed(1) #not doing this lends an interesting speckled pattern to the lighting,
               #as each proc that's started up (max of 4 atm) has different
               #values for the arealight vectors (i think the work is done as threads
               #on those procs, so the vectors are *not* changed w/ each call to 
               #apply_sync)
numArealightVectors = 20
arealightVectors = []
for i in range(numArealightVectors):
    temp = []
    for i in range(3):
        temp.append(2.0*(random.random() - 0.5))
    arealightVectors.append( [temp[0],temp[1],temp[2]] )
 
def rayTrace(eye, dir, k):
    hitDist = 1e6
    hitType, hitIndex, hitDist, hitLoc, hitNormal = rayFindObstacle(eye,dir,hitDist)

    if hitType != -1:

        # Found an obstacle. Next, find out how it is illuminated.
        # Shoot a ray to each lightsource, and determine if there
        # is an obstacle behind it. This is called "diffuse light".
        # To smooth out the infinitely sharp shadows caused by 
        # infinitely small point-lightsources, assume the lightsource
        # is actually a cloud of small lightsources around its center.
        diffuseLight = [0.0,0.0,0.0]
        specularLight = [0.0,0.0,0.0]
        pigment = [1.0, .98, .94] #default pigment
        for i in range(numLights):
            for j in range(numArealightVectors):
                v = vSub(vAdd(lights[i].where,arealightVectors[j]),hitLoc)
                lightDist = vLen(v)
                v = vNorm(v)
                diffuseEffect = vDot(hitNormal,v) / (numArealightVectors*1.0)
                attenuation = (1 + ((lightDist/34.0)**2.0))
                diffuseEffect /= attenuation
                if diffuseEffect > 1e-3:
                    shadowDist = lightDist - 1e-4
                    t,hi,hd,hl,hn = rayFindObstacle(vAdd(hitLoc,vMulD(v,1e-4)),v,shadowDist)
                    if t == -1: #no obstacle occluding the light
                        diffuseLight = vAdd(diffuseLight,vMulD(lights[i].color,diffuseEffect))

        if k > 1:
            # add specular light/reflection, unless recursion depth is maxed
            v = vNeg(dir)
            v = vMirrorAround(v, hitNormal)
            specularLight = rayTrace(vAdd(hitLoc, vMulD(v,1e-4)),v,k-1)

        if hitType == 0: #plane
            diffuseLight = vMulD(diffuseLight,0.9)
            specularLight = vMulD(specularLight,0.5)
            # color the different walls differently
            idx = hitIndex % 3
            if   idx == 0: pigment = [0.9,0.7,0.6]
            elif idx == 1: pigment = [0.6,0.7,0.7]
            elif idx == 2: pigment = [0.5,0.8,0.3]

        elif hitType == 1: #sphere
            diffuseLight = vMulD(diffuseLight,1.0)
            specularLight = vMulD(specularLight, 0.34)
        
        return vMul(vAdd(diffuseLight,specularLight),pigment)

    #didn't hit anything, return black
    return [0.0,0.0,0.0]

def getPix(x, y, w, h, zoom, camlookmatrix, campos, MAXTRACE):

    camray = [ x/float(w) - 0.5,
               y/float(h) - 0.5,
               zoom ]
    camray[0] *= 4.0/3 # Aspect Ratio Correction
    camray = vNorm(camray)
    camray = mTransform(camlookmatrix, camray)
    campix = rayTrace(campos, camray, MAXTRACE)

    campix = vMulD(campix, 0.5) #Adjust brightness
    
    return campix

def main():

    pygame.display.init()
    screenw, screenh = 640, 480
    w, h = int(sys.argv[1]),int(sys.argv[2])
    screen = pygame.display.set_mode((screenw, screenh), pygame.DOUBLEBUF)
    frame = pygame.Surface((w,h))
        
    camangle =      [  0.0,  0.0,  0.0]
    camangledelta = [-.005,-.011,-.017]
    camlook =       [  0.0,  0.0,  0.0]
    camlookdelta =  [-.001, .005, .004]

    zoom = 46.0
    zoomdelta = 0.99

    contrast = 32
    contrast_offset = -0.17

    frames = []

    #numFrames = 9300
    numFrames = int(sys.argv[3])

    MAXTRACE = int(sys.argv[4])
    
    pool = Pool(processes=4)

    for frameNo in range(numFrames):
        # Put camera between central sphere and walls
        campos = [0.0,0.0,16.0]
        camrotatematrix = vGetRotMatrix(camangle)
        campos = mTransform(camrotatematrix, campos)
        camlookmatrix = vGetRotMatrix(camlook)

        # Determine contrast ratio for this frame's pixels
        thisframe_min = 100
        thisframe_max = -100

        pixels = pygame.PixelArray(frame)
                
        for y in range(h):
        
            results = []

            for x in range(w):
                result = pool.apply_async(getPix, [x, y, w, h, zoom, camlookmatrix, campos, MAXTRACE])
                results.append(result)

            for x in range(w):
            
                campix = results[x].get()
                
                # update frame luma info for automatic contrast adjuster
                lum = vLuma(campix)
                if lum < thisframe_min: thisframe_min = lum
                if lum > thisframe_max: thisframe_max = lum
                
                # exagerate the colors to bring contrast better forth
                campix = vMulD(vAddD(campix,contrast_offset),contrast)

                # Clamp, and compensate for display gamma (for dithering)
                # ...But don't actually dither. Because that shit is bananas.
                # Maybe later, I can look up the EGA palette, etc.
                campix = vClampWithDesaturation(campix)

                # Draw pixel (use pygame)
                r = int(campix[0] * 255)
                g = int(campix[1] * 255)
                b = int(campix[2] * 255) 

                pixels[x][y] = pygame.Color(r,g,b)

                #del pixels?
                pygame.transform.scale(frame, (screenw,screenh), screen)
                pygame.display.flip() 

                for e in pygame.event.get():
                    if e.type == pygame.QUIT:
                        pygame.quit()
                        sys.exit()

        frames.append(frame.copy()) 

        print 'frame ',frameNo
        sys.stdout.flush()

        # Tweak coordinates/camera params for next frame
        much = 1.0

        # In the beginning, do some camera action (play with zoom)
        if zoom <= 1.1:
            zoom = 1.1
        else: 
            if zoom > 40:
                if zoomdelta > 0.95:
                    zoomdelta -= 0.001
            elif zoom < 3:
                if zoomdelta < 0.99:
                    zoomdelta += 0.001
            zoom *= zoomdelta
            much = 1.1 / ((zoom/1.1)**3.0)

        # Update rotation angle
        camlook =  vAdd(camlook,  vMulD(camlookdelta,much))
        camangle = vAdd(camangle, vMulD(camangledelta,much))

        # dynamically adjust the contrast based on the contents of the
        # last frame

        middle = (thisframe_min + thisframe_max) * 0.5
        span = (thisframe_max - thisframe_min)
        thisframe_min = middle - span*0.60 # Avoid dark tones
        thisframe_max = middle + span*0.37 # Emphasize bright tones

        new_contrast_offset = -thisframe_min
        new_contrast = 1 / (thisframe_max - thisframe_min)

        # Avoid abrupt changes, though
        l = 0.85
        if frameNo == 0: l = 0.7
        contrast_offset = (contrast_offset*l + new_contrast_offset*(1.0-l))
        contrast        = (contrast       *l + new_contrast       *(1.0-l))
		
	sys.exit()
    while True:
        for f in frames:
            pygame.transform.scale(f, (screenw,screenh), screen)
            pygame.display.flip() 
            pygame.time.wait(33)
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
            
if __name__ == '__main__':
    main()