/**
 * 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 <signal.h>

#include <OpenSim/OpenSim.h>
#include <OpenSim/Common/IO.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-6;

static Logger &logger = Logger::getInstance();
static LocomotorPrimitivesManager *manager = NULL;

static bool yes = false;

static void signal_handler(int signr)
{
	switch (signr) {
	case SIGINT:
		logger.log("Got signal SIGINT. Quitting...\n");
		if (manager)
			manager->halt();
	default:
		break;
	}
}

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.");
		if (!yes)
			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.01);
		m->setDefaultFiberLength(0.1);
	}

//	logger.log("+ Adding Actuation Analysis\n");
//	OpenSim::Actuation *actuation = new OpenSim::Actuation(&model);
//	actuation->setStepInterval(1);
//	model.addAnalysis(actuation);

	// Create the force reporter for obtaining the forces applied to the model
	// during a forward simulation
	logger.log("+ Adding Force Reporter\n");
	OpenSim::ForceReporter *reporter = new OpenSim::ForceReporter(&model);
	model.addAnalysis(reporter);

	// Muscle analysis for recording and computting basic quantities
	// (length, shortening velocity, tendon length, ...) for muscles during
	// a simulation
//	logger.log("+ Adding Muscle Analysis\n");
//	OpenSim::MuscleAnalysis *analysis = new OpenSim::MuscleAnalysis(&model);
//	model.addAnalysis(analysis);

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

void simulateModel(OpenSim::Model &model, const double initial_time,
		   const double final_time, double fixed_step_size,
		   std::string output_dir)
{
	if (NO_SIM) {
		logger.log("Skipping simulation as per NO_SIM=%d\n", NO_SIM);
		return;
	}

	std::string cwd = OpenSim::IO::getCwd();
	if (OpenSim::IO::chDir(output_dir) != 0) {
		logger.err("Failed to change to output directory %s", output_dir.c_str());
		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 and step size %f (0.0 means dynamic)\n",
			INTEGRATOR_ACCURACY, fixed_step_size);
	// Create the integrator and manager for the simulation
	SimTK::RungeKuttaMersonIntegrator integrator(model.getMultibodySystem());
	if (fixed_step_size > 0)
		integrator.setFixedStepSize(fixed_step_size);
	integrator.setAccuracy(INTEGRATOR_ACCURACY);

	logger.log("+ Creating simulation manager\n");
	manager = new LocomotorPrimitivesManager(model, integrator);
	// Integrate from initial time to final time
	manager->setInitialTime(initial_time);
	manager->setFinalTime(final_time);
	logger.log("+ Integrating from time %f to %f\n", initial_time, final_time);
	try {
		manager->integrate(si);
	} catch (...) {
		logger.log("Got exception during integration\n");
	}

	// 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");

	logger.log("+ Writing states (degrees) to %s_states_degrees.mot\n", MODEL_NAME.c_str());
	model.updSimbodyEngine().convertRadiansToDegrees(statesDegrees);
	statesDegrees.setWriteSIMMHeader(true);
	statesDegrees.print(MODEL_NAME + "_states_degrees.mot");

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

	// Save the actuation
//	logger.log("+ Writing actuation to %s_actuation.sto\n", MODEL_NAME.c_str());
	// TODO

//	logger.log("+ Writing muscle analysis to %s_muscles\n", MODEL_NAME.c_str());
	// TODO

	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
		  << "  -o DIR   set output directory to DIR (default: current directory)" << std::endl
		  << "  -s N     set fixed step size for integrator to N (default: dynamic)" << std::endl
		  << "  -y       don't prompt for any input on errors and finishing" << std::endl
		  << "  -h       show this help and exit" << std::endl;
}

int main(int argc, char **argv)
{
	clock_t ts_start = clock();
	double fixed_step_size = 0.0;
	/* set default values (to work on windows) */
	std::string model_file = FIXED_IN_SPACE ? "../../locomotor-primitives-fixed.osim" : "../../locomotor-primitives.osim";
	std::string data_file = "../../data/adults.sto";
	std::string output_dir = ".";

	sigset_t block_mask;
	struct sigaction saction;

	for (int i = 1; (i < argc) && (argv[i][0] == '-'); i++) {
		switch (argv[i][1]) {
		case 'd':
			data_file = argv[++i];
			break;
		case 'm':
			model_file = argv[++i];
			break;
		case 'o':
			output_dir = argv[++i];
			break;
		case 's':
			fixed_step_size = strtod(argv[++i], NULL);
			break;
		case 'y':
			yes = true;
			break;
		case 'h':
			usage();
			exit(0);
		default:
			std::cerr << "unknown option: -" << argv[i][1] << std::endl;
			usage();
			exit(-1);
		}
	}

	OpenSim::Storage actData(data_file);
	const double initial_time = actData.getFirstTime();
	const double final_time = actData.getLastTime();

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

	/* Register signal handler */
	sigfillset(&block_mask);
	saction.sa_handler = signal_handler;
	saction.sa_mask = block_mask;
	saction.sa_flags = SA_RESTART;
	sigaction(SIGINT, &saction, NULL);

	logger.log("Starting simulation %s with model file %s, actuation data file %s\n, fixed step size: %f\n", MODEL_NAME.c_str(), model_file.c_str(), data_file.c_str(), fixed_step_size);

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

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

		constructModel(osimModel, actData);
		simulateModel(osimModel, initial_time, final_time, fixed_step_size, output_dir);
	} catch (OpenSim::Exception ex) {
		logger.err("Exception: %s. Press ENTER to quit", ex.getMessage());
		if (!yes)
			std::cin.get();
		return 1;
	} catch (std::exception ex) {
		logger.err("Exception: %s. Press ENTER to quit", ex.what());
		if (!yes)
			std::cin.get();
		return 1;
	} catch (...) {
		logger.err("Unrecognized Exception. Press ENTER to quit");
		if (!yes)
			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());
	if (!yes)
		std::cin.get();
	return 0;
}