/**
 * Simulation for the Locomotor Primitives (Dominici et. al 2012) in a Human
 * Leg model with 14 Muscles.
 *
 * Author: Tobias Klauser <tobias.klauser@uzh.ch>
 *
 * Copyright (C) 2013 Tobias Klauser <tobias.klauser@uzh.ch>
 */

#include <OpenSim/OpenSim.h>

#include "LocomotorPrimitivesManager.h"
#include "LocomotorPrimitivesController.h"
#include "Logger.h"

#define FIXED_IN_SPACE	1
#define NO_SIM		0

static const std::string MODEL_NAME = "LocomotorPrimitives";
static const double INTEGRATOR_ACCURACY = 1.0e-3;

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

static void constructModel(OpenSim::Model &model, OpenSim::Storage actData)
{
	logger.log("Constructing model %s\n", MODEL_NAME.c_str());

	// Define controller for the model
	LocomotorPrimitivesController *control = new LocomotorPrimitivesController(actData, 0.01);
	control->setActuators(model.updActuators());
	if (control->checkControls() != 0) {
		logger.err("Control data for some muscles in the model is missing. Press ENTER to continue, ^-C to quit.");
		std::cin.get();
	}
	model.addController(control);

	// Set default activation and fiber length on all muscles of the model
	const OpenSim::Set<OpenSim::Actuator> &actSet = model.getActuators();
	int sz = actSet.getSize();
	logger.log("+ Defining initial muscle states for %d actuators:\n", sz);

	// Set default activation and fiber length of all muscles
	for (int i = 0; i < sz; i++) {
		OpenSim::ActivationFiberLengthMuscle *m = dynamic_cast<OpenSim::ActivationFiberLengthMuscle *>(&actSet.get(i));
		logger.log("  => %s\n", m->getName().c_str());
//		m->setDefaultActivation(0.5);
//		m->setDefaultFiberLength(0.1);
	}

	// Create the force reporter for obtaining the forces applied to the model
	// during a forward simulation
	OpenSim::ForceReporter *reporter = new OpenSim::ForceReporter(&model);
	model.addAnalysis(reporter);

	logger.log("Finished constructing model\n");
}

void simulateModel(OpenSim::Model &model, const double initialTime, const double finalTime)
{
	if (NO_SIM) {
		logger.log("Skipping simulation as per NO_SIM=%d\n", NO_SIM);
		return;
	}

	logger.log("Simulating model %s\n", MODEL_NAME.c_str());

	logger.log("+ Initializing system\n");
	SimTK::State &si = model.initSystem();

//	std::cout << "=== MODEL SUMMARY ===" << std::endl;
//	model.printDetailedInfo(si, std::cout);
//	std::cout << "=== END MODEL SUMMARY ===" << std::endl << std::endl;

	logger.log("+ Creating integrator (RungeKuttaMerson) with accuracy %e\n", INTEGRATOR_ACCURACY);
	// Create the integrator and manager for the simulation
	SimTK::RungeKuttaMersonIntegrator integrator(model.getMultibodySystem());
	integrator.setAccuracy(INTEGRATOR_ACCURACY);

	logger.log("+ Creating simulation manager\n");
	LocomotorPrimitivesManager manager(model, integrator);
	// Integrate from initial time to final time
	manager.setInitialTime(initialTime);
	manager.setFinalTime(finalTime);
	logger.log("+ Integrating from time %f to %f\n", initialTime, finalTime);
	manager.integrate(si);

	// Save the model states from forward integration
	logger.log("+ Writing states to %s_states.sto\n", MODEL_NAME.c_str());
	OpenSim::Storage statesDegrees(manager.getStateStorage());
	statesDegrees.print(MODEL_NAME + "_states.sto");

	// Save the forces
	logger.log("+ Writing forces to %s_forces.mot\n", MODEL_NAME.c_str());
	OpenSim::Analysis &reporter = model.getAnalysisSet()[0];
	((OpenSim::ForceReporter &)reporter).getForceStorage().print(MODEL_NAME + "_forces.mot");

	logger.log("Finished simulating model\n");
}

static void usage()
{
	std::cout << "usage: " << MODEL_NAME << " [OPTION...]" << std::endl << std::endl
		  << "  -d FILE  specify .sto datafile to use for muscle activation data" << std::endl
		  << "  -m FILE  specify .osim to use as model" << std::endl
		  << "  -h       show this help and exit" << std::endl;
}

int main(int argc, char **argv)
{
	clock_t ts_start = clock();

	/* set default values (to work on windows) */
	std::string modelFile = FIXED_IN_SPACE ? "../../locomotor-primitives-fixed.osim" : "../../locomotor-primitives.osim";
	std::string dataFile = "../../data/adults.sto";

	for (int i = 1; (i < argc) && (argv[i][0] == '-'); i++) {
		switch (argv[i][1]) {
		case 'd':
			dataFile = argv[++i];
			break;
		case 'm':
			modelFile = argv[++i];
			break;
		case 'h':
			usage();
			exit(0);
		default:
			std::cerr << "unknown option: -" << argv[i][1] << std::endl;
			usage();
			exit(-1);
		}
	}

	OpenSim::Storage actData(dataFile);
	const double initialTime = actData.getFirstTime();
	const double finalTime = actData.getLastTime();

	if (initialTime < 0 || finalTime < 0 || finalTime <= initialTime) {
		logger.err("invalid time range in data file: %f - %f\n", initialTime, finalTime);
		exit(-1);
	}

	logger.log("Starting simulation %s with model file %s, actuation data file %s\n", MODEL_NAME.c_str(), modelFile.c_str(), dataFile.c_str());

	try {
		// Create an OpenSim model and set its name
		OpenSim::Model osimModel(modelFile);

		osimModel.setName(MODEL_NAME);
		osimModel.setUseVisualizer(false);

		constructModel(osimModel, actData);
		simulateModel(osimModel, initialTime, finalTime);
	} catch (OpenSim::Exception ex) {
		logger.err("Exception: %s. Press ENTER to quit", ex.getMessage());
		std::cin.get();
		return 1;
	} catch (std::exception ex) {
		logger.err("Exception: %s. Press ENTER to quit", ex.what());
		std::cin.get();
		return 1;
	} catch (...) {
		logger.err("Unrecognized Exception. Press ENTER to quit");
		std::cin.get();
		return 1;
	}

	logger.log("Total runtime = %fms\n", 1.e3 * (clock() - ts_start) / CLOCKS_PER_SEC);
	logger.log("Simulation %s completed successfully. Press ENTER to quit", MODEL_NAME.c_str());
	std::cin.get();
	return 0;
}