// MIT License

// Copyright (c) 2019 Erin Catto

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include "box2d/b2_draw.h"
#include "box2d/b2_rope.h"

#include <stdio.h>

struct b2RopeStretch
{
	int32 i1, i2;
	float invMass1, invMass2;
	float L;
	float lambda;
	float spring;
	float damper;
};

struct b2RopeBend
{
	int32 i1, i2, i3;
	float invMass1, invMass2, invMass3;
	float invEffectiveMass;
	float lambda;
	float L1, L2;
	float alpha1, alpha2;
	float spring;
	float damper;
};

b2Rope::b2Rope()
{
	m_position.SetZero();
	m_count = 0;
	m_stretchCount = 0;
	m_bendCount = 0;
	m_stretchConstraints = nullptr;
	m_bendConstraints = nullptr;
	m_bindPositions = nullptr;
	m_ps = nullptr;
	m_p0s = nullptr;
	m_vs = nullptr;
	m_invMasses = nullptr;
	m_gravity.SetZero();
}

b2Rope::~b2Rope()
{
	b2Free(m_stretchConstraints);
	b2Free(m_bendConstraints);
	b2Free(m_bindPositions);
	b2Free(m_ps);
	b2Free(m_p0s);
	b2Free(m_vs);
	b2Free(m_invMasses);
}

void b2Rope::Create(const b2RopeDef& def)
{
	b2Assert(def.count >= 3);
	m_position = def.position;
	m_count = def.count;
	m_bindPositions = (b2Vec2*)b2Alloc(m_count * sizeof(b2Vec2));
	m_ps = (b2Vec2*)b2Alloc(m_count * sizeof(b2Vec2));
	m_p0s = (b2Vec2*)b2Alloc(m_count * sizeof(b2Vec2));
	m_vs = (b2Vec2*)b2Alloc(m_count * sizeof(b2Vec2));
	m_invMasses = (float*)b2Alloc(m_count * sizeof(float));

	for (int32 i = 0; i < m_count; ++i)
	{
		m_bindPositions[i] = def.vertices[i];
		m_ps[i] = def.vertices[i] + m_position;
		m_p0s[i] = def.vertices[i] + m_position;
		m_vs[i].SetZero();

		float m = def.masses[i];
		if (m > 0.0f)
		{
			m_invMasses[i] = 1.0f / m;
		}
		else
		{
			m_invMasses[i] = 0.0f;
		}
	}

	m_stretchCount = m_count - 1;
	m_bendCount = m_count - 2;

	m_stretchConstraints = (b2RopeStretch*)b2Alloc(m_stretchCount * sizeof(b2RopeStretch));
	m_bendConstraints = (b2RopeBend*)b2Alloc(m_bendCount * sizeof(b2RopeBend));

	for (int32 i = 0; i < m_stretchCount; ++i)
	{
		b2RopeStretch& c = m_stretchConstraints[i];

		b2Vec2 p1 = m_ps[i];
		b2Vec2 p2 = m_ps[i+1];

		c.i1 = i;
		c.i2 = i + 1;
		c.L = b2Distance(p1, p2);
		c.invMass1 = m_invMasses[i];
		c.invMass2 = m_invMasses[i + 1];
		c.lambda = 0.0f;
		c.damper = 0.0f;
		c.spring = 0.0f;
	}

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		b2RopeBend& c = m_bendConstraints[i];

		b2Vec2 p1 = m_ps[i];
		b2Vec2 p2 = m_ps[i + 1];
		b2Vec2 p3 = m_ps[i + 2];

		c.i1 = i;
		c.i2 = i + 1;
		c.i3 = i + 2;
		c.invMass1 = m_invMasses[i];
		c.invMass2 = m_invMasses[i + 1];
		c.invMass3 = m_invMasses[i + 2];
		c.invEffectiveMass = 0.0f;
		c.L1 = b2Distance(p1, p2);
		c.L2 = b2Distance(p2, p3);
		c.lambda = 0.0f;

		// Pre-compute effective mass (TODO use flattened config)
		b2Vec2 e1 = p2 - p1;
		b2Vec2 e2 = p3 - p2;
		float L1sqr = e1.LengthSquared();
		float L2sqr = e2.LengthSquared();

		if (L1sqr * L2sqr == 0.0f)
		{
			continue;
		}

		b2Vec2 Jd1 = (-1.0f / L1sqr) * e1.Skew();
		b2Vec2 Jd2 = (1.0f / L2sqr) * e2.Skew();

		b2Vec2 J1 = -Jd1;
		b2Vec2 J2 = Jd1 - Jd2;
		b2Vec2 J3 = Jd2;

		c.invEffectiveMass = c.invMass1 * b2Dot(J1, J1) + c.invMass2 * b2Dot(J2, J2) + c.invMass3 * b2Dot(J3, J3);
	
		b2Vec2 r = p3 - p1;

		float rr = r.LengthSquared();
		if (rr == 0.0f)
		{
			continue;
		}

		// a1 = h2 / (h1 + h2)
		// a2 = h1 / (h1 + h2)
		c.alpha1 = b2Dot(e2, r) / rr;
		c.alpha2 = b2Dot(e1, r) / rr;
	}

	m_gravity = def.gravity;

	SetTuning(def.tuning);
}

void b2Rope::SetTuning(const b2RopeTuning& tuning)
{
	m_tuning = tuning;

	// Pre-compute spring and damper values based on tuning

	const float bendOmega = 2.0f * b2_pi * m_tuning.bendHertz;

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		b2RopeBend& c = m_bendConstraints[i];

		float L1sqr = c.L1 * c.L1;
		float L2sqr = c.L2 * c.L2;

		if (L1sqr * L2sqr == 0.0f)
		{
			c.spring = 0.0f;
			c.damper = 0.0f;
			continue;
		}

		// Flatten the triangle formed by the two edges
		float J2 = 1.0f / c.L1 + 1.0f / c.L2;
		float sum = c.invMass1 / L1sqr + c.invMass2 * J2 * J2 + c.invMass3 / L2sqr;
		if (sum == 0.0f)
		{
			c.spring = 0.0f;
			c.damper = 0.0f;
			continue;
		}

		float mass = 1.0f / sum;

		c.spring = mass * bendOmega * bendOmega;
		c.damper = 2.0f * mass * m_tuning.bendDamping * bendOmega;
	}
	
	const float stretchOmega = 2.0f * b2_pi * m_tuning.stretchHertz;

	for (int32 i = 0; i < m_stretchCount; ++i)
	{
		b2RopeStretch& c = m_stretchConstraints[i];

		float sum = c.invMass1 + c.invMass2;
		if (sum == 0.0f)
		{
			continue;
		}

		float mass = 1.0f / sum;

		c.spring = mass * stretchOmega * stretchOmega;
		c.damper = 2.0f * mass * m_tuning.stretchDamping * stretchOmega;
	}
}

void b2Rope::Step(float dt, int32 iterations, const b2Vec2& position)
{
	if (dt == 0.0)
	{
		return;
	}

	const float inv_dt = 1.0f / dt;
	float d = expf(- dt * m_tuning.damping);

	// Apply gravity and damping
	for (int32 i = 0; i < m_count; ++i)
	{
		if (m_invMasses[i] > 0.0f)
		{
			m_vs[i] *= d;
			m_vs[i] += dt * m_gravity;
		}
		else
		{
			m_vs[i] = inv_dt * (m_bindPositions[i] + position - m_p0s[i]);
		}
	}

	// Apply bending spring
	if (m_tuning.bendingModel == b2_springAngleBendingModel)
	{
		ApplyBendForces(dt);
	}

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		m_bendConstraints[i].lambda = 0.0f;
	}

	for (int32 i = 0; i < m_stretchCount; ++i)
	{
		m_stretchConstraints[i].lambda = 0.0f;
	}

	// Update position
	for (int32 i = 0; i < m_count; ++i)
	{
		m_ps[i] += dt * m_vs[i];
	}

	// Solve constraints
	for (int32 i = 0; i < iterations; ++i)
	{
		if (m_tuning.bendingModel == b2_pbdAngleBendingModel)
		{
			SolveBend_PBD_Angle();
		}
		else if (m_tuning.bendingModel == b2_xpbdAngleBendingModel)
		{
			SolveBend_XPBD_Angle(dt);
		}
		else if (m_tuning.bendingModel == b2_pbdDistanceBendingModel)
		{
			SolveBend_PBD_Distance();
		}
		else if (m_tuning.bendingModel == b2_pbdHeightBendingModel)
		{
			SolveBend_PBD_Height();
		}
		else if (m_tuning.bendingModel == b2_pbdTriangleBendingModel)
		{
			SolveBend_PBD_Triangle();
		}

		if (m_tuning.stretchingModel == b2_pbdStretchingModel)
		{
			SolveStretch_PBD();
		}
		else if (m_tuning.stretchingModel == b2_xpbdStretchingModel)
		{
			SolveStretch_XPBD(dt);
		}
	}

	// Constrain velocity
	for (int32 i = 0; i < m_count; ++i)
	{
		m_vs[i] = inv_dt * (m_ps[i] - m_p0s[i]);
		m_p0s[i] = m_ps[i];
	}
}

void b2Rope::Reset(const b2Vec2& position)
{
	m_position = position;

	for (int32 i = 0; i < m_count; ++i)
	{
		m_ps[i] = m_bindPositions[i] + m_position;
		m_p0s[i] = m_bindPositions[i] + m_position;
		m_vs[i].SetZero();
	}

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		m_bendConstraints[i].lambda = 0.0f;
	}

	for (int32 i = 0; i < m_stretchCount; ++i)
	{
		m_stretchConstraints[i].lambda = 0.0f;
	}
}

void b2Rope::SolveStretch_PBD()
{
	const float stiffness = m_tuning.stretchStiffness;

	for (int32 i = 0; i < m_stretchCount; ++i)
	{
		const b2RopeStretch& c = m_stretchConstraints[i];

		b2Vec2 p1 = m_ps[c.i1];
		b2Vec2 p2 = m_ps[c.i2];

		b2Vec2 d = p2 - p1;
		float L = d.Normalize();

		float sum = c.invMass1 + c.invMass2;
		if (sum == 0.0f)
		{
			continue;
		}

		float s1 = c.invMass1 / sum;
		float s2 = c.invMass2 / sum;

		p1 -= stiffness * s1 * (c.L - L) * d;
		p2 += stiffness * s2 * (c.L - L) * d;

		m_ps[c.i1] = p1;
		m_ps[c.i2] = p2;
	}
}

void b2Rope::SolveStretch_XPBD(float dt)
{
	b2Assert(dt > 0.0f);

	for (int32 i = 0; i < m_stretchCount; ++i)
	{
		b2RopeStretch& c = m_stretchConstraints[i];

		b2Vec2 p1 = m_ps[c.i1];
		b2Vec2 p2 = m_ps[c.i2];

		b2Vec2 dp1 = p1 - m_p0s[c.i1];
		b2Vec2 dp2 = p2 - m_p0s[c.i2];

		b2Vec2 u = p2 - p1;
		float L = u.Normalize();

		b2Vec2 J1 = -u;
		b2Vec2 J2 = u;

		float sum = c.invMass1 + c.invMass2;
		if (sum == 0.0f)
		{
			continue;
		}

		const float alpha = 1.0f / (c.spring * dt * dt);	// 1 / kg
		const float beta = dt * dt * c.damper;				// kg * s
		const float sigma = alpha * beta / dt;				// non-dimensional
		float C = L - c.L;

		// This is using the initial velocities
		float Cdot = b2Dot(J1, dp1) + b2Dot(J2, dp2);

		float B = C + alpha * c.lambda + sigma * Cdot;
		float sum2 = (1.0f + sigma) * sum + alpha;

		float impulse = -B / sum2;

		p1 += (c.invMass1 * impulse) * J1;
		p2 += (c.invMass2 * impulse) * J2;

		m_ps[c.i1] = p1;
		m_ps[c.i2] = p2;
		c.lambda += impulse;
	}
}

void b2Rope::SolveBend_PBD_Angle()
{
	const float stiffness = m_tuning.bendStiffness;

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		const b2RopeBend& c = m_bendConstraints[i];

		b2Vec2 p1 = m_ps[c.i1];
		b2Vec2 p2 = m_ps[c.i2];
		b2Vec2 p3 = m_ps[c.i3];

		b2Vec2 d1 = p2 - p1;
		b2Vec2 d2 = p3 - p2;
		float a = b2Cross(d1, d2);
		float b = b2Dot(d1, d2);

		float angle = b2Atan2(a, b);

		float L1sqr, L2sqr;
		
		if (m_tuning.isometric)
		{
			L1sqr = c.L1 * c.L1;
			L2sqr = c.L2 * c.L2;
		}
		else
		{
			L1sqr = d1.LengthSquared();
			L2sqr = d2.LengthSquared();
		}

		if (L1sqr * L2sqr == 0.0f)
		{
			continue;
		}

		b2Vec2 Jd1 = (-1.0f / L1sqr) * d1.Skew();
		b2Vec2 Jd2 = (1.0f / L2sqr) * d2.Skew();

		b2Vec2 J1 = -Jd1;
		b2Vec2 J2 = Jd1 - Jd2;
		b2Vec2 J3 = Jd2;

		float sum;
		if (m_tuning.fixedEffectiveMass)
		{
			sum = c.invEffectiveMass;
		}
		else
		{
			sum = c.invMass1 * b2Dot(J1, J1) + c.invMass2 * b2Dot(J2, J2) + c.invMass3 * b2Dot(J3, J3);
		}

		if (sum == 0.0f)
		{
			sum = c.invEffectiveMass;
		}

		float impulse = -stiffness * angle / sum;

		p1 += (c.invMass1 * impulse) * J1;
		p2 += (c.invMass2 * impulse) * J2;
		p3 += (c.invMass3 * impulse) * J3;

		m_ps[c.i1] = p1;
		m_ps[c.i2] = p2;
		m_ps[c.i3] = p3;
	}
}

void b2Rope::SolveBend_XPBD_Angle(float dt)
{
	b2Assert(dt > 0.0f);

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		b2RopeBend& c = m_bendConstraints[i];

		b2Vec2 p1 = m_ps[c.i1];
		b2Vec2 p2 = m_ps[c.i2];
		b2Vec2 p3 = m_ps[c.i3];

		b2Vec2 dp1 = p1 - m_p0s[c.i1];
		b2Vec2 dp2 = p2 - m_p0s[c.i2];
		b2Vec2 dp3 = p3 - m_p0s[c.i3];

		b2Vec2 d1 = p2 - p1;
		b2Vec2 d2 = p3 - p2;

		float L1sqr, L2sqr;

		if (m_tuning.isometric)
		{
			L1sqr = c.L1 * c.L1;
			L2sqr = c.L2 * c.L2;
		}
		else
		{
			L1sqr = d1.LengthSquared();
			L2sqr = d2.LengthSquared();
		}

		if (L1sqr * L2sqr == 0.0f)
		{
			continue;
		}

		float a = b2Cross(d1, d2);
		float b = b2Dot(d1, d2);

		float angle = b2Atan2(a, b);

		b2Vec2 Jd1 = (-1.0f / L1sqr) * d1.Skew();
		b2Vec2 Jd2 = (1.0f / L2sqr) * d2.Skew();

		b2Vec2 J1 = -Jd1;
		b2Vec2 J2 = Jd1 - Jd2;
		b2Vec2 J3 = Jd2;

		float sum;
		if (m_tuning.fixedEffectiveMass)
		{
			sum = c.invEffectiveMass;
		}
		else
		{
			sum = c.invMass1 * b2Dot(J1, J1) + c.invMass2 * b2Dot(J2, J2) + c.invMass3 * b2Dot(J3, J3);
		}

		if (sum == 0.0f)
		{
			continue;
		}

		const float alpha = 1.0f / (c.spring * dt * dt);
		const float beta = dt * dt * c.damper;
		const float sigma = alpha * beta / dt;
		float C = angle;

		// This is using the initial velocities
		float Cdot = b2Dot(J1, dp1) + b2Dot(J2, dp2) + b2Dot(J3, dp3);

		float B = C + alpha * c.lambda + sigma * Cdot;
		float sum2 = (1.0f + sigma) * sum + alpha;

		float impulse = -B / sum2;

		p1 += (c.invMass1 * impulse) * J1;
		p2 += (c.invMass2 * impulse) * J2;
		p3 += (c.invMass3 * impulse) * J3;

		m_ps[c.i1] = p1;
		m_ps[c.i2] = p2;
		m_ps[c.i3] = p3;
		c.lambda += impulse;
	}
}

void b2Rope::ApplyBendForces(float dt)
{
	// omega = 2 * pi * hz
	const float omega = 2.0f * b2_pi * m_tuning.bendHertz;

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		const b2RopeBend& c = m_bendConstraints[i];

		b2Vec2 p1 = m_ps[c.i1];
		b2Vec2 p2 = m_ps[c.i2];
		b2Vec2 p3 = m_ps[c.i3];

		b2Vec2 v1 = m_vs[c.i1];
		b2Vec2 v2 = m_vs[c.i2];
		b2Vec2 v3 = m_vs[c.i3];

		b2Vec2 d1 = p2 - p1;
		b2Vec2 d2 = p3 - p2;

		float L1sqr, L2sqr;

		if (m_tuning.isometric)
		{
			L1sqr = c.L1 * c.L1;
			L2sqr = c.L2 * c.L2;
		}
		else
		{
			L1sqr = d1.LengthSquared();
			L2sqr = d2.LengthSquared();
		}

		if (L1sqr * L2sqr == 0.0f)
		{
			continue;
		}

		float a = b2Cross(d1, d2);
		float b = b2Dot(d1, d2);

		float angle = b2Atan2(a, b);

		b2Vec2 Jd1 = (-1.0f / L1sqr) * d1.Skew();
		b2Vec2 Jd2 = (1.0f / L2sqr) * d2.Skew();

		b2Vec2 J1 = -Jd1;
		b2Vec2 J2 = Jd1 - Jd2;
		b2Vec2 J3 = Jd2;

		float sum;
		if (m_tuning.fixedEffectiveMass)
		{
			sum = c.invEffectiveMass;
		}
		else
		{
			sum = c.invMass1 * b2Dot(J1, J1) + c.invMass2 * b2Dot(J2, J2) + c.invMass3 * b2Dot(J3, J3);
		}

		if (sum == 0.0f)
		{
			continue;
		}

		float mass = 1.0f / sum;

		const float spring = mass * omega * omega;
		const float damper = 2.0f * mass * m_tuning.bendDamping * omega;

		float C = angle;
		float Cdot = b2Dot(J1, v1) + b2Dot(J2, v2) + b2Dot(J3, v3);

		float impulse = -dt * (spring * C + damper * Cdot);

		m_vs[c.i1] += (c.invMass1 * impulse) * J1;
		m_vs[c.i2] += (c.invMass2 * impulse) * J2;
		m_vs[c.i3] += (c.invMass3 * impulse) * J3;
	}
}

void b2Rope::SolveBend_PBD_Distance()
{
	const float stiffness = m_tuning.bendStiffness;

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		const b2RopeBend& c = m_bendConstraints[i];

		int32 i1 = c.i1;
		int32 i2 = c.i3;

		b2Vec2 p1 = m_ps[i1];
		b2Vec2 p2 = m_ps[i2];

		b2Vec2 d = p2 - p1;
		float L = d.Normalize();

		float sum = c.invMass1 + c.invMass3;
		if (sum == 0.0f)
		{
			continue;
		}

		float s1 = c.invMass1 / sum;
		float s2 = c.invMass3 / sum;

		p1 -= stiffness * s1 * (c.L1 + c.L2 - L) * d;
		p2 += stiffness * s2 * (c.L1 + c.L2 - L) * d;

		m_ps[i1] = p1;
		m_ps[i2] = p2;
	}
}

// Constraint based implementation of:
// P. Volino: Simple Linear Bending Stiffness in Particle Systems
void b2Rope::SolveBend_PBD_Height()
{
	const float stiffness = m_tuning.bendStiffness;

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		const b2RopeBend& c = m_bendConstraints[i];

		b2Vec2 p1 = m_ps[c.i1];
		b2Vec2 p2 = m_ps[c.i2];
		b2Vec2 p3 = m_ps[c.i3];

		// Barycentric coordinates are held constant
		b2Vec2 d = c.alpha1 * p1 + c.alpha2 * p3 - p2;
		float dLen = d.Length();

		if (dLen == 0.0f)
		{
			continue;
		}

		b2Vec2 dHat = (1.0f / dLen) * d;

		b2Vec2 J1 = c.alpha1 * dHat;
		b2Vec2 J2 = -dHat;
		b2Vec2 J3 = c.alpha2 * dHat;

		float sum = c.invMass1 * c.alpha1 * c.alpha1 + c.invMass2 + c.invMass3 * c.alpha2 * c.alpha2;

		if (sum == 0.0f)
		{
			continue;
		}

		float C = dLen;
		float mass = 1.0f / sum;
		float impulse = -stiffness * mass * C;

		p1 += (c.invMass1 * impulse) * J1;
		p2 += (c.invMass2 * impulse) * J2;
		p3 += (c.invMass3 * impulse) * J3;

		m_ps[c.i1] = p1;
		m_ps[c.i2] = p2;
		m_ps[c.i3] = p3;
	}
}

// M. Kelager: A Triangle Bending Constraint Model for PBD
void b2Rope::SolveBend_PBD_Triangle()
{
	const float stiffness = m_tuning.bendStiffness;

	for (int32 i = 0; i < m_bendCount; ++i)
	{
		const b2RopeBend& c = m_bendConstraints[i];

		b2Vec2 b0 = m_ps[c.i1];
		b2Vec2 v = m_ps[c.i2];
		b2Vec2 b1 = m_ps[c.i3];

		float wb0 = c.invMass1;
		float wv = c.invMass2;
		float wb1 = c.invMass3;

		float W = wb0 + wb1 + 2.0f * wv;
		float invW = stiffness / W;

		b2Vec2 d = v - (1.0f / 3.0f) * (b0 + v + b1);

		b2Vec2 db0 = 2.0f * wb0 * invW * d;
		b2Vec2 dv = -4.0f * wv * invW * d;
		b2Vec2 db1 = 2.0f * wb1 * invW * d;

		b0 += db0;
		v += dv;
		b1 += db1;

		m_ps[c.i1] = b0;
		m_ps[c.i2] = v;
		m_ps[c.i3] = b1;
	}
}

void b2Rope::Draw(b2Draw* draw) const
{
	b2Color c(0.4f, 0.5f, 0.7f);
	b2Color pg(0.1f, 0.8f, 0.1f);
	b2Color pd(0.7f, 0.2f, 0.4f);

	for (int32 i = 0; i < m_count - 1; ++i)
	{
		draw->DrawSegment(m_ps[i], m_ps[i+1], c);

		const b2Color& pc = m_invMasses[i] > 0.0f ? pd : pg;
		draw->DrawPoint(m_ps[i], 5.0f, pc);
	}

	const b2Color& pc = m_invMasses[m_count - 1] > 0.0f ? pd : pg;
	draw->DrawPoint(m_ps[m_count - 1], 5.0f, pc);
}