(* Demonstrates (affine) abstract types.  Correct. *)

(*
    This program uses a read-write lock array library in, where the
    locks are statically checked capabilities.  It runs an algorithm
    that copies an array a bunch of times, mutates the copies in
    parallel, and recombines the results.  The locks ensure that we
    don't read an array at the same time someone is writing it.
*)

#load "libarraycap"

open AArray

module SkewnessExample = struct
  let sum (a: (float, 't) array) (c: ('t, 'c) readcap) =
    fold (+.) 0.0 a c
  
  let mean (a: (float, 't) array) (c: ('t, 'c) readcap) =
    let (total, c) = sum a c in
      (total /. float_of_int (size a), c)
  
  let stdDev (a: (float, 't) array) (c: ('t, 'c) readcap) =
    let (mean, c) = mean a c in
    let (num, c)  = fold
                      (fun (x: float) (acc: float) ->
                         acc +. (x -. mean) ** 2.0)
                      0.0 a c in
      (sqrt (num /. float_of_int (size a)), c)
  
  let skewness (a: (float, 't) array) (c: ('t, 'c) readcap) =
    let n         = float_of_int (size a) in
    let (m, c)    = mean a c in
    let (s, c)    = stdDev a c in
    let (devs, c) = fold
                      (fun (x: float) (acc: float) ->
                         (x -. m) ** 3.0 +.  acc)
                      0.0 a c in
      (devs /. ((n -. 1.0) *. s ** 3.0), c)
  
  type transformation = T of string * (float -> float)
  
  let reduceSkewness (ts: transformation list)
                     (a: (float, 't) array)
                     (c0: 't writecap) =
    let rec replace (i: int)
                    (T(_, ft) as t: transformation)
                    (c: 't writecap)
                    : 't writecap =
      if i < size a
        then let (x, c) = get a i c in
             let c      = set a i (ft x) c in
               replace (i + 1) t c
        else c in
    (* Need type annotation for polymorphic recursion: *)
    let rec find :
        ∀ 'd. int → transformation list → ('t, 'd) readcap →
              float * transformation * ('t, 'd) readcap =
      λ ix ts c →
        match ts with
        | []  -> let (sk, c) = skewness a c in
                   (sk, T("identity", fun f -> f), c)
        | (T(_, ft) as t) ∷ ts
              -> let ((sk1, t1), (sk2, t2), c) =
                   par
                     (fun (c: ('t, 'c) readcap) -> find (ix + 1) ts c)
                     (fun (c: ('t, 'c) readcap) ->
                       let ((b, d), c) = map ft a c in
                       let (sk, d) = skewness b d in
                         (sk, t, c))
                     c
                  in if absf sk2 <. absf sk1
                       then (sk2, t2, c)
                       else (sk1, t1, c) in
    let (sk, t, c) = find 0 ts c0 in
      (sk, t, replace 0 t c)
  
  let newDistribution (n: int) (T(_, gen): transformation)
                         : ex 't. (float, 't) array * 't writecap =
    let (a, c) = new n 0.0 in
      let rec loop (i: int) c =
        if i < n
          then loop (i + 1) (set a i (gen (float_of_int (i + 1))) c)
          else c in
        (a, loop 0 c)
  
  let functions (n: int) =
    [T("1",         fun (ix: float) -> 1.0)
    ,T("x",         fun (ix: float) -> ix)
    ,T("x^2",       flip ( ** ) 2.0)
    ,T("sqrt x",    sqrt)
    ,T("x^5",       flip ( ** ) 5.0)
    ,T("x^1/5",     flip ( ** ) 0.2)
    ,T("e^x",       ( ** ) 2.718)
    ,T("log x",     log)
    ,T("1/x",       (/.) 1.0)
    ,T("-x",        (-.) (float_of_int n))
    ]
  
  let testCase (n: int) (T(name, _) as t: transformation) =
    let (a, c)               = newDistribution n t in
    let (sk0, c)             = skewness a c in
    let (sk, T(name', _), c) = reduceSkewness (functions n) a c in
      putStrLn ("Distribution:      " ^ name);
      putStrLn ("Original skewness: " ^ string_of sk0);
      putStrLn ("Improved skewness: " ^ string_of sk);
      putStrLn ("Winning function:  " ^ name');
      putStrLn ""
  
  let tests (n: int) =
    foldl (fun (t: transformation) () -> testCase n t)
          () (functions n)
end

in
  SkewnessExample.tests 30