structure LayerTree =
struct
  type square = {r: int, g: int, b: int, a: int}

  type grid = square vector vector

  datatype tree =
    NODE of {key: int, value: grid, left: tree, right: tree}
  | LEAF

  val emptyPixel = {r = 0, g = 0, b = 0, a = 0}

  val emptyTree = LEAF

  fun isBlank ({a, ...}: square) = a = 0

  fun insert (newKey, newValue, tree) =
    case tree of
      LEAF => NODE {key = newKey, value = newValue, left = LEAF, right = LEAF}
    | NODE {key, value, left, right} =>
        if newKey < key then
          let val left = insert (newKey, newValue, left)
          in NODE {key = key, value = value, left = left, right = right}
          end
        else if newKey > key then
          let val right = insert (newKey, newValue, right)
          in NODE {key = key, value = value, left = left, right = right}
          end
        else
          NODE {key = key, value = newValue, left = left, right = right}

  fun get (searchKey, tree) =
    case tree of
      LEAF => NONE
    | NODE {key, value, left, right} =>
        if searchKey < key then get (searchKey, left)
        else if searchKey > key then get (searchKey, right)
        else SOME value

  fun foldl (f, tree, acc) =
    case tree of
      LEAF => acc
    | NODE {value, left, right, ...} =>
        let
          val acc = foldl (f, left, acc)
          val acc = f (value, acc)
        in
          foldl (f, right, acc)
        end

  fun map (f, tree) =
    case tree of
      LEAF => LEAF
    | NODE {key, value, left, right} =>
        let
          val left = map (f, left)
          val right = map (f, right)
          val newValue = f value
        in
          NODE {key = key, value = value, left = left, right = right}
        end

  (* copies non-blank pixels in value vector into acc *)
  fun helpFlatten (value, acc) =
    Vector.mapi
      (fn (xIdx, valueYAxis) =>
         Vector.mapi
           (fn (yIdx, valuePixel) =>
              if isBlank valuePixel then
                let val accYAxis = Vector.sub (acc, xIdx)
                in Vector.sub (accYAxis, yIdx)
                end
              else
                valuePixel) valueYAxis) value

  fun makeEmptyGrid maxSide =
    Vector.tabulate (maxSide, fn _ =>
      Vector.tabulate (maxSide, fn _ => emptyPixel))

  fun flatten (maxSide, tree) =
    foldl (helpFlatten, tree, makeEmptyGrid maxSide)

  fun helpChangeGridSize maxSide squares =
    Vector.tabulate (maxSide, fn i =>
      if i < Vector.length squares then
        let
          val yAxis = Vector.sub (squares, i)
        in
          Vector.tabulate (maxSide, fn ii =>
            if ii < Vector.length yAxis then Vector.sub (yAxis, ii)
            else emptyPixel)
        end
      else
        Vector.tabulate (maxSide, fn _ => emptyPixel))

  fun changeGridSize (maxSide, tree) =
    let val f = helpChangeGridSize maxSide
    in map (f, tree)
    end
end