package ru.dbotthepony.mc.otm.core.math

import it.unimi.dsi.fastutil.objects.ObjectArrayList
import ru.dbotthepony.mc.otm.core.random
import java.util.random.RandomGenerator
import kotlin.math.min

interface Cluster<V : Any> {
	val values: List<V>
	val center: V

	data class Impl<V : Any>(override val values: List<V>, override val center: V) : Cluster<V>
}

private class ClusterValue<V : Comparable<V>>(val value: V, var cluster: MutableCluster<V>, var error: V) {
	fun updateError(abs: (V) -> V, minus: (V, V) -> V) {
		error = abs(minus(cluster.center, value))
	}

	fun maybeSwitchCluster(cluster: MutableCluster<V>, abs: (V) -> V, minus: (V, V) -> V): Boolean {
		if (cluster == this.cluster) return false

		val newError = abs(minus(cluster.center, value))

		if (newError < error) {
			error = newError
			this.cluster.values.remove(this)
			cluster.values.add(this)
			cluster.generation++
			this.cluster.generation++
			this.cluster = cluster
			return true
		} else {
			return false
		}
	}
}

private class MutableCluster<V : Comparable<V>>(var center: V) {
	val values = ObjectArrayList<ClusterValue<V>>()
	var generation = 0
	var centerCalculatedAtGeneration = -1

	fun calculateCenter(identity: V, plus: (V, V) -> V, minus: (V, V) -> V, divInt: (V, Int) -> V, abs: (V) -> V): Boolean {
		if (values.isEmpty) return false
		if (centerCalculatedAtGeneration == generation) return true

		centerCalculatedAtGeneration = generation
		var value = identity
		values.forEach { value = plus(value, it.value) }
		val old = center
		center = divInt(value, values.size)

		if (old != center) {
			values.forEach { it.updateError(abs, minus) }
		}

		return true
	}

	fun pullFrom(clusters: MutableList<MutableCluster<V>>, identity: V): Boolean {
		var candidate: ClusterValue<V>? = null

		for (eCluser in clusters) {
			if (eCluser.values.size > 1) {
				for (value in eCluser.values) {
					if (candidate == null || candidate.error < value.error) {
						candidate = value
					}
				}
			}
		}

		if (candidate != null) {
			values.add(candidate)
			candidate.cluster.values.remove(candidate)
			this.generation++
			candidate.cluster.generation++
			candidate.cluster = this
			center = candidate.value
			candidate.error = identity
			clusters.add(this)
			return true
		} else {
			return false
		}
	}
}

private fun <V : Comparable<V>> Iterable<V>.clusterize(
	random: RandomGenerator,
	initialClusters: Int = 1,
	zeroBound: Boolean = false,
	identity: V,
	plus: (V, V) -> V,
	minus: (V, V) -> V,
	divInt: (V, Int) -> V,
	abs: (V) -> V,
	heuristics: (min: V, max: V, error: V) -> Boolean,
): List<Cluster<V>> {
	require(initialClusters > 0) { "Invalid amount of initial clusters: $initialClusters" }
	val itr = iterator()

	if (!itr.hasNext())
		return listOf()

	val clusters = ObjectArrayList<MutableCluster<V>>(initialClusters)

	for (i in 0 until initialClusters) {
		clusters.add(MutableCluster(identity))
	}

	val values = ObjectArrayList<ClusterValue<V>>()
	var min = itr.next()
	var max = min

	values.add(ClusterValue(min, clusters.random(random), min))
	values[0].cluster.values.add(values[0])

	for (value in itr) {
		min = minOf(min, value)
		max = maxOf(max, value)

		val cluster = clusters.random(random)
		val wrapped = ClusterValue(value, cluster, value)
		values.add(wrapped)
		cluster.values.add(wrapped)
	}

	if (zeroBound) {
		min = minOf(min, identity)
		max = maxOf(max, identity)
	}

	if (min == max) {
		return listOf(Cluster.Impl(listOf(min), min))
	} else {
		// if "clusters" already satisfy constraints...
		clusters.forEach { it.calculateCenter(identity, plus, minus, divInt, abs) }
		values.forEach { it.updateError(abs, minus) }

		val maxError = values.maxOf { it.error }

		if (!heuristics(min, max, maxError)) {
			return clusters.stream()
				.filter { it.values.isNotEmpty() }
				.map { Cluster.Impl(it.values.map { it.value }, it.center) }
				.sorted { o1, o2 -> o2.values.size.compareTo(o1.values.size) } // большие кластеры должны идти первыми
				.toList()
		}
	}

	while (true) {
		var converged = false
		var oversaturation = false

		clusters.forEach { it.calculateCenter(identity, plus, minus, divInt, abs) }

		while (!converged) {
			converged = true

			for (value in values) {
				clusters.forEach { converged = !value.maybeSwitchCluster(it, abs, minus) && converged }
			}

			if (!converged) {
				val emptyClusters = ObjectArrayList<MutableCluster<V>>()
				val citr = clusters.iterator()

				for (cluster in citr) {
					if (!cluster.calculateCenter(identity, plus, minus, divInt, abs)) {
						emptyClusters.add(cluster)
						citr.remove()
					}
				}

				for (cluster in emptyClusters) {
					if (!cluster.pullFrom(clusters, identity)) {
						oversaturation = true
						break
					}
				}
			}
		}

		val maxError = values.maxOf { it.error }

		if (!oversaturation && clusters.size < values.size && heuristics(min, max, maxError)) {
			val cluster = MutableCluster(identity)
			check(cluster.pullFrom(clusters, identity)) { "Newly created cluster couldn't pull first value from other clusters (expected at least one cluster with 2 values in it)" }
			clusters.add(cluster)
		} else {
			return clusters.stream()
				.filter { it.values.isNotEmpty() }
				.map { Cluster.Impl(it.values.map { it.value }, it.center) }
				.sorted { o1, o2 -> o2.values.size.compareTo(o1.values.size) } // большие кластеры должны идти первыми
				.toList()
		}
	}
}

// TODO: could use some tweaking
private val DECIMAL_ERROR_TOLERANCE = Decimal("0.02")

fun Iterable<Decimal>.clusterize(random: RandomGenerator, clusters: Int? = null, zeroBound: Boolean = false): List<Cluster<Decimal>> {
	return clusterize(random, clusters ?: 1, zeroBound, Decimal.ZERO, Decimal::plus, Decimal::minus, Decimal::div, Decimal::absoluteValue) { min, max, error ->
		if (clusters != null)
			false
		else
			error / (max - min) >= DECIMAL_ERROR_TOLERANCE
	}
}
