import sys, time
from math import sqrt
from math import sin
from math import cos
from math import tan
from math import pi
from math import floor

# use PIL if not using java
if 'java' in sys.platform:
    from java.util import Random
    randomObject = Random()
    random = randomObject.nextFloat
else:
    from random import random

def DOT(a, b):
    ax,ay,az = a.val
    bx,by,bz = b.val
    return ax*bx + ay*by + az*bz

class Image:
    def __init__(self, width, height):
        self.size = (width, height)
        self.grid = []
        for i in xrange(0, width*height):
            self.grid.append((0,0,0))
    def put(self, pos, color):
        x = pos[0]
        y = pos[1]
        self.grid[x + y*self.size[0]] = color

    def update(self):
        from PIL import Image as pilImage
        im = pilImage.new("RGB", self.size, "white")
        for x in xrange(0, self.size[0]):
            for y in xrange(0, self.size[1]):
                color = self.grid[x + y*self.size[1]]
#                print("(%i, %i) - %f %f %f" % (x, y, color[0], color[1], color[2]))
                r = int(255*color[0])
                g = int(255*color[1])
                b = int(255*color[2])
                im.putpixel((x,y), (r,g,b))
        im.show()

class Tuple3(object):
    # make Tuple3 immutable
    def __setattr__(self, *args):
        raise TypeError("Tuple3 immutable")
    __delattr__ = __setattr__

#    def __mul__(self, b):
#        ax,ay,az = self.val
#        bx,by,bz = b.val
#        return ax*bx + ay*by + az*bz
    def __init__(self, x, y, z):
        super(Tuple3, self).__setattr__('val', (x,y,z))

    def x(self):
        return self.val[0]
    def y(self):
        return self.val[1]
    def z(self):
        return self.val[2]

class Vector3(Tuple3):
    def __init__(self, x, y, z):
        Tuple3.__init__(self, x,y,z)
    def length(self):
        x,y,z = self.val
        return sqrt(x*x + y*y + z*z)
    def normalize(self):
        x,y,z = self.val
        inv_l = 1.0 / self.length()
        return Vector3(x*inv_l, y*inv_l, z*inv_l)
    def cross(self, b):
        a = self
        a1,a2,a3 = a.val
        b1,b2,b3 = b.val
        return Vector3(a2*b3 - a3*b2, a3*b1 - a1*b3, a1*b2 - a2*b1)
    def __mul__(self, s):
        x,y,z = self.val
        return Vector3(s*x,s*y,s*z)
    def __sub__(self, b):
        ax,ay,az = self.val
        bx,by,bz = b.val
        return Vector3(ax-bx,ay-by,az-bz)
    def __str__(self):
        x,y,z = self.val
        return "<%f,%f,%f>" % (x,y,z)

class Point3(Tuple3):
    def __init__(self, x, y, z):
        Tuple3.__init__(self, x,y,z)
    def __sub__(self, b):
        ax,ay,az = self.val
        bx,by,bz = b.val
        return Vector3(ax-bx,ay-by,az-bz)
    def __add__(self, vec):
        ax,ay,az = self.val
        vx,vy,vz = vec.val
        return Point3(ax+vx, ay+vy, az+vz)
    def __str__(self):
        return "(%f,%f,%f)" % self.val

class Color(Tuple3):
    def __init__(self, x, y, z):
        Tuple3.__init__(self, x,y,z)
    def tuple(self):
        return self.val
    def toneMapTuple(self, gamma=1.0):
        x,y,z = self.val
        return (x/(1.0+x), y/(1.0+y), z/(1.0+z))
    def scale(self, s):
        x,y,z = self.val
        return Color(s*x,s*y,s*z)
    def modulate(self, b):
        ax,ay,az = self.val
        bx,by,bz = b.val
        return Color(ax*bx, ay*by, az*bz)
    def __add__(self, b):
        ax,ay,az = self.val
        bx,by,bz = b.val
        return Color(ax+bx, ay+by, az+bz)


    def __str__(self):
        return "(%f,%f,%f)" % self.val

class Ray(object):
    # make Ray immutable
    def __setattr__(self, *args):
        raise TypeError("Ray immutable")
    __delattr__ = __setattr__

    def __init__(self, ray_orig, ray_dir, color, t=sys.maxint, hit=None, prim=None, scratch=None, depth=0):
        super(Ray, self).__setattr__('orig', ray_orig)
        super(Ray, self).__setattr__('dir', ray_dir)
        super(Ray, self).__setattr__('color', color)
        super(Ray, self).__setattr__('t', t)
        super(Ray, self).__setattr__('hit', hit)
        super(Ray, self).__setattr__('primitive', prim)
        super(Ray, self).__setattr__('scratch', scratch)
        super(Ray, self).__setattr__('depth', depth)
#    def updateT(self, time, primitive, scr):
#        ray = Ray(self.orig, self.dir, self.color, t=time, prim=primitive, scratch=scr)
#        return ray
    def updateT(self, t, primitive, scr):
        ray = Ray(self.orig, self.dir, self.color, t, primitive, scr)
        return ray
    def surfOffset(self):
        # move ray over to surface and offset a bit
        new_orig = self.orig + self.dir*self.t + self.hit.hitNormal(self)*.000001
        return Ray(new_orig, self.dir, self.color, self.t, self.hit, self.primitive, self.scratch, self.depth+1)
    def reflect(self, normal):
        refl_dir = self.dir - normal * (2*DOT(self.dir, normal))
        ray = Ray(self.orig, refl_dir, self.color)
        return ray

class Material:
    def __init__(self):
        pass
    def shade(self):
        raise NotImplementedError("Material shade is virtual and must be overridden.")

class CheckerMat3D(Material):
    def __init__(self, color1, color2, scale_x=1, scale_y=1, scale_z=1):
        self.color1 = color1
        self.color2 = color2
    def shade(self, ray):
        pos = ray.hit.hitPosition(ray)
        x,y,z = pos.val
        fx,fy,fz = floor(x), floor(y), floor(z)
        if (fx + fy + fz) % 2 == 0:
            return self.color1
        else:
            return self.color2

class MetalMat(Material):
    def shade(self, ray):
        normal = ray.hit.hitNormal(ray)
        ray = ray.surfOffset()
        r_ray = ray.reflect(normal)
        ray = None # to avoid using it for now


#        print("start")
#        print(ray.orig)
#        print(r_ray.orig)
#        print("end")

        if r_ray.depth > 5:
            return r_ray.hit.error_mat.shade(r_ray)

        for primitive in scene.geometry:
            r_ray = primitive.intersect(r_ray)

            if r_ray.hit is not None:
                color = r_ray.hit.material.shade(r_ray)
                return color
        else:
            return Color(0,0,0)

class DiffuseMat(Material):
    def __init__(self, color):
        self.color = color
    def shade(self, ray):
        ray = ray.surfOffset()

        total_color = Color(0,0,0)
        for light in scene.lights:
            for plight in light.samples():
                light_dir = plight.position - ray.orig
                light_dist = light_dir.length()
                light_dir = light_dir.normalize()
                light_ray = Ray(ray.orig, light_dir, Color(0,0,0))

                for primitive in scene.geometry:
                    light_ray = primitive.intersect(light_ray)

                # not occluded (in shadow)
                if light_ray.t > light_dist:
#                if ray.hit is None:
                    light_contrib = plight.color.scale(DOT(light_dir, ray.hit.hitNormal(ray)) * plight.intensity)
                    filter = light_contrib.modulate(self.color)
                    total_color = total_color + filter
#                    color = r_ray.hit.material.shade(r_ray)
        return total_color

class Light:
    def __init__(self):
        pass
    def generateRays(self):
        raise NotImplementedError("Light generateRays is virtual and must be overridden.")

class PointLight(Light):
    def __init__(self, color, intensity, position):
        self.color = color
        self.intensity = intensity
        self.position = position
    def generateRays(self, num_rays):
        rays = []
        for i in xrange(num_rays):
            # generate a random point on a sphere
            # taken from http://demonstrations.wolfram.com/RandomPointsOnASphere/
            u = 2*random() - 1
            theta = 2*pi*random()
            x = cos(theta)*sqrt(1-(u*u))
            y = sin(theta)*sqrt(1-(u*u))
            z = u
            dir = Vector3(x,y,z)
            rays.append(Ray(self.position, dir.normalize(), self.color))

        return rays
    def samples(self):
        return [self]

class Primitive:
    def __init__(self):
        self.error_mat = CheckerMat3D(Color(1,0,1), Color(1,1,0))
        self.material = CheckerMat3D(Color(1,0,1), Color(0,1,1))
    def setMaterial(self, mat):
        self.material = mat
    def intersect(self, ray):
        raise NotImplementedError("Primitive intersect is virtual and must be overridden.")
    def hitNormal(self, ray):
        raise NotImplementedError("Primitive hitNormal is virtual and must be overridden.")
    def hitPosition(self, ray):
        return ray.orig + ray.dir*ray.t

class Plane(Primitive):
    def __init__(self, normal, d):
        Primitive.__init__(self)
        self.normal = normal.normalize()
        self.d = d
    def intersect(self, ray):
        t = -(DOT(self.normal, ray.orig) + self.d) / (DOT(self.normal, ray.dir))
        if t > 0 and t < ray.t:
            ray = ray.updateT(t, self, t)
        return ray
    def hitNormal(self, ray):
        return self.normal

class Sphere(Primitive):
    def __init__(self, center, radius):
        Primitive.__init__(self)
        self.center = center
        self.radius = radius
        self.radiusSq = radius*radius
    def intersect(self, ray):
        CO = self.center - ray.orig
        L2_co = DOT(CO, CO)
        t_ca = DOT(CO, ray.dir)
        if L2_co < self.radiusSq:
            t2_hc = self.radiusSq - L2_co + t_ca*t_ca
            root = t_ca + sqrt(t2_hc)
        else:
            if t_ca < 0:
                return ray # miss
            else:
                t2_hc = self.radiusSq - L2_co + t_ca*t_ca
                if t2_hc < 0:
                    return ray # miss
                else:
                    root = t_ca - sqrt(t2_hc)
        if root > 0 and root < ray.t:
            ray = ray.updateT(root, self, root)
        return ray

    def hitNormal(self, ray):
        return (self.hitPosition(ray) - self.center).normalize()

class Box(Primitive):
    def __init__(self, min, max):
        Primitive.__init__(self)
        self.bounds = (min, max)
    def intersect(self, ray):
        # taken from "An Efficient and Robust Ray-Box Intersection Algorithm" by Williams et al
        bounds = self.bounds
        normal = Vector3(0,0,1)
        if ray.dir.x() >= 0:
            tmin = (bounds[0].x() - ray.orig.x()) / ray.dir.x()
            tmax = (bounds[1].x() - ray.orig.x()) / ray.dir.x()
        else:
            tmin = (bounds[1].x() - ray.orig.x()) / ray.dir.x()
            tmax = (bounds[0].x() - ray.orig.x()) / ray.dir.x()

        if (ray.dir.y() >= 0):
            tymin = (bounds[0].y() - ray.orig.y()) / ray.dir.y()
            tymax = (bounds[1].y() - ray.orig.y()) / ray.dir.y()
        else:
            tymin = (bounds[1].y() - ray.orig.y()) / ray.dir.y()
            tymax = (bounds[0].y() - ray.orig.y()) / ray.dir.y()

        if (tmin > tymax or tymin > tmax):
            return ray

        if (tymin > tmin):
            tmin = tymin
            normal = Vector3(0,1,0)
        if (tymax < tmax):
            tmax = tymax
            normal = Vector3(1,0,0)
        if (ray.dir.z() >= 0):
            tzmin = (bounds[0].z() - ray.orig.z()) / ray.dir.z()
            tzmax = (bounds[1].z() - ray.orig.z()) / ray.dir.z()
        else:
            tzmin = (bounds[1].z() - ray.orig.z()) / ray.dir.z()
            tzmax = (bounds[0].z() - ray.orig.z()) / ray.dir.z()

        if (tmin > tzmax or tzmin > tmax):
            return ray
        if (tzmin > tmin):
            tmin = tzmin
            normal = Vector3(0,0,1)
        if (tzmax < tmax):
            tmax = tzmax
            normal = Vector3(0,0,1)

        if tmin > 0 and tmin < ray.t:
            ray = ray.updateT(tmin, self, (tmin, normal))

        return ray

    def hitNormal(self, ray):
        return ray.scratch[1]
        return Vector3(1,1,1).normalize()

class Camera:
    def __init__(self):
        pass

class PinholeCamera(Camera):
    def __init__(self, eye, lookat, up):
        self.eye = eye
        self.lookat = lookat
        self.fov = 45
        self.lookdir = (self.lookat - self.eye).normalize()
        self.right = self.lookdir.cross(up).normalize()

        # recalculate eye to be orthogonal to look direction
        self.up = self.right.cross(self.lookdir)

    def preprocess(self, dims):
        width, height = dims
        aspect = float(width) / float(height)
        fov_length = 2*tan(self.fov*.5)
        self.u = self.right * fov_length * aspect
        self.v = self.up * fov_length

    def computeRay(self, i, j):
        screen = self.eye + self.lookdir + self.u*i + self.v*j
        ray_dir = screen - self.eye
        return Ray(self.eye, ray_dir.normalize(), Tuple3(0,0,0))

class PhotonMap:
    def __init__(self):
        self.photons = []
        pass
    def storePhoton(self, photon):
        self.photons.append(photon)

class Scene:
    def __init__(self):
        self.lights = []
        self.geometry = []
        self.image = Image(512,512)
#        self.image = Image(256,256)
        self.photon_map = PhotonMap()
    def addLight(self, light):
        self.lights.append(light)
    def addGeometry(self, geometry):
        self.geometry.append(geometry)
    def setCamera(self, camera):
        self.camera = camera
    def render(self):
        # do some preprocessing
        self.camera.preprocess(self.image.size)

#        self.image.show()
#        draw = ImageDraw.Draw(self.image)
#        draw.line((0, 0) + self.image.size, fill=128)
#        draw.line((0, self.image.size[1], self.image.size[0], 0), fill=128)
#        del draw
#        self.image.show()

        # start the clock
        start = time.time()

        # send out photon rays


        # render from screen
        print("Rendering image")
        width = self.image.size[0]
        height = self.image.size[1]
        dx = 2.0 / width
        dy = 2.0 / height
        hw = width / 2.0
        hh = height / 2.0
        aspect = float(width) / float(height)
        total_pixels = float(width*height)
        for i in xrange(0,width):
            for j in xrange(0,height):
                percent_done = 100*(i*height+j)/total_pixels
                sys.stdout.write('\r working: %.0f%%' % percent_done)
                sys.stdout.flush()
                ray = self.camera.computeRay(i*dx-1,j*dy-1)

                for primitive in self.geometry:
                    ray = primitive.intersect(ray)

                if ray.hit is not None:
                    color = ray.hit.material.shade(ray)
                    color = color.toneMapTuple()
                    self.image.put((i,self.image.size[1]-j-1), color)

        elapsed = int(time.time() - start)
        hours = elapsed / 360
        elapsed -= 360*hours
        minutes = elapsed / 60
        elapsed -= 60*minutes
        seconds = elapsed
        print("")
        print("Done in %ih %im %is" % (hours, minutes, seconds))

        # display on screen
        self.image.update()



# create the scene
scene = Scene()

scene.addLight(PointLight(Color(1,1,1), 100, Point3(0,10,0)))

plane1 = Plane(Vector3(0,1,0), 0.001)
plane1.setMaterial(DiffuseMat(Color(1,1,0)))
scene.addGeometry(plane1)

sphere1 = Sphere(Point3(0,2,0), 2)
sphere1.setMaterial(MetalMat())
scene.addGeometry(sphere1)

sphere2 = Sphere(Point3(-3,2,3), 2)
sphere2.setMaterial(CheckerMat3D(Color(1,0,0), Color(0,0,1)))
scene.addGeometry(sphere2)

scene.setCamera(PinholeCamera(Point3(8,8,8),
                              Point3(0,0,0),
                              Vector3(0,1,0)))
#cProfile.run('scene.render()', 'profile')

scene.render()