# nicked from https://github.com/jamesbowman/raytrace/blob/master/rt2.py

from functools import reduce
import numpy as np
import time
from pixelflut import Pixelflut

class vec3():
    def __init__(self, x, y, z):
        (self.x, self.y, self.z) = (x, y, z)
    def __mul__(self, other):
        return vec3(self.x * other, self.y * other, self.z * other)
    def __add__(self, other):
        return vec3(self.x + other.x, self.y + other.y, self.z + other.z)
    def __sub__(self, other):
        return vec3(self.x - other.x, self.y - other.y, self.z - other.z)
    def dot(self, other):
        return (self.x * other.x) + (self.y * other.y) + (self.z * other.z)
    def __abs__(self):
        return self.dot(self)
    def norm(self):
        mag = np.sqrt(abs(self))
        return self * (1.0 / np.where(mag == 0, 1, mag))
    def components(self):
        return (self.x, self.y, self.z)
        
def html(s):
    r = int(s[1:3], 16) / 255.0
    g = int(s[3:5], 16) / 255.0
    b = int(s[5:7], 16) / 255.0
    return vec3(r, g, b)
        
(w, h) = (256, 144)         # Screen size
L = vec3(5, 5., -10)        # Point light position
E = vec3(0., 0.35, -1.)     # Eye position
FARAWAY = 1.0e39            # an implausibly huge distance

class Sphere:
    def __init__(self, center, r, diffuse, mirror = 0.5):
        self.c = center
        self.r = r
        self.diffuse = diffuse
        self.mirror = mirror

    def intersect(self, O, D):
        b = 2 * D.dot(O - self.c)
        c = abs(self.c) + abs(O) - 2 * self.c.dot(O) - (self.r * self.r)
        disc = (b ** 2) - (4 * c)
        sq = np.sqrt(np.maximum(0, disc))
        h0 = (-b - sq) / 2
        h1 = (-b + sq) / 2
        h = np.where((h0 > 0) & (h0 < h1), h0, h1)

        pred = (disc > 0) & (h > 0)
        return np.where(pred, h, FARAWAY)

    def diffusecolor(self, M):
        return self.diffuse

    def light(self, O, D, d, scene, bounce):
        M = (O + D * d)                         # intersection point
        N = (M - self.c) * (1. / self.r)        # normal
        toL = (L - M).norm()                    # direction to light
        toO = (E - M).norm()                    # direction to ray origin
        nudged = M + N * .0001                  # M nudged to avoid itself

        # Shadow: find if the point is shadowed or not.
        light_distances = [s.intersect(nudged, toL) for s in scene]
        light_nearest = reduce(np.minimum, light_distances)
        seelight = light_distances[scene.index(self)] == light_nearest

        # Ambient
        color = vec3(0.05, 0.05, 0.05)

        # Lambert shading (diffuse)
        lv = np.maximum(N.dot(toL), 0)
        color += self.diffusecolor(M) * lv * seelight

        # Reflection
        if bounce < 2:
            rayD = (D - N * 2 * D.dot(N)).norm()
            color += raytrace(nudged, rayD, scene, bounce + 1) * self.mirror

        # Blinn-Phong shading (specular)
        phong = N.dot((toL + toO).norm())
        color += vec3(1, 1, 1) * np.power(np.clip(phong, 0, 1), 50) * seelight
        return color

class CheckeredSphere(Sphere):
    def diffusecolor(self, M):
        checker = ((M.x * 2 + 1000).astype(int) % 2) == ((M.z * 2 + 1000).astype(int) % 2)
        return self.diffuse * checker

def raytrace(O, D, scene, bounce = 0):
    # O is the ray origin, D is the normalized ray direction
    # scene is a list of Sphere objects (see below)
    # bounce is the number of the bounce, starting at zero for camera rays

    distances = [s.intersect(O, D) for s in scene]
    nearest = reduce(np.minimum, distances)
    color = vec3(0, 0, 0)
    for (s, d) in zip(scene, distances):
        color += s.light(O, D, d, scene, bounce) * (nearest != FARAWAY) * (d == nearest)
    return color

class Task(Pixelflut):
    

    def run(self):
        
        scene = [
            Sphere(vec3(.75, .1, 1.), .6, html('#238acc'), 0.7),
            Sphere(vec3(-.75, .1, 2.25), .6, html('#e5185d'), 0.5),
            Sphere(vec3(-2.75, .1, 3.5), .6, html('#fad31c'), 0.3),
            CheckeredSphere(vec3(0,-99999.5, 0), 99999, html('#e9d4a7'), 0.0),
            ]
        
        r = float(w) / h
        # Screen coordinates: x0, y0, x1, y1.
        S = (-1., 1. / r + .25, 1., -1. / r + .25)
        x = np.tile(np.linspace(S[0], S[2], w), h)
        y = np.repeat(np.linspace(S[1], S[3], h), w)
        
        t0 = time.time()
        Q = vec3(x, y, 0)
        color = raytrace(E, (Q - E).norm(), scene)
        print(f"Rendered image in {(time.time() - t0):.3f} seconds.")
        for y in range(144):
            for x in range(256):
                offset = y * 256 + x
                self.set_pixel(x, y, color.x[offset] * 255, 
                                     color.y[offset] * 255,
                                     color.z[offset] * 255)