import taichi as ti                                                                # of course, we need taichi
from taichi.math import *                                                # need common mathematical operations

ti.init(arch=ti.gpu, default_ip=ti.i32, default_fp=ti.f32)        # initialize, use GPU, set default ip and fp

image_resolution = (512, 512)                                         # resolution of the image, not too large
image_buffer = ti.Vector.field(4, float, image_resolution)    # image buffer field for recording sample counts
image_pixels = ti.Vector.field(3, float, image_resolution)           # for output display pixels to the screen

Ray = ti.types.struct(origin=vec3, direction=vec3, color=vec3)      # the struct representing camera light ray
Material = ti.types.struct(albedo=vec3, emission=vec3)             # Cornell Box need only albedo and emission
Transform = ti.types.struct(position=vec3, rotation=vec3, scale=vec3)          # Transformation of SDF objects
SDFObject = ti.types.struct(distance=float, transform=Transform, material=Material)              # SDF objects
HitRecord = ti.types.struct(object=SDFObject, position=vec3, distance=float, hit=bool)   # for ray-hit-surface

objects    = SDFObject.field(shape=8)                           # field for storing SDF objects with 8 objects
objects[0] = SDFObject(transform=Transform(vec3(0, 0, -1), vec3(0, 0, 0), vec3(1, 1, 0.2)),
                material=Material(vec3(1, 1, 1)*0.4, vec3(1)))                                        # wall 1
objects[1] = SDFObject(transform=Transform(vec3(0, 1, 0), vec3(90, 0, 0), vec3(1, 1, 0.2)),
                material=Material(vec3(1, 1, 1)*0.4, vec3(1)))                                        # wall 2
objects[2] = SDFObject(transform=Transform(vec3(0, -1, 0), vec3(90, 0, 0), vec3(1, 1, 0.2)),
                material=Material(vec3(1, 1, 1)*0.4, vec3(1)))                                        # wall 3
objects[3] = SDFObject(transform=Transform(vec3(-1, 0, 0), vec3(0, 90, 0), vec3(1, 1, 0.2)),
                material=Material(vec3(1, 0, 0)*0.5, vec3(1)))                                        # wall 4
objects[4] = SDFObject(transform=Transform(vec3(1, 0, 0), vec3(0, 90, 0), vec3(1, 1, 0.2)),
                material=Material(vec3(0, 1, 0)*0.5, vec3(1)))                                        # wall 5
objects[5] = SDFObject(transform=Transform(vec3(-0.275, -0.3, -0.2), vec3(0, 112, 0), vec3(0.25, 0.5, 0.25)),
                material=Material(vec3(1, 1, 1)*0.4, vec3(1)))                                    # taller box
objects[6] = SDFObject(transform=Transform(vec3(0.275,-0.55, 0.2), vec3(0, -197, 0), vec3(0.25, 0.25, 0.25)),
                material=Material(vec3(1, 1, 1)*0.4, vec3(1)))                                           # box
objects[7] = SDFObject(transform=Transform(vec3(0, 0.809, 0), vec3(90, 0, 0), vec3(0.2, 0.2, 0.01)),
                material=Material(vec3(1, 1, 1)*1, vec3(100)))                                         # light

@ti.func
def angle(a: vec3) -> mat3:                                          # convert Euler angles to rotation matrix
    s, c = sin(a), cos(a)                                          # first calculate the two axial projections
    return mat3(c.z, s.z, 0, -s.z, c.z, 0, 0, 0, 1) @ \
           mat3(c.y, 0, -s.y, 0, 1, 0, s.y, 0, c.y) @ \
           mat3(1, 0, 0, 0, c.x, s.x, 0, -s.x, c.x)      # convert to rotation matrix in XYZ and multiply left

@ti.func
def signed_distance(obj: SDFObject, pos: vec3) -> float:     # calc the signed distance from pos to SDF object
    p = angle(radians(obj.transform.rotation)) @ (pos - obj.transform.position)    # translate and then rotate
    q = abs(p) - obj.transform.scale
    return length(max(q, 0)) + min(max(q.x, max(q.y, q.z)), 0)               # return the sdf value of the box

@ti.func
def nearest_object(p: vec3) -> SDFObject:                                        # find the nearest sdf object
    o = objects[0]; o.distance = abs(signed_distance(o, p))                   # we start with the first object
    for i in range(1, 8):                                                                  # for all 8 objects
        oi = objects[i]; oi.distance = abs(signed_distance(oi, p))     # handling the interior of SDF with abs
        if oi.distance < o.distance: o = oi                  # we need the nearest object to step into the ray
    return o                                             # this can also be seen as a concatenation of the SDF

@ti.func
def calc_normal(obj: SDFObject, p: vec3) -> vec3: # representing the surface normal by the gradient of the SDF
    e = vec2(1, -1) * 0.5773 * 0.005                # calculation of gradients using the Tetrahedron technique
    return normalize(e.xyy * signed_distance(obj, p + e.xyy) + \
                     e.yyx * signed_distance(obj, p + e.yyx) + \
                     e.yxy * signed_distance(obj, p + e.yxy) + \
                     e.xxx * signed_distance(obj, p + e.xxx) )

@ti.func
def raycast(ray: Ray) -> HitRecord:                 # ray marching to obtain the intersection with the surface
    record = HitRecord(distance=0.0005)                                  # step a little off the surface first
    for _ in range(256):                                                      # need a maximum number of steps
        record.position  = ray.origin + record.distance * ray.direction
        record.object    = nearest_object(record.position)    # according to the nearest distance ray marching
        record.distance += record.object.distance                                             # sphere tracing
        record.hit       = record.object.distance < 0.00001         # less than the surface thickness is a hit
        if record.distance > 2000.0 or record.hit: break                        # no need to continue stepping
    return record

@ti.func
def hemispheric_sampling(normal: vec3) -> vec3:           # choose a random direction in the normal hemisphere
    z = 2.0 * ti.random() - 1.0
    a = ti.random() * 2.0 * pi
    xy = sqrt(1.0 - z*z) * vec2(sin(a), cos(a))
    return normalize(normal + vec3(xy, z))

@ti.func
def raytrace(ray: Ray) -> Ray:                                                                  # Path Tracing
    for i in range(3):                   # 3 times is already enough to bring Global Illumination to the scene
        inv_pdf = exp(float(i) / 128.0)
        roulette_prob = 1.0 - (1.0 / inv_pdf)  # Russian Roulette for spreading the computation between frames
        if ti.random() < roulette_prob: ray.color *= roulette_prob; break                     # end of tracing

        record = raycast(ray)                           # calculate the intersection of the ray with the scene
        if not record.hit: ray.color = vec3(0); break                           # not hitting the light source

        normal  = calc_normal(record.object, record.position)     # calc the normal of the intersection points
        ray.direction = hemispheric_sampling(normal)                # approximate diffuse reflection direction
        ray.color *= record.object.material.albedo                  # ray needs to be multiplied by the albedo
        ray.origin = record.position                                         # update light departure position

        intensity  = dot(ray.color, vec3(0.299, 0.587, 0.114))              # calculating the intensity of ray
        ray.color *= record.object.material.emission                                # multiplying the emission
        visible    = dot(ray.color, vec3(0.299, 0.587, 0.114))                                # new brightness
        if intensity < visible or visible < 0.000001: break           # too dark or arrive at the light source
    return ray

@ti.kernel
def render(camera_position: vec3, camera_lookat: vec3, camera_up: vec3):
    for i, j in image_pixels:                                         # iterate through all pixels in parallel
        buffer = image_buffer[i, j]                                                     # current buffer color

        z = normalize(camera_position - camera_lookat)
        x = normalize(cross(camera_up, z))                          # calculating the camera coordinate system
        y = cross(z, x)
        
        half_width = half_height = tan(radians(35) * 0.5)           # calculate camera frame position and size
        lower_left_corner = camera_position - half_width * x - half_height * y - z
        horizontal = 2.0 * half_width  * x
        vertical   = 2.0 * half_height * y

        uv = (vec2(i, j) + vec2(ti.random(), ti.random())) / vec2(image_resolution)             # oversampling
        po = lower_left_corner + uv.x * horizontal + uv.y * vertical
        rd = normalize(po - camera_position)                          # calculation of ray direction by camera

        ray = raytrace(Ray(camera_position, rd, vec3(1)))                                       # Path Tracing
        buffer += vec4(ray.color, 1.0)              # accumulate colors and record the number of accumulations
        image_buffer[i, j] = buffer                                                      # updating the buffer

        color = buffer.rgb / buffer.a                                  # calculate the average value of colors
        color = pow(color, vec3(1.0 / 2.2))                     # Gamma correction, then use ACES tone mapping
        color = mat3(0.597190, 0.35458, 0.04823, 0.07600, 0.90834, 0.01566, 0.02840, 0.13383, 0.83777) @ color
        color = (color * (color + 0.024578) - 0.0000905) / (color * (0.983729 * color + 0.4329510) + 0.238081)
        color = mat3(1.60475, -0.531, -0.0736, -0.102, 1.10813, -0.00605, -0.00327, -0.07276, 1.07602) @ color
        image_pixels[i, j] = clamp(color, 0, 1)  # write pixels, clamp the brightness that cannot be displayed

def main():
    window = ti.ui.Window("Cornell Box", image_resolution)                                     # create window
    canvas = window.get_canvas()
    while window.running:                                                            # main loop of the window
        render(vec3(0, 0, 3.5), vec3(0, 0, -1), vec3(0, 1, 0))        # set the camera parameters, then render
        canvas.set_image(image_pixels)                                              # writing pixels to canvas
        window.show()                                                                # continue to show window

if __name__ == '__main__':
    main()                                                                             # By HK-SHAO 2023/01/30
# ------------------------------------------------------------------------------------------------------------