@file:OptIn(FlowPreview::class, ExperimentalCoroutinesApi::class)

package network.rs485.ben.computervision

import kotlinx.coroutines.*
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ProducerScope
import kotlinx.coroutines.flow.*
import org.opencv.core.*
import org.opencv.imgproc.Imgproc
import java.io.FileNotFoundException
import java.nio.file.FileSystemNotFoundException
import java.nio.file.Path
import java.nio.file.Paths
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.Executors
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger
import kotlin.io.path.*
import kotlin.math.acos
import kotlin.system.exitProcess
import kotlin.time.Duration
import kotlin.time.Duration.Companion.seconds

internal const val HIST_SIZE = 256
internal const val FOREGROUND_THRESHOLD = 120.0f
internal const val GAUSSIAN_KERNEL_SIZE = 3
internal val COROUTINE_BUFFER_SIZE = System.getProperties().getProperty("cv.coroutines.buffer.size", "4").toInt()
internal val MAXIMUM_PARALLEL_IMAGES = System.getProperties().getProperty("cv.images.buffer.size", "1").toInt()
//internal val ANIMATION_DELAY: Duration = 100.milliseconds
internal val ANIMATION_DELAY: Duration = Duration.ZERO
//internal val WINDOW_MANAGER_FACTORY = DEFAULT_WINDOW_MANAGER_FACTORY
internal val WINDOW_MANAGER_FACTORY = NOOP_WINDOW_MANAGER_FACTORY
private val SINGLE_DISPATCHER = Executors.newSingleThreadExecutor().asCoroutineDispatcher()

internal object Resource {
    fun getPath(name: String): Path {
        val path = try {
            javaClass.getResource("/$name")?.toURI()?.let(Paths::get)
        } catch (error: FileSystemNotFoundException) {
            System.err.println("Could not read resource $name, error:\n  $error\n  using fallback..")
            null
        } ?: Paths.get(name).apply {
            if (!exists()) {
                throw FileNotFoundException("Cannot find $this on filesystem $fileSystem")
            }
        }
        return path
    }
}

typealias Candidate = Triple<Pixel, Vector2D, Double>

private suspend fun findGoldenPoI(
    image: Mat,
    identity: IdentifiableEntity,
    golden: Golden,
    centrumMask: Mat,
): List<Pair<Golden, Pair<Pixel, Candidate>>> {
    val goldenOrientation = golden.orientation.vec2D()
    val goldenDirection = Vector2D(goldenOrientation).normalize()
    val searchMag = golden.orientationRadius
        .let { (it - DISTANCE_DEVIATION..it + DISTANCE_DEVIATION) }
    val bestCenters = channelFlow bestCenters@{
        for (center in centrumMask.filterPixels { it[0] > 0.0 }) {
            launch findCenterDeviation@{
                val intoGoldSpace = golden.minMaxDistance.maxLoc.rounded() - center
                val rect = golden.rect.apply {
                    x += identity.rect.x - intoGoldSpace.x
                    y += identity.rect.y - intoGoldSpace.y
                }
                if (
                    rect.x < 0 ||
                    rect.y < 0 ||
                    rect.x + rect.width >= image.width() ||
                    rect.y + rect.height >= image.height()
                ) {
                    if (ANIMATION_DELAY.isPositive()) println("Not considering $center for $identity on $golden")
                    return@findCenterDeviation
                }
                image.withSubmatSafe(rect) { subImage ->
                    val deviationFlow = channelFlow {
                        suspend fun ProducerScope<Candidate>.emitDeviation(startPixel: Pixel) {
                            val centerStart = center - startPixel
                            if (centerStart == Pixel(0, 0)) return
                            val dirVec = centerStart.vec2D().normalize()
                            val endpix = try {
                                identity.entityMask.searchGreatest(
                                    center = center,
                                    direction = dirVec,
                                    searchRange = searchMag.progression(step = DISTANCE_STEP)
                                )
                            } catch (error: IllegalArgumentException) {
                                System.err.println(
                                    "Error when finding greatest magnitude of $center for $identity.\n" +
                                            "  Values: center = $center, dir = $dirVec, searchMag = $searchMag\n" +
                                            "  $error"
                                )
                                null
                            } ?: return
                            val angle = acos(dirVec.dot(goldenDirection))
                            val deviation = CvContext.withMat(MatInfo.from(subImage)) { tmp ->
                                subImage.copyTo(tmp)
                                CvContext.withRotationMatrix(
                                    angle = angle,
                                    center = golden.minMaxDistance.maxLoc,
                                ) { intoGoldSpaceRot ->
                                    Imgproc.warpAffine(tmp, tmp, intoGoldSpaceRot, tmp.size())
                                }
                                Core.absdiff(golden.image, tmp, tmp)
                                Imgproc.filter2D(tmp, tmp, -1, CvContext.gaussianKernel)
                                CvContext.withGrayCalcHist(
                                    image = tmp,
                                    mask = golden.entityMaskDilated,
                                ) { hist ->
                                    hist.sumOfHist(
                                        indices = golden.deviationParameters.deviationRange,
                                        maximum = golden.deviationParameters.maximumDeviation,
                                        mapFunc = golden.deviationParameters.deviationFunc::calculate,
                                    )
                                }
                            }
                            if (deviation <= golden.deviationParameters.maximumDeviation) {
                                send(Candidate(endpix, dirVec, deviation))
                            }
                        }

                        identity.transformedLines.forEach {
                            launch {
                                emitDeviation(it.first)
                            }
                            launch {
                                emitDeviation(it.second)
                            }
                        }
                    }
                    deviationFlow.buffer(COROUTINE_BUFFER_SIZE)
                        .fold(null) { prev: Candidate?, other: Candidate ->
                            when {
                                prev === null -> other
                                prev.third > other.third -> other
                                else -> prev
                            }
                        }
                        ?.also { this@bestCenters.send(center to it) }
                }
            }
        }
    }.toList().sortedBy { (_, candidate) -> candidate.third }

    // filter all candidates that are too close to each other
    return bestCenters
        .filterIndexed { index, (center, _) ->
            (0 until index).map { bestCenters[it] }
                .none { (betterCenter, _) ->
                    (betterCenter - center).vec2D().magnitude < golden.orientationRadius
                }
        }.map { golden to it }
}

private suspend fun findGoldenPoIAnimate(
    golden: Golden,
    distanceMultiplier: Scalar,
    identity: IdentifiableEntity,
    centrumMask: Mat,
    image: Mat,
    displayImage: Mat,
    windowMgr: IWindowManager,
    parentScope: CoroutineScope,
): List<Pair<Golden, Pair<Pixel, Candidate>>> {
    withTmpMat { tmp ->
        identity.distanceMat.convertTo(tmp, CvType.CV_8U)
        Core.multiply(tmp, distanceMultiplier, tmp)
        displayImage.withSubmatSafe(identity.rect) { target ->
            Imgproc.cvtColor(tmp, tmp, Imgproc.COLOR_GRAY2BGR)
            tmp.copyTo(target, identity.entityMask)
            tmp.setTo(COLOR_BLACK)
            tmp.copyTo(target, centrumMask)
            Imgproc.cvtColor(centrumMask, tmp, Imgproc.COLOR_GRAY2BGR)
            Core.multiply(tmp, golden.goldenColor, tmp, 1.0)
            Core.add(tmp, target, target, centrumMask, CvType.CV_8UC3)
        }
    }
    windowMgr.updateImage(displayImage)
    delay(ANIMATION_DELAY)
    return findGoldenPoI(
        image = image,
        identity = identity,
        golden = golden,
        centrumMask = centrumMask
    ).also { goldenCandidates ->
        val goldenDirection = Vector2D(golden.orientation.vec2D()).normalize()
        withContext(SINGLE_DISPATCHER) {
            withTmpMat { tmp ->
                withTmpMat { displayTmp ->
                    goldenCandidates.forEach { (golden, centerToCandidate) ->
                        val (center, candidate) = centerToCandidate
                        val intoGoldSpace = golden.minMaxDistance.maxLoc.rounded() - center
                        val rect = golden.rect.apply {
                            x += identity.rect.x - intoGoldSpace.x
                            y += identity.rect.y - intoGoldSpace.y
                        }
                        image.withSubmatSafe(rect) { subImage ->
                            subImage.copyTo(displayTmp)
                            val (endpix, dirVec, deviation) = candidate
                            val wmg = WINDOW_MANAGER_FACTORY.createWindowManager("$golden on ${identity.rect} with $deviation")
                            wmg.setDimension(HIST_SIZE, HIST_SIZE)
                            parentScope.launch {
                                wmg.display(displayTmp)
                            }
                            val angle = acos(dirVec.dot(goldenDirection))
                            subImage.copyTo(tmp)
                            Imgproc.cvtColor(tmp, displayTmp, Imgproc.COLOR_GRAY2BGR)
                            displayTmp.drawLine(
                                golden.minMaxDistance.maxLoc.rounded() to intoGoldSpace + endpix,
                                COLOR_GREEN,
                            )
                            launch {
                                wmg.updateImage(displayTmp)
                            }
                            delay(ANIMATION_DELAY * 1000)
                            CvContext.withRotationMatrix(
                                angle = angle,
                                center = golden.minMaxDistance.maxLoc,
                            ) { intoGoldSpaceRot ->
                                Imgproc.warpAffine(tmp, tmp, intoGoldSpaceRot, tmp.size())
                            }
                            Imgproc.cvtColor(tmp, displayTmp, Imgproc.COLOR_GRAY2BGR)
                            displayTmp.drawLine(
                                golden.orientation,
                                COLOR_YELLOW,
                            )
                            launch {
                                wmg.updateImage(displayTmp)
                            }
                            delay(ANIMATION_DELAY * 1000)
                            Core.absdiff(golden.image, tmp, tmp)
                            Imgproc.filter2D(tmp, displayTmp, -1, CvContext.gaussianKernel)
                            launch {
                                wmg.updateImage(displayTmp)
                            }
                            delay(ANIMATION_DELAY)
                            wmg.updateImage(displayTmp)
                            delay(10.seconds)
                            wmg.dispose()
                        }
                    }
                }
            }
        }
    }
}

suspend fun main(args: Array<String>) {
    CvContext.initialize()
    val imageDirectory = Path(args.getOrElse(0) { "." })
    val comparisonData = imageDirectory.resolve("export_data.csv").bufferedReader().lineSequence().drop(1)
        .map { line -> line.trimEnd().split(",") }
        .groupingBy { parts -> parts[0] }
        .aggregateTo(ConcurrentHashMap()) { _, accumulator: MutableList<Pair<String, Rect>>?, parts, _ ->
            val descriptionToRect = parts[5] to Rect(
                parts[1].toInt(),
                parts[2].toInt(),
                parts[3].toInt() - parts[1].toInt() + 1,
                parts[4].toInt() - parts[2].toInt() + 1,
            )
            (accumulator ?: mutableListOf()).apply { add(descriptionToRect) }
        }
    val hits = AtomicInteger(0)
    val falseHit = AtomicInteger(0)
    val count = AtomicInteger(0)
    val imagePaths = imageDirectory.listDirectoryEntries("*.bmp")
    val perImageWindow = WINDOW_MANAGER_FACTORY.createWindowManager("Current Image")
    val scanner = EntityScanner(animate = false, scanBarSpacing = 6, foregroundThreshold = FOREGROUND_THRESHOLD)

    coroutineScope {
        val chaetomium = Golden.createGolden(
            parentScope = this,
            scanner = scanner,
            goldenName = "Chaetomium",
            goldenColor = COLOR_GREEN,
            goldenPath = Resource.getPath("single_chaetomium.bmp"),
            deviationParameters = DeviationParameters(
                deviationMin = 40,
                deviationFunc = { level, value ->
                    level * (value / 40f)
                },
                maximumDeviation = 100.0
            ),
        )
        val stachybotry = Golden.createGolden(
            parentScope = this,
            scanner = scanner,
            goldenName = "Stachybotrys",
            goldenPath = Resource.getPath("single_stachybotry.bmp"),
            goldenColor = COLOR_RED,
            deviationParameters = DeviationParameters(
                deviationMin = 20,
                deviationFunc = { level, value ->
                    level * (value / 12f)
                },
                maximumDeviation = 200.0
            ),
        )
        val goldenList = listOf(chaetomium, stachybotry)
        val localBufferInfo = when (COROUTINE_BUFFER_SIZE) {
            Channel.RENDEZVOUS -> "RENDEZVOUS"
            Channel.CONFLATED -> "CONFLATED"
            Channel.BUFFERED -> "BUFFERED"
            Channel.UNLIMITED -> "UNLIMITED"
            else -> toString()
        }
        println("Coroutine buffer: $localBufferInfo, Parallel images: $MAXIMUM_PARALLEL_IMAGES")
        val deferred = async {
            val maxDistance = goldenList.maxOf { it.minMaxDistance.maxVal }
            val distanceMultiplier = COLOR_WHITE / maxDistance
            if (ANIMATION_DELAY.isPositive()) {
                goldenList.forEach {
                    Core.multiply(it.displayImage, distanceMultiplier, it.displayImage)
                    launch {
                        it.updateImage()
                    }
                    if (ANIMATION_DELAY.isPositive()) delay(ANIMATION_DELAY)
                    Imgproc.cvtColor(it.image, it.displayImage, Imgproc.COLOR_GRAY2BGR)
                    it.displayImage.drawLine(it.orientation, COLOR_YELLOW)
                    it.updateImage()
                    if (ANIMATION_DELAY.isPositive()) delay(ANIMATION_DELAY)
                }
            }

            val windowInitialized = AtomicBoolean(false)
            val imageFlow = flow {
                for (imagePath in imagePaths.shuffled()) {
                    try {
                        val (displayImage, image) = CvContext.readAndConvertImage(imagePath)
                        if (!windowInitialized.getAndSet(true) && ANIMATION_DELAY.isPositive()) {
                            this@coroutineScope.launch {
                                try {
                                    perImageWindow.display(displayImage)
                                } catch (e: CancellationException) {
                                    this@async.cancel("Window was closed", e)
                                }
                            }
                        }
                        val comparisonSet: MutableList<Pair<String, Rect>> = comparisonData[imagePath.name] ?: mutableListOf()
                        emit(Triple(displayImage, image, comparisonSet))
                    } catch (error: CvException) {
                        System.err.println("Cannot read $imagePath: ${error.message}")
                    }
                }
            }
            imageFlow
                .buffer(if (ANIMATION_DELAY.isPositive()) Channel.RENDEZVOUS else MAXIMUM_PARALLEL_IMAGES)
                .collect { (displayImage, image, comparison) ->
                    scanner.scanImage(image, displayImage, perImageWindow).asFlow()
                        .buffer(COROUTINE_BUFFER_SIZE)
                        .mapNotNull entityMatching@{ entity ->
                            withIdEntity(entity) { identity ->
                                var goldenCandidates = goldenList.flatMap { golden ->
                                    golden.withCentrumMask(identity) { centrumMask ->
                                        if (ANIMATION_DELAY.isPositive()) {
                                            findGoldenPoIAnimate(
                                                golden = golden,
                                                distanceMultiplier = distanceMultiplier,
                                                identity = identity,
                                                centrumMask = centrumMask,
                                                image = image,
                                                displayImage = displayImage,
                                                windowMgr = perImageWindow,
                                                parentScope = this@coroutineScope,
                                            )
                                        } else {
                                            findGoldenPoI(
                                                image = image,
                                                identity = identity,
                                                golden = golden,
                                                centrumMask = centrumMask,
                                            )
                                        }
                                    }
                                }.sortedBy { it.second.second.third } // deviation
                                goldenCandidates = goldenCandidates
                                    .filterIndexed { index, (_, it) ->
                                        val (center, _) = it
                                        (0 until index).map { goldenCandidates[it] }
                                            .none { (betterGold, other) ->
                                                (other.first - center).vec2D().magnitude < betterGold.orientationRadius
                                            }
                                    }
                                if (ANIMATION_DELAY.isPositive()) {
                                    goldenCandidates = goldenCandidates.onEach { (gold, it) ->
                                        val (center, candidate) = it
                                        val (_, _, deviation) = candidate
                                        println("Found minimum for ${identity.rect} on $gold: $center (deviation: $deviation)")
                                        Imgproc.drawMarker(
                                            /* img = */ displayImage,
                                            /* position = */ center.vec2D() + identity.rect.tl(),
                                            /* color = */ gold.goldenColor,
                                            /* markerType = */ Imgproc.MARKER_CROSS,
                                            /* markerSize = */ 3,
                                            /* thickness = */ 3,
                                        )
                                        Imgproc.rectangle(displayImage, identity.rect, gold.goldenColor)
                                        perImageWindow.updateImage(displayImage)
                                        if (ANIMATION_DELAY.isPositive()) delay(ANIMATION_DELAY)
                                    }
                                }
                                return@entityMatching goldenCandidates.groupingBy { (gold, _) -> gold.goldenName }
                                    .eachCount()
                                    .maxByOrNull { it.value }
                                    ?.let { goldenAmount ->
                                        val description = if (goldenAmount.value > 1) {
                                            "${goldenAmount.key}_Agglomerate_Candidate"
                                        } else {
                                            "${goldenAmount.key}_Einzelspore_Candidate"
                                        }
                                        Triple(description, identity.rect, comparison)
                                    }
                            }
                        }
                        .collect { result ->
                            count.incrementAndGet()
                            val matchingResults = result.third.filter { (description, comparisonRect) ->
                                description == result.first &&
                                        (comparisonRect.contains(result.second.tl()) &&
                                                comparisonRect.contains(result.second.br())) ||
                                        (result.second.contains(comparisonRect.tl()) &&
                                                result.second.contains(comparisonRect.br()))
                            }
                            result.third.removeAll(matchingResults)
                            when (matchingResults.size) {
                                0 -> falseHit.incrementAndGet().also { println("False hit ${result.first}") }
                                1 -> hits.incrementAndGet().also { println("Hit ${result.first}") }
                                else -> falseHit.incrementAndGet().also { println("Hit too many ${result.first}") }
                            }
                        }
                    if (ANIMATION_DELAY.isPositive()) delay(ANIMATION_DELAY * 2000)
                    println("Releasing image $displayImage\n and $image")
                    displayImage.release()
                    CvContext.recycle(image)
                }
        }
        deferred.invokeOnCompletion {
            perImageWindow.dispose()
            goldenList.forEach { it.release() }
            if (deferred.isCancelled) {
                val error = deferred.getCompletionExceptionOrNull()!!
                if (error is CancellationException) {
                    val cause = error.cause?.takeIf { it.message != error.message }?.let { ". Cause: ".plus(it.message) }.orEmpty()
                    System.err.println("Cancelled program. Cause: ${error.message}" + cause)
                } else {
                    error.printStackTrace()
                }
                exitProcess(1)
            }
        }
        deferred.join()
        val misses = comparisonData.values.sumOf { it.size }
        println("Hits ${hits.get()}, False-hits ${falseHit.get()}, Total ${count.get()}, Misses $misses")
    }
}

internal inline fun <T> Mat.withSubmatSafe(rect: Rect, block: (mat: Mat) -> T): T? {
    val submat = try {
        submat(rect)
    } catch (error: CvException) {
        System.err.println("Cannot create submatrix for $this, rect $rect, error:\n  $error")
        null
    }
    return submat?.let { mat ->
        kotlin.runCatching { block(mat) }.also { mat.release() }.getOrThrow()
    }
}

private fun ClosedFloatingPointRange<Double>.progression(step: Number): List<Double> = sequence {
    var x = start
    while (lessThanOrEquals(x, endInclusive)) {
        yield(x)
        x += step.toDouble()
    }
}.toList()