structure LayerTree =
struct
  datatype t = NODE of {key: int, value: Grid.t, left: t, right: t} | LEAF

  val minKey = 1

  fun init maxSide =
    let val grid = Grid.makeEmpty maxSide
    in NODE {key = minKey, value = grid, left = LEAF, right = LEAF}
    end

  fun singleton grid =
    NODE {key = minKey, value = grid, left = LEAF, right = LEAF}

  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
          NODE
            { key = key
            , value = value
            , left = insert (newKey, newValue, left)
            , right = right
            }
        else if newKey > key then
          NODE
            { key = key
            , value = value
            , left = left
            , right = insert (newKey, newValue, right)
            }
        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 foldr (f, tree, acc) =
    case tree of
      LEAF => acc
    | NODE {value, left, right, ...} =>
        let
          val acc = foldr (f, right, acc)
          val acc = f (value, acc)
        in
          foldr (f, left, 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 = newValue, 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 Grid.isBlank valuePixel then
                let val accYAxis = Vector.sub (acc, xIdx)
                in Vector.sub (accYAxis, yIdx)
                end
              else
                valuePixel) valueYAxis) value

  fun flatten (maxSide, tree) =
    foldl (helpFlatten, tree, Grid.makeEmpty maxSide)

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

  fun addPixel (key, newX, newY, maxSide, pixel, tree) =
    let
      val grid =
        case get (key, tree) of
          SOME grid => grid
        | NONE => Grid.makeEmpty maxSide

      val grid = Grid.updateGrid (grid, newX, newY, pixel)
    in
      insert (key, grid, tree)
    end

  fun flipHorizontally tree = map (Grid.flipHorizontally, tree)
end