package ru.dbotthepony.kbox2d.dynamics.joint

import ru.dbotthepony.kbox2d.api.*
import ru.dbotthepony.kbox2d.api.B2SolverData
import ru.dbotthepony.kbox2d.api.b2Mul
import ru.dbotthepony.kstarbound.math.MutableMatrix2d
import ru.dbotthepony.kstarbound.math.Vector2d
import ru.dbotthepony.kstarbound.math.times
import ru.dbotthepony.kstarbound.util.Color
import kotlin.math.cos
import kotlin.math.sin

// Point-to-point constraint
// C = p2 - p1
// Cdot = v2 - v1
//      = v2 + cross(w2, r2) - v1 - cross(w1, r1)
// J = [-I -r1_skew I r2_skew ]
// Identity used:
// w k % (rx i + ry j) = w * (-ry i + rx j)

// Motor constraint
// Cdot = w2 - w1
// J = [0 0 -1 0 0 1]
// K = invI1 + invI2

class RevoluteJoint(def: RevoluteJointDef) : AbstractJoint(def) {
	// Solver shared
	val localAnchorA: Vector2d = def.localAnchorA
	val localAnchorB: Vector2d = def.localAnchorB
	private var impulse: Vector2d = Vector2d.ZERO
	private var motorImpulse: Double = 0.0
	private var lowerImpulse: Double = 0.0
	private var upperImpulse: Double = 0.0
	val referenceAngle: Double = def.referenceAngle

	// Solver temp
	private var indexA: Int = 0
	private var indexB: Int = 0
	private var rA: Vector2d = Vector2d.ZERO
	private var rB: Vector2d = Vector2d.ZERO
	private var localCenterA: Vector2d = Vector2d.ZERO
	private var localCenterB: Vector2d = Vector2d.ZERO
	private var invMassA: Double = 0.0
	private var invMassB: Double = 0.0
	private var invIA: Double = 0.0
	private var invIB: Double = 0.0
	private var K: MutableMatrix2d = MutableMatrix2d().also { it.zero() }
	private var angle: Double = 0.0
	private var axialMass: Double = 0.0

	override fun initVelocityConstraints(data: B2SolverData) {
		indexA = bodyA.islandIndex
		indexB = bodyB.islandIndex
		localCenterA = bodyA.sweep.localCenter
		localCenterB = bodyB.sweep.localCenter
		invMassA = bodyA.invMass
		invMassB = bodyB.invMass
		invIA = bodyA.rotInertiaInv
		invIB = bodyB.rotInertiaInv

		val aA = data.positions[indexA].a
		var vA = data.velocities[indexA].v
		var wA = data.velocities[indexA].w

		val aB = data.positions[indexB].a
		var vB = data.velocities[indexB].v
		var wB = data.velocities[indexB].w

		val qA = Rotation(aA)
		val qB = Rotation(aB)

		rA = b2Mul(qA, localAnchorA - localCenterA)
		rB = b2Mul(qB, localAnchorB - localCenterB)

		// J = [-I -r1_skew I r2_skew]
		// r_skew = [-ry; rx]

		// Matlab
		// K = [ mA+r1y^2*iA+mB+r2y^2*iB,  -r1y*iA*r1x-r2y*iB*r2x]
		//     [  -r1y*iA*r1x-r2y*iB*r2x, mA+r1x^2*iA+mB+r2x^2*iB]

		val mA = invMassA
		val mB = invMassB
		val iA = invIA
		val iB = invIB

		K.m00 = mA + mB + rA.y * rA.y * iA + rB.y * rB.y * iB
		K.m01 = -rA.y * rA.x * iA - rB.y * rB.x * iB
		K.m10 = K.m01
		K.m11 = mA + mB + rA.x * rA.x * iA + rB.x * rB.x * iB

		axialMass = iA + iB
		val fixedRotation: Boolean

		if (axialMass > 0.0f) {
			axialMass = 1.0f / axialMass
			fixedRotation = false
		} else {
			fixedRotation = true
		}

		angle = aB - aA - referenceAngle

		if (!enableLimit || fixedRotation) {
			lowerImpulse = 0.0
			upperImpulse = 0.0
		}

		if (!enableMotor || fixedRotation) {
			motorImpulse = 0.0
		}

		if (data.step.warmStarting) {
			// Scale impulses to support a variable time step.
			impulse *= data.step.dtRatio
			motorImpulse *= data.step.dtRatio
			lowerImpulse *= data.step.dtRatio
			upperImpulse *= data.step.dtRatio

			val axialImpulse = motorImpulse + lowerImpulse - upperImpulse
			val P = Vector2d(impulse.x, impulse.y)

			vA -= mA * P
			wA -= iA * (b2Cross(rA, P) + axialImpulse)

			vB += mB * P
			wB += iB * (b2Cross(rB, P) + axialImpulse)
		} else {
			impulse = Vector2d.ZERO
			motorImpulse = 0.0
			lowerImpulse = 0.0
			upperImpulse = 0.0
		}

		data.velocities[indexA].v = vA
		data.velocities[indexA].w = wA
		data.velocities[indexB].v = vB
		data.velocities[indexB].w = wB
	}

	override fun solveVelocityConstraints(data: B2SolverData) {
		var vA = data.velocities[indexA].v
		var wA = data.velocities[indexA].w
		var vB = data.velocities[indexB].v
		var wB = data.velocities[indexB].w

		val mA = invMassA
		val mB = invMassB
		val iA = invIA
		val iB = invIB

		val fixedRotation = iA + iB == 0.0

		// Solve motor constraint.
		if (enableMotor && !fixedRotation) {
			val Cdot = wB - wA - motorSpeed
			var impulse = -axialMass * Cdot
			val oldImpulse = motorImpulse
			val maxImpulse = data.step.dt * maxMotorTorque
			motorImpulse = b2Clamp(motorImpulse + impulse, -maxImpulse, maxImpulse)
			impulse = motorImpulse - oldImpulse

			wA -= iA * impulse
			wB += iB * impulse
		}

		if (enableLimit && !fixedRotation) {
			// Lower limit
			run {
				val C = this.angle - this.lowerAngle
				val Cdot = wB - wA
				var impulse = -this.axialMass * (Cdot + b2Max(C, 0.0) * data.step.inv_dt)
				val oldImpulse = this.lowerImpulse
				this.lowerImpulse = b2Max(this.lowerImpulse + impulse, 0.0)
				impulse = this.lowerImpulse - oldImpulse

				wA -= iA * impulse
				wB += iB * impulse
			}

			// Upper limit
			// Note: signs are flipped to keep C positive when the constraint is satisfied.
			// This also keeps the impulse positive when the limit is active.
			run {
				val C = this.upperAngle - this.angle
				val Cdot = wA - wB
				var impulse = -this.axialMass * (Cdot + b2Max(C, 0.0) * data.step.inv_dt)
				val oldImpulse = this.upperImpulse
				this.upperImpulse = b2Max(this.upperImpulse + impulse, 0.0)
				impulse = this.upperImpulse - oldImpulse

				wA += iA * impulse
				wB -= iB * impulse
			}
		}

		// Solve point-to-point constraint
		run {
			val Cdot = vB + b2Cross(wB, this.rB) - vA - b2Cross(wA, this.rA)
			val impulse = this.K.solve(-Cdot)

			this.impulse += impulse

			vA -= mA * impulse
			wA -= iA * b2Cross(this.rA, impulse)

			vB += mB * impulse
			wB += iB * b2Cross(this.rB, impulse)
		}

		data.velocities[indexA].v = vA
		data.velocities[indexA].w = wA
		data.velocities[indexB].v = vB
		data.velocities[indexB].w = wB
	}

	override fun solvePositionConstraints(data: B2SolverData): Boolean {
		var cA = data.positions[indexA].c
		var aA = data.positions[indexA].a
		var cB = data.positions[indexB].c
		var aB = data.positions[indexB].a

		val qA = Rotation(aA)
		val qB = Rotation(aB)

		var angularError = 0.0
		var positionError = 0.0

		val fixedRotation = invIA + invIB == 0.0

		// Solve angular limit constraint
		if (enableLimit && !fixedRotation) {
			val angle = aB - aA - referenceAngle
			var C = 0.0

			if (b2Abs(upperAngle - lowerAngle) < 2.0f * b2_angularSlop) {
				// Prevent large angular corrections
				C = b2Clamp(angle - lowerAngle, -b2_maxAngularCorrection, b2_maxAngularCorrection)
			} else if (angle <= lowerAngle) {
				// Prevent large angular corrections and allow some slop.
				C = b2Clamp(angle - lowerAngle + b2_angularSlop, -b2_maxAngularCorrection, 0.0)
			} else if (angle >= upperAngle) {
				// Prevent large angular corrections and allow some slop.
				C = b2Clamp(angle - upperAngle - b2_angularSlop, 0.0, b2_maxAngularCorrection)
			}

			val limitImpulse = -axialMass * C
			aA -= invIA * limitImpulse
			aB += invIB * limitImpulse
			angularError = b2Abs(C)
		}

		// Solve point-to-point constraint.
		run {
			qA.set(aA)
			qB.set(aB)
			val rA = b2Mul(qA, this.localAnchorA - this.localCenterA)
			val rB = b2Mul(qB, this.localAnchorB - this.localCenterB)

			val C = cB + rB - cA - rA
			positionError = C.length

			val mA = this.invMassA
			val mB = this.invMassB
			val iA = this.invIA
			val iB = this.invIB

			val K = MutableMatrix2d()
			K.m00 = mA + mB + iA * rA.y * rA.y + iB * rB.y * rB.y
			K.m10 = -iA * rA.x * rA.y - iB * rB.x * rB.y
			K.m01 = K.m10
			K.m11 = mA + mB + iA * rA.x * rA.x + iB * rB.x * rB.x

			val impulse = -K.solve(C)

			cA -= mA * impulse
			aA -= iA * b2Cross(rA, impulse)

			cB += mB * impulse
			aB += iB * b2Cross(rB, impulse)
		}

		data.positions[indexA].c = cA
		data.positions[indexA].a = aA
		data.positions[indexB].c = cB
		data.positions[indexB].a = aB

		return positionError <= b2_linearSlop && angularError <= b2_angularSlop
	}

	/**
	 * Get the reaction force given the inverse time step.
	 *
	 * Unit is N.
	 */
	override fun getReactionForce(inv_dt: Double): Vector2d {
		return inv_dt * Vector2d(impulse.x, impulse.y)
	}

	/**
	 * Get the reaction torque due to the joint limit given the inverse time step.
	 *
	 * Unit is N*m.
	 */
	override fun getReactionTorque(inv_dt: Double): Double {
		return inv_dt * (motorImpulse + lowerImpulse - upperImpulse)
	}

	override val anchorA: Vector2d
		get() = bodyA.getWorldPoint(localAnchorA)
	override val anchorB: Vector2d
		get() = bodyB.getWorldPoint(localAnchorB)

	/**
	 * Get the current joint angle in radians.
	 */
	val jointAngle: Double get() {
		return bodyB.sweep.a - bodyA.sweep.a - referenceAngle
	}

	/**
	 * Get the current joint angle speed in radians per second.
	 */
	val jointSpeed: Double get() {
		return bodyB.angularVelocity - bodyA.angularVelocity
	}

	var enableMotor: Boolean = def.enableMotor
		set(value) {
			if (value != field) {
				field = value
				bodyA.isAwake = true
				bodyB.isAwake = true
			}
		}

	/**
	 * Get the current motor torque given the inverse time step.
	 *
	 * Unit is N*m.
	 */
	fun getMotorTorque(inv_dt: Double): Double {
		return inv_dt * motorImpulse
	}

	/**
	 *  Set the motor speed in radians per second.
	 */
	var motorSpeed: Double = def.motorSpeed
		set(value) {
			if (field != value) {
				if (enableMotor) {
					bodyA.isAwake = true
					bodyB.isAwake = true
				}

				field = value
			}
		}

	/**
	 * Set the maximum motor torque, usually in N-m.
	 */
	var maxMotorTorque: Double = def.maxMotorTorque
		set(value) {
			if (field != value) {
				if (enableMotor) {
					bodyA.isAwake = true
					bodyB.isAwake = true
				}

				field = value
			}
		}

	/**
	 * Enable/disable the joint limit.
	 */
	var enableLimit: Boolean = def.enableLimit
		set(value) {
			if (field != value) {
				bodyA.isAwake = true
				bodyB.isAwake = true
				field = value
				lowerImpulse = 0.0
				upperImpulse = 0.0
			}
		}

	/**
	 * Get the lower joint limit in radians.
	 */
	var lowerAngle: Double = def.lowerAngle
		private set

	/**
	 * Get the upper joint limit in radians.
	 */
	var upperAngle: Double = def.upperAngle
		private set

	init {
		// KBox2D: Overlooked check in constructor
		// TODO: Notify Erin
		require(lowerAngle <= upperAngle) { "$lowerAngle !<= $upperAngle" }
	}

	/**
	 * Set the joint limits in radians.
	 */
	fun setLimits(lower: Double, upper: Double) {
		require(lower <= upper) { "$lower !<= $upper" }

		if (lower != lowerAngle || upper != upperAngle) {
			bodyA.isAwake = true
			bodyB.isAwake = true
			lowerImpulse = 0.0
			upperImpulse = 0.0
			lowerAngle = lower
			upperAngle = upper
		}
	}

	override fun draw(draw: IDebugDraw) {
		val xfA = bodyA.transform
		val xfB = bodyB.transform
		val pA = b2Mul(xfA, localAnchorA)
		val pB = b2Mul(xfB, localAnchorB)

		draw.drawPoint(pA, 5.0, c4)
		draw.drawPoint(pB, 5.0, c5)

		val aA = bodyA.angle
		val aB = bodyB.angle
		val angle = aB - aA - referenceAngle

		val r = L * Vector2d(cos(angle), sin(angle))
		draw.drawSegment(pB, pB + r, c1)
		draw.drawCircle(pB, L, c1)

		if (enableLimit) {
			val rlo = L * Vector2d(cos(lowerAngle), sin(lowerAngle))
			val rhi = L * Vector2d(cos(upperAngle), sin(upperAngle))

			draw.drawSegment(pB, pB + rlo, c2)
			draw.drawSegment(pB, pB + rhi, c3)
		}

		draw.drawSegment(xfA.p, pA, color)
		draw.drawSegment(pA, pB, color)
		draw.drawSegment(xfB.p, pB, color)
	}

	companion object {
		private const val L = 0.5
		private val color = Color(0.5f, 0.8f, 0.8f)
		private val c1 = Color(0.7f, 0.7f, 0.7f)
		private val c2 = Color(0.3f, 0.9f, 0.3f)
		private val c3 = Color(0.9f, 0.3f, 0.3f)
		private val c4 = Color(0.3f, 0.3f, 0.9f)
		private val c5 = Color(0.4f, 0.4f, 0.4f)
	}
}