(*
This file is part of licht.

licht is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

licht is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with licht.  If not, see <http://www.gnu.org/licenses/>.
*)

open OUnit2

(** Module for the keys *)
module K = struct

  type 'a t = Int : int -> int t [@@unboxed]

  let comp:type a b. a t -> b t -> (a, b) Tools.cmp = fun a b -> begin
    match a, b with Int a, Int b ->

    if a < b then
      Tools.Lt
    else if a > b then
      Tools.Gt
    else
      Tools.Eq
  end

  let repr: type a. Format.formatter -> a t -> unit = fun formatter (Int a) ->
    Format.fprintf formatter "%d" a

end

module SplayTest = Splay.Make(K)

let test_add ctx = begin

  let tree =
     SplayTest.empty
  |> SplayTest.add (K.Int 1) (-1)
  |> SplayTest.add (K.Int 3) 2
  |> SplayTest.add (K.Int 5) 2
  |> SplayTest.add (K.Int 4) 1 in

  let formatter = Format.str_formatter in
  SplayTest.repr formatter tree;
  let str = Format.flush_str_formatter () in

  let expected = "digraph G {\n\
    \"4\" -> \"3\"\n\
    \"3\" -> \"1\"\n\
    \"4\" -> \"5\"\n\
    }" in
  let result = str in

  let msg = Printf.sprintf "Expected %s, but got %s." expected result in

  assert_equal
    ~msg
    expected
    result
end

let test_removeLeaf ctx = begin

  let tree =
     SplayTest.empty
  |> SplayTest.remove (K.Int 4)
  in

  let formatter = Format.str_formatter in
  SplayTest.repr formatter tree;
  let str = Format.flush_str_formatter () in

  let expected = "digraph G {}" in
  let result = str in

  let msg = Printf.sprintf "Expected %s, but got %s." expected result in

  assert_equal
    ~msg
    expected
    result
end

let test_removeTop ctx = begin

  let tree =
     SplayTest.empty
  |> SplayTest.add (K.Int 1) (-1)
  |> SplayTest.add (K.Int 3) 2
  |> SplayTest.add (K.Int 5) 2
  |> SplayTest.add (K.Int 4) 1
  |> SplayTest.remove (K.Int 4)
  in

  let formatter = Format.str_formatter in
  SplayTest.repr formatter tree;
  let str = Format.flush_str_formatter () in

  let expected = "digraph G {\n\
    \"3\" -> \"1\"\n\
    \"3\" -> \"5\"\n\
    }" in
  let result = str in

  let msg = Printf.sprintf "Expected %s, but got %s." expected result in

  assert_equal
    ~msg
    expected
    result
end

let test_removeLeft ctx = begin
  let tree =
     SplayTest.empty
  |> SplayTest.add (K.Int 1) (-1)
  |> SplayTest.add (K.Int 3) 2
  |> SplayTest.add (K.Int 5) 2
  |> SplayTest.add (K.Int 4) 1
  |> SplayTest.remove (K.Int 3)
  in

  let formatter = Format.str_formatter in
  SplayTest.repr formatter tree;
  let str = Format.flush_str_formatter () in

  let expected = "digraph G {\n\
    \"1\" -> \"4\"\n\
    \"4\" -> \"5\"\n\
    }" in
  let result = str in

  let msg = Printf.sprintf "Expected %s, but got %s." expected result in

  assert_equal
    ~msg
    expected
    result
end

(** Remove a value not present in the tree.
  The tree shall be rebalanced. *)
let test_removeOutMax ctx = begin
  let tree =
     SplayTest.empty
  |> SplayTest.add (K.Int 1) (-1)
  |> SplayTest.add (K.Int 3) 2
  |> SplayTest.add (K.Int 5) 2
  |> SplayTest.add (K.Int 4) 1
  |> SplayTest.remove (K.Int 6)
  in

  let formatter = Format.str_formatter in
  SplayTest.repr formatter tree;
  let str = Format.flush_str_formatter () in

  let expected = "digraph G {\n\
    \"5\" -> \"4\"\n\
    \"4\" -> \"3\"\n\
    \"3\" -> \"1\"\n\
    }" in
  let result = str in

  let msg = Printf.sprintf "Expected %s, but got %s." expected result in

  assert_equal
    ~msg
    expected
    result
end

(** Remove the maximum value in the tree *)
let test_removeMax ctx = begin
  let tree =
     SplayTest.empty
  |> SplayTest.add (K.Int 1) (-1)
  |> SplayTest.add (K.Int 3) 2
  |> SplayTest.add (K.Int 5) 2
  |> SplayTest.add (K.Int 4) 1
  |> SplayTest.remove (K.Int 5)
  in

  let formatter = Format.str_formatter in
  SplayTest.repr formatter tree;
  let str = Format.flush_str_formatter () in

  let expected = "digraph G {\n\
    \"4\" -> \"3\"\n\
    \"3\" -> \"1\"\n\
    }" in
  let result = str in

  let msg = Printf.sprintf "Expected %s, but got %s." expected result in

  assert_equal
    ~msg
    expected
    result
end

let tests = "splay_test">::: [
  "add"           >:: test_add;
  "removeLeaf"    >:: test_removeLeaf;
  "removeTop"     >:: test_removeTop;
  "removeLeft"    >:: test_removeLeft;
  "removeMax"     >:: test_removeMax;
  "removeOutMax"  >:: test_removeOutMax;
]