computer-vision-project/src/main/kotlin/network/rs485/ben/computervision/EntityScanner.kt

package network.rs485.ben.computervision

import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.*
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import org.opencv.core.Mat

class EntityScanner(private val animate: Boolean, private val scanBarSpacing: Int, private val foregroundThreshold: Float) {

    private fun Mat.isValid(row: Int, column: Int): Boolean =
        (this[row, column]?.getOrNull(0) ?: 255.0) < foregroundThreshold

    private fun Mat.isValid(point: Pixel): Boolean =
        (this[point]?.getOrNull(0) ?: 255.0) < foregroundThreshold

    suspend fun scanImage(image: Mat, displayImage: Mat, windowMgr: IWindowManager) = coroutineScope {
        val entities = mutableListOf<Entity>()
        val mergeMutex = Mutex()
        lineScan(image, displayImage, windowMgr)
            .map { line ->
                lineGrowToEntityDown(line, image, displayImage, windowMgr).also { entity ->
                    // add missing upper parts of all entities
                    lineGrowToEntityUp(line, image, displayImage, windowMgr, entity)
                }
            }
            .collect {
                mergeMutex.withLock {
                    entities.merge(it)
                }
            }
        entities.removeIf { it.lines.size < 6 }
        if (animate) {
            entities.colorRanged(COLOR_GREEN, COLOR_RED, displayImage, windowMgr)
            windowMgr.updateImage(displayImage)
        }
        return@coroutineScope entities
    }

    @OptIn(ExperimentalCoroutinesApi::class)
    private suspend fun lineScan(
        image: Mat,
        displayImage: Mat,
        windowMgr: IWindowManager,
    ) = channelFlow {
        coroutineScope {
            for (i in 0 until image.height().div(scanBarSpacing)) {
                launch {
                    val lineNumber = i * scanBarSpacing
                    var start = Pixel(0, lineNumber)
                    var end = Pixel(image.width() - 1, lineNumber)
                    if (animate) {
                        displayImage.drawLine(Line(start, end), COLOR_RED)
                        launch {
                            windowMgr.updateImage(displayImage)
                        }
                    }
                    var onEntity = false
                    for (x in 0 until image.width()) {
                        val color = if (image.isValid(lineNumber, x)) {
                            COLOR_GREEN.also {
                                if (!onEntity) {
                                    onEntity = true
                                    start = Pixel(x, start.y)
                                }
                            }
                        } else {
                            COLOR_BLUE.also {
                                if (onEntity) {
                                    onEntity = false
                                    send(start to Pixel(x, start.y))
                                }
                            }
                        }
                        end = Pixel(x, end.y)
                        if (animate) {
                            // color the current dot on the line
                            displayImage.put(lineNumber, x, *color.asArray(3))
                            launch {
                                windowMgr.updateImage(displayImage)
                            }
                            if (ANIMATION_DELAY.isPositive()) delay(ANIMATION_DELAY)
                        }
                    }
                    if (onEntity) {
                        send(start to Pixel(image.width() - 1, start.y))
                    }
                }
            }
        }
    }

    private suspend fun lineGrowToEntityDown(
        line: Line,
        image: Mat,
        displayImage: Mat,
        windowMgr: IWindowManager,
        entity: Entity = Entity(),
    ) = lineGrowToEntity(
        line = line,
        image = image,
        displayImage = displayImage,
        windowMgr = windowMgr,
        entity = entity,
        last = (line.first.y + scanBarSpacing - 1).coerceIn(0, image.height() - 1),
        step = 1,
    )

    private suspend fun lineGrowToEntityUp(
        line: Line,
        image: Mat,
        displayImage: Mat,
        windowMgr: IWindowManager,
        entity: Entity = Entity(),
    ) = lineGrowToEntity(
        line = line,
        image = image,
        displayImage = displayImage,
        windowMgr = windowMgr,
        entity = entity,
        last = (line.first.y - scanBarSpacing + 1).coerceIn(0, image.height() - 1),
        step = -1,
    )

    private suspend fun lineGrowToEntity(
        line: Line,
        image: Mat,
        displayImage: Mat,
        windowMgr: IWindowManager,
        entity: Entity,
        last: Int,
        step: Int,
    ): Entity = coroutineScope {
        assert(step != 0)

        suspend fun lineFound(foundLine: Line) {
            assert(foundLine.first.x >= 0)
            assert(foundLine.first.y >= 0)
            assert(foundLine.first.x <= foundLine.second.x)
            assert(foundLine.first.y == foundLine.second.y)
            entity.lines.add(foundLine)
            if (animate) {
                displayImage.drawLine(foundLine, COLOR_YELLOW)
                launch {
                    windowMgr.updateImage(displayImage)
                }
                if (ANIMATION_DELAY.isPositive()) delay(ANIMATION_DELAY)
            }
        }

        lineFound(line)

        suspend fun lineGrowRecursion(foundLine: Line) {
            lineFound(foundLine)
            if (step > 0 && foundLine.first.y <= last || step < 0 && foundLine.first.y >= last) {
                lineGrowToEntity(
                    line = foundLine,
                    image = image,
                    displayImage = displayImage,
                    windowMgr = windowMgr,
                    entity = entity,
                    last = last,
                    step = step,
                )
            }
        }

        var start = Pixel(line.first.x, line.first.y + step)
        if (start.y < 0 || start.y >= image.height()) return@coroutineScope entity
        val row = start.y
        fun IntProgression.lastValid() = takeWhile { image.isValid(row, it) }.last()
        var end = start
        if (image.isValid(start)) {
            start = Pixel(line.first.x.downTo(0).lastValid(), start.y)
            end = Pixel(line.first.x.until(image.width()).lastValid(), end.y)
            lineGrowRecursion(start to end)
        }
        var next = end.x + 2
        while (next <= line.second.x) {
            if (image.isValid(row, next)) {
                start = Pixel(end.x + 2, start.y)
                end = Pixel(next.until(image.width()).lastValid(), start.y)
                lineGrowRecursion(start to end)
                next = end.x + 2
            } else {
                ++next
            }
        }
        return@coroutineScope entity
    }
}