type interval = int * int

type set = Empty | Node of set * interval * set * int

type t = {cmp: interval -> interval -> int; set: set}

(* returns height of interval set *)
let height = function Node (_, _, _, h) -> h | Empty -> 0

(* Interval set constructor *)
let make l i r = Node (l, i, r, max (height l) (height r) + 1)

let bal l k r =
  let hl = height l in
  let hr = height r in
  if hl > hr + 2 then
    match l with
    | Node (ll, lk, lr, _) -> (
        if height ll >= height lr then make ll lk (make lr k r)
        else
          match lr with
          | Node (lrl, lrk, lrr, _) -> make (make ll lk lrl) lrk (make lrr k r)
          | Empty -> assert false )
    | Empty -> assert false
  else if hr > hl + 2 then
    match r with
    | Node (rl, rk, rr, _) -> (
        if height rr >= height rl then make (make l k rl) rk rr
        else
          match rl with
          | Node (rll, rlk, rlr, _) -> make (make l k rll) rlk (make rlr rk rr)
          | Empty -> assert false )
    | Empty -> assert false
  else Node (l, k, r, max hl hr + 1)

let rec min_elt = function
  | Node (Empty, k, _, _) -> k
  | Node (l, _, _, _) -> min_elt l
  | Empty -> raise Not_found

let rec remove_min_elt = function
  | Node (Empty, _, r, _) -> r
  | Node (l, k, r, _) -> bal (remove_min_elt l) k r
  | Empty -> invalid_arg "PSet.remove_min_elt"

let merge t1 t2 =
  match (t1, t2) with
  | Empty, _ -> t2
  | _, Empty -> t1
  | _ ->
      let i = min_elt t2 in
      bal t1 i (remove_min_elt t2)

(* TODO check if default compare works the same. If you are my reviewer please reminde me about this ;) *)
let interval_compare (a, b) (c, d) = if a = c then compare b d else compare a c

let empty = {cmp= interval_compare; set= Empty}

let are_overlapping a b =
  let check (a1, a2) (b1, b2) =
    (a1 <= b1 && b1 <= a2) || (a1 <= b2 && b2 <= a2)
  in
  check a b || check b a

let are_neighboring a b =
  let check (_, a2) (b1, _) = a2 + 1 = b1 in
  are_overlapping a b || check a b || check b a

let union (a1, a2) (b1, b2) = (min a1 b1, max a2 b2)

let checked (a, b) =
  if a <= b then (a, b)
  else raise (Invalid_argument "invalid interval values order")

let rec add_one cmp x s =
  let rec loop i = function
    | Empty -> (i, Empty)
    | Node (l, v, r, _) ->
        if are_neighboring i v then loop (union i v) (merge l r)
        else
          let c = cmp i v in
          if 0 < c then
            let ir, nr = loop i r in
            (ir, make l v nr)
          else
            let il, nl = loop i l in
            (il, make nl v r)
  in
  let x = checked x in
  match s with
  | Empty -> make Empty x Empty
  | Node (l, i, r, h) ->
      let c = cmp x i in
      let overlap = are_neighboring x i in
      if c = 0 then make l x r
      else if not overlap then
        if c < 0 then
          let nl = add_one cmp x l in
          bal nl i r
        else
          let nr = add_one cmp x r in
          bal l i nr
      else
        let u = union x i in
        let il, nl = loop u l in
        let ir, nr = loop u r in
        bal nl (union u (union il ir)) nr

(*
let rec add_one cmp x = function
  | Node (l, i, r, h) ->
      let c = cmp x i in
      if c = 0 then Node (l, x, r, h)
      else if c < 0 then
        let nl = add_one cmp x l in
        bal nl i r
      else
        let nr = add_one cmp x r in
        bal l i nr
  | Empty -> make Empty x Empty *)

let add x {cmp; set} = {cmp; set= add_one cmp x set}

let rec join cmp l v r =
  match (l, r) with
  | Empty, _ -> add_one cmp v r
  | _, Empty -> add_one cmp v l
  | Node (ll, lv, lr, lh), Node (rl, rv, rr, rh) ->
      if lh > rh + 2 then bal ll lv (join cmp lr v r)
      else if rh > lh + 2 then bal (join cmp l v rl) rv rr
      else make l v r

let split x {cmp; set} =
  let rec loop x = function
    | Empty -> (Empty, false, Empty)
    | Node (l, v, r, _) ->
        let c = cmp x v in
        if c = 0 then (l, true, r)
        else if c < 0 then
          let ll, pres, rl = loop x l in
          (ll, pres, join cmp rl v r)
        else
          let lr, pres, rr = loop x r in
          (join cmp l v lr, pres, rr)
  in
  let setl, pres, setr = loop x set in
  ({cmp; set= setl}, pres, {cmp; set= setr})

let intersection (a, b) (c, d) =
  if a > d || c > b then None else Some (max a c, min b d)

let difference a b =
  match intersection a b with
  | None -> (Some a, None)
  | Some i ->
      if a = i then (None, None)
      else
        let (a1, a2), (i1, i2) = (a, i) in
        if a1 = i1 then (None, Some (i2 + 1, a2))
        else if a2 = i2 then (Some (a1, i1 - 1), None)
        else (Some (a1, i1 - 1), Some (i2 + 1, a2))

let remove x {cmp; set} =
  let rec loop = function
    | Empty -> Empty
    | Node (l, i, r, _) -> (
        let overlap = are_overlapping x i in
        if not overlap then
          let c = cmp x i in
          if c < 0 then bal (loop l) i r else bal l i (loop r)
        else
          match difference i x with
          | Some a, Some b -> add_one cmp a (bal l b r)
          | Some a, None -> bal l a (loop r)
          | None, Some b -> bal (loop l) b r
          | None, None -> merge (loop l) (loop r) )
  in
  {cmp; set= loop set}

(* let remove x { cmp = cmp; set = set } =
  let rec loop = function
    | Node (l, k, r, _) ->
        let c = cmp x k in
        if c = 0 then merge l r else
        if c < 0 then bal (loop l) k r else bal l k (loop r)
    | Empty -> Empty in
  { cmp = cmp; set = loop set } *)

let mem x {cmp; set} =
  let x = (x, x) in
  let rec loop = function
    | Node (l, k, r, _) ->
        let c = cmp x k in
        are_overlapping x k || loop (if c < 0 then l else r)
    | Empty -> false
  in
  loop set

let iter f {set} =
  let rec loop = function
    | Empty -> ()
    | Node (l, k, r, _) -> loop l ; f k ; loop r
  in
  loop set

let fold f {cmp; set} acc =
  let rec loop acc = function
    | Empty -> acc
    | Node (l, k, r, _) -> loop (f k (loop acc l)) r
  in
  loop acc set

let elements {set} =
  let rec loop acc = function
    | Empty -> acc
    | Node (l, k, r, _) -> loop (k :: loop acc r) l
  in
  loop [] set

let is_empty s = s = Empty

let lenght (a, b) =
  let l = 1 + b - a in
  if l < 0 then max_int else l

let below x t =
  let below_interval (a, b) x =
    if are_overlapping (x, x) (a, b) then lenght (a, x) else 0
  in
  let helper i acc =
    if acc = max_int then max_int else acc + below_interval i x
  in
  fold helper t 0

let split x ({cmp} as t) =
  let helper (a, b) (l, p, r) =
    if are_overlapping (x, x) (a, b) then
      (add_one cmp (a, x - 1) l, true, add_one cmp (x + 1, b) r)
    else if x < a then (l, p, add_one cmp (a, b) r)
    else (add_one cmp (a, b) l, p, r)
  in
  let l, p, r = fold helper t (Empty, false, Empty) in
  ({cmp; set= l}, p, {cmp; set= r})