icegps-geotools/android/src/main/java/com/icegps/geotools/marchingsquares/MarchingSquares.kt

package com.icegps.geotools.marchingsquares

import com.icegps.math.geometry.Rectangle
import com.icegps.math.geometry.Vector2D
import com.icegps.math.geometry.Vector2I
import com.icegps.geotools.ktx.mix
import kotlin.math.max
import kotlin.math.min

private const val closeEpsilon = 1E-6

data class Segment2D(
    val start: Vector2D,
    val control: List<Vector2D>,
    val end: Vector2D,
    val corner: Boolean = false
)

fun Segment2D(start: Vector2D, end: Vector2D, corner: Boolean = true) =
    Segment2D(start, emptyList(), end, corner)

fun Segment2D(start: Vector2D, c0: Vector2D, c1: Vector2D, end: Vector2D, corner: Boolean = true) =
    Segment2D(start, listOf(c0, c1), end, corner)

data class ShapeContour(
    val segments: List<Segment2D>,
    val closed: Boolean,
) {
    companion object {
        val EMPTY = ShapeContour(
            segments = emptyList(),
            closed = false,
        )

        /**
         * Creates a ShapeContour from a list of points, specifying whether the contour is closed and its y-axis polarity.
         *
         * @param points A list of points (Vector2) defining the vertices of the contour.
         * @param closed Boolean indicating whether the contour should be closed (forms a loop).
         * @return A ShapeContour object representing the resulting contour.
         */
        fun fromPoints(
            points: List<Vector2D>,
            closed: Boolean,
        ): ShapeContour = if (points.isEmpty()) {
            EMPTY
        } else {
            if (!closed) {
                ShapeContour((0 until points.size - 1).map {
                    Segment2D(
                        points[it],
                        points[it + 1]
                    )
                }, false)
            } else {
                val d = (points.last() - points.first()).lengthSquared
                val usePoints = if (d > closeEpsilon) points else points.dropLast(1)
                ShapeContour((usePoints.indices).map {
                    Segment2D(
                        usePoints[it],
                        usePoints[(it + 1) % usePoints.size]
                    )
                }, true)
            }
        }
    }
}

data class LineSegment(val start: Vector2D, val end: Vector2D)

/**
 * Find contours for a function [f] using the marching squares algorithm. A contour is found when f(x) crosses zero.
 * @param f the function
 * @param area a rectangular area in which the function should be evaluated
 * @param cellSize the size of the cells, smaller size gives higher resolution
 * @param useInterpolation intersection points will be interpolated if true, default true
 * @return a list of [ShapeContour] instances
 */
fun findContours(
    f: (Vector2D) -> Double,
    area: Rectangle,
    cellSize: Double,
    useInterpolation: Boolean = true
): List<ShapeContour> {
    val segments = mutableListOf<LineSegment>()
    val values = mutableMapOf<Vector2I, Double>()
    val segmentsMap = mutableMapOf<Vector2D, MutableList<LineSegment>>()

    for (y in 0 until (area.height / cellSize).toInt()) {
        for (x in 0 until (area.width / cellSize).toInt()) {
            values[Vector2I(x, y)] = f(Vector2D(x * cellSize + area.x, y * cellSize + area.y))
        }
    }

    val zero = 0.0
    for (y in 0 until (area.height / cellSize).toInt()) {
        for (x in 0 until (area.width / cellSize).toInt()) {

            // Here we check if we are at a right or top border. This is to ensure we create closed contours
            // later on in the process.
            val v00 = if (x == 0 || y == 0) zero else (values[Vector2I(x, y)] ?: zero)
            val v10 = if (y == 0) zero else (values[Vector2I(x + 1, y)] ?: zero)
            val v01 = if (x == 0) zero else (values[Vector2I(x, y + 1)] ?: zero)
            val v11 = (values[Vector2I(x + 1, y + 1)] ?: zero)

            val p00 = Vector2D(x.toDouble(), y.toDouble()) * cellSize + area.topLeft
            val p10 = Vector2D((x + 1).toDouble(), y.toDouble()) * cellSize + area.topLeft
            val p01 = Vector2D(x.toDouble(), (y + 1).toDouble()) * cellSize + area.topLeft
            val p11 = Vector2D((x + 1).toDouble(), (y + 1).toDouble()) * cellSize + area.topLeft

            val index = (if (v00 >= 0.0) 1 else 0) +
                    (if (v10 >= 0.0) 2 else 0) +
                    (if (v01 >= 0.0) 4 else 0) +
                    (if (v11 >= 0.0) 8 else 0)

            fun blend(v1: Double, v2: Double): Double {
                if (useInterpolation) {
                    require(!v1.isNaN() && !v2.isNaN()) {
                        "Input values v1=$v1 or v2=$v2 are NaN, which is not allowed."
                    }
                    val f1 = min(v1, v2)
                    val f2 = max(v1, v2)
                    val v = (-f1) / (f2 - f1)

                    require(v == v && v in 0.0..1.0) {
                        "Invalid value calculated during interpolation: v=$v"
                    }

                    return if (f1 == v1) {
                        v
                    } else {
                        1.0 - v
                    }
                } else {
                    return 0.5
                }
            }

            fun emitLine(
                p00: Vector2D, p01: Vector2D, v00: Double, v01: Double,
                p10: Vector2D, p11: Vector2D, v10: Double, v11: Double
            ) {
                val r0 = blend(v00, v01)
                val r1 = blend(v10, v11)

                val v0 = p00.mix(p01, r0)
                val v1 = p10.mix(p11, r1)
                val l0 = LineSegment(v0, v1)
                segmentsMap.getOrPut(v1) { mutableListOf() }.add(l0)
                segmentsMap.getOrPut(v0) { mutableListOf() }.add(l0)
                segments.add(l0)
            }

            when (index) {
                0, 15 -> {}
                1, 15 xor 1 -> {
                    emitLine(p00, p01, v00, v01, p00, p10, v00, v10)
                }

                2, 15 xor 2 -> {
                    emitLine(p00, p10, v00, v10, p10, p11, v10, v11)
                }

                3, 15 xor 3 -> {
                    emitLine(p00, p01, v00, v01, p10, p11, v10, v11)
                }

                4, 15 xor 4 -> {
                    emitLine(p00, p01, v00, v01, p01, p11, v01, v11)
                }

                5, 15 xor 5 -> {
                    emitLine(p00, p10, v00, v10, p01, p11, v01, v11)
                }

                6, 15 xor 6 -> {
                    emitLine(p00, p01, v00, v01, p00, p10, v00, v10)
                    emitLine(p01, p11, v01, v11, p10, p11, v10, v11)
                }

                7, 15 xor 7 -> {
                    emitLine(p01, p11, v01, v11, p10, p11, v10, v11)
                }
            }
        }
    }

    val processedSegments = mutableSetOf<LineSegment>()
    val contours = mutableListOf<ShapeContour>()
    for (segment in segments) {
        if (segment in processedSegments) {
            continue
        } else {
            val collected = mutableListOf<Vector2D>()
            var current: LineSegment? = segment
            var closed = true
            var lastVertex = Vector2D.INFINITY
            do {
                current!!
                if (lastVertex.squaredDistanceTo(current.start) > 1E-5) {
                    collected.add(current.start)
                }
                lastVertex = current.start
                processedSegments.add(current)
                if (segmentsMap[current.start]!!.size < 2) {
                    closed = false
                }
                val hold = current
                current = segmentsMap[current.start]?.firstOrNull { it !in processedSegments }
                if (current == null) {
                    current = segmentsMap[hold.end]?.firstOrNull { it !in processedSegments }
                }
            } while (current != segment && current != null)

            contours.add(ShapeContour.fromPoints(collected, closed = closed))
        }
    }
    return contours
}