(*
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/>.
*)

type cell = int * int

module type T_DEFAULT = sig

  type t

  val default : t

end

module MapArray(T:T_DEFAULT) = struct

  (** The type is composed by the number of defined cell in the page, and the page itself *)
  type t = int * (T.t array array)

  let find (x:int) (y:int) (t:t) : T.t = begin
    let block = snd t in
    block.(y).(x)
  end

  let add (x:int) (y:int) (value:T.t) (t:t) : t = begin
    let n, block = t in
    let n' =
      if (block.(y).(x) == T.default) then
        n + 1
      else
        n in
    block.(y).(x) <- value;
    n', block
  end

  let remove (x:int) (y:int) (t:t) : t = begin
    let n, block = t in
    if (block.(y).(x) = T.default) then
     t
    else (
      if n = 1 then
        (* Do not keep empty block in memory *)
        raise Not_found
      else (
        block.(y).(x) <- T.default;
        (n -1, block)
      )
    )
  end

  let create array_size =  begin
    0, Array.make_matrix array_size array_size T.default
  end

  let fold_line f y init t = begin
    let n, block = t
    and res = ref init in

    let array_size = Array.length block in
    for x = 0 to (array_size - 1) do
      let value = block.(y).(x) in
      if value != T.default then
        res := f x value !res;
    done;
    !res
  end

end

module SplayMap(T:T_DEFAULT) = struct

  let array_size = 8

  module PageMap = MapArray(T)

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

    type 'a t = K : (int * int) -> PageMap.t t [@@unboxed]

    let comp:type a b. a t -> b t -> (a, b) Tools.cmp = fun a b -> begin
      match a, b with K (x1, y1), K (x2, y2) ->
        let res = Pervasives.compare (y1, x1) (y2, x2) in
        if res < 0 then
          Tools.Lt
        else if res > 0 then
          Tools.Gt
        else
          Tools.Eq
    end

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

  end

  module Map = Splay.Make(K)

  type t = Map.t

  (* Values are always positive *)
  let get_bounded_values (x, y) = (max 0 x), (max 0 y)

  let find (id:cell) (t:Map.t) : T.t = begin
    let x, y = get_bounded_values id in
    let block_x = x / array_size
    and block_y = y / array_size in
    try
      let block = Map.find (K (block_x, block_y)) t in
      PageMap.find (x mod array_size) (y mod array_size) block
    with Not_found -> T.default
  end

  let add (id:cell) (value:T.t) (t:Map.t) : Map.t = begin
    let x, y = get_bounded_values id in
    let block_x = x / array_size
    and block_y = y / array_size in
    let block =
      try Map.find (K (block_x, block_y)) t
      with Not_found -> PageMap.create array_size in
    let page = PageMap.add (x mod array_size) (y mod array_size) value block in
    Map.add (K (block_x, block_y)) page t
  end

  let remove (id:cell) (t:Map.t) : Map.t = begin
    let x, y = get_bounded_values id in
    let block_x = x / array_size
    and block_y = y / array_size in
    try
      let block = Map.find (K (block_x, block_y)) t in
      try
        let block' = PageMap.remove (x mod array_size) (y mod array_size) block in
        Map.add (K (block_x, block_y)) block' t
      with Not_found ->
        Map.remove (K (block_x, block_y)) t
    with Not_found -> t
  end

  (** Empty map *)
  let empty = Map.empty

  (** Fold over the elements in the Map.*)
  let fold f (t:Map.t) init = begin
    let res = ref init in

    let call_function column row x value acc = begin
      f (column + x, row) value acc
    end in

    (* Call process_line for each block on the same row *)
    let process_pages block_y acc = begin
      let blocks = List.rev acc
      and row_index = block_y * array_size in
      for y = 0 to (array_size - 1) do
        let row = row_index + y in
        res := List.fold_left (fun init (column, block) ->
          PageMap.fold_line (call_function column row) y init block
        ) !res blocks;

      done
    end in

    let fold_blocks (current_row, acc) (Map.C key_val) = begin
      match key_val with ((K.K (block_x, block_y)), (block:PageMap.t)) ->
      (* As long as the page lay in the same row, accumulate it *)
      if current_row = block_y then
        current_row, (block_x * array_size, block)::acc
      else (
        (* We apply the function for each accumulated block in the row *)
        process_pages current_row acc;
        block_y, (block_x, block)::[]
      )
    end in

    let row_number, acc = Map.fold fold_blocks (1, []) t in
    (* Apply the function to the last row *)
    process_pages row_number acc;
    !res
  end


end