structure CollisionTree =
struct
  structure BinTree =
  struct
    datatype 'a bintree =
      NODE of
        { x: int
        , y: int
        , ex: int
        , ey: int
        , data: 'a
        , left: 'a bintree
        , right: 'a bintree
        }
    | LEAF

    val empty = LEAF

    fun insert (newItem as {x, y, ex, ey, data}, tree) =
      case tree of
        LEAF =>
          NODE
            { x = x
            , y = y
            , ex = ex
            , ey = ey
            , data = data
            , left = LEAF
            , right = LEAF
            }
      | NODE {x = ox, y = oy, ex = oex, ey = oey, data = oldData, left, right} =>
          let
            val dir =
              if x < ox then
                LESS
              else if x > ox then
                GREATER
              else
                (if y < oy then
                   LESS
                 else if y > oy then
                   GREATER
                 else
                   (if ex < oex then
                      LESS
                    else if ex > oex then
                      GREATER
                    else
                      (if ey < oey then LESS
                       else if ey > oey then GREATER
                       else EQUAL)))
          in
            case dir of
              LESS =>
                NODE
                  { left = insert (newItem, left)
                  , right = right
                  , x = ox
                  , y = oy
                  , ex = oex
                  , ey = oey
                  , data = oldData
                  }
            | GREATER =>
                NODE
                  { right = insert (newItem, right)
                  , left = left
                  , x = ox
                  , y = oy
                  , ex = oex
                  , ey = oey
                  , data = oldData
                  }
            | EQUAL =>
                NODE
                  { left = left
                  , right = right
                  , x = x
                  , y = y
                  , ex = ex
                  , ey = ey
                  , data = data
                  }
          end

    fun foldr (f, tree, acc) =
      case tree of
        NODE {x, y, ex, ey, data, left, right} =>
          let
            val acc = foldr (f, right, acc)
            val item = {x = x, y = y, ex = ex, ey = ey, data = data}
            val acc = f (item, acc)
          in
            foldr (f, left, acc)
          end
      | LEAF => acc

    fun toList tree =
      foldr (fn (item, acc) => item :: acc, tree, [])
  end

  fun shouldIgnoreData {a, r = _, g = _, b = _} = a = 0

  local
    fun loopYAxis (x, y, eX, eY, yAxis, col) =
      if y > eY orelse y >= Vector.length yAxis then
        true
      else
        let
          val newCol = Vector.sub (yAxis, y)
        in
          if col = newCol then loopYAxis (x, y + 1, eX, eY, yAxis, col)
          else false
        end

    fun loopColour (x, y, eX, eY, grid, col) =
      if x > eX orelse x >= Vector.length grid then
        true
      else
        let
          val yAxis = Vector.sub (grid, x)
        in
          if loopYAxis (x, y, eX, eY, yAxis, col) then
            loopColour (x + 1, y, eX, eY, grid, col)
          else
            false
        end
  in
    fun quadHasSameColour (startX, startY, endX, endY, grid) =
      let
        val yAxis = Vector.sub (grid, startX)
        val col = Vector.sub (yAxis, startY)
      in
        loopColour (startX, startY, endX, endY, grid, col)
      end
  end

  (* tree creation *)
  fun build (x, y, size, grid, bintree) =
    if x >= Vector.length grid orelse y >= Vector.length grid then
      bintree
    else if quadHasSameColour (x, y, x + size, y + size, grid) then
      let
        val yAxis = Vector.sub (grid, x)
        val data = Vector.sub (yAxis, y)
      in
        if shouldIgnoreData data then
          bintree
        else
          let
            val ex = x + size
            val ey = y + size
            val ex = Int.min (ex, Vector.length grid - 1)
            val ey = Int.min (ey, Vector.length grid - 1)
            val item = {x = x, y = y, ex = ex, ey = ey, data = data}
          in
            BinTree.insert (item, bintree)
          end
      end
    else
      (if size mod 2 = 0 orelse size = 1 then
         let
           val halfSize = size div 2
           val bintree = build (x, y, halfSize, grid, bintree)
           val bintree = build (x + halfSize, y, halfSize, grid, bintree)
           val bintree = build (x, y + halfSize, halfSize, grid, bintree)
         in
           build (x + halfSize, y + halfSize, halfSize, grid, bintree)
         end
       else
         (* handles odd-number divisions.
          * For example, `7 div 2` is 3 because of integer division. 
          * We would not cover every pixel unless we handle odd numbers specially. *)
         let
           val halfSizeBefore = size div 2
           val halfSizeAfter = (size + 1) div 2
           val bintree = build (x, y, halfSizeAfter, grid, bintree)
           val bintree = build
             (x + halfSizeBefore, y, halfSizeAfter, grid, bintree)
           val bintree = build
             (x, y + halfSizeBefore, halfSizeAfter, grid, bintree)
         in
           build
             ( x + halfSizeBefore
             , y + halfSizeBefore
             , halfSizeAfter
             , grid
             , bintree
             )
         end)

  fun getClickPoint (clickPoints, pos) =
    let val idx = Int.min (pos, Vector.length clickPoints - 1)
    in Vector.sub (clickPoints, idx)
    end

  fun folder
    ( windowWidth
    , windowHeight
    , canvasWidth
    , canvasHeight
    , xClickPoints
    , yClickPoints
    ) ({x, ex, y, ey, data}, acc) =
    let
      val ex = if ex = x then x + 1 else ex + 1
      val ey = if ey = y then y + 1 else ey + 1

      val x = getClickPoint (xClickPoints, x)
      val y = getClickPoint (yClickPoints, y)
      val ex = getClickPoint (xClickPoints, ex)
      val ey = getClickPoint (yClickPoints, ey)

      val startX = Ndc.fromPixelX (x, windowWidth, windowHeight)
      val endX = Ndc.fromPixelX (ex, windowWidth, windowHeight)
      val startY = Ndc.fromPixelY (y, windowWidth, windowHeight)
      val endY = Ndc.fromPixelY (ey, windowWidth, windowHeight)

      val {r, g, b, a} = data
      val r = Real32.fromInt r / 255.0
      val g = Real32.fromInt g / 255.0
      val b = Real32.fromInt b / 255.0
      val a = Real32.fromInt a / 255.0
    in
      Ndc.ltrbToVertexRgb (startX, startY, endX, endY, r, g, b) :: acc
    end

  fun toTriangles
    ( windowWidth
    , windowHeight
    , squares
    , size
    , canvasWidth
    , canvasHeight
    , xClickPoints
    , yClickPoints
    ) =
    let
      val bintree = build (0, 0, size, squares, BinTree.empty)

      val f = folder
        ( windowWidth
        , windowHeight
        , canvasWidth
        , canvasHeight
        , xClickPoints
        , yClickPoints
        )
      val vec = BinTree.foldr (f, bintree, [])
    in
      Vector.concat vec
    end

  (* building and querying quad tree, plus compression *)
  datatype quad_tree =
    LEAF of {x: int, y: int, ex: int, ey: int, data: AppType.square}
  | NODE of {tl: quad_tree, tr: quad_tree, bl: quad_tree, br: quad_tree}
  | EMPTY

  fun foldWithDuplicates (f, tree, acc) =
    case tree of
      EMPTY => acc
    | LEAF item => f (item, acc)
    | NODE {tl, tr, bl, br} =>
        let
          val acc = foldWithDuplicates (f, tl, acc)
          val acc = foldWithDuplicates (f, tr, acc)
          val acc = foldWithDuplicates (f, bl, acc)
        in
          foldWithDuplicates (f, br, acc)
        end

  fun toBintree qtree =
    foldWithDuplicates (BinTree.insert, qtree, BinTree.empty)

  fun buildTree (x, y, size, grid) =
    if x >= Vector.length grid orelse y >= Vector.length grid then
      EMPTY
    else if quadHasSameColour (x, y, x + size, y + size, grid) then
      let
        val yAxis = Vector.sub (grid, x)
        val data = Vector.sub (yAxis, y)
      in
        if shouldIgnoreData data then
          EMPTY
        else
          let
            val ex = x + size
            val ex = Int.min (ex, Vector.length grid - 1)
            val ey = y + size
            val ey = Int.min (ey, Vector.length grid - 1)
          in
            LEAF {x = x, y = y, ex = ex, ey = ey, data = data}
          end
      end
    else
      (if size mod 2 = 0 orelse size = 1 then
         let
           val halfSize = size div 2
           val tl = buildTree (x, y, halfSize, grid)
           val tr = buildTree (x + halfSize, y, halfSize, grid)
           val bl = buildTree (x, y + halfSize, halfSize, grid)
           val br = buildTree (x + halfSize, y + halfSize, halfSize, grid)
         in
           NODE {tl = tl, tr = tr, bl = bl, br = br}
         end
       else
         (* handles odd-number divisions.
          * For example, `7 div 2` is 3 because of integer division. 
          * We would not cover every pixel unless we handle odd numbers specially. *)
         let
           val halfSizeBefore = size div 2
           val halfSizeAfter = (size + 1) div 2
           val tl = buildTree (x, y, halfSizeAfter, grid)
           val tr = buildTree (x + halfSizeBefore, y, halfSizeAfter, grid)
           val bl = buildTree (x, y + halfSizeBefore, halfSizeAfter, grid)
           val br =
             buildTree
               (x + halfSizeBefore, y + halfSizeBefore, halfSizeAfter, grid)
         in
           NODE {tl = tl, tr = tr, bl = bl, br = br}
         end)

  local
    fun loop (x, y, ex, ey, grid) =
      if x < 0 then
        0
      else if quadHasSameColour (x, y, ex, ey, grid) then
        loop (x - 1, y, x, ey, grid)
      else
        ex
  in
    fun getLeftmostX ({x, y, ex, ey, data}, grid) =
      loop (x - 1, y, x, ey, grid)
  end

  local
    fun loop (x, y, ex, ey, grid) =
      if x < Vector.length grid andalso ex < Vector.length grid then
        if quadHasSameColour (x, y, ex, ey, grid) then
          loop (ex, y, ex + 1, ey, grid)
        else
          x
      else
        Vector.length grid - 1
  in
    fun getRightmostX ({x, y, ex, ey, data}, grid) =
      loop (ex, y, ex + 1, ey, grid)
  end

  local
    fun loop (x, y, ex, ey, grid) =
      if y < 0 then
        0
      else if quadHasSameColour (x, y, ex, ey, grid) then
        loop (x, y - 1, ex, y, grid)
      else
        ey
  in
    fun getTopmostY ({x, y, ex, ey, data}, grid) =
      if y < 0 orelse ey <= 0 then
        0
      else if quadHasSameColour (x, y, ex, ey, grid) then
        loop (x, y - 1, ex, y, grid)
      else
        y
  end

  local
    fun loop (x, y, ex, ey, grid) =
      if y < Vector.length grid andalso ey < Vector.length grid then
        if quadHasSameColour (x, y, ex, ey, grid) then
          loop (x, ey, ex, ey + 1, grid)
        else
          y
      else
        Vector.length grid
  in
    fun getBottomY ({x, y, ex, ey, data}, grid) =
      if quadHasSameColour (x, y, ex, ey, grid) then loop (x, y, ex, ey, grid)
      else y
  end

  local
    fun loop (tree, grid) =
      case tree of
        EMPTY => (EMPTY, false)
      | LEAF (oldItem as {x, y, ex, ey, data}) =>
          let
            val topY = getTopmostY (oldItem, grid)
            val bottomY = getBottomY (oldItem, grid)
            val newItem = {y = topY, ey = bottomY, x = x, ex = ex, data = data}
            val didItemChange = newItem <> oldItem
          in
            (LEAF newItem, didItemChange)
          end
      | NODE {tl, tr, bl, br} =>
          let
            val (tl, didTlChange) = loop (tl, grid)
            val (tr, didTrChange) = loop (tr, grid)
            val (bl, didBlChange) = loop (bl, grid)
            val (br, didBrChange) = loop (br, grid)

            val node = NODE {tl = tl, tr = tr, bl = bl, br = br}
            val didChange =
              didTlChange orelse didTrChange orelse didBlChange
              orelse didBrChange
          in
            (node, false)
          end
  in
    fun mergeVertical (tree, grid) =
      let val (newTree, didChange) = loop (tree, grid)
      in if didChange then mergeVertical (newTree, grid) else newTree
      end
  end

  local
    fun loop (tree, grid) =
      case tree of
        EMPTY => (EMPTY, false)
      | LEAF (oldItem as {x, y, ex, ey, data}) =>
          let
            val leftX = getLeftmostX (oldItem, grid)
            val rightX = getRightmostX (oldItem, grid)
            val newItem = {x = leftX, ex = rightX, y = y, ey = ey, data = data}
            val didItemChange = newItem <> oldItem
          in
            (LEAF newItem, didItemChange)
          end
      | NODE {tl, tr, bl, br} =>
          let
            val (tl, didTlChange) = loop (tl, grid)
            val (tr, didTrChange) = loop (tr, grid)
            val (bl, didBlChange) = loop (bl, grid)
            val (br, didBrChange) = loop (br, grid)

            val node = NODE {tl = tl, tr = tr, bl = bl, br = br}
            val didChange =
              didTlChange orelse didTrChange orelse didBlChange
              orelse didBrChange
          in
            (node, didChange)
          end
  in
    fun mergeHorizontal (tree, grid) =
      let val (newTree, didChange) = loop (tree, grid)
      in if didChange then mergeHorizontal (newTree, grid) else newTree
      end
  end

  fun merge (tree, grid) =
    let
      val tree = mergeVertical (tree, grid)
      val tree = mergeHorizontal (tree, grid)
    in
      toBintree tree
    end

  fun toSaveStringFolder ({x, ex, y, ey, data = {r, g, b, a}}, acc) =
    let
      val item = String.concat
        [ "{"
        , Int.toString x
        , " "
        , Int.toString y
        , " "
        , Int.toString ex
        , " "
        , Int.toString ey
        , " "
        , Int.toString r
        , " "
        , Int.toString g
        , " "
        , Int.toString b
        , " "
        , Int.toString a
        , " } "
        ]
    in
      item :: acc
    end

  fun toSaveString (squares, canvasWidth, canvasHeight) =
    let
      val size = Int.max (canvasWidth, canvasHeight)
      val qtree = buildTree (0, 0, size, squares)
      val bintree = merge (qtree, squares)

      val initial = ["}"]
      val acc = BinTree.foldr (toSaveStringFolder, bintree, initial)
      val acc =
        String.concat
          [Int.toString canvasWidth, " ", Int.toString canvasHeight, " { "]
        :: acc
    in
      String.concat acc
    end

  fun intToRealString num =
    let
      val result = Real.fromInt num
      val result = Real.fmt (StringCvt.FIX (SOME 15)) result
    in
      if String.isSubstring "." result then result else result ^ ".0"
    end

  fun colToRealString col =
    let
      val result = Real.fromInt col / 255.0
      val result = Real.fmt (StringCvt.FIX (SOME 15)) result
    in
      if String.isSubstring "." result then result else result ^ ".0"
    end

  fun makeXString x =
    let val x = intToRealString x
    in "xToNdc (xOffset, " ^ x ^ ", scale, halfWidth)"
    end

  fun makeEndXString (startX, endX) =
    let val endX = intToRealString endX
    in "endXToNdc (xOffset, " ^ startX ^ ", " ^ endX ^ ", scale, halfWidth)"
    end

  fun makeYString y =
    let val y = intToRealString y
    in "yToNdc (yOffset, " ^ y ^ ", scale, halfHeight)"
    end

  fun makeEndYString (startY, endY) =
    let val endY = intToRealString endY
    in "endYToNdc (yOffset, " ^ startY ^ ", " ^ endY ^ ", scale, halfHeight)"
    end

  fun toExportStringFolder ({x, ex, y, ey, data = {r, g, b, a}}, acc) =
    let
      val ey = if ey = y then y + 1 else ey

      val x = makeXString x
      val y = makeYString y
      val ex = makeEndXString (x, ex)
      val ey = makeYString ey
      val r = colToRealString r
      val g = colToRealString g
      val b = colToRealString b

      (* based on triangle order formed by `Ndc.ltrbToVertexRgb` function *)
      val item = String.concatWith ",\n"
        [ x
        , ey
        , r
        , g
        , b
        , ex
        , ey
        , r
        , g
        , b
        , x
        , y
        , r
        , g
        , b
        , x
        , y
        , r
        , g
        , b
        , ex
        , ey
        , r
        , g
        , b
        , ex
        , y
        , r
        , g
        , b
        ]
    in
      item :: acc
    end

  fun toExportString (squares, canvasWidth, canvasHeight) =
    let
      val size = Int.max (canvasWidth, canvasHeight)
      val qtree = buildTree (0, 0, size, squares)
      val bintree = merge (qtree, squares)

      val coords = BinTree.foldr (toExportStringFolder, bintree, [])
      val coords = String.concatWith ",\n" coords
    in
      String.concat
        [ "structure AAA = \nstruct\n"

        , "  fun xToNdc (xOffset, xpos, scale, halfWidth) =\n"
        , "    ((xpos * scale + xOffset) - halfWidth) / halfWidth\n\n"

        , "  fun endXToNdc (xOffset, startX, endX, scale, halfWidth) =\n"
        , "    (((endX - startX) * scale + xOffset) - halfWidth) / halfWidth\n\n"

        , "  fun yToNdc (yOffset, ypos, scale, halfHeight) =\n"
        , "   ~(((ypos * scale + yOffset) - halfHeight) / halfHeight)\n\n"

        , "  fun endYToNdc (yOffset, startY, endY, scale, halfHeight) =\n"
        , "   ~((((endY - startY) * scale + yOffset) - halfHeight) / halfHeight)\n\n"

        , "  fun lerp (xOffset, yOffset, scale, windowWidth, windowHeight) =\n"
        , "    let\n"
        , "      val windowWidth = Real32.fromInt windowWidth\n"
        , "      val halfWidth = windowWidth / 2.0\n"
        , "      val windowHeight = Real32.fromInt windowHeight\n"
        , "      val halfHeight = windowHeight / 2.0\n"
        , "    in\n"
        , "     #[\n"
        , coords
        , "\n"
        , "      ]\n"
        , "    end\n"
        , "end\n"
        ]
    end

  (* functions for exporting a collision detection string *)
  fun mapItem (item as {r, g, b, a}) =
    if shouldIgnoreData item then item else {r = 1, g = 1, b = 1, a = 1}

  fun mapGrid grid =
    Vector.map (fn yAxis => Vector.map (fn item => mapItem item) yAxis) grid

  fun toCollisionStringFolder ({x, ex, y, ey, data = _}, acc) =
    let
      val ex = if x = ex then ex + 1 else ex
      val ey = if y = ey then ey + 1 else ey

  fun toCollisionString (squares, canvasWidth, canvasHeight) =
    let
      val size = Int.max (canvasWidth, canvasHeight)
      val qtree = buildTree (0, 0, size, squares)
      val bintree = merge (qtree, squares)

      val collisions = BinTree.foldr (toExportStringFolder, bintree, [])
      val collisions = String.concatWith ",\n" coords
end