/// Ray tracer in F#
/// (C) Flying Frog Consultancy Ltd., 2007
/// http://www.ffconsultancy.com

open Math
open Math.Notation
open System.Drawing
open System.Windows.Forms
open System.Threading

let pi = 4. * atan 1.
let delta = sqrt epsilon_float
let sqr x : float = x * x
let clamp l u x = if x < l then l else if x > u then u else x
let vec x y z = vector [x; y; z; 1.]
let zero = vec 0. 0. 0.
let ( .+ ) s r = Vector.create 3 s + r
let dot a b = a.[0] * b.[0] + a.[1] * b.[1] + a.[2] * b.[2]
let length2 r = sqr r.[0] + sqr r.[1] + sqr r.[2]
let length r = sqrt(length2 r)
let unitise r = 1. / length r $* r
let init = Matrix.init 4 4
let translate (r : vector) = init (fun i j -> if i = j then 1. else if j = 3 then r.[i] else 0.)
let identity = Matrix.identity 4
let rot_x t =
  let c = cos t and s = sin t in
  matrix [[ 1.; 0.; 0.; 0.];
          [ 0.;  c;  s; 0.];
          [ 0.; -s;  c; 0.];
          [ 0.; 0.; 0.; 1.]]
let rot_y t =
  let c = cos t and s = sin t in
  matrix [[  c; 0.;  s; 0.];
          [ 0.; 1.; 0.; 0.];
          [ -s; 0.;  c; 0.];
          [ 0.; 0.; 0.; 1.]]
let rot_z t =
  let c = cos t and s = sin t in
  matrix [[  c;  s; 0.; 0.];
          [ -s;  c; 0.; 0.];
          [ 0.; 0.; 1.; 0.];
          [ 0.; 0.; 0.; 1.]]

// Ray tracing primitives
type material = { color: vector; shininess: float }
type sphere = { center: vector; radius: float }
type obj = Sphere of sphere * material | Group of sphere * obj list
type ray = { origin: vector; direction: vector }

// Direction of the light
let light = unitise (vec -1. -3. 2.)

// Levels of spheres
let level = 6

// Ratio of one level's radius to the next
let s = 1. /. 3.

// Sky material
let sky_material = { color = vector [0.; 0.1; 0.2]; shininess = 0. }

// Floor material as a function of coordinate
let floor_material x y =
  let z = Complex.mkRect(x, y) in
  let arg = Complex.phase z and s = Complex.magnitude z in
  if int_of_float(s + arg / pi) % 2 = 0 then
    { color = vector [0.8; 1.; 0.7]; shininess = 0.15 }
  else
    { color = vector [0.; 0.; 0.]; shininess = 0. }

// Sphere material as a function of level
let sphere_material n =
  let t = pi / 180. * 15. * float n in
  { color = vector[sin t; sin(t + 2. * pi / 3.); sin(t + 4. * pi / 3.)]; shininess = 0.5 }

// Ray-sphere intersection
let ray_sphere ray sphere =
  let v = sphere.center - ray.origin in
  let b = dot v ray.direction in
  let disc = sqr b - length2 v + sqr sphere.radius in
  if disc < 0. then infinity else
    let disc = sqrt disc in
    let t2 = b + disc in
    if t2 < 0. then infinity else
      let t1 = b - disc in
      if t1 > 0. then t1 else t2

// Intersect a ray with the tree of spheres
let rec intersect_spheres ray (lambda, _, _ as hit) = function
  | Sphere (sphere, material) ->
      let lambda' = ray_sphere ray sphere in
      if lambda' >= lambda then hit else
        let normal = unitise (ray.origin + (lambda' $* ray.direction) - sphere.center) in
        lambda', normal, material
  | Group (bound, scenes) ->
      let lambda' = ray_sphere ray bound in
      if lambda' >= lambda then hit else List.fold_left (intersect_spheres ray) hit scenes

// Intersect a ray with the floor
let intersect_floor ray (lambda, _, _ as hit) =
  let lambda' = -ray.origin.[1] / ray.direction.[1] in
  if ray.direction.[1] > 0. || lambda' >= lambda then hit else
    let r = ray.origin + (lambda' $* ray.direction) in
    lambda', vec 0. 1. 0., floor_material r.[0] r.[2]

// Find the first intersection of the given ray with scene and floor
let intersect ray scene =
  intersect_floor ray (intersect_spheres ray (infinity, zero, sky_material) scene)

(* Trace a single ray *)
let rec ray_trace weight light ray scene =
  let lambda, normal, material = intersect ray scene in
  if lambda = infinity then vector [0.; sqrt(ray.direction.[1]); 1.] else
    let o = ray.origin + (lambda $* ray.direction) + (delta $* normal) in
    (* Recursively examine specular reflections *)
    let color =
      if weight < 0.1 then Vector.create 3 0.5 else
        let d = ray.direction in
        let ray = { origin = o; direction = ray.direction - (2. * dot ray.direction normal $* normal) } in
        material.shininess $* ray_trace (weight * material.shininess) light ray scene in
    (* Calculate the final color, taking account of shadows, specular and
       diffuse reflection. *)
    let s = match intersect { origin = o; direction = zero - light } scene with
        slambda, _, _ when slambda = infinity -> max 0. (-dot normal light)
      | _ -> 0. in
    (s .+ color) .* material.color

// Find the bounding sphere of a given scene
let rec bound b = function
  | Sphere(s, _) -> { b with radius = max b.radius (length(b.center - s.center) + s.radius) }
  | Group(_, scenes) -> List.fold_left bound b scenes

// Build the scene
let scene =
    let rec aux level r m =
      let sphere = { center = m * zero; radius = r } in
      let material = sphere_material level in
      let obj = Sphere (sphere, material) in
      if level = 1 then obj else begin
        let make_top i =
          m * rot_y(pi / 6. + 2. * pi / 3. * float i) * rot_z(pi / 4.) * translate(vec 0. ((1. + s) * r) 0.) in
        let make_bottom i =
          m * rot_y(pi / 3. * float i) * rot_z(110. * pi / 180.) * translate(vec 0. ((1. + s) * r) 0.) in
        
        let objects =
          List.map make_top [0 .. 2] @ List.map make_bottom [0 .. 5] |>
            List.map (aux (level - 1) (s * r)) in
        
        Group (List.fold_left bound sphere objects, obj :: objects)
      end in
    aux level 1. (translate (vec 0. 1. 0.))

// Render a pixel
let pixel w h x y =
  let ray =
    let o = vec 0. 2. -6. in
    let x, y = float x - 0.5 * float w, float y - 0.5 * float h in
    let d = unitise(vec x y (float (max w h))) in
    let d = rot_x -0.1 * d in
    { origin = o ; direction = d } in
  let c = ray_trace 1. light ray scene in
  let scale x = clamp 0 255 (int_of_float(0.5 + 255. * x)) in
  Color.FromArgb(scale c.[0], scale c.[1], scale c.[2])

// As threads complete rasters they are pushed onto this stack
let rasters = ref (ref [])

// Render a raster and then pop it on the raster stack
let raster r w h y =
  // If this image becomes obsolete then we raise and catch an exception to return immediately
  try
    let data = Array.init w (fun x -> if !rasters != r then raise Exit; pixel w h x y) in
    Idioms.lock !rasters (fun () -> r := (h - 1 - y, data) :: !r)
  with Exit -> ()

// Spawn a thread for each raster
let render (form : #Form) =
  for y = 0 to form.Height - 1 do
    assert(ThreadPool.QueueUserWorkItem(fun _ -> raster !rasters form.Width form.Height y; form.Invalidate()));
  done

type Form1 = class
  inherit Form

  val mutable bitmap : Bitmap

  new() as form = { bitmap = null } then
    // Start rendering rasters
    render form;
    
    // Stop Windows from drawing the background for this form
    form.SetStyle(Enum.combine[ControlStyles.AllPaintingInWmPaint; ControlStyles.Opaque], true);
    
    form.Text <- "Ray tracer";
    form.bitmap <- new Bitmap(form.Width, form.Height, Imaging.PixelFormat.Format24bppRgb);
    form.Show()

  override form.OnPaint e =
    // Copy any rasters that have been rendered into this form's bitmap
    // (the bitmap is only accessible from inside this rendering thread)
    Idioms.lock !rasters (fun () ->
      let draw(y, data) = Array.iteri (fun x c -> try form.bitmap.SetPixel(x, y, c) with _ -> ()) data in
      List.iter draw (! !rasters);
      !rasters := []);
    
    // Draw the bitmap on the form
    e.Graphics.DrawImage(form.bitmap, form.ClientRectangle, new Rectangle(0, 0, form.Width, form.Height), GraphicsUnit.Pixel)

  override form.OnResize e =
    // Reset the raster and replace the bitmap
    rasters := ref [];
    form.bitmap <- new Bitmap(form.Width, form.Height, Imaging.PixelFormat.Format24bppRgb);
    render form;
    form.Invalidate()

  override form.OnKeyDown e = if e.KeyCode = Keys.Escape then form.Close()
end

do let form = new Form1() in while form.Created do Thread.Sleep(1); Application.DoEvents() done