#include "orbits.h"


const static uint ELLIPSE_VERT_COUNT = 256;


void
systemInit(TwoBodySystem& system, GravBody gb, OrbitalElements el)
{
	system.body = gb;
	system.elements = el;
	system.ep = ellipseInitAE(el.a, el.e);
	system.epsilon = orbitGetSpecificEnergy(system.ep.a, gb.mu);
	system.h = orbitGetAngularMomentum(system.ep.p, gb.mu);
	system.orbital_period = orbitGetPeriod(system.ep.a, gb.mu);
	system.r_periapsis = system.ep.a - system.ep.c;
	system.r_apoapsis = 2 * system.ep.a - system.r_periapsis;
	system.rotation = orbitGetXForm(el);

	system.sat.theta = el.nu;
	system.sat.r = orbitGetRadialDistance(system.ep.e, system.ep.p, el.nu);
	system.sat.v = orbitGetVelocity(system.epsilon, gb.mu, system.sat.r);
	system.sat.position =
		system.rotation * orbitGetPositionVector(system.sat.r, el.nu);
}

GravBody
gravBodyInit(double mu, double r)
{
	GravBody gb = {0};
	gb.mu = mu;
	gb.radius = r;
	return gb;
}

EllipseParameters
ellipseInitAB(double a, double b)
{
	assert(a > 0 && b > 0 && a >= b);
	EllipseParameters ep = { a, b };
	ep.c = sqrt(a * a - b * b);
	ep.e = ep.c / ep.a;
	ep.p = ep.a * (1 - pow(ep.e, 2));
	ep.f1.x = ep.c;
	ep.f2.x = -1 * ep.c;
	return ep;
}

// FIXME: we should avoid calling ellipseInitAB, and recalculate the properties
//	to avoid floating point errors in the known quantity 'e'
EllipseParameters
ellipseInitAE(double a, double e)
{
	assert(e >= 0 && e < 1);
	double b = a * sqrt(1 - pow(e, 2.0));
	return ellipseInitAB(a, b);
}

OrbitalElements
orbitInit(double a, double e, double iota, double ohm, double omega, double nu)
{
	OrbitalElements o = {0};
	o.a = a;
	o.e = e;
	o.iota = iota;
	o.ohm = ohm;
	o.omega = omega;
	o.nu = nu;
	return o;
}

OrbitalElements
orbitGetElementsFromStateVectors(glm::dvec3 r, glm::dvec3 v, double mu)
{
	OrbitalElements el = {0};
	const glm::dvec3 I = glm::dvec3(1, 0, 0);
	const glm::dvec3 J = glm::dvec3(0, 1, 0);
	const glm::dvec3 K = glm::dvec3(0, 0, 1);
	double r_mag = orbitGetVectorMagnitude(r);
	double v_mag = orbitGetVectorMagnitude(v);
	double epsilon = orbitGetSpecificEnergyFromStateVectors(r_mag, v_mag, mu);
	// TODO: orbits other than ellipses
	assert(epsilon < 0);
	glm::dvec3 ecc_v = orbitGetEccentricityVector(r, v, mu);

	el.a = orbitGetSemiMajorAxis(epsilon, mu);
	el.e = fabs(orbitGetVectorMagnitude(ecc_v));

	glm::dvec3 h = glm::cross(r, v);
	double cosi = glm::dot(K, h) / orbitGetVectorMagnitude(h);
	el.iota = acos(cosi);

	if (el.iota == 0) { // prograde equatorial orbit
		el.ohm = 0;
		double i_dot_e = glm::dot(I, ecc_v);
		el.omega = acos(i_dot_e / el.e);
		if (ecc_v.y < 0) el.omega *= -1; // quadrant check
	} else if (el.iota == M_PI) { // retrograde equatorial orbit
		// FIXME: retrograde equatorial orbit case
		assert(0);
	} else {
		glm::dvec3 n = glm::cross(K, h); // ascending node vector
		double n_mag = orbitGetVectorMagnitude(n);
		double cos_ohm = glm::dot(I, n) / n_mag;
		double sin_ohm = glm::dot(J, n) / n_mag;
		el.ohm = atan2(sin_ohm, cos_ohm);

		double cos_omega = glm::dot(n, ecc_v) / (n_mag * el.e);
		el.omega = acos(cos_omega);
		if (ecc_v.z < 0) el.omega = 2 * M_PI - el.omega; // quadrant check
	}

	if (el.e == 0) {
		el.nu = 0;
	} else {
		double cos_theta = glm::dot(r, ecc_v) / (el.e * r_mag);

		// NOTE: clamp to vaild range for acos function
		// 	we were getting slight floating point errors out of the glm::dot()
		// 	function above
		if (cos_theta > 1.0)
			cos_theta = 1.0;
		else if (cos_theta < - 1.0)
			cos_theta = -1.0;

		el.nu= acos(cos_theta);
	}

	// FIXME: breaks test case
	//if (glm::dot(r, v) < 0) el.nu *= -1; //quadrant check
	if (glm::dot(r, v) < 0) el.nu = 2 * M_PI - el.nu; //quadrant check

	return el;
}

StateVectors
orbitGetStateVectorsFromElements(const OrbitalElements& el, double mu)
{
	double p = 0;

	// FIXME: need a helper here
	if (el.e < 1) { // ellipse or circle
		p = el.a * (1 - pow(el.e, 2));
	} else if (el.e == 1) { // parabola
		assert(0);
	} else { // hyperbola
		assert(0);
	}

	// FIXME: we need to ensure that we update el.nu instead of sat.theta
	double r = orbitGetRadialDistance(el.e, p, el.nu);
	double h = orbitGetAngularMomentum(p, mu);
	StateVectors sv;
	// TODO: getPos/getVel don't need to be interface functions
	sv.position = orbitGetPositionVector(r, el.nu);
	sv.velocity = orbitGetVelocityVector(mu, h, el.e, el.nu);
	glm::dmat3 M = orbitGetXForm(el);
	sv.position = M * sv.position;
	sv.velocity = M * sv.velocity;

	return sv;
}

glm::dvec3
orbitGetEccentricityVector(glm::dvec3 r, glm::dvec3 v, double mu)
{
	double v_mag = orbitGetVectorMagnitude(v);
	double r_mag = orbitGetVectorMagnitude(r);
	return 1 / mu * ((pow(v_mag, 2) - mu / r_mag) * r - (glm::dot(r, v) * v));
}

glm::dvec3
orbitGetPositionVector(double r, double theta)
{
	return glm::dvec3(r * cos(theta), r * sin(theta), 0);
}

glm::dvec3
orbitGetVelocityVector(double mu, double h, double e, double theta)
{
	return glm::dvec3(-1 * (mu / h) * sin(theta), mu / h * (e + cos(theta)), 0);
}

glm::dmat3
orbitGetXForm(OrbitalElements elements)
{
	const OrbitalElements& el = elements;
	glm::mat3 M(1.0);
	M[0][0] = cos(el.ohm) * cos(el.omega) - sin(el.ohm) * sin(el.omega) * cos(el.iota);
	M[1][0] = -cos(el.ohm) * sin(el.omega) - sin(el.ohm) * cos(el.omega) * cos(el.iota);
	M[2][0] = sin(el.ohm) * sin(el.iota);
	M[0][1] = sin(el.ohm) * cos(el.omega) + cos(el.ohm) * sin(el.omega) * cos(el.iota);
	M[1][1] = -sin(el.ohm) * sin(el.omega) + cos(el.ohm) * cos(el.omega) * cos(el.iota);
	M[2][1] = -cos(el.ohm) * sin(el.iota);
	M[0][2] = sin(el.omega) * sin(el.iota);
	M[1][2] = cos(el.omega) * sin(el.iota);
	M[2][2] = cos(el.iota);
	return M;
}

double
orbitGetVectorMagnitude(glm::dvec3 v)
{
	return(sqrt(pow(v.x, 2) + pow(v.y, 2) + pow(v.z, 2)));
}


//
// NOTE: propagate anomaly functions:
//

// FIXME: variable names 'ecc' should be just 'e'
inline double
getEccAnomFromTrueAnom(double ecc, double true_anom)
{
	return 2 * atan(sqrt((1 - ecc) / (1 + ecc)) * tan(true_anom / 2));
}

inline double
getTrueAnomFromEccAnom(double ecc, double ecc_anom)
{
	return 2 * atan(sqrt((1 + ecc) / (1 - ecc)) * tan(ecc_anom / 2));
}

inline double
getMeanAnomFromEccAnom(double ecc_anom, double ecc)
{
	return ecc_anom - ecc * sin(ecc_anom);
}

inline double
getMeanMotion(double mu, double a)
{
	return sqrt(mu / pow(a, 3));
}

inline double
getPropagatedMeanAnom(double mean_anom, double mean_motion, double time_step)
{
	return mean_anom + mean_motion * (time_step);
}

inline double
getInitialTrialValue(double mean_anom, double ecc)
{
	return mean_anom + ecc * sin(mean_anom)
		+ ((pow(ecc, 2) / 2) * sin(2 * mean_anom));
}

inline double
getTrialError(double ecc, double test_anom, double mean_anom)
{
	return test_anom - ecc * sin(test_anom) - mean_anom;
}

inline double
getNextTrialValue(double err, double ecc, double test_anom, double mean_anom)
{
	// compute derivative of the error function
	double derr = 1 - ecc * cos(test_anom);

	// use Newton's method to compute next trial value of E2
	return test_anom - (err / derr);
}

double
getPropagatedEccAnomaly(TwoBodySystem sys,
	double initial_anom,
	double time_step)
{
	double e = sys.ep.e;
	double E1 = getEccAnomFromTrueAnom(e, initial_anom);
	double M1 = getMeanAnomFromEccAnom(E1, e);
	double n = getMeanMotion(sys.body.mu, sys.elements.a);
	double M2 = getPropagatedMeanAnom(M1, n, time_step);
	double E2_1 = getInitialTrialValue(M2, e);

	// test if guess is a solution to kepler's equation
	const double ACCEPTABLE_ERROR = 0.00000001;
	double E2_test = E2_1;

	for (uint i = 0; i < 10; i++) {
		double err = getTrialError(e, E2_test, M2);

		if (fabs(err) < ACCEPTABLE_ERROR)
			break;

		E2_test = getNextTrialValue(err, e, E2_test, M2);
	}

	return E2_test;
}

double
orbitGetPropagatedTrueAnomaly(TwoBodySystem sys,
	double initial_anom,
	double time_step)
{
	// FIXME: I don't think we need this now that we have gs->running?
	// NOTE: 'pause' simulation when time_step is set close to 0
	if (time_step < 1e-8 && time_step > -1e-8)
		return initial_anom;

	double ecc_anom = getPropagatedEccAnomaly(sys, initial_anom, time_step);
	return getTrueAnomFromEccAnom(sys.ep.e, ecc_anom);
}


// FIXME: organize into interface/internal functions

double
orbitTimeSincePeriapsis(TwoBodySystem sys, double theta);

double
orbitGetTimeOfFlight(TwoBodySystem sys, double theta_begin, double theta_end)
{
	double e = sys.ep.e;

	double n = getMeanMotion(sys.body.mu, sys.ep.a);

	double ecc_begin = getEccAnomFromTrueAnom(sys.ep.e, theta_begin);
	double ecc_end = getEccAnomFromTrueAnom(sys.ep.e, theta_end);

	// NOTE: test if flight passes through perisapsis
	if (ecc_begin > ecc_end)
		ecc_end += 2 * M_PI;

	double M1 = getMeanAnomFromEccAnom(ecc_begin, e);
	double M2 = getMeanAnomFromEccAnom(ecc_end, e);

	// NOTE: Kepler's equation for time of flight
	double tof_begin = 1 / n * M1;
	double tof_end = 1 / n * M2;

	return tof_end - tof_begin;
}

double
orbitGetTransferVelocity(const TwoBodySystem& sys,
	const OrbitalElements& target)
{
	const OrbitalElements& el1 = sys.elements;
	assert(el1.iota == target.iota);
	assert(el1.e == 0 && target.e == 0);

	const double mu = sys.body.mu;
	double a_t = (el1.a + target.a) / 2;
	double transfer_energy = orbitGetSpecificEnergy(a_t, mu);
	double r_periapse = el1.a; // NOTE: circular orbit

	double v1 = orbitGetVelocity(orbitGetSpecificEnergy(el1.a, mu), mu, el1.a);
	double vt = orbitGetVelocity(transfer_energy, mu, r_periapse);
	return vt - v1;
}

double
orbitGetCircVelocity(const TwoBodySystem& sys, bool raise_apoapse)
{
	double r = (raise_apoapse) ? sys.r_apoapsis : sys.r_periapsis;
	double mu = sys.body.mu;
	double epsilon_target = orbitGetSpecificEnergy(r, mu);
	double v = orbitGetVelocity(epsilon_target, mu, r);

	return v - sys.sat.v;
}


// internal

double
getApoapsis(double a, double e)
{
	return a * (1 + e);
}

double
getPeriapsis(double a, double e)
{
	return a * (1 - e);
}