(*
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
  type t = int * T.t array array
  (** The type is composed by the number of defined cell in the page, and the page itself *)

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

  let add (x : int) (y : int) (value : T.t) (t : t) : t =
    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)

  let remove (x : int) (y : int) (t : t) : t =
    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))

  let create array_size = (0, Array.make_matrix array_size array_size T.default)

  let fold_line f y init t =
    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

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 ->
      match (a, b) with
      | K (x1, y1), K (x2, y2) ->
          let res = Stdlib.compare (y1, x1) (y2, x2) in
          if res < 0 then Tools.Lt else if res > 0 then Tools.Gt else Tools.Eq

    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 =
    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

  let add (id : cell) (value : T.t) (t : Map.t) : Map.t =
    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

  let remove (id : cell) (t : Map.t) : Map.t =
    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

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

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

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

    (* Call process_line for each block on the same row *)
    let process_pages block_y acc =
      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
    in

    let fold_blocks (current_row, acc) (Map.C key_val) =
      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) ]))
    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