/*
 * Author: Tobias Klauser <tobias.klauser@uzh.ch>
 *
 * Copyright (C) 2013 Tobias Klauser <tobias.klauser@uzh.ch>
 */

#include <OpenSim/OpenSim.h>

#include "LocomotorPrimitivesController.h"

static const double TIME_DAMP = 0.1;
static const unsigned int N_PRINT = 1000;

enum {
	BIFEMLH,
	BIFEMSH,
	SEMIMEM,
	SEMITEN,
	GLUT_MAX1,
	GLUT_MAX2,
	GLUT_MAX3,
	TFL,
};

/* TODO: Make the whole scaling business a bit more sophisticated, e.g. by
 * looking up a mapping name -> enum value */
static const double SCALE_MACT[] = {
	[BIFEMLH]	= 1.0,
	[BIFEMSH]	= 1.0,
	[SEMIMEM]	= 1.0,
	[SEMITEN]	= 1.0,
	[GLUT_MAX1]	= 1.0,
	[GLUT_MAX2]	= 1.0,
	[GLUT_MAX3]	= 1.0,
	[TFL]		= 1.0,
};

static Logger &logger = Logger::getInstance();

int LocomotorPrimitivesController::checkControls()
{
	const int act_set_siz = getActuatorSet().getSize();
	int ret = 0;

	/* Iterate over all controls in the model and check whether they're
	 * available in the loaded storage */
	for (int i = 0; i < act_set_siz; i++) {
		const OpenSim::Muscle *muscle = dynamic_cast<const OpenSim::Muscle *>(&getActuatorSet().get(i));
		std::string muscle_name = muscle->getName();

		/* If it is one of the 'special' combined muscles, check the
		 * combination instead. See comment in computeControls() for
		 * details. */
		if (muscle_name.compare("bifemlh_r") == 0) {
			muscle_name = "hs_r";
		} else if (muscle_name.compare("bifemlh_l") == 0) {
			muscle_name = "hs_l";
		} else if (muscle_name.compare("bifemsh_r") == 0) {
			muscle_name = "hs_r";
		} else if (muscle_name.compare("bifemsh_l") == 0) {
			muscle_name = "hs_l";
		} else if (muscle_name.compare("semimem_r") == 0) {
			muscle_name = "hs_r";
		} else if (muscle_name.compare("semimem_l") == 0) {
			muscle_name = "hs_l";
		} else if (muscle_name.compare("semiten_r") == 0) {
			muscle_name = "hs_r";
		} else if (muscle_name.compare("semiten_l") == 0) {
			muscle_name = "hs_l";
		} else if (muscle_name.compare("glut_max1_r") == 0) {
			muscle_name = "gm_r";
		} else if (muscle_name.compare("glut_max1_l") == 0) {
			muscle_name = "gm_l";
		} else if (muscle_name.compare("glut_max2_r") == 0) {
			muscle_name = "gm_r";
		} else if (muscle_name.compare("glut_max2_l") == 0) {
			muscle_name = "gm_l";
		} else if (muscle_name.compare("glut_max3_r") == 0) {
			muscle_name = "gm_r";
		} else if (muscle_name.compare("glut_max3_l") == 0) {
			muscle_name = "gm_l";
		}

		const OpenSim::Array<int> indices = _act.getColumnIndicesForIdentifier(muscle_name);
		if (indices.getSize() == 0) {
			logger.err("no actuation data for muscle '%s' found\n", muscle->getName().c_str());
			ret--;
		}
	}

	return ret;
}

double LocomotorPrimitivesController::getMuscleActivation(const std::string &muscle_name) const
{
	double mact;
	const OpenSim::Array<int> indices = _act.getColumnIndicesForIdentifier(muscle_name);

	if (indices.getSize() != 0) {
		int idx = indices.get(0) - 1;
		if (_muscle_act[idx] > 1.0) {
			logger.err("%s (%d): muscle actuation larger than 1.0 (%f), truncating to 1.0\n", muscle_name.c_str(), idx, _muscle_act[idx]);
			_muscle_act[idx] = 1.0;
		}

		mact = _muscle_act[idx];
	} else {
		mact = 0.0;
	}

	return mact;
}

/* Initialize the muscle actuation data array for all muscles, at a given
 * timestep */
int LocomotorPrimitivesController::getAndInitActuationData(const double t) const
{
	/* Extract muscle activation for each of the muscles */
	const int act_set_siz = getActuatorSet().getSize();
	const int nc = _act.getSmallestNumberOfStates();

	/* Get the first nc states at a specified time (_NOT_ the state at nc as
	 * with getData()). */
	_act.getDataAtTime(t, nc, _muscle_act);

	return act_set_siz;
}

void LocomotorPrimitivesController::computeControls(const SimTK::State &s,
						    SimTK::Vector &controls) const
{
	double t = s.getTime();
	static double last_twitch = t;
	static unsigned int n = 0;
	const int act_set_siz = getAndInitActuationData(t);

	for (int i = 0; i < act_set_siz; i++) {
		const OpenSim::Muscle *muscle = dynamic_cast<const OpenSim::Muscle *>(&getActuatorSet().get(i));
		std::string muscle_name = muscle->getName();
		double mact;

		/* Special case: In the paper some muscles were measured
		 * in combination but they're present as individual
		 * muscles in the model (see list below). Thus we
		 * linearly distribute the activation value among the
		 * individual muscles of the combination.
		 *
		 * TODO: Introduce scale factor for each muscle.
		 *
		 * measurement           muscle model
		 * -----------------------------------------------------
		 * HS (hamstring)        bifemlh bifemsh semimem semiten
		 * GM (gluteus maximus)  glut_max1 glut_max2 glut_max3
		 */
		if (muscle_name.compare("bifemlh_r") == 0) {
			mact = getMuscleActivation("hs_r");
			mact /= 4.0;
			mact *= SCALE_MACT[BIFEMLH];
		} else if (muscle_name.compare("bifemlh_l") == 0) {
			mact = getMuscleActivation("hs_l");
			mact /= 4.0;
			mact *= SCALE_MACT[BIFEMLH];
		} else if (muscle_name.compare("bifemsh_r") == 0) {
			mact = getMuscleActivation("hs_r");
			mact /= 4.0;
			mact *= SCALE_MACT[BIFEMSH];
		} else if (muscle_name.compare("bifemsh_l") == 0) {
			mact = getMuscleActivation("hs_l");
			mact /= 4.0;
			mact *= SCALE_MACT[BIFEMSH];
		} else if (muscle_name.compare("semimem_r") == 0) {
			mact = getMuscleActivation("hs_r");
			mact /= 4.0;
			mact *= SCALE_MACT[SEMIMEM];
		} else if (muscle_name.compare("semimem_l") == 0) {
			mact = getMuscleActivation("hs_l");
			mact /= 4.0;
			mact *= SCALE_MACT[SEMIMEM];
		} else if (muscle_name.compare("semiten_r") == 0) {
			mact = getMuscleActivation("hs_r");
			mact /= 4.0;
			mact *= SCALE_MACT[SEMITEN];
		} else if (muscle_name.compare("semiten_l") == 0) {
			mact = getMuscleActivation("hs_l");
			mact /= 4.0;
			mact *= SCALE_MACT[SEMITEN];
		} else if (muscle_name.compare("glut_max1_r") == 0) {
			mact = getMuscleActivation("gm_r");
			mact /= 3.0;
			mact *= SCALE_MACT[GLUT_MAX1];
		} else if (muscle_name.compare("glut_max1_l") == 0) {
			mact = getMuscleActivation("gm_l");
			mact /= 3.0;
			mact *= SCALE_MACT[GLUT_MAX1];
		} else if (muscle_name.compare("glut_max2_r") == 0) {
			mact = getMuscleActivation("gm_r");
			mact /= 3.0;
			mact *= SCALE_MACT[GLUT_MAX2];
		} else if (muscle_name.compare("glut_max2_l") == 0) {
			mact = getMuscleActivation("gm_l");
			mact /= 3.0;
			mact *= SCALE_MACT[GLUT_MAX2];
		} else if (muscle_name.compare("glut_max3_r") == 0) {
			mact = getMuscleActivation("gm_r");
			mact /= 3.0;
			mact *= SCALE_MACT[GLUT_MAX3];
		} else if (muscle_name.compare("glut_max3_l") == 0) {
			mact = getMuscleActivation("gm_l");
			mact /= 3.0;
			mact *= SCALE_MACT[GLUT_MAX3];
		} else {
			/* Otherwise just get the actuaction data for
			 * the plain name, it must be there! */
			mact = getMuscleActivation(muscle_name);

			/* Allow to scale some muscles afterwards */
			if (muscle_name.compare("tfl_l") == 0)
				mact *= SCALE_MACT[TFL];
			if (muscle_name.compare("tfl_r") == 0)
				mact *= SCALE_MACT[TFL];
		}

		if (mact > 1.0) {
			logger.warn("Actuation data for '%s' larger than 1.0, truncating.\n", muscle_name.c_str());
			mact = 1.0;
		}

		if (n % N_PRINT == 0)
			logger.log("timestep %f: actuation data for "
				   "muscle '%s': %f\n",
				   t, muscle_name.c_str(), mact);

		SimTK::Vector ctrl(1, mact);
		muscle->addInControls(ctrl, controls);
	}

	n++;
}