package ru.dbotthepony.mc.otm.core

import net.minecraft.world.entity.Entity
import net.minecraft.world.level.Level
import net.minecraft.world.phys.AABB
import ru.dbotthepony.mc.otm.core.math.Vector
import ru.dbotthepony.mc.otm.core.math.minus
import java.util.LinkedList
import java.util.function.Predicate
import kotlin.math.pow

data class EntityDistance<T : Entity>(val entity: T, val distance: Double) : Comparable<EntityDistance<*>> {
	override fun compareTo(other: EntityDistance<*>): Int {
		return distance.compareTo(other.distance)
	}
}

/**
 * Pair of entity and distance fraction to center of ellipsoid
 */
fun Level.getEntitiesInEllipsoid(pos: Vector, dimensions: Vector, except: Entity?, predicate: Predicate<in Entity>): MutableList<EntityDistance<Entity>> {
	val entities = getEntities(except, AABB(
		pos.x - dimensions.x,
		pos.y - dimensions.y,
		pos.z - dimensions.z,

		pos.x + dimensions.x,
		pos.y + dimensions.y,
		pos.z + dimensions.z,
	), predicate)

	val result = ArrayList<EntityDistance<Entity>>()

	for (it in entities) {
		val a = (it.position - pos).let { vec ->
			vec.x.pow(2.0) / dimensions.x.pow(2.0) +
			vec.y.pow(2.0) / dimensions.y.pow(2.0) +
			vec.z.pow(2.0) / dimensions.z.pow(2.0)
		}

		val b = it.boundingBox.center.let { vec ->
			(vec.x - pos.x).pow(2.0) / dimensions.x.pow(2.0) +
			(vec.y - pos.y).pow(2.0) / dimensions.y.pow(2.0) +
			(vec.z - pos.z).pow(2.0) / dimensions.z.pow(2.0)
		}

		val min = a.coerceAtMost(b)

		if (min <= 1.0) {
			result.add(EntityDistance(it, min))
		}
	}

	return result
}

/**
 * Pair of entity and distance fraction to center of ellipsoid
 */
fun <T : Entity> Level.getEntitiesInEllipsoid(type: Class<out T>, pos: Vector, dimensions: Vector, predicate: Predicate<in Entity>): MutableList<EntityDistance<T>> {
	val entities = getEntitiesOfClass(type, AABB(
		pos.x - dimensions.x,
		pos.y - dimensions.y,
		pos.z - dimensions.z,

		pos.x + dimensions.x,
		pos.y + dimensions.y,
		pos.z + dimensions.z,
	), predicate)

	val result = ArrayList<EntityDistance<T>>()

	for (it in entities) {
		val a = (it.position - pos).let { vec ->
			vec.x.pow(2.0) / dimensions.x.pow(2.0) +
			vec.y.pow(2.0) / dimensions.y.pow(2.0) +
			vec.z.pow(2.0) / dimensions.z.pow(2.0)
		}

		val b = it.boundingBox.center.let { vec ->
			(vec.x - pos.x).pow(2.0) / dimensions.x.pow(2.0) +
			(vec.y - pos.y).pow(2.0) / dimensions.y.pow(2.0) +
			(vec.z - pos.z).pow(2.0) / dimensions.z.pow(2.0)
		}

		val min = a.coerceAtMost(b)

		if (min <= 1.0) {
			result.add(EntityDistance(it, min))
		}
	}

	return result
}

/**
 * Pair of entity and distance fraction to center of ellipsoid
 */
fun Level.getEntitiesInEllipsoid(pos: Vector, dimensions: Vector, predicate: Predicate<in Entity>): MutableList<EntityDistance<Entity>> {
	return getEntitiesInEllipsoid(pos, dimensions, null, predicate)
}

/**
 * Pair of entity and distance fraction to center of ellipsoid
 */
fun Level.getEntitiesInSphere(pos: Vector, radius: Double, except: Entity?, predicate: Predicate<in Entity>): MutableList<EntityDistance<Entity>> {
	return getEntitiesInEllipsoid(pos, Vector(radius, radius, radius), except, predicate)
}

/**
 * Pair of entity and distance fraction to center of ellipsoid
 */
fun Level.getEntitiesInSphere(pos: Vector, radius: Double, predicate: Predicate<in Entity>): MutableList<EntityDistance<Entity>> {
	return getEntitiesInEllipsoid(pos, Vector(radius, radius, radius), null, predicate)
}

/**
 * Pair of entity and distance fraction to center of ellipsoid
 */
fun <T : Entity> Level.getEntitiesInSphere(type: Class<out T>, pos: Vector, radius: Double, predicate: Predicate<in Entity>): MutableList<EntityDistance<T>> {
	return getEntitiesInEllipsoid(type, pos, Vector(radius, radius, radius), predicate)
}
