package ru.dbotthepony.kbox2d.dynamics.joint

import ru.dbotthepony.kbox2d.api.*
import ru.dbotthepony.kbox2d.api.B2SolverData
import ru.dbotthepony.kbox2d.api.b2Cross
import ru.dbotthepony.kbox2d.api.b2Mul
import ru.dbotthepony.kstarbound.math.*
import ru.dbotthepony.kstarbound.util.Color

// Linear constraint (point-to-line)
// d = p2 - p1 = x2 + r2 - x1 - r1
// C = dot(perp, d)
// Cdot = dot(d, cross(w1, perp)) + dot(perp, v2 + cross(w2, r2) - v1 - cross(w1, r1))
//      = -dot(perp, v1) - dot(cross(d + r1, perp), w1) + dot(perp, v2) + dot(cross(r2, perp), v2)
// J = [-perp, -cross(d + r1, perp), perp, cross(r2,perp)]
//
// Angular constraint
// C = a2 - a1 + a_initial
// Cdot = w2 - w1
// J = [0 0 -1 0 0 1]
//
// K = J * invM * JT
//
// J = [-a -s1 a s2]
//     [0  -1  0  1]
// a = perp
// s1 = cross(d + r1, a) = cross(p2 - x1, a)
// s2 = cross(r2, a) = cross(p2 - x2, a)

// Motor/Limit linear constraint
// C = dot(ax1, d)
// Cdot = -dot(ax1, v1) - dot(cross(d + r1, ax1), w1) + dot(ax1, v2) + dot(cross(r2, ax1), v2)
// J = [-ax1 -cross(d+r1,ax1) ax1 cross(r2,ax1)]

// Predictive limit is applied even when the limit is not active.
// Prevents a constraint speed that can lead to a constraint error in one time step.
// Want C2 = C1 + h * Cdot >= 0
// Or:
// Cdot + C1/h >= 0
// I do not apply a negative constraint error because that is handled in position correction.
// So:
// Cdot + max(C1, 0)/h >= 0

// Block Solver
// We develop a block solver that includes the angular and linear constraints. This makes the limit stiffer.
//
// The Jacobian has 2 rows:
// J = [-uT -s1 uT s2] // linear
//     [0   -1   0  1] // angular
//
// u = perp
// s1 = cross(d + r1, u), s2 = cross(r2, u)
// a1 = cross(d + r1, v), a2 = cross(r2, v)

class PrismaticJoint(def: PrismaticJointDef) : AbstractJoint(def) {
	internal val localAnchorA: Vector2d = def.localAnchorA
	internal val localAnchorB: Vector2d = def.localAnchorB
	internal val localXAxisA: Vector2d = def.localAxisA.normalized
	internal val localYAxisA: Vector2d = b2Cross(1.0, localXAxisA)
	internal val referenceAngle: Double = def.referenceAngle

	private var impulse: Vector2d = Vector2d.ZERO
	private var motorImpulse: Double = 0.0
	private var lowerImpulse: Double = 0.0
	private var upperImpulse: Double = 0.0

	// Solver temp
	private var indexA: Int = 0
	private var indexB: Int = 0
	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 axis: Vector2d = Vector2d.ZERO
	private var perp: Vector2d = Vector2d.ZERO
	private var s1: Double = 0.0
	private var s2: Double = 0.0
	private var a1: Double = 0.0
	private var a2: Double = 0.0
	private var K: MutableMatrix2d = MutableMatrix2d()
	private var translation: 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 cA = data.positions[indexA].c
		val aA = data.positions[indexA].a
		var vA = data.velocities[indexA].v
		var wA = data.velocities[indexA].w

		val cB = data.positions[indexB].c
		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)

		// Compute the effective masses.
		val rA = b2Mul(qA, localAnchorA - localCenterA)
		val rB = b2Mul(qB, localAnchorB - localCenterB)
		val d = (cB - cA) + rB - rA

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

		// Compute motor Jacobian and effective mass.
		run {
			this.axis = b2Mul(qA, this.localXAxisA)
			this.a1 = b2Cross(d + rA, this.axis)
			this.a2 = b2Cross(rB, this.axis)

			this.axialMass = mA + mB + iA * this.a1 * this.a1 + iB * this.a2 * this.a2
			if (this.axialMass > 0.0f)
			{
				this.axialMass = 1.0f / this.axialMass
			}
		}

		// Prismatic constraint.
		run {
			this.perp = b2Mul(qA, this.localYAxisA)

			this.s1 = b2Cross(d + rA, this.perp)
			this.s2 = b2Cross(rB, this.perp)

			val k11 = mA + mB + iA * this.s1 * this.s1 + iB * this.s2 * this.s2
			val k12 = iA * this.s1 + iB * this.s2
			var k22 = iA + iB

			if (k22 == 0.0) {
				// For bodies with fixed rotation.
				k22 = 1.0
			}

			this.K.m00 = k11
			this.K.m10 = k12

			this.K.m01 = k12
			this.K.m11 = k22
		}

		if (enableLimit) {
			translation = b2Dot(axis, d)
		} else {
			lowerImpulse = 0.0
			upperImpulse = 0.0
		}

		if (!enableMotor) {
			motorImpulse = 0.0
		}

		if (data.step.warmStarting) {
			// Account for 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 = impulse.x * perp + axialImpulse * axis
			val LA = impulse.x * s1 + impulse.y + axialImpulse * a1
			val LB = impulse.x * s2 + impulse.y + axialImpulse * a2

			vA -= mA * P
			wA -= iA * LA

			vB += mB * P
			wB += iB * LB
		} 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

		// Solve linear motor constraint
		if (enableMotor) {
			val Cdot = b2Dot(axis, vB - vA) + a2 * wB - a1 * wA
			var impulse = axialMass * (motorSpeed - Cdot)
			val oldImpulse = motorImpulse
			val maxImpulse = data.step.dt * maxMotorForce
			motorImpulse = b2Clamp(motorImpulse + impulse, -maxImpulse, maxImpulse)
			impulse = motorImpulse - oldImpulse

			val P = impulse * axis
			val LA = impulse * a1
			val LB = impulse * a2

			vA -= mA * P
			wA -= iA * LA
			vB += mB * P
			wB += iB * LB
		}

		if (enableLimit) {
			// Lower limit
			run {
				val C = this.translation - this.lowerTranslation
				val Cdot = b2Dot(this.axis, vB - vA) + this.a2 * wB - this.a1 * 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

				val P = impulse * this.axis
				val LA = impulse * this.a1
				val LB = impulse * this.a2

				vA -= mA * P
				wA -= iA * LA
				vB += mB * P
				wB += iB * LB
			}

			// 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.upperTranslation - this.translation
				val Cdot = b2Dot(this.axis, vA - vB) + this.a1 * wA - this.a2 * 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

				val P = impulse * this.axis
				val LA = impulse * this.a1
				val LB = impulse * this.a2

				vA += mA * P
				wA += iA * LA
				vB -= mB * P
				wB -= iB * LB
			}
		}

		// Solve the prismatic constraint in block form.
		run {
			val Cdot = Vector2d(
				x = b2Dot(this.perp, vB - vA) + this.s2 * wB - this.s1 * wA,
				y = wB - wA
			)

			val df = this.K.solve(-Cdot)
			this.impulse += df

			val P = df.x * this.perp
			val LA = df.x * this.s1 + df.y
			val LB = df.x * this.s2 + df.y

			vA -= mA * P
			wA -= iA * LA

			vB += mB * P
			wB += iB * LB
		}

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

	/**
	 * A velocity based solver computes reaction forces(impulses) using the velocity constraint solver.Under this context,
	 * the position solver is not there to resolve forces.It is only there to cope with integration error.
	 *
	 * Therefore, the pseudo impulses in the position solver do not have any physical meaning.Thus it is okay if they suck.
	 *
	 * We could take the active state from the velocity solver.However, the joint might push past the limit when the velocity
	 * solver indicates the limit is inactive.
	 */
	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)

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

		// Compute fresh Jacobians
		val rA = b2Mul(qA, localAnchorA - localCenterA)
		val rB = b2Mul(qB, localAnchorB - localCenterB)
		val d = cB + rB - cA - rA

		val axis = b2Mul(qA, localXAxisA)
		val a1 = b2Cross(d + rA, axis)
		val a2 = b2Cross(rB, axis)
		val perp = b2Mul(qA, localYAxisA)

		val s1 = b2Cross(d + rA, perp)
		val s2 = b2Cross(rB, perp)

		val impulse: Vector3d

		val C1 = Vector2d(
			x = b2Dot(perp, d),
			y = aB - aA - referenceAngle
		)

		var linearError = b2Abs(C1.x)
		val angularError = b2Abs(C1.y)

		var active = false
		var C2 = 0.0

		if (enableLimit) {
			val translation = b2Dot(axis, d)
			if (b2Abs(upperTranslation - lowerTranslation) < 2.0 * b2_linearSlop) {
				C2 = translation
				linearError = b2Max(linearError, b2Abs(translation))
				active = true
			} else if (translation <= lowerTranslation) {
				C2 = b2Min(translation - lowerTranslation, 0.0)
				linearError = b2Max(linearError, lowerTranslation - translation)
				active = true
			} else if (translation >= upperTranslation) {
				C2 = b2Max(translation - upperTranslation, 0.0)
				linearError = b2Max(linearError, translation - upperTranslation)
				active = true
			}
		}

		if (active) {
			val k11 = mA + mB + iA * s1 * s1 + iB * s2 * s2
			val k12 = iA * s1 + iB * s2
			val k13 = iA * s1 * a1 + iB * s2 * a2
			var k22 = iA + iB

			if (k22 == 0.0) {
				// For fixed rotation
				k22 = 1.0
			}

			val k23 = iA * a1 + iB * a2
			val k33 = mA + mB + iA * a1 * a1 + iB * a2 * a2

			val K = MutableMatrix3d()

			K.m00 = k11
			K.m10 = k12
			K.m20 = k13

			K.m01 = k12
			K.m11 = k22
			K.m21 = k23

			K.m02 = k13
			K.m12 = k23
			K.m22 = k33

			val C = Vector3d(
				x = C1.x,
				y = C1.y,
				z = C2,
			)

			impulse = K.solve(-C)
		} else {
			val k11 = mA + mB + iA * s1 * s1 + iB * s2 * s2
			val k12 = iA * s1 + iB * s2
			var k22 = iA + iB

			if (k22 == 0.0) {
				k22 = 1.0
			}

			val K = MutableMatrix2d()

			K.m00 = k11
			K.m10 = k12

			K.m01 = k12
			K.m11 = k22

			val impulse1 = K.solve(-C1)
			impulse = Vector3d(impulse1.x, impulse1.y)
		}

		val P = impulse.x * perp + impulse.z * axis
		val LA = impulse.x * s1 + impulse.y + impulse.z * a1
		val LB = impulse.x * s2 + impulse.y + impulse.z * a2

		cA -= mA * P
		aA -= iA * LA
		cB += mB * P
		aB += iB * LB

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

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

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

	override fun getReactionForce(inv_dt: Double): Vector2d {
		return inv_dt * (impulse.x * perp + (motorImpulse + lowerImpulse - upperImpulse) * axis)
	}

	override fun getReactionTorque(inv_dt: Double): Double {
		return inv_dt * impulse.y
	}

	/**
	 * Get the current joint translation, usually in meters.
	 */
	val jointTranslation: Double get() {
		val pA = bodyA.getWorldPoint(localAnchorA)
		val pB = bodyB.getWorldPoint(localAnchorB)
		val d = pB - pA
		val axis = bodyA.getWorldVector(localXAxisA)

		return b2Dot(d, axis)
	}

	/**
	 * Get the current joint translation speed, usually in meters per second.
	 */
	val jointSpeed: Double get() {
		val bA = bodyA
		val bB = bodyB

		val rA = b2Mul(bA.transform.q, localAnchorA - bA.sweep.localCenter)
		val rB = b2Mul(bB.transform.q, localAnchorB - bB.sweep.localCenter)
		val p1 = bA.sweep.c + rA
		val p2 = bB.sweep.c + rB
		val d = p2 - p1
		val axis = b2Mul(bA.transform.q, localXAxisA)

		val vA = bA.linearVelocity
		val vB = bB.linearVelocity
		val wA = bA.angularVelocity
		val wB = bB.angularVelocity

		return b2Dot(d, b2Cross(wA, axis)) + b2Dot(axis, vB + b2Cross(wB, rB) - vA - b2Cross(wA, rA))
	}

	/**
	 * Get the lower joint limit, usually in meters.
	 */
	var lowerTranslation: Double = def.lowerTranslation
		private set

	/**
	 * Get the upper joint limit, usually in meters.
	 */
	var upperTranslation: Double = def.upperTranslation
		private set

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

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

	/**
	 * Set the maximum motor force, usually in N.
	 */
	var maxMotorForce: Double = def.maxMotorForce
		set(value) {
			if (value != field) {
				field = value

				if (enableMotor) {
					bodyA.isAwake = true
					bodyB.isAwake = true
				}
			}
		}

	/**
	 * Set the motor speed, usually in meters per second.
	 */
	var motorSpeed: Double = def.motorSpeed
		set(value) {
			if (value != field) {
				field = value

				if (enableMotor) {
					bodyA.isAwake = true
					bodyB.isAwake = true
				}
			}
		}

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

	/**
	 * Enable/disable the joint motor.
	 */
	var enableMotor: Boolean = def.enableMotor
		set(value) {
			if (value != field) {
				field = value
				bodyA.isAwake = true
				bodyB.isAwake = true
			}
		}

	init {
		require(lowerTranslation <= upperTranslation) { "$lowerTranslation !<= $upperTranslation" }
	}

	fun getMotorForce(inv_dt: Double): Double {
		return inv_dt * motorImpulse
	}

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

		val axis = b2Mul(xfA.q, localXAxisA)

		draw.drawSegment(pA, pB, c5)

		if (enableLimit) {
			val lower = pA + lowerTranslation * axis
			val upper = pA + upperTranslation * axis
			val perp = b2Mul(xfA.q, localYAxisA)
			draw.drawSegment(lower, upper, c1)
			draw.drawSegment(lower - 0.5 * perp, lower + 0.5 * perp, c2)
			draw.drawSegment(upper - 0.5 * perp, upper + 0.5 * perp, c3)
		} else {
			draw.drawSegment(pA - 1.0 * axis, pA + 1.0 * axis, c1)
		}

		draw.drawPoint(pA, 5.0, c1)
		draw.drawPoint(pB, 5.0, c4)
	}

	companion object {
		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)
	}
}