// Source: http://www.tryfsharp.org/create/cpoulain/shared/raytracer.fsx
module RayTrace

open System
open System.Text

type Vector(x: float, y: float, z: float) =
    member this.X = x
    member this.Y = y
    member this.Z = z
    static member (*)(k, (v: Vector)) = Vector(k * v.X, k * v.Y, k * v.Z)

    static member (-)(v1: Vector, v2: Vector) =
        Vector(v1.X - v2.X, v1.Y - v2.Y, v1.Z - v2.Z)

    static member (+)(v1: Vector, v2: Vector) =
        Vector(v1.X + v2.X, v1.Y + v2.Y, v1.Z + v2.Z)

    static member Dot(v1: Vector, v2: Vector) = v1.X * v2.X + v1.Y * v2.Y + v1.Z * v2.Z

    static member Mag(v: Vector) =
        sqrt (v.X * v.X + v.Y * v.Y + v.Z * v.Z)

    static member Norm(v: Vector) =
        let mag = Vector.Mag v

        let div =
            if mag = 0.0 then
                infinity
            else
                1.0 / mag

        div * v

    static member Cross(v1: Vector, v2: Vector) =
        Vector(v1.Y * v2.Z - v1.Z * v2.Y, v1.Z * v2.X - v1.X * v2.Z, v1.X * v2.Y - v1.Y * v2.X)

type Color(r: float, g: float, b: float) =
    static let clamp v = Math.Floor(255.0 * Math.Min(v, 1.0))
    member this.R = r
    member this.G = g
    member this.B = b
    static member Scale(k, v: Color) = Color(k * v.R, k * v.G, k * v.B)

    static member (+)(v1: Color, v2: Color) =
        Color(v1.R + v2.R, v1.G + v2.G, v1.B + v2.B)

    static member (*)(v1: Color, v2: Color) =
        Color(v1.R * v2.R, v1.G * v2.G, v1.B * v2.B)

    static member White = Color(1.0, 1.0, 1.0)
    static member Grey = Color(0.5, 0.5, 0.5)
    static member Black = Color(0.0, 0.0, 0.0)
    static member Background = Color.Black
    static member DefaultColor = Color.Black

    static member ToDrawingColor(c: Color) =
        new Color(clamp c.R, clamp c.G, clamp c.B)

type Camera(pos: Vector, lookAt: Vector) =
    let forward = Vector.Norm(lookAt - pos)
    let down = Vector(0.0, -1.0, 0.0)
    let right = 1.5 * Vector.Norm(Vector.Cross(forward, down))
    let up = 1.5 * Vector.Norm(Vector.Cross(forward, right))
    member c.Pos = pos
    member c.Forward = forward
    member c.Up = up
    member c.Right = right

type Ray =
    {
        Start: Vector
        Dir: Vector
    }

type Surface =
    abstract Diffuse: Vector -> Color
    abstract Specular: Vector -> Color
    abstract Reflect: Vector -> double
    abstract Roughness: double

type Intersection =
    {
        Thing: SceneObject
        Ray: Ray
        Dist: double
    }

and SceneObject =
    abstract Surface: Surface
    abstract Intersect: Ray -> Intersection option
    abstract Normal: Vector -> Vector

let Sphere (center, radius, surface) =
    let radius2 = radius * radius

    { new SceneObject with
        member this.Surface = surface
        member this.Normal pos = Vector.Norm(pos - center)

        member this.Intersect(ray: Ray) =
            let eo = center - ray.Start
            let v = Vector.Dot(eo, ray.Dir)

            let dist =
                if (v < 0.0) then
                    0.0
                else
                    let disc = radius2 - (Vector.Dot(eo, eo) - (v * v))

                    if disc < 0.0 then
                        0.0
                    else
                        v - (sqrt (disc))

            if dist = 0.0 then
                None
            else
                Some
                    {
                        Thing = this
                        Ray = ray
                        Dist = dist
                    }
    }

let Plane (norm, offset, surface) =
    { new SceneObject with
        member this.Surface = surface
        member this.Normal pos = norm

        member this.Intersect(ray) =
            let denom = Vector.Dot(norm, ray.Dir)

            if denom > 0.0 then
                None
            else
                let dist = (Vector.Dot(norm, ray.Start) + offset) / (-denom)

                Some
                    {
                        Thing = this
                        Ray = ray
                        Dist = dist
                    }
    }

type Light =
    {
        Pos: Vector
        Color: Color
    }

type Scene =
    {
        Things: SceneObject list
        Lights: Light list
        Camera: Camera
    }

type RayTracer(screenWidth, screenHeight) =

    let maxDepth = 5

    let Intersections ray scene =
        scene.Things
        |> List.choose (fun sceneObj -> sceneObj.Intersect(ray))
        |> List.sortBy (fun intersection -> intersection.Dist)

    let TestRay (ray, scene) =
        match Intersections ray scene with
        | [] -> None
        | isect :: _ -> Some isect.Dist

    let rec TraceRay (ray, scene, depth: int) =
        match Intersections ray scene with
        | [] -> Color.Background
        | isect :: _ -> Shade isect scene depth

    and Shade isect scene depth =
        let d = isect.Ray.Dir
        let pos = isect.Dist * d + isect.Ray.Start
        let normal = isect.Thing.Normal(pos)
        let reflectDir = d - 2.0 * Vector.Dot(normal, d) * normal

        let naturalcolor =
            Color.DefaultColor
            + GetNaturalColor(isect.Thing, pos, normal, reflectDir, scene)

        let reflectedColor =
            if depth >= maxDepth then
                Color(0.5, 0.5, 0.5)
            else
                GetReflectionColor(isect.Thing, pos + (0.001 * reflectDir), normal, reflectDir, scene, depth)

        naturalcolor + reflectedColor

    and GetReflectionColor (thing: SceneObject, pos, normal: Vector, rd: Vector, scene: Scene, depth: int) =
        Color.Scale(
            thing.Surface.Reflect(pos),
            TraceRay(
                {
                    Start = pos
                    Dir = rd
                },
                scene,
                depth + 1
            )
        )

    and GetNaturalColor (thing, pos, norm, rd, scene) =
        let addLight col (light: Light) =
            let ldis = light.Pos - pos
            let livec = Vector.Norm(ldis)

            let neatIsect =
                TestRay(
                    {
                        Start = pos
                        Dir = livec
                    },
                    scene
                )

            let isInShadow =
                match neatIsect with
                | None -> false
                | Some d -> not (d > Vector.Mag(ldis))

            if isInShadow then
                col
            else
                let illum = Vector.Dot(livec, norm)

                let lcolor =
                    if illum > 0.0 then
                        Color.Scale(illum, light.Color)
                    else
                        Color.DefaultColor

                let specular = Vector.Dot(livec, Vector.Norm(rd))

                let scolor =
                    if specular > 0.0 then
                        Color.Scale(System.Math.Pow(specular, thing.Surface.Roughness), light.Color)
                    else
                        Color.DefaultColor

                col + thing.Surface.Diffuse(pos) * lcolor + thing.Surface.Specular(pos) * scolor

        List.fold addLight Color.DefaultColor scene.Lights

    let GetPoint x y (camera: Camera) =
        let RecenterX x =
            (float x - (float screenWidth / 2.0)) / (2.0 * float screenWidth)

        let RecenterY y =
            -(float y - (float screenHeight / 2.0)) / (2.0 * float screenHeight)

        Vector.Norm(camera.Forward + RecenterX(x) * camera.Right + RecenterY(y) * camera.Up)

    member this.Render(scene, rgb: Color[]) =
        for y = 0 to screenHeight - 1 do
            let stride = y * screenWidth

            for x = 0 to screenWidth - 1 do
                let color =
                    TraceRay(
                        {
                            Start = scene.Camera.Pos
                            Dir = GetPoint x y scene.Camera
                        },
                        scene,
                        0
                    )

                rgb.[x + stride] <- Color.ToDrawingColor color

        rgb

module Surfaces =
    let Shiny =
        { new Surface with
            member s.Diffuse pos = Color.White
            member s.Specular pos = Color.Grey
            member s.Reflect pos = 0.7
            member s.Roughness = 250.0
        }

    let Checkerboard =
        { new Surface with
            member s.Diffuse pos =
                // if (int (Math.Floor(pos.Z) + Math.Floor(pos.X))) % 2 <> 0
                if ((int pos.Z) + (int pos.X)) % 2 <> 0 then
                    Color.White
                else
                    Color.Black

            member s.Specular pos = Color.White

            member s.Reflect pos =
                // if (int (Math.Floor(pos.Z) + Math.Floor(pos.X))) % 2 <> 0
                if ((int pos.Z) + (int pos.X)) % 2 <> 0 then
                    0.1
                else
                    0.7

            member s.Roughness = 150.0
        }

let scene =
    {
        Things =
            [
                Plane(Vector(0.0, 1.0, 0.0), 0.0, Surfaces.Checkerboard)
                Sphere(Vector(0.0, 1.0, -0.25), 1.0, Surfaces.Shiny)
                Sphere(Vector(-1.0, 0.5, 1.5), 0.5, Surfaces.Shiny)
            ]
        Lights =
            [
                {
                    Pos = Vector(-2.0, 2.5, 0.0)
                    Color = Color(0.49, 0.07, 0.07)
                }
                {
                    Pos = Vector(1.5, 2.5, 1.5)
                    Color = Color(0.07, 0.07, 0.49)
                }
                {
                    Pos = Vector(1.5, 2.5, -1.5)
                    Color = Color(0.07, 0.49, 0.071)
                }
                {
                    Pos = Vector(0.0, 3.5, 0.0)
                    Color = Color(0.21, 0.21, 0.35)
                }
            ]
        Camera = Camera(Vector(3.0, 2.0, 4.0), Vector(-1.0, 0.5, 0.0))
    }

// Compute the scene
let computeScene width height =
    let raytracer = RayTracer(width, height)
    let rgbBuffer = Array.zeroCreate (width * height)
    let colors = raytracer.Render(scene, rgbBuffer)
    colors
// |> Array.fold (fun (sb:StringBuilder) c -> sb.AppendFormat("#{0:x2}{1:x2}{2:x2}", int c.R, int c.G, int c.B)) (new StringBuilder())
// |> fun s -> TryFSharp.Canvas.JavaScriptFunction("RayTracer.render").Invoke(s.ToString())
// |> ignore

let measure f x y =
    let dtStart = DateTime.UtcNow
    let res = f x y
    let elapsed = DateTime.UtcNow - dtStart
    res, elapsed.TotalSeconds

let x, y = 100, 100

[ 1..3 ]
|> List.iter (fun i ->
    let colors, elapsed = measure computeScene x y

    printfn "run %d: Ray tracing scene size (%d,%d), elapsed %f sec" i x y elapsed
)
