#include "Controller.h"
#include "Character.h"
#include "MultilevelSpline.h"
#include <boost/filesystem.hpp>
#include <Eigen/QR>
#include <fstream>
#include <numeric>
#include <algorithm>
namespace DPhy
{

Controller::Controller(ReferenceManager* ref, bool adaptive, bool parametric, bool record, int id)
	:mControlHz(30),mSimulationHz(150),mCurrentFrame(0),
	w_p(0.35),w_v(0.1),w_ee(0.3),w_com(0.25),
	terminationReason(-1),mIsNanAtTerminal(false), mIsTerminal(false)
{
	this->mRescaleParameter = std::make_tuple(1.0, 1.0, 1.0);
	this->isAdaptive = adaptive;
	this->isParametric = parametric;
	this->mRecord = record;
	this->mReferenceManager = ref;
	this->id = id;
	this->mParamGoal = mReferenceManager->GetParamGoal();
	this->mCurrentFrameOnPhase = 0;

	this->mSimPerCon = mSimulationHz / mControlHz;
	this->mWorld = std::make_shared<dart::simulation::World>();

	this->mBaseGravity = Eigen::Vector3d(0,-9.81, 0);
	this->mWorld->setGravity(this->mBaseGravity);

	this->mWorld->setTimeStep(1.0/(double)mSimulationHz);
	this->mWorld->getConstraintSolver()->setCollisionDetector(dart::collision::DARTCollisionDetector::create());
	dynamic_cast<dart::constraint::BoxedLcpConstraintSolver*>(mWorld->getConstraintSolver())->setBoxedLcpSolver(std::make_shared<dart::constraint::PgsBoxedLcpSolver>());

	this->mGround = DPhy::SkeletonBuilder::BuildFromFile(std::string(CAR_DIR)+std::string("/character/ground.xml")).first;
	this->mGround->getBodyNode(0)->setFrictionCoeff(1.0);
	this->mWorld->addSkeleton(this->mGround);

	std::string path = std::string(CAR_DIR)+std::string("/character/") + std::string(CHARACTER_TYPE) + std::string(".xml");
	this->mCharacter = new DPhy::Character(path);
	this->mWorld->addSkeleton(this->mCharacter->GetSkeleton());

	this->mBaseMass = mCharacter->GetSkeleton()->getMass();
	this->mMass = mBaseMass;

	Eigen::VectorXd kp(this->mCharacter->GetSkeleton()->getNumDofs()), kv(this->mCharacter->GetSkeleton()->getNumDofs());

	kp.setZero();
	kv.setZero();
	this->mCharacter->SetPDParameters(kp,kv);
	mContacts.clear();
	mContacts.push_back("RightToe");
	mContacts.push_back("RightFoot");
	mContacts.push_back("LeftToe");
	mContacts.push_back("LeftFoot");

	mInterestedDof = mCharacter->GetSkeleton()->getNumDofs() - 6;
	mRewardDof = mCharacter->GetSkeleton()->getNumDofs();

	auto collisionEngine = mWorld->getConstraintSolver()->getCollisionDetector();
	this->mCGL = collisionEngine->createCollisionGroup(this->mCharacter->GetSkeleton()->getBodyNode("LeftFoot"));
	this->mCGR = collisionEngine->createCollisionGroup(this->mCharacter->GetSkeleton()->getBodyNode("RightFoot"));
	this->mCGEL = collisionEngine->createCollisionGroup(this->mCharacter->GetSkeleton()->getBodyNode("LeftToe"));
	this->mCGER = collisionEngine->createCollisionGroup(this->mCharacter->GetSkeleton()->getBodyNode("RightToe"));
	this->mCGHL = collisionEngine->createCollisionGroup(this->mCharacter->GetSkeleton()->getBodyNode("LeftHand"));
	this->mCGHR = collisionEngine->createCollisionGroup(this->mCharacter->GetSkeleton()->getBodyNode("RightHand"));
	this->mCGG = collisionEngine->createCollisionGroup(this->mGround.get());

	int num_body_nodes = mInterestedDof / 3;
	int dof = this->mCharacter->GetSkeleton()->getNumDofs();

	mActions = Eigen::VectorXd::Zero(mInterestedDof + 1);
	mActions.setZero();

	mEndEffectors.clear();
	mEndEffectors.push_back("RightFoot");
	mEndEffectors.push_back("LeftFoot");
	mEndEffectors.push_back("LeftHand");
	mEndEffectors.push_back("RightHand");
	mEndEffectors.push_back("Head");

	this->mTargetPositions = Eigen::VectorXd::Zero(dof);
	this->mTargetVelocities = Eigen::VectorXd::Zero(dof);

	this->mPDTargetPositions = Eigen::VectorXd::Zero(dof);
	this->mPDTargetVelocities = Eigen::VectorXd::Zero(dof);

	//temp
	this->mRewardParts.resize(7, 0.0);
	mParamCur.resize(mReferenceManager->GetParamGoal().rows());
	this->mNumState = this->GetState().rows();

	this->mNumAction = mActions.size();

	ClearRecord();

	mRewardLabels.clear();
	if(isAdaptive) {
		mRewardLabels.push_back("total_d");
		mRewardLabels.push_back("total_s");
		mRewardLabels.push_back("tracking");
		mRewardLabels.push_back("time");
		mRewardLabels.push_back("similarity");
	} else {
		mRewardLabels.push_back("total");
		mRewardLabels.push_back("p");
		mRewardLabels.push_back("com");
		mRewardLabels.push_back("ee");
		mRewardLabels.push_back("v");
		mRewardLabels.push_back("time");
	}

}
const dart::dynamics::SkeletonPtr&
Controller::GetSkeleton() {
	return this->mCharacter->GetSkeleton();
}
void
Controller::
Step()
{
	if(IsTerminalState())
		return;

	Eigen::VectorXd s = this->GetState();

	Eigen::VectorXd a = mActions;

	// set action target pos
	int num_body_nodes = mInterestedDof / 3;
	int dof = this->mCharacter->GetSkeleton()->getNumDofs();

	for(int i = 0; i < mInterestedDof; i++){
		mActions[i] = dart::math::clip(mActions[i]*0.2, -0.7*M_PI, 0.7*M_PI);
	}
	int sign = 1;
	if(mActions[mInterestedDof] < 0)
		sign = -1;

	mActions[mInterestedDof] = dart::math::clip(mActions[mInterestedDof]*1.5, -3.0, 1.0);
	mActions[mInterestedDof] = exp(mActions[mInterestedDof]);
	mAdaptiveStep = mActions[mInterestedDof];
	// /////BASELINE
	// mAdaptiveStep = 1;
	// if(!isAdaptive)
	// 	mAdaptiveStep = 1;

	mPrevFrameOnPhase = this->mCurrentFrameOnPhase;
	this->mCurrentFrame += mAdaptiveStep;
	this->mCurrentFrameOnPhase += mAdaptiveStep;
	nTotalSteps += 1;
	int n_bnodes = mCharacter->GetSkeleton()->getNumBodyNodes();
	// if(mRecord)
	// 	std::cout << mCurrentFrameOnPhase << " "<< mAdaptiveStep << " "<< mReferenceManager->GetTimeStep(mPrevFrameOnPhase, true) << std::endl;

	Motion* p_v_target = mReferenceManager->GetMotion(mCurrentFrame, isAdaptive);
	this->mTargetPositions = p_v_target->GetPosition();
	this->mTargetVelocities = mCharacter->GetSkeleton()->getPositionDifferences(mTargetPositions, mPrevTargetPositions) / 0.033 * (mCurrentFrame - mPrevFrame);
	delete p_v_target;

	p_v_target = mReferenceManager->GetMotion(mCurrentFrame, false);
	this->mPDTargetPositions = p_v_target->GetPosition();
	this->mPDTargetVelocities = p_v_target->GetVelocity();
	delete p_v_target;

	int count_dof = 0;

	for(int i = 1; i <= num_body_nodes; i++){
		int idx = mCharacter->GetSkeleton()->getBodyNode(i)->getParentJoint()->getIndexInSkeleton(0);
		int dof = mCharacter->GetSkeleton()->getBodyNode(i)->getParentJoint()->getNumDofs();
		mPDTargetPositions.block(idx, 0, dof, 1) += mActions.block(count_dof, 0, dof, 1);
		count_dof += dof;

	}

	mSumTorque.resize(dof);
	mSumTorque.setZero();
	Eigen::VectorXd torque;
	Eigen::Vector3d d = Eigen::Vector3d(0, 0, 1);
	double end_f_sum = 0;

	for(int i = 0; i < this->mSimPerCon; i += 2){

		for(int j = 0; j < 2; j++) {
			mCharacter->GetSkeleton()->setSPDTarget(mPDTargetPositions, 600, 49);
			//Eigen::VectorXd torque = mCharacter->GetSkeleton()->getSPDForces(mPDTargetPositions, 600, 49, mWorld->getConstraintSolver());
			// for(int j = 0; j < num_body_nodes; j++) {
			// 	int idx = mCharacter->GetSkeleton()->getBodyNode(j)->getParentJoint()->getIndexInSkeleton(0);
			// 	int dof = mCharacter->GetSkeleton()->getBodyNode(j)->getParentJoint()->getNumDofs();
			// 	std::string name = mCharacter->GetSkeleton()->getBodyNode(j)->getName();
			// 	double torquelim = mCharacter->GetTorqueLimit(name) * 1.5;
			// 	double torque_norm = torque.block(idx, 0, dof, 1).norm();

			// 	torque.block(idx, 0, dof, 1) = std::max(-torquelim, std::min(torquelim, torque_norm)) * torque.block(idx, 0, dof, 1).normalized();
			// }

			//mCharacter->GetSkeleton()->setForces(torque);
			mWorld->step(false);
			//mSumTorque += torque.cwiseAbs();

		}
		if(mCurrentFrameOnPhase >= 18 && mControlFlag[0] == 0) {
			Eigen::Vector3d c_vel = mCharacter->GetSkeleton()->getCOMLinearVelocity();
			double rf = mCharacter->GetSkeleton()->getBodyNode("LeftFoot")->getWorldTransform().translation()(1);
			double lf = mCharacter->GetSkeleton()->getBodyNode("RightFoot")->getWorldTransform().translation()(1);
			if(mVelocity < c_vel(1)) {
				mVelocity = c_vel(1);
				mMomentum = mCharacter->GetSkeleton()->getMass() * c_vel;
			}
		}
		mTimeElapsed += 2 * mAdaptiveStep;
	}
	if(this->mCurrentFrameOnPhase > mReferenceManager->GetPhaseLength()){
		this->mCurrentFrameOnPhase -= mReferenceManager->GetPhaseLength();
		mRootZero = mCharacter->GetSkeleton()->getPositions().segment<6>(0);

		if(isAdaptive) {
			mTrackingRewardTrajectory /= mCountTracking;
			mFitness.sum_pos /= mCountTracking;
			mFitness.sum_vel /= mCountTracking;
			mFitness.sum_pos_threshold /= mCountTracking;
			mFitness.sum_vel_threshold /= mCountTracking;
			mFitness.sum_slide /= mCountSlide;

			mReferenceManager->SaveTrajectories(data_raw, std::tuple<double, double, Fitness>(mTrackingRewardTrajectory, mParamRewardTrajectory, mFitness), mParamCur);
			data_raw.clear();

			mFitness.sum_contact = 0;
			mFitness.sum_pos = 0;
			mFitness.sum_vel = 0;
			mFitness.sum_pos_threshold = 0;
			mFitness.sum_vel_threshold = 0;
			mFitness.sum_slide = 0;
			mFitness.sum_reward = 0;

			mTrackingRewardTrajectory = 0;
			mParamRewardTrajectory = 0;

			mControlFlag.setZero();
			mCountParam = 0;
			mCountTracking = 0;
			mCondiff = 0;
			mCountContact = 0;
			mCountSlide = 0;

			mVelocity = 0;
			mMomentum.setZero();
			mMaxCOM.setZero();
			mPrevHeight = 0;

			////BASELINE
			mParamRewardMax = 0;
		}
	}
	if(isAdaptive) {
		this->UpdateAdaptiveReward();
	}
	else
		this->UpdateReward();

	this->UpdateTerminalInfo();

	if(mRecord) {
		SaveStepInfo();
	}

	if(isAdaptive)
	{
		data_raw.push_back(std::pair<Eigen::VectorXd,double>(mCharacter->GetSkeleton()->getPositions(), mCurrentFrameOnPhase));
	}

	mPrevTargetPositions = mTargetPositions;
	mPrevFrame = mCurrentFrame;

	if(mPosQueue.size() >= 3)
		mPosQueue.pop();
	if(mTimeQueue.size() >= 3)
		mTimeQueue.pop();
	mPosQueue.push(mCharacter->GetSkeleton()->getPositions());
	mTimeQueue.push(mCurrentFrame);

	if(isAdaptive && mIsTerminal)
		data_raw.clear();

}
void
Controller::
SaveStepInfo()
{
	mRecordBVHPosition.push_back(mReferenceManager->GetPosition(mCurrentFrame, false));
	mRecordTargetPosition.push_back(mTargetPositions);
	mRecordPosition.push_back(mCharacter->GetSkeleton()->getPositions());
	mRecordVelocity.push_back(mCharacter->GetSkeleton()->getVelocities());
	mRecordCOM.push_back(mCharacter->GetSkeleton()->getCOM());
	mRecordPhase.push_back(mCurrentFrame);

	bool rightContact = CheckCollisionWithGround("RightFoot") || CheckCollisionWithGround("RightToe");
	bool leftContact = CheckCollisionWithGround("LeftFoot") || CheckCollisionWithGround("LeftToe");

	mRecordFootContact.push_back(std::make_pair(rightContact, leftContact));
}
void
Controller::
ClearRecord()
{
	this->mRecordVelocity.clear();
	this->mRecordPosition.clear();
	this->mRecordCOM.clear();
	this->mRecordTargetPosition.clear();
	this->mRecordBVHPosition.clear();
	this->mRecordObjPosition.clear();
	this->mRecordPhase.clear();
	this->mRecordFootContact.clear();

	this->mControlFlag.resize(4);
	this->mControlFlag.setZero();

	mCountParam = 0;
	mCountTracking = 0;

	while(!mPosQueue.empty())
		mPosQueue.pop();
	while(!mTimeQueue.empty())
		mTimeQueue.pop();

	data_raw.clear();
	mVelocity = 0;
	mMomentum.setZero();
	mMaxCOM.setZero();
	mCondiff = 0;
	mCountContact = 0;

	////BASELINE
	mParamRewardMax = 0;
	mPrevHeight = 0;
}

std::vector<double>
Controller::
GetTrackingReward(Eigen::VectorXd position, Eigen::VectorXd position2,
	Eigen::VectorXd velocity, Eigen::VectorXd velocity2, std::vector<std::string> list, bool useVelocity)
{
	auto& skel = this->mCharacter->GetSkeleton();
	int dof = skel->getNumDofs();
	int num_body_nodes = skel->getNumBodyNodes();

	Eigen::VectorXd p_save = skel->getPositions();
	Eigen::VectorXd v_save = skel->getVelocities();

	Eigen::VectorXd p_diff = skel->getPositionDifferences(position, position2);
	Eigen::VectorXd p_diff_reward;

	p_diff_reward = p_diff;
	//OURS
	// if(isAdaptive) {
	// 	p_diff_reward.segment<6>(0) *= 3;

	// }
	Eigen::VectorXd v_diff, v_diff_reward;

	if(useVelocity) {
		v_diff = skel->getVelocityDifferences(velocity, velocity2);
		v_diff_reward = v_diff;
	}

	skel->setPositions(position);
	skel->computeForwardKinematics(true,false,false);

	std::vector<Eigen::Isometry3d> ee_transforms;
	Eigen::VectorXd ee_diff(mEndEffectors.size()*3);
	ee_diff.setZero();
	for(int i=0;i<mEndEffectors.size(); i++){
		ee_transforms.push_back(skel->getBodyNode(mEndEffectors[i])->getWorldTransform());
	}

	Eigen::Vector3d com_diff = skel->getCOM();

	skel->setPositions(position2);
	skel->computeForwardKinematics(true,false,false);

	for(int i=0;i<mEndEffectors.size();i++){
		Eigen::Isometry3d diff = ee_transforms[i].inverse() * skel->getBodyNode(mEndEffectors[i])->getWorldTransform();
		ee_diff.segment<3>(3*i) = diff.translation();
	}
	com_diff -= skel->getCOM();

	double scale = 1.0;

	double sig_p = 0.4 * scale;
	double sig_v = 3 * scale;
	double sig_com = 0.2 * scale;
	double sig_ee = 0.2 * scale;

	double r_p = exp_of_squared(p_diff_reward,sig_p);
	double r_v;
	if(useVelocity)
	{
		r_v = exp_of_squared(v_diff_reward,sig_v);
	}
	double r_ee = exp_of_squared(ee_diff,sig_ee);
	double r_com = exp_of_squared(com_diff,sig_com);

	std::vector<double> rewards;
	rewards.clear();

	rewards.push_back(r_p);
	rewards.push_back(r_com);
	rewards.push_back(r_ee);

	if(useVelocity) {
		rewards.push_back(r_v);
	}

	skel->setPositions(p_save);
	skel->setVelocities(v_save);
	skel->computeForwardKinematics(true,true,false);

	return rewards;

}
std::vector<std::pair<bool, Eigen::Vector3d>>
Controller::
GetContactInfo(Eigen::VectorXd pos)
{
	auto& skel = this->mCharacter->GetSkeleton();
	Eigen::VectorXd p_save = skel->getPositions();
	Eigen::VectorXd v_save = skel->getVelocities();

	skel->setPositions(pos);
	skel->computeForwardKinematics(true,false,false);

	std::vector<std::string> contact;
	contact.push_back("RightFoot");
	contact.push_back("RightToe");
	contact.push_back("LeftFoot");
	contact.push_back("LeftToe");

	std::vector<std::pair<bool, Eigen::Vector3d>> result;
	result.clear();
	for(int i = 0; i < contact.size(); i++) {
		Eigen::Vector3d p = skel->getBodyNode(contact[i])->getWorldTransform().translation();
		if(p[1] < 0.07) {
			result.push_back(std::pair<bool, Eigen::Vector3d>(true, p));
		} else {
			result.push_back(std::pair<bool, Eigen::Vector3d>(false, p));
		}
	}

	skel->setPositions(p_save);
	skel->setVelocities(v_save);
	skel->computeForwardKinematics(true,true,false);

	return result;
}
double
Controller::
GetSimilarityReward()
{

	auto& skel = this->mCharacter->GetSkeleton();
	Eigen::VectorXd p_save = skel->getPositions();
	Eigen::VectorXd v_save = skel->getVelocities();

	auto p_v_target = mReferenceManager->GetMotion(mCurrentFrameOnPhase, false);
	Eigen::VectorXd pos = p_v_target->GetPosition();
	Eigen::VectorXd vel = p_v_target->GetVelocity();
	delete p_v_target;

	std::vector<std::pair<bool, Eigen::Vector3d>> contacts_ref = GetContactInfo(pos);
	std::vector<std::pair<bool, Eigen::Vector3d>> contacts_cur = GetContactInfo(skel->getPositions());

	double con_diff = 0;

	for(int i = 0; i < contacts_cur.size(); i++) {
		if(contacts_ref[i].first && !contacts_cur[i].first) {
			con_diff += abs(std::max(0.0, (contacts_cur[i].second)(1) - 0.07));
		} else if(!contacts_ref[i].first && contacts_cur[i].first) {
			con_diff += abs(std::max(0.0, (contacts_ref[i].second)(1) - 0.07));
		}
	}

	//double r_con = exp(-con_diff);
	Eigen::VectorXd p_aligned = skel->getPositions();
	std::vector<Eigen::VectorXd> p_with_zero;
	p_with_zero.push_back(mRootZero);
	p_with_zero.push_back(p_aligned.segment<6>(0));
	p_with_zero = Align(p_with_zero, mReferenceManager->GetPosition(0, false));
	p_aligned.segment<6>(0) = p_with_zero[1];

	Eigen::VectorXd v = skel->getPositionDifferences(skel->getPositions(), mPosQueue.front()) / (mCurrentFrame - mTimeQueue.front() + 1e-10) / 0.033;

	for(auto& jn : skel->getJoints()){
		if(dynamic_cast<dart::dynamics::RevoluteJoint*>(jn)!=nullptr){
			double v_ = v[jn->getIndexInSkeleton(0)];
			if(v_ > M_PI){
				v_ -= 2*M_PI;
			}
			else if(v_ < -M_PI){
				v_ += 2*M_PI;
			}
			v[jn->getIndexInSkeleton(0)] = v_;
		}
	}

	Eigen::VectorXd p_diff = skel->getPositionDifferences(pos, p_aligned);
	Eigen::VectorXd p_diff_th = p_diff;

	Eigen::VectorXd v_diff = skel->getVelocityDifferences(vel, v);

	int num_body_nodes = skel->getNumBodyNodes();
	for(int i =0 ; i < vel.rows(); i++) {
		v_diff(i) = v_diff(i) / std::max(0.5, vel(i));
	}
	Eigen::VectorXd v_diff_th = v_diff;

	for(int i = 0; i < num_body_nodes; i++) {
		std::string name = mCharacter->GetSkeleton()->getBodyNode(i)->getName();
		int idx = mCharacter->GetSkeleton()->getBodyNode(i)->getParentJoint()->getIndexInSkeleton(0);

		if(name.compare("Hips") == 0 ) {
			p_diff.segment<3>(idx) *= 3;
			p_diff.segment<3>(idx + 3) *= 5;

			p_diff_th.segment<3>(idx) *= 3;
			p_diff_th.segment<3>(idx + 3) *= 5;

			v_diff.segment<3>(idx + 3) *= 5;
			v_diff_th.segment<3>(idx + 3) *= 5;
		}
	}

	double footSlide = 0;
	if(mCurrentFrameOnPhase >= 41){
		Eigen::Vector3d lf = skel->getBodyNode("LeftFoot")->getWorldTransform().translation();
		lf += skel->getBodyNode("LeftToe")->getWorldTransform().translation();
		lf /= 2.0;

		Eigen::Vector3d rf = skel->getBodyNode("RightFoot")->getWorldTransform().translation();
		rf += skel->getBodyNode("RightToe")->getWorldTransform().translation();
		rf /= 2.0;

		Eigen::VectorXd foot_diff (6);
		foot_diff << (lf-stickLeftFoot), (rf- stickRightFoot);
		foot_diff(1) = 0;
		foot_diff(4) = 0;

		footSlide += foot_diff.segment<3>(0).norm() + foot_diff.segment<3>(3).norm();
		mCountSlide += 1;

	} else if(mCurrentFrameOnPhase >= 35) {
		Eigen::Vector3d lf = skel->getBodyNode("LeftFoot")->getWorldTransform().translation();
		lf += skel->getBodyNode("LeftToe")->getWorldTransform().translation();
		lf /= 2.0;

		Eigen::Vector3d rf = skel->getBodyNode("RightFoot")->getWorldTransform().translation();
		rf += skel->getBodyNode("RightToe")->getWorldTransform().translation();
		rf /= 2.0;

		stickRightFoot = rf;
		stickLeftFoot = lf;
	}

	double r_con = exp(-con_diff);
	double r_ee = exp_of_squared(v_diff, 3);
	double r_p = exp_of_squared(p_diff, 0.3);

	mPrevFrame = mCurrentFrame;

	mFitness.sum_contact += con_diff;
	mFitness.sum_pos += p_diff.dot(p_diff) / p_diff.rows();
	mFitness.sum_vel += v_diff.dot(v_diff) / v_diff.rows();

	mFitness.sum_pos_threshold += p_diff_th.dot(p_diff_th) / p_diff_th.rows();
	mFitness.sum_vel_threshold += v_diff_th.dot(v_diff_th) / v_diff_th.rows();

	mFitness.sum_slide += footSlide;
	return exp(-r_con)  * r_p * r_ee;
}
double
Controller::
GetParamReward()
{
	double r_param = 0;
	auto& skel = this->mCharacter->GetSkeleton();
	if(mCurrentFrameOnPhase >= 26 && mControlFlag[0] == 0) {
		Eigen::Vector3d momentumGoal = Eigen::Vector3d(0, 165, 0);
		Eigen::Vector3d m_diff = momentumGoal - mMomentum;
		m_diff *= 0.1;
		m_diff(1) *= 4;

		double r_m = exp_of_squared(m_diff, 1.5);
		if(r_m < 0.4) {
			mParamCur(0) = -5;
		} else {
			mParamCur(0) = mParamGoal(0);
		}
		mFitness.sum_reward = (0.5 + 0.5 * r_m);
		r_param = 0.6 * r_m;

		mControlFlag[0] = 1;
		if(mRecord) {
			std::cout << "momentum: " << mMomentum.transpose() << " / " << m_diff.transpose() << " / " << r_m << std::endl;
		}
	} if(mCurrentFrameOnPhase >= 30 && mControlFlag[0] == 1) {
		std::vector<std::pair<bool, Eigen::Vector3d>> contacts_ref = GetContactInfo(mReferenceManager->GetPosition(mCurrentFrameOnPhase, false));
		std::vector<std::pair<bool, Eigen::Vector3d>> contacts_cur = GetContactInfo(skel->getPositions());

		for(int i = 0; i < contacts_cur.size(); i++) {
			if(contacts_ref[i].first && !contacts_cur[i].first) {
				mCondiff += pow(std::max(0.0, (contacts_cur[i].second)(1) - 0.07), 2);
			} else if(!contacts_ref[i].first && contacts_cur[i].first) {
				mCondiff += pow(std::max(0.0, (contacts_ref[i].second)(1) - 0.07), 2);
			}
		}
		mCountContact += 1;

		if(mCurrentFrameOnPhase >= 44) {

			// // /////BASELINE
			// // Eigen::Vector3d momentumGoal = Eigen::Vector3d(0, 165, 0);
			// // Eigen::Vector3d m_diff = momentumGoal - mMomentum;
			// // m_diff *= 0.1;
			// // m_diff(1) *= 4;

			// // double r_m = exp_of_squared(m_diff, 1.5);
			// // if(r_m < 0.4) {
			// // 	mParamCur(0) = -5;
			// // } else {
			// // 	mParamCur(0) = mParamGoal(0);
			// // }
			// // r_param = r_m;
			// // mParamRewardMax = r_param;

			double r_c = exp(-mCondiff * 7.5);
			if(r_c < 0.5) {
				mParamCur(0) = -5;
			}
			// else {
			// 	mParamCur(0) = mParamGoal(0);
			// }
			r_param = 0.6 * r_c;
			mFitness.sum_reward *= r_c;

			mControlFlag[0] = 2;

			if(mRecord) {
				std::cout << "contact: " << mCondiff << " / " << r_c << std::endl;
				std::cout << "final parameter: " <<  mParamCur.transpose() << std::endl;
			}
		}
	}
	return r_param;

}
void
Controller::
UpdateAdaptiveReward()
{

	auto& skel = this->mCharacter->GetSkeleton();

	std::vector<double> tracking_rewards_bvh = this->GetTrackingReward(skel->getPositions(), mTargetPositions,
								 skel->getVelocities(), mTargetVelocities, mRewardBodies, true);
	double accum_bvh = std::accumulate(tracking_rewards_bvh.begin(), tracking_rewards_bvh.end(), 0.0) / tracking_rewards_bvh.size();
	double time_diff = mAdaptiveStep  - mReferenceManager->GetTimeStep(mPrevFrameOnPhase, true);
	double r_time = exp(-pow(time_diff, 2)*75);

	double r_tracking = 0.85 * accum_bvh + 0.15 * r_time;

	double r_similarity = this->GetSimilarityReward();
	double r_param = this->GetParamReward();

	double r_tot = r_tracking;

	mRewardParts.clear();

	if(dart::math::isNan(r_tot)){
		mRewardParts.resize(mRewardLabels.size(), 0.0);
	}
	else {
		mRewardParts.push_back(r_tot);
		mRewardParts.push_back(20 * r_param);
		mRewardParts.push_back(accum_bvh);
		mRewardParts.push_back(r_time);
		mRewardParts.push_back(r_similarity);
	}
	if(r_param != 0) {
		if(mParamRewardTrajectory == 0) {
			mParamRewardTrajectory = r_param;
		}
		else {
			mParamRewardTrajectory *= r_param;
		}
	}
	mTrackingRewardTrajectory += accum_bvh;
	mCountTracking += 1;
}
void
Controller::
UpdateReward()
{
	auto& skel = this->mCharacter->GetSkeleton();
	std::vector<double> tracking_rewards_bvh = this->GetTrackingReward(skel->getPositions(), mTargetPositions,
								 skel->getVelocities(), mTargetVelocities, mRewardBodies, true);
	double accum_bvh = std::accumulate(tracking_rewards_bvh.begin(), tracking_rewards_bvh.end(), 0.0) / tracking_rewards_bvh.size();

	double r_time = exp(-pow((mActions[mInterestedDof] - 1),2)*40);

	mSumTorque /= mSimPerCon;
	double r_torque = exp_of_squared(mSumTorque, 50);

	mRewardParts.clear();
	double r_tot = 0.9 * (0.5 * tracking_rewards_bvh[0] + 0.1 * tracking_rewards_bvh[1] + 0.3 * tracking_rewards_bvh[2] + 0.1 * tracking_rewards_bvh[3] ) + 0.1 * r_time;
	// r_tot = 0.99 * r_tot + 0.01 * r_torque;
	if(dart::math::isNan(r_tot)){
		mRewardParts.resize(mRewardLabels.size(), 0.0);
	}
	else {
		mRewardParts.push_back(r_tot);
		mRewardParts.push_back(tracking_rewards_bvh[0]);
		mRewardParts.push_back(tracking_rewards_bvh[1]);
		mRewardParts.push_back(tracking_rewards_bvh[2]);
		mRewardParts.push_back(tracking_rewards_bvh[3]);
		mRewardParts.push_back(r_torque);
	}
}
void
Controller::
UpdateTerminalInfo()
{
	Eigen::VectorXd p_ideal = mTargetPositions;
	auto& skel = this->mCharacter->GetSkeleton();

	Eigen::VectorXd p = skel->getPositions();
	Eigen::VectorXd v = skel->getVelocities();
	Eigen::Vector3d root_pos = skel->getPositions().segment<3>(3);
	Eigen::Isometry3d cur_root_inv = skel->getRootBodyNode()->getWorldTransform().inverse();
	double root_y = skel->getBodyNode(0)->getTransform().translation()[1];

	Eigen::VectorXd p_save = skel->getPositions();
	Eigen::VectorXd v_save = skel->getVelocities();

	skel->setPositions(mTargetPositions);
	skel->computeForwardKinematics(true,false,false);

	Eigen::Isometry3d root_diff = cur_root_inv * skel->getRootBodyNode()->getWorldTransform();

	Eigen::AngleAxisd root_diff_aa(root_diff.linear());
	double angle = RadianClamp(root_diff_aa.angle());
	Eigen::Vector3d root_pos_diff = root_diff.translation();

	// check nan
	if(dart::math::isNan(p)){
		mIsNanAtTerminal = true;
		mIsTerminal = true;
		terminationReason = 3;
	} else if(dart::math::isNan(v)){
		mIsNanAtTerminal = true;
		mIsTerminal = true;
		terminationReason = 4;
	}
	//characterConfigration
	else if(!mRecord && root_pos_diff.norm() > TERMINAL_ROOT_DIFF_THRESHOLD){
		mIsTerminal = true;
		terminationReason = 2;
	} else if(!mRecord && root_y<TERMINAL_ROOT_HEIGHT_LOWER_LIMIT || root_y > TERMINAL_ROOT_HEIGHT_UPPER_LIMIT){
		mIsTerminal = true;
		terminationReason = 1;
	} else if(!mRecord && std::abs(angle) > TERMINAL_ROOT_DIFF_ANGLE_THRESHOLD){
		mIsTerminal = true;
		terminationReason = 5;
	} else if(isAdaptive && mCurrentFrame > mReferenceManager->GetPhaseLength()* 3 + 10) { // this->mBVH->GetMaxFrame() - 1.0){
		mIsTerminal = true;
		terminationReason =  8;
	} else if(!isAdaptive && mCurrentFrame > mReferenceManager->GetPhaseLength()* 5 + 10) { // this->mBVH->GetMaxFrame() - 1.0){
		mIsTerminal = true;
		terminationReason =  8;
	}
	if(mRecord) {
		if(mIsTerminal) std::cout << terminationReason << std::endl;
	}

	skel->setPositions(p_save);
	skel->setVelocities(v_save);
	skel->computeForwardKinematics(true,true,false);

}
bool
Controller::
FollowBvh()
{
	if(IsTerminalState())
		return false;
	auto& skel = mCharacter->GetSkeleton();

	Motion* p_v_target = mReferenceManager->GetMotion(mCurrentFrame);
	mTargetPositions = p_v_target->GetPosition();
	mTargetVelocities = p_v_target->GetVelocity();
	delete p_v_target;

	for(int i=0;i<this->mSimPerCon;i++)
	{
		skel->setPositions(mTargetPositions);
		skel->setVelocities(mTargetVelocities);
		skel->computeForwardKinematics(true, true, false);
	}
	this->mCurrentFrame += 1;
	this->nTotalSteps += 1;
	return true;
}
void
Controller::
SetGoalParameters(Eigen::VectorXd tp)
{
	mParamGoal = tp;
	this->mWorld->setGravity(exp(mParamGoal(0))*mBaseGravity);
	// this->SetSkeletonWeight(mParamGoal(1)*mBaseMass);
}

void
Controller::
SetSkeletonWeight(double mass)
{

	double m_new = mass / mMass;

	std::vector<std::tuple<std::string, Eigen::Vector3d, double>> deform;
	int n_bnodes = mCharacter->GetSkeleton()->getNumBodyNodes();

	for(int i = 0; i < n_bnodes; i++){
		std::string name = mCharacter->GetSkeleton()->getBodyNode(i)->getName();
		deform.push_back(std::make_tuple(name, Eigen::Vector3d(1, 1, 1), m_new));
	}
	DPhy::SkeletonBuilder::DeformSkeleton(mCharacter->GetSkeleton(), deform);
	mMass = mCharacter->GetSkeleton()->getMass();
}
void
Controller::
Reset(bool RSI)
{
	this->mWorld->reset();
	auto& skel = mCharacter->GetSkeleton();
	skel->clearConstraintImpulses();
	skel->clearInternalForces();
	skel->clearExternalForces();

	//RSI
	if(RSI && !isAdaptive) {
		this->mCurrentFrame = (int) dart::math::Random::uniform(0.0, mReferenceManager->GetPhaseLength()-5.0);
	}
	else {
		this->mCurrentFrame = 0; // 0;
		this->mParamRewardTrajectory = 0;
		this->mTrackingRewardTrajectory = 0;

		mFitness.sum_contact = 0;
		mFitness.sum_pos = 0;
		mFitness.sum_vel = 0;
		mFitness.sum_pos_threshold = 0;
		mFitness.sum_pos_threshold = 0;
		mFitness.sum_reward = 0;

	}

	this->mCurrentFrameOnPhase = this->mCurrentFrame;
	this->mStartFrame = this->mCurrentFrame;
	this->nTotalSteps = 0;
	this->mTimeElapsed = 0;

	Motion* p_v_target;
	p_v_target = mReferenceManager->GetMotion(mCurrentFrame, isAdaptive);
	this->mTargetPositions = p_v_target->GetPosition();
	this->mTargetVelocities = p_v_target->GetVelocity();
	delete p_v_target;

	this->mPDTargetPositions = mTargetPositions;
	this->mPDTargetVelocities = mTargetVelocities;
	int num_body_nodes = skel->getNumBodyNodes();

	skel->setPositions(mTargetPositions);
	skel->setVelocities(mTargetVelocities);
	skel->computeForwardKinematics(true,true,false);

	this->mIsNanAtTerminal = false;
	this->mIsTerminal = false;

	ClearRecord();
	SaveStepInfo();

	mRootZero = mCharacter->GetSkeleton()->getPositions().segment<6>(0);

	mPrevPositions = mCharacter->GetSkeleton()->getPositions();
	mPrevTargetPositions = mTargetPositions;

	mPrevFrame = mCurrentFrame;
	mPrevFrame2 = mPrevFrame;

	mPosQueue.push(mCharacter->GetSkeleton()->getPositions());
	mTimeQueue.push(0);
	mAdaptiveStep = 1;
	if(isAdaptive)
	{
		data_raw.push_back(std::pair<Eigen::VectorXd,double>(mCharacter->GetSkeleton()->getPositions(), mCurrentFrame));
	}

}
int
Controller::
GetNumState()
{
	return this->mNumState;
}
int
Controller::
GetNumAction()
{
	return this->mNumAction;
}
void
Controller::
SetAction(const Eigen::VectorXd& action)
{
	this->mActions = action;
}
Eigen::VectorXd
Controller::
GetEndEffectorStatePosAndVel(const Eigen::VectorXd pos, const Eigen::VectorXd vel) {
	Eigen::VectorXd ret;
	auto& skel = mCharacter->GetSkeleton();
	dart::dynamics::BodyNode* root = skel->getRootBodyNode();
	Eigen::Isometry3d cur_root_inv = root->getWorldTransform().inverse();

	int num_ee = mEndEffectors.size();
	Eigen::VectorXd p_save = skel->getPositions();
	Eigen::VectorXd v_save = skel->getVelocities();

	skel->setPositions(pos);
	skel->setVelocities(vel);
	skel->computeForwardKinematics(true, true, false);

	ret.resize((num_ee)*12+15);
//	ret.resize((num_ee)*9+12);

	for(int i=0;i<num_ee;i++)
	{
		Eigen::Isometry3d transform = cur_root_inv * skel->getBodyNode(mEndEffectors[i])->getWorldTransform();
		//Eigen::Quaterniond q(transform.linear());
		// Eigen::Vector3d rot = QuaternionToDARTPosition(Eigen::Quaterniond(transform.linear()));
		ret.segment<9>(9*i) << transform.linear()(0,0), transform.linear()(0,1), transform.linear()(0,2),
							   transform.linear()(1,0), transform.linear()(1,1), transform.linear()(1,2),
							   transform.translation();
//		ret.segment<6>(6*i) << rot, transform.translation();
	}

	for(int i=0;i<num_ee;i++)
	{
	    int idx = skel->getBodyNode(mEndEffectors[i])->getParentJoint()->getIndexInSkeleton(0);
		ret.segment<3>(9*num_ee + 3*i) << vel.segment<3>(idx);
//	    ret.segment<3>(6*num_ee + 3*i) << vel.segment<3>(idx);

	}

	// root diff with target com
	Eigen::Isometry3d transform = cur_root_inv * skel->getRootBodyNode()->getWorldTransform();
	//Eigen::Quaterniond q(transform.linear());

	Eigen::Vector3d rot = QuaternionToDARTPosition(Eigen::Quaterniond(transform.linear()));
	Eigen::Vector3d root_angular_vel_relative = cur_root_inv.linear() * skel->getRootBodyNode()->getAngularVelocity();
	Eigen::Vector3d root_linear_vel_relative = cur_root_inv.linear() * skel->getRootBodyNode()->getCOMLinearVelocity();

	ret.tail<15>() << transform.linear()(0,0), transform.linear()(0,1), transform.linear()(0,2),
					  transform.linear()(1,0), transform.linear()(1,1), transform.linear()(1,2),
					  transform.translation(), root_angular_vel_relative, root_linear_vel_relative;
//	ret.tail<12>() << rot, transform.translation(), root_angular_vel_relative, root_linear_vel_relative;

	// restore
	skel->setPositions(p_save);
	skel->setVelocities(v_save);
	skel->computeForwardKinematics(true, true, false);

	return ret;
}
bool
Controller::
CheckCollisionWithGround(std::string bodyName){
	auto collisionEngine = mWorld->getConstraintSolver()->getCollisionDetector();
	dart::collision::CollisionOption option;
	dart::collision::CollisionResult result;
	if(bodyName == "RightFoot"){
		bool isCollide = collisionEngine->collide(this->mCGR.get(), this->mCGG.get(), option, &result);
		return isCollide;
	}
	else if(bodyName == "LeftFoot"){
		bool isCollide = collisionEngine->collide(this->mCGL.get(), this->mCGG.get(), option, &result);
		return isCollide;
	}
	else if(bodyName == "RightToe"){
		bool isCollide = collisionEngine->collide(this->mCGER.get(), this->mCGG.get(), option, &result);
		return isCollide;
	}
	else if(bodyName == "LeftToe"){
		bool isCollide = collisionEngine->collide(this->mCGEL.get(), this->mCGG.get(), option, &result);
		return isCollide;
	}
	else if(bodyName == "RightHand"){
		bool isCollide = collisionEngine->collide(this->mCGHR.get(), this->mCGG.get(), option, &result);
		return isCollide;
	}
	else if(bodyName == "LeftHand"){
		bool isCollide = collisionEngine->collide(this->mCGHL.get(), this->mCGG.get(), option, &result);
		return isCollide;
	}
	else{ // error case
		std::cout << "check collision : bad body name" << std::endl;
		return false;
	}
}
Eigen::VectorXd
Controller::
GetState()
{
	if(mIsTerminal && terminationReason != 8){
		return Eigen::VectorXd::Zero(mNumState);
	}
	auto& skel = mCharacter->GetSkeleton();

	double root_height = skel->getRootBodyNode()->getCOM()[1];

	Eigen::VectorXd p_save = skel->getPositions();
	Eigen::VectorXd v_save = skel->getVelocities();
	Eigen::VectorXd p,v;
	// p.resize(p_save.rows()-6);
	// p = p_save.tail(p_save.rows()-6);

	int n_bnodes = mCharacter->GetSkeleton()->getNumBodyNodes();
	int num_p = (n_bnodes - 1) * 6;
	p.resize(num_p);

	for(int i = 1; i < n_bnodes; i++){
		Eigen::Isometry3d transform = skel->getBodyNode(i)->getRelativeTransform();
		// Eigen::Quaterniond q(transform.linear());
		p.segment<6>(6*(i-1)) << transform.linear()(0,0), transform.linear()(0,1), transform.linear()(0,2),
								 transform.linear()(1,0), transform.linear()(1,1), transform.linear()(1,2);
	}

	v = v_save;

	dart::dynamics::BodyNode* root = skel->getRootBodyNode();
	Eigen::Isometry3d cur_root_inv = root->getWorldTransform().inverse();
	Eigen::VectorXd ee;
	ee.resize(mEndEffectors.size()*3);
	for(int i=0;i<mEndEffectors.size();i++)
	{
		Eigen::Isometry3d transform = cur_root_inv * skel->getBodyNode(mEndEffectors[i])->getWorldTransform();
		ee.segment<3>(3*i) << transform.translation();
	}
	double t = mReferenceManager->GetTimeStep(mCurrentFrameOnPhase, isAdaptive);

	Motion* p_v_target = mReferenceManager->GetMotion(mCurrentFrame+t, isAdaptive);
	Eigen::VectorXd p_next = GetEndEffectorStatePosAndVel(p_v_target->GetPosition(), p_v_target->GetVelocity()*t);

	delete p_v_target;

	Eigen::Vector3d up_vec = root->getTransform().linear()*Eigen::Vector3d::UnitY();
	double up_vec_angle = atan2(std::sqrt(up_vec[0]*up_vec[0]+up_vec[2]*up_vec[2]),up_vec[1]);
	double phase = ((int) mCurrentFrame % mReferenceManager->GetPhaseLength()) / (double) mReferenceManager->GetPhaseLength();
	Eigen::VectorXd state;

	double com_diff = 0;

	std::vector<std::string> contact;
	contact.push_back("RightFoot");
	contact.push_back("RightToe");
	contact.push_back("LeftFoot");
	contact.push_back("LeftToe");

	Eigen::Vector4d foot_height;
	for(int i = 0; i < contact.size(); i++) {
		Eigen::Vector3d p = skel->getBodyNode(contact[i])->getWorldTransform().translation();
		foot_height[i] = p[1];
	}

	if(isParametric) {
		state.resize(p.rows()+v.rows()+1+1+p_next.rows()+ee.rows()+2+mParamGoal.rows()+4);
		state<< p, v, up_vec_angle, root_height, p_next, mAdaptiveStep, ee, mCurrentFrameOnPhase, mParamGoal, foot_height;
	}
	else {
		state.resize(p.rows()+v.rows()+1+1+p_next.rows()+ee.rows()+2);
		state<< p, v, up_vec_angle, root_height, p_next, mAdaptiveStep, ee, mCurrentFrameOnPhase;
	}

	return state;
}
void
Controller::SaveTimeData(std::string directory) {
	std::string path = std::string(CAR_DIR) + std::string("/") +  directory;
	std::cout << "save results to" << path << std::endl;

	std::ofstream ofs(path);
	ofs << mReferenceManager->GetPhaseLength() << std::endl;

	for(int i = 0; i < mRecordPhase.size() - 1; i++) {
		ofs << i << " " << mRecordPhase[i] << " " << mRecordPhase[i+1] - mRecordPhase[i] << std::endl;
	}
	ofs.close();

}
void
Controller::SaveDisplayedData(std::string directory, bool normalized) {
	std::string path = directory;
	std::cout << "save results to" << path << std::endl;
	std::vector<std::string> HIERARCHY = mReferenceManager->GetHierarchyStr();

	std::ofstream ofs(path);
	std::vector<std::string> bvh_order;
	bvh_order.push_back("Hips");
	bvh_order.push_back("Spine");
	bvh_order.push_back("Spine1");
	bvh_order.push_back("Spine2");
	bvh_order.push_back("Neck");
	bvh_order.push_back("Head");
	bvh_order.push_back("LeftShoulder");
	bvh_order.push_back("LeftArm");
	bvh_order.push_back("LeftForeArm");
	bvh_order.push_back("LeftHand");
	bvh_order.push_back("RightShoulder");
	bvh_order.push_back("RightArm");
	bvh_order.push_back("RightForeArm");
	bvh_order.push_back("RightHand");
	bvh_order.push_back("LeftUpLeg");
	bvh_order.push_back("LeftLeg");
	bvh_order.push_back("LeftFoot");
	bvh_order.push_back("LeftToe");
	bvh_order.push_back("RightUpLeg");
	bvh_order.push_back("RightLeg");
	bvh_order.push_back("RightFoot");
	bvh_order.push_back("RightToe");

	std::vector<Eigen::VectorXd> normalizedPosition;
	std::vector<double> normalizedDPhase;

	if(normalized){
		int count = 0;
		for(int i = 0; i < mReferenceManager->GetPhaseLength(); i++) {
			while(count + 1 < mRecordPosition.size() && i >= mRecordPhase[count+1])
				count += 1;
			Eigen::VectorXd p(mCharacter->GetSkeleton()->getNumDofs());
			double dp = 0;
			if(i < mRecordPhase[count]) {
				p = mRecordPosition[count];
				dp = mRecordPhase[count + 1] - mRecordPhase[count - 1];
			} else if(count == data_raw.size() - 1 && i > mRecordPhase[count]) {
				p = mRecordPosition[count];
				dp = mRecordPhase[count] - mRecordPhase[count - 1];
			} else if(i == mRecordPhase[count]) {
				p = mRecordPosition[count];
				dp = mRecordPhase[count+1] - mRecordPhase[count];
			} else {
				double weight = 1.0 - (mRecordPhase[count+1] - i) / (mRecordPhase[count+1] - mRecordPhase[count]);
				double dp0 = mRecordPhase[count+1] - mRecordPhase[count];
				double dp1 = mRecordPhase[count+2] - mRecordPhase[count+1];
				p = DPhy::BlendPosition(mRecordPosition[count], mRecordPosition[count+1], weight);
				dp = (1 - weight) * dp0 + weight * dp1;
			}
			normalizedPosition.push_back(p);
			normalizedDPhase.push_back(dp);
		}
	}

	ofs << "HIERARCHY" << std::endl;
	ofs << "ROOT Hips" << std::endl;
	for(int i = 0; i < HIERARCHY.size(); i++) {
		ofs << HIERARCHY[i] << std::endl;
	}
	ofs << "MOTION" << std::endl;
	if(normalized) {
		ofs << "Frames: " << std::to_string(normalizedPosition.size()) << std::endl;
	} else {
		ofs << "Frames: " << std::to_string(mRecordPosition.size()) << std::endl;
	}
	ofs << "Frame Time:	0.0333333" << std::endl;
	if(normalized) {
		for(auto t: normalizedPosition) {

			ofs << t.segment<3>(3).transpose() * 100 << " ";

			for(int i = 0; i < bvh_order.size(); i++) {
				int idx = mCharacter->GetSkeleton()->getBodyNode(bvh_order[i])->getParentJoint()->getIndexInSkeleton(0);
				Eigen::AngleAxisd aa(t.segment<3>(idx).norm(), t.segment<3>(idx).normalized());
				Eigen::Matrix3d m;
				m = aa;
				Eigen::Vector3d v = dart::math::matrixToEulerZXY(m);
				ofs << v.transpose() * 180 / M_PI << " ";
			}
			ofs << std::endl;
		}
	} else {
		for(auto t: mRecordPosition) {
			ofs << t.segment<3>(3).transpose() * 100 << " ";

			for(int i = 0; i < bvh_order.size(); i++) {
				int idx = mCharacter->GetSkeleton()->getBodyNode(bvh_order[i])->getParentJoint()->getIndexInSkeleton(0);
				Eigen::AngleAxisd aa(t.segment<3>(idx).norm(), t.segment<3>(idx).normalized());
				Eigen::Matrix3d m;
				m = aa;
				Eigen::Vector3d v = dart::math::matrixToEulerZXY(m);
				ofs << v.transpose() * 180 / M_PI << " ";
			}
			ofs << std::endl;
		}

	}
	std::cout << "saved position: " << mRecordPosition.size() << ", "<< mReferenceManager->GetPhaseLength() << ", " << mRecordPosition[0].rows() << std::endl;
	ofs.close();

	ofs.open(path+"time");
	for(auto t: normalizedDPhase) {
		ofs << t << std::endl;
	}
	ofs.close();

}
}


#include <gflags/gflags.h>
#include <functional>
#include <iostream>
#include <fstream>
#include <random>
#include <memory>
#include <chrono>
#include <vector>
#include <string>
#include <utility>
#include <algorithm>
#include <iterator>
#include <map>
#include <thread>
#include <mutex>
#include <math.h>

#include <inference_engine.hpp>

#include <samples/slog.hpp>

#include "customflags.hpp"
#include "detectors.hpp"
#include "cnn.hpp"
#include "face_reid.hpp"
#include "tracker.hpp"
#include "classes.hpp"

#include <ie_iextension.h>
#include <ext_list.hpp>

#include <opencv2/opencv.hpp>

#include <dlib/image_processing/frontal_face_detector.h>
#include <dlib/image_processing/render_face_detections.h>
#include <dlib/image_processing.h>
#include <dlib/gui_widgets.h>
#include <dlib/image_io.h>
#include <dlib/opencv.h>

#include <boost/circular_buffer.hpp>

#ifdef SIMULATOR
#include "rclcpp/rclcpp.hpp"
#include "ets_msgs/msg/truck.hpp"

void ros_client(Truck *truck)
{
    auto node = rclcpp::Node::make_shared("ets_client");

    auto sub = node->create_subscription<ets_msgs::msg::Truck>(
        "truck", std::bind(&Truck::ros_callback, truck, std::placeholders::_1), rmw_qos_profile_default);

    rclcpp::spin(node);
}
#endif

using namespace InferenceEngine;

static dlib::rectangle openCVRectToDlib(cv::Rect r)
{
    return dlib::rectangle((long)r.tl().x, (long)r.tl().y, (long)r.br().x - 1, (long)r.br().y - 1);
}

float distanceAtoB(cv::Point2f A, cv::Point2f B)
{
    float distance_l = sqrt((A.x - B.x) * (A.x - B.x) + (A.y - B.y) * (A.y - B.y));
    return distance_l;
}

bool ParseAndCheckCommandLine(int argc, char *argv[])
{
    // ---------------------------Parsing and validation of input args--------------------------------------
    gflags::ParseCommandLineNonHelpFlags(&argc, &argv, true);
    if (FLAGS_h)
    {
        showUsage();
        return false;
    }
    slog::info << "Parsing input parameters" << slog::endl;

    if (FLAGS_i.empty())
    {
        throw std::logic_error("Parameter -i is not set");
    }

    if (FLAGS_m.empty())
    {
        throw std::logic_error("Parameter -m is not set");
    }

    if (FLAGS_n_ag < 1)
    {
        throw std::logic_error("Parameter -n_ag cannot be 0");
    }

    if (FLAGS_n_hp < 1)
    {
        throw std::logic_error("Parameter -n_hp cannot be 0");
    }

    // no need to wait for a key press from a user if an output image/video file is not shown.
    FLAGS_no_wait |= FLAGS_no_show;

    return true;
}

enum distractionLevel
{
    NOT_DISTRACTED = 0,
    DISTRACTED,
    PHONE,
};

int isDistracted(float y, float p, float r)
{
    int result = 0;
    if (abs(y) > 30 || abs(p) > 30)
    {
        if (abs(y) > 20 && p > 10 && r < 0)
            result = PHONE;
        else
            result = DISTRACTED;
    }
    return result;
}

bool identify_driver(cv::Mat frame, std::vector<FaceDetection::Result> *results, VectorCNN *landmarks_detector,
                     VectorCNN *face_reid, EmbeddingsGallery *face_gallery, std::string *driver_name)
{
    bool ret = false;
    std::vector<cv::Mat> face_rois, landmarks, embeddings;

    if (results->empty())
        return ret;

    for (const auto &face : *results)
    {
        cv::Rect rect = face.location;
        float scale_factor_x = 0.15;
        float scale_factor_y = 0.20;
        double aux_x = (rect.x > 3 ? rect.x : 3);
        double aux_y = (rect.y > 3 ? rect.y : 3);
        double aux_width = (rect.width + aux_x < frame.cols ? rect.width : frame.cols - aux_x);
        double aux_height = (rect.height + aux_y < frame.rows ? rect.height : frame.rows - aux_y);
        aux_x += scale_factor_x * aux_width;
        aux_y += scale_factor_y * aux_height;
        aux_width = aux_width * (1 - 2 * scale_factor_x);
        aux_height = aux_height * (1 - scale_factor_y);
        cv::Rect aux_rect = cv::Rect(aux_x, aux_y, aux_width, aux_height);
        face_rois.push_back(frame(aux_rect));
    }

    if (!face_rois.empty())
    {
        landmarks_detector->Compute(face_rois, &landmarks, cv::Size(2, 5));
        AlignFaces(&face_rois, &landmarks);
        face_reid->Compute(face_rois, &embeddings);
        auto ids = face_gallery->GetIDsByEmbeddings(embeddings);

        if (!ids.empty() && ids[0] != EmbeddingsGallery::unknown_id)
        {
            ret = true;
            *driver_name = face_gallery->GetLabelByID(ids[0]);
        }
        else
            *driver_name = "Unknown";
    }

    return ret;
}

// Global variables declaration
int timer_danger = 150; // N frames for DANGER sign
int timer_off = 0;      // N frames for Welcome sign
bool face_identified = false;
bool first_stage_completed = (FLAGS_d_recognition ? false : true);
int biggest_head = 0;
bool alarmDistraction = false;
std::string driver_name = "";
Timer timer;
int firstTime = 0;
Truck truck;
bool fSim = false;

// Alarm Logic Function
std::string labelAlarm = "";

std::string alarmDetection(int is_dist, int yawn_total, int blinl_total)
{
    if (is_dist)
    {
        switch (is_dist)
        {
        case DISTRACTED:
            if (truck.getSpeed() * 3.6 >= 5 || !fSim)
            {
                labelAlarm = "EYES OUT OF ROAD";
                alarmDistraction = true;
            }
            break;
        case PHONE:
            if (truck.getSpeed() * 3.6 >= 2 || !fSim)
            {
                labelAlarm = "LOOKING AT THE PHONE";
                alarmDistraction = true;
            }
            break;
        default:
            alarmDistraction = false;
            break;
        }
    }
    else
    {
        labelAlarm = "";
        alarmDistraction = false;
    }

    return labelAlarm;
}

int maxNormal = 3;
int maxWarning = 7;
int maxCritical = 10;
int x_vum = 5;
int y_vum = 10;
int tDrowsiness = 0 ;

void alarmDrowsiness(cv::Mat prev_frame, int yawn_total, int blinl_total,int width, int height, int x_alarm, int y_alarm)
{

    tDrowsiness = (yawn_total + blinl_total);

    if (tDrowsiness <= maxNormal)
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195, height - y_alarm + 15, tDrowsiness * x_vum, y_vum), cv::Scalar(0, 255, 0), -1);

    else if (tDrowsiness > maxNormal && tDrowsiness <= maxWarning)
    {
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195, height - y_alarm + 15, maxNormal * x_vum, y_vum), cv::Scalar(0, 255, 0), -1);
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195 + (maxNormal * x_vum), height - y_alarm + 15, (tDrowsiness - maxNormal) * x_vum, y_vum), cv::Scalar(0, 255, 255), -1);
    }

    else if (tDrowsiness > maxWarning && tDrowsiness <= maxCritical)
    {
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195, height - y_alarm + 15, maxNormal * x_vum, y_vum), cv::Scalar(0, 255, 0), -1);
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195 + (maxNormal * x_vum), height - y_alarm + 15, (maxWarning - maxNormal) * x_vum, y_vum), cv::Scalar(0, 255, 255), -1);
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195 + (maxWarning * x_vum), height - y_alarm + 15, (tDrowsiness - maxWarning) * x_vum, y_vum), cv::Scalar(0, 0, 255), -1);
    }
    else
    {
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195, height - y_alarm + 15, maxNormal * x_vum, y_vum), cv::Scalar(0, 255, 0), -1);
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195 + (maxNormal * x_vum), height - y_alarm + 15, (maxWarning - maxNormal) * x_vum, y_vum), cv::Scalar(0, 255, 255), -1);
        cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195 + (maxWarning * x_vum), height - y_alarm + 15, (maxCritical - maxWarning) * x_vum, y_vum), cv::Scalar(0, 0, 255), -1);
    }

    // VUmeter - White Rectangle
    cv::rectangle(prev_frame, cv::Rect(width - x_alarm + 195, height - y_alarm + 15, maxCritical * x_vum, y_vum), cv::Scalar(255, 255, 255), 1);

}

// Thread 1: Driver Recognition
void driver_recognition(cv::Mat prev_frame, std::vector<FaceDetection::Result> prev_detection_results, VectorCNN landmarks_detector, VectorCNN face_reid, EmbeddingsGallery face_gallery, std::string *driver_name)
{
    if (timer["face_identified"].getSmoothedDuration() > 60000.0 && face_identified && firstTime == 1 ||
        timer["face_identified"].getSmoothedDuration() > 1000.0 && !face_identified && firstTime == 1 ||
        firstTime == 0)
    {
        cv::Mat aux_prev_frame = prev_frame.clone();
        face_identified = identify_driver(aux_prev_frame, &prev_detection_results, &landmarks_detector, &face_reid, &face_gallery, driver_name);
        if (!prev_detection_results.empty())
            cv::rectangle(prev_frame, prev_detection_results[0].location, cv::Scalar(255, 255, 255), 1);
        firstTime = 1;
        timer.start("face_identified");
    }
}

// Thread 2: Driver Behavior
void driver_behavior(cv::Mat frame, cv::Mat prev_frame, std::vector<FaceDetection::Result> prev_detection_results, std::ostringstream out, dlib::shape_predictor sp, HeadPoseDetection headPoseDetector, size_t width, size_t height)
{
    float EYE_AR_THRESH = 0.195;
    float MOUTH_EAR_THRESH = 0.65;
    float EYE_AR_CONSEC_FRAMES = 3;
    float MOUTH_EAR_CONSEC_FRAMES = 5;

    bool eye_closed = false;

    int blink_counter = 0;
    int yawn_counter = 0;
    int last_blink_counter = 0;
    int last_yawn_counter = 0;
    int blinl_total = 0;
    int yawn_total = 0;
    boost::circular_buffer<float> ear_5(5);
    boost::circular_buffer<float> ear_5_mouth(5);

    int i = 0;
    std::vector<cv::Point2f> left_eye, right_eye, mouth;
    for (auto &result : prev_detection_results)
    {
        // Complete this thread.
    }
}

void beeping (Player* beep, bool* finished) {

    while (!(*finished)) {
        if (tDrowsiness <= 10)
            beep->setGain(exp((float)(tDrowsiness-10)));
        if(tDrowsiness > 3) {
            if (!beep->isPlaying())
                beep->play();
        }
    }
}

int headbuttDetection(boost::circular_buffer<double>* angle_p)
{
	boost::circular_buffer<double>& pitch = *angle_p;
	bool ret = false;
	const int full = 5;
	double delta = 0;

	int lim = std::min(full,(int)pitch.size()-1); // Wait for one frame to calculate speed
	if (lim > 1) {
		for (int i = 0; i < lim; i++)
			delta = delta + (pitch[i] - pitch[i+1]);
		delta = delta / lim;
		if (delta > 6) // magic threshold (adjust)
			ret=true;
	}
	return ret;
}

int main(int argc, char *argv[])
{

#ifdef SIMULATOR
    rclcpp::init(argc, argv);
    std::thread truck_data(ros_client, &truck);
    fSim = true;
#endif

    try
    {
        timer.start("face_identified");
        dlib::shape_predictor sp;
        dlib::deserialize("../data/shape_predictor_68_face_landmarks.dat") >> sp;
        std::vector<dlib::full_object_detection> shapes;

        std::chrono::high_resolution_clock::time_point slp1, slp2;

        float EYE_AR_THRESH = 0.195;
        float MOUTH_EAR_THRESH = 0.65;
        float EYE_AR_CONSEC_FRAMES = 3;
        float MOUTH_EAR_CONSEC_FRAMES = 5;

        bool eye_closed = false;

        int blink_counter = 0;
        int yawn_counter = 0;
        int last_blink_counter = 0;
        int last_yawn_counter = 0;
        int blinl_total = 0;
        int yawn_total = 0;
        boost::circular_buffer<float> ear_5(5);
        boost::circular_buffer<float> ear_5_mouth(5);

        std::cout << "InferenceEngine: " << GetInferenceEngineVersion() << std::endl;

        // ------------------------------ Parsing and validation of input args ---------------------------------
        if (!ParseAndCheckCommandLine(argc, argv))
        {
            return 0;
        }

        slog::info << "Reading input" << slog::endl;
        cv::VideoCapture cap;
        const bool isCamera = FLAGS_i == "cam";
        if (FLAGS_i == "cam")
        {
            if (!cap.open(0))
                throw std::logic_error("Cannot open input file or camera: " + FLAGS_i);
        }
        else if (FLAGS_i == "cam1")
        {
            if (!cap.open(1))
                throw std::logic_error("Cannot open input file or camera: " + FLAGS_i);
            cap.set(cv::CAP_PROP_FOURCC, cv::VideoWriter::fourcc('M', 'J', 'P', 'G'));
            //cap.set(cv::CAP_PROP_FRAME_HEIGHT, 480);
            //cap.set(cv::CAP_PROP_FRAME_WIDTH, 640);
            //cap.set(cv::CAP_PROP_FPS, 30);
        }
        else if (!cap.open(FLAGS_i))
        {
            throw std::logic_error("Cannot open input file or camera: " + FLAGS_i);
        }

        // Size force
        cap.set(cv::CAP_PROP_FRAME_HEIGHT, 720);
        cap.set(cv::CAP_PROP_FRAME_WIDTH, 1280);
        // cap.set(cv::CAP_PROP_FPS, 30);

        const size_t width = (size_t)cap.get(cv::CAP_PROP_FRAME_WIDTH);
        const size_t height = (size_t)cap.get(cv::CAP_PROP_FRAME_HEIGHT);

        int x = 200;
        int y = 140;
        int x_truck_i = width - (x + 10);
        int y_driver_i = y + 30;

        // read input (video) frame
        cv::Mat frame;
        if (!cap.read(frame))
        {
            throw std::logic_error("Failed to get frame from cv::VideoCapture");
        }
        // -----------------------------------------------------------------------------------------------------
        // --------------------------- 1. Load Plugin for inference engine -------------------------------------
        std::map<std::string, InferencePlugin> pluginsForDevices;
        std::vector<std::pair<std::string, std::string>> cmdOptions = {
            {FLAGS_d, FLAGS_m}, {FLAGS_d_ag, FLAGS_m_ag}, {FLAGS_d_hp, FLAGS_m_hp}, {FLAGS_d_em, FLAGS_m_em}};
        FaceDetection faceDetector(FLAGS_m, FLAGS_d, 1, false, FLAGS_async, FLAGS_t, FLAGS_r);
        HeadPoseDetection headPoseDetector(FLAGS_m_hp, FLAGS_d_hp, FLAGS_n_hp, FLAGS_dyn_hp, FLAGS_async);
        //	FacialLandmarksDetection facialLandmarksDetector(FLAGS_m_lm, FLAGS_d_lm, FLAGS_n_lm, FLAGS_dyn_lm, FLAGS_async);

        auto fr_model_path = FLAGS_m_reid;
        std::cout << fr_model_path << std::endl;
        auto fr_weights_path = fileNameNoExt(FLAGS_m_reid) + ".bin";
        auto lm_model_path = FLAGS_m_lm;
        auto lm_weights_path = fileNameNoExt(FLAGS_m_lm) + ".bin";

        std::map<std::string, InferencePlugin> plugins_for_devices;
        std::vector<std::string> devices = {FLAGS_d_lm, FLAGS_d_reid};

        for (const auto &device : devices)
        {
            if (plugins_for_devices.find(device) != plugins_for_devices.end())
            {
                continue;
            }
            slog::info << "Loading plugin " << device << slog::endl;
            InferencePlugin plugin = PluginDispatcher({"../../../lib/intel64", ""}).getPluginByDevice(device);
            printPluginVersion(plugin, std::cout);
            /** Load extensions for the CPU plugin **/
            if ((device.find("CPU") != std::string::npos))
            {
                plugin.AddExtension(std::make_shared<Extensions::Cpu::CpuExtensions>());
                if (!FLAGS_l.empty())
                {
                    // CPU(MKLDNN) extensions are loaded as a shared library and passed as a pointer to base extension
                    auto extension_ptr = make_so_pointer<IExtension>(FLAGS_l);
                    plugin.AddExtension(extension_ptr);
                    slog::info << "CPU Extension loaded: " << FLAGS_l << slog::endl;
                }
            }
            else if (!FLAGS_c.empty())
            {
                // Load Extensions for other plugins not CPU
                plugin.SetConfig({{PluginConfigParams::KEY_CONFIG_FILE, FLAGS_c}});
            }
            plugin.SetConfig({{PluginConfigParams::KEY_DYN_BATCH_ENABLED, PluginConfigParams::YES}});
            if (FLAGS_pc)
                plugin.SetConfig({{PluginConfigParams::KEY_PERF_COUNT, PluginConfigParams::YES}});
            plugins_for_devices[device] = plugin;
        }

        CnnConfig reid_config(fr_model_path, fr_weights_path);
        reid_config.max_batch_size = 16;
        reid_config.enabled = /*face_config.enabled*/ true && !fr_model_path.empty() && !lm_model_path.empty();
        reid_config.plugin = plugins_for_devices[FLAGS_d_reid];
        VectorCNN face_reid(reid_config);

        // Load landmarks detector
        CnnConfig landmarks_config(lm_model_path, lm_weights_path);
        landmarks_config.max_batch_size = 16;
        landmarks_config.enabled = /*face_config.enabled*/ true && reid_config.enabled && !lm_model_path.empty();
        landmarks_config.plugin = plugins_for_devices[FLAGS_d_lm];
        VectorCNN landmarks_detector(landmarks_config);

        double t_reid = 0.4; // Cosine distance threshold between two vectors for face reidentification.
        EmbeddingsGallery face_gallery(FLAGS_fg, t_reid, landmarks_detector, face_reid);

        for (auto &&option : cmdOptions)
        {
            auto deviceName = option.first;
            auto networkName = option.second;

            if (deviceName == "" || networkName == "")
            {
                continue;
            }

            if (pluginsForDevices.find(deviceName) != pluginsForDevices.end())
            {
                continue;
            }
            slog::info << "Loading plugin " << deviceName << slog::endl;
            InferencePlugin plugin = PluginDispatcher({"../../../lib/intel64", ""}).getPluginByDevice(deviceName);

            /** Printing plugin version **/
            printPluginVersion(plugin, std::cout);

            /** Load extensions for the CPU plugin **/
            if ((deviceName.find("CPU") != std::string::npos))
            {
                plugin.AddExtension(std::make_shared<Extensions::Cpu::CpuExtensions>());

                if (!FLAGS_l.empty())
                {
                    // CPU(MKLDNN) extensions are loaded as a shared library and passed as a pointer to base extension
                    auto extension_ptr = make_so_pointer<IExtension>(FLAGS_l);
                    plugin.AddExtension(extension_ptr);
                    slog::info << "CPU Extension loaded: " << FLAGS_l << slog::endl;
                }
            }
            else if (!FLAGS_c.empty())
            {
                // Load Extensions for other plugins not CPU
                plugin.SetConfig({{PluginConfigParams::KEY_CONFIG_FILE, FLAGS_c}});
            }
            pluginsForDevices[deviceName] = plugin;
        }

        /** Per layer metrics **/
        if (FLAGS_pc)
        {
            for (auto &&plugin : pluginsForDevices)
            {
                plugin.second.SetConfig({{PluginConfigParams::KEY_PERF_COUNT, PluginConfigParams::YES}});
            }
        }
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 2. Read IR models and load them to plugins ------------------------------
        // Disable dynamic batching for face detector as long it processes one image at a time.
        Load(faceDetector).into(pluginsForDevices[FLAGS_d], false);
        Load(headPoseDetector).into(pluginsForDevices[FLAGS_d_hp], FLAGS_dyn_hp);
        // -----------------------------------------------------------------------------------------------------

        // --------------------------- 3. Do inference ---------------------------------------------------------
        // Start inference & calc performance.
        slog::info << "Start inference " << slog::endl;
        if (!FLAGS_no_show)
        {
            std::cout << "Press any key to stop" << std::endl;
        }

        bool isFaceAnalyticsEnabled = headPoseDetector.enabled();

        //Timer timer;
        timer.start("total");

        std::ostringstream out;
        size_t framesCounter = 0;
        bool frameReadStatus;
        bool isLastFrame;
        cv::Mat prev_frame, next_frame;

        // Detect all faces on the first frame and read the next one.
        timer.start("detection");
        faceDetector.enqueue(frame);
        faceDetector.submitRequest();
        timer.finish("detection");

        prev_frame = frame.clone();

        // Read next frame.
        timer.start("video frame decoding");
        frameReadStatus = cap.read(frame);
        timer.finish("video frame decoding");

        //dlib
        //dlib::image_window win, win_faces;

        FaceDetection::Result big_head;
        big_head.label = 0;
        big_head.confidence = 0;
        big_head.location = cv::Rect(0, 0, 0, 0);

        boost::circular_buffer<double> pitch = boost::circular_buffer<double>(5);
        bool headbutt = false;

        bool processing_finished = false;
        Player beep("../data/beep.ogg");
        std::thread beep_thread(beeping, &beep, &processing_finished);

        while (true)
        {
            framesCounter++;
            isLastFrame = !frameReadStatus;

            timer.start("detection");
            // Retrieve face detection results for previous frame.
            faceDetector.wait();
            faceDetector.fetchResults();
            auto prev_detection_results = faceDetector.results;
            if (!prev_detection_results.empty())
            {
                for (int i = 0; i < prev_detection_results.size(); i++)
                {
                    if (big_head.location.area() < prev_detection_results[i].location.area())
                    {
                        big_head = prev_detection_results[i];
                        biggest_head = i;
                    }
                }
                prev_detection_results.clear();
                prev_detection_results.push_back(big_head);
                big_head.label = 0;
                big_head.confidence = 0;
                big_head.location = cv::Rect(0, 0, 0, 0);
            }
            // No valid frame to infer if previous frame is last.
            if (!isLastFrame)
            {
                faceDetector.enqueue(frame);
                faceDetector.submitRequest();
            }
            timer.finish("detection");

            timer.start("data preprocessing");
            // Fill inputs of face analytics networks.
            for (auto &&face : prev_detection_results)
            {
                if (isFaceAnalyticsEnabled)
                {
                    auto clippedRect = face.location & cv::Rect(0, 0, width, height);
                    cv::Mat face = prev_frame(clippedRect);
                    headPoseDetector.enqueue(face);
                }
            }
            timer.finish("data preprocessing");

            // Run age-gender recognition, head pose estimation and emotions recognition simultaneously.
            timer.start("face analytics call");
            if (isFaceAnalyticsEnabled)
            {
                headPoseDetector.submitRequest();
            }
            timer.finish("face analytics call");

            // Read next frame if current one is not last.
            if (!isLastFrame)
            {
                timer.start("video frame decoding");
                frameReadStatus = cap.read(next_frame);
                timer.finish("video frame decoding");
            }

            timer.start("face analytics wait");
            if (isFaceAnalyticsEnabled)
            {
                headPoseDetector.wait();
            }
            timer.finish("face analytics wait");

            // Visualize results.
            if (!FLAGS_no_show)
            {
                TrackedObjects tracked_face_objects;
                timer.start("visualization");
                out.str("");
                out << "OpenCV cap/render time: " << std::fixed << std::setprecision(2)
                    << (timer["video frame decoding"].getSmoothedDuration() +
                        timer["visualization"].getSmoothedDuration())
                    << " ms";
                cv::putText(prev_frame, out.str(), cv::Point2f(10, 25), cv::FONT_HERSHEY_TRIPLEX, 0.4,
                            cv::Scalar(255, 0, 0));

                out.str("");
                out << "Face detection time: " << std::fixed << std::setprecision(2)
                    << timer["detection"].getSmoothedDuration()
                    << " ms ("
                    << 1000.f / (timer["detection"].getSmoothedDuration())
                    << " fps)";
                cv::putText(prev_frame, out.str(), cv::Point2f(10, 45), cv::FONT_HERSHEY_TRIPLEX, 0.4,
                            cv::Scalar(255, 0, 0));

                if (isFaceAnalyticsEnabled)
                {
                    out.str("");
                    out << "Face Analysics Networks "
                        << "time: " << std::fixed << std::setprecision(2)
                        << timer["face analytics call"].getSmoothedDuration() +
                               timer["face analytics wait"].getSmoothedDuration()
                        << " ms ";
                    if (!prev_detection_results.empty())
                    {
                        out << "("
                            << 1000.f / (timer["face analytics call"].getSmoothedDuration() +
                                         timer["face analytics wait"].getSmoothedDuration())
                            << " fps)";
                    }
                    cv::putText(prev_frame, out.str(), cv::Point2f(10, 65), cv::FONT_HERSHEY_TRIPLEX, 0.4,
                                cv::Scalar(255, 0, 0));
                }

                if ((truck.getEngine() && fSim) || !fSim)
                { // Detect if Engine = ON and Simulator Flag
                    // Thread 1: Driver Recognition
                    std::thread thread_recognition(driver_recognition, prev_frame, prev_detection_results, landmarks_detector, face_reid, face_gallery, &driver_name);

                    // Tread 2: Diver Behavior
                    // std::thread thread_behavior(driver_behavior, frame, prev_frame, prev_detection_results, out, sp, headPoseDetector, width, height);

                    // Driver Label (CHECK! -> Not here)
                    cv::rectangle(prev_frame, cv::Rect(width - (x + 20), y_driver_i, x, y), cv::Scalar(0, 0, 0), -1);
                    cv::rectangle(prev_frame, cv::Rect(width - (x + 20), y_driver_i, x, y), cv::Scalar(255, 255, 255), 2);

                    // For every detected face.
                    int i = 0;
                    std::vector<cv::Point2f> left_eye, right_eye, mouth;
                    for (auto &result : prev_detection_results)
                    {
                        cv::Rect rect = result.location;

                        out.str("");
                        cv::rectangle(prev_frame, rect, cv::Scalar(255, 255, 255), 1);
                        if (FLAGS_dlib_lm)
                        {
                            float scale_factor_x = 0.15;
                            float scale_factor_y = 0.20;
                            cv::Rect aux_rect = cv::Rect(rect.x + scale_factor_x * rect.width, rect.y + scale_factor_y * rect.height, rect.width * (1 - 2 * scale_factor_x), rect.height * (1 - scale_factor_y));
                            //dlib facial landmarks
                            dlib::array2d<dlib::rgb_pixel> img;
                            dlib::assign_image(img, dlib::cv_image<dlib::bgr_pixel>(prev_frame));
                            dlib::rectangle det = openCVRectToDlib(aux_rect);
                            dlib::full_object_detection shape = sp(img, det);
                            for (int i = 0; i < shape.num_parts(); i++)
                            {
                                if (i >= 36 && i <= 41)
                                {
                                    left_eye.push_back(cv::Point2l(shape.part(i).x(), shape.part(i).y()));
                                    cv::circle(prev_frame, cv::Point2l(shape.part(i).x(), shape.part(i).y()), 1 + static_cast<int>(0.0012 * rect.width), cv::Scalar(0, 255, 255), -1);
                                }
                                if (i >= 42 && i <= 47)
                                {
                                    right_eye.push_back(cv::Point2l(shape.part(i).x(), shape.part(i).y()));
                                    cv::circle(prev_frame, cv::Point2l(shape.part(i).x(), shape.part(i).y()), 1 + static_cast<int>(0.0012 * rect.width), cv::Scalar(0, 255, 255), -1);
                                }
                                //48 - 54. 50 - 58. 52 - 56.

                                if (i == 48 || i == 54 || i == 50 || i == 58 || i == 52 || i == 56)
                                {
                                    mouth.push_back(cv::Point2l(shape.part(i).x(), shape.part(i).y()));
                                    cv::circle(prev_frame, cv::Point2l(shape.part(i).x(), shape.part(i).y()), 1 + static_cast<int>(0.0012 * rect.width), cv::Scalar(0, 255, 255), -1);
                                }
                            }
                            float ear_left = 0;
                            float ear_right = 0;
                            float ear = 0;
                            ear_left = (distanceAtoB(left_eye[1], left_eye[5]) + distanceAtoB(left_eye[2], left_eye[4])) / (2 * distanceAtoB(left_eye[0], left_eye[3]));
                            ear_right = (distanceAtoB(right_eye[1], right_eye[5]) + distanceAtoB(right_eye[2], right_eye[4])) / (2 * distanceAtoB(right_eye[0], right_eye[3]));
                            ear = (ear_left + ear_right) / 2;
                            ear_5.push_front(ear);
                            float ear_avg = 0;
                            for (auto &&i : ear_5)
                            {
                                ear_avg = ear_avg + i;
                            }
                            ear_avg = ear_avg / ear_5.size();
                            if (ear_avg < EYE_AR_THRESH)
                            {
                                blink_counter += 1;
                                if (blink_counter >= 90)
                                    eye_closed = true;
                            }
                            else
                            {
                                if (blink_counter >= EYE_AR_CONSEC_FRAMES)
                                {
                                    blinl_total += 1;
                                    last_blink_counter = blink_counter;
                                }
                                blink_counter = 0;
                            }
                            if (eye_closed && timer_danger > 0)
                            {
                                cv::putText(frame, "DANGER", cv::Point2f(50, 250), cv::FONT_HERSHEY_SIMPLEX, 5, cv::Scalar(0, 0, 255), 5);
                                cv::putText(frame, "Blink time: " + std::to_string(last_blink_counter) + " frames", cv::Point2f(250, 100), cv::FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2);
                                timer_danger--;
                            }
                            if (timer_danger == 0)
                            {
                                eye_closed = false;
                                timer_danger = 150;
                            }

                            cv::putText(prev_frame, "Blinks: " + std::to_string(blinl_total), cv::Point2f(x_truck_i, y_driver_i + 60), cv::FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1);
                            //cv::putText(frame, "EAR: " + std::to_string(ear_avg), cv::Point2f(300, 100), cv::FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2);

                            //Yawn detection
                            float ear_mouth = (distanceAtoB(mouth[1], mouth[5]) + distanceAtoB(mouth[2], mouth[4])) / (2 * distanceAtoB(mouth[0], mouth[3]));
                            ear_5_mouth.push_front(ear_mouth);
                            float ear_avg_mouth = 0;
                            for (auto &&i : ear_5_mouth)
                            {
                                ear_avg_mouth = ear_avg_mouth + i;
                            }
                            ear_avg_mouth = ear_avg_mouth / ear_5_mouth.size();
                            if (ear_avg_mouth > MOUTH_EAR_THRESH)
                            {
                                yawn_counter += 1;
                            }
                            else
                            {
                                if (yawn_counter >= MOUTH_EAR_CONSEC_FRAMES)
                                {
                                    yawn_total += 1;
                                    last_yawn_counter = yawn_counter;
                                }
                                yawn_counter = 0;
                            }
                            cv::putText(prev_frame, "Yawn time: " + std::to_string(last_yawn_counter) + " frames", cv::Point2f(x_truck_i, y_driver_i + 100), cv::FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1);
                            cv::putText(prev_frame, "Yawns: " + std::to_string(yawn_total), cv::Point2f(x_truck_i, y_driver_i + 80), cv::FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1);
                            // cv::putText(frame, "EAR: " + std::to_string(ear_avg_mouth), cv::Point2f(10, 160), cv::FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2);
                        }

                        cv::putText(prev_frame,
                                    out.str(),
                                    cv::Point2f(result.location.x, result.location.y - 15),
                                    cv::FONT_HERSHEY_COMPLEX_SMALL,
                                    0.8,
                                    cv::Scalar(0, 0, 255));

                        if (headPoseDetector.enabled() && i < headPoseDetector.maxBatch)
                        {
                            if (FLAGS_r)
                            {
                                std::cout << "Head pose results: yaw, pitch, roll = "
                                          << headPoseDetector[i].angle_y << ";"
                                          << headPoseDetector[i].angle_p << ";"
                                          << headPoseDetector[i].angle_r << std::endl;
                            }
                            cv::Point3f center(rect.x + rect.width / 2, rect.y + rect.height / 2, 0);
                            headPoseDetector.drawAxes(prev_frame, center, headPoseDetector[i], 50);
                            pitch.push_front(headPoseDetector[i].angle_p);
                            headbutt = headbuttDetection(&pitch);
                            int is_dist = isDistracted(headPoseDetector[i].angle_y, headPoseDetector[i].angle_p, headPoseDetector[i].angle_r);

                            // Alarm Label
                            int x_alarm = 300;
                            int y_alarm = 100;
                            cv::Rect rect(width - (x_alarm + 20), height - (y_alarm + 20), x_alarm, y_alarm);
                            cv::rectangle(prev_frame, rect, cv::Scalar(0, 0, 0), -1);
                            cv::rectangle(prev_frame, rect, cv::Scalar(255, 255, 255), 2);

                            cv::putText(prev_frame, "Alarms", cv::Point2f(width - x_alarm, height - y_alarm + 4), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 2);
                            cv::putText(prev_frame, "Distraction", cv::Point2f(width - x_alarm + 100, height - y_alarm + 4), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1);
                            cv::putText(prev_frame, "Drowsiness", cv::Point2f(width - x_alarm + 100, height - y_alarm + 24), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1);
                            cv::putText(prev_frame, "Description", cv::Point2f(width - x_alarm, height - y_alarm + 44), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 2);

                            // Call alarmDetection Function
                            cv::putText(prev_frame, alarmDetection(is_dist, yawn_total, blinl_total), cv::Point2f(width - x_alarm, height - y_alarm + 64), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255), 2);

                            if (alarmDistraction)
                                cv::circle(prev_frame, cv::Point2f(width - x_alarm + 200, height - y_alarm), 7, cv::Scalar(0, 0, 255), -1);
                            else
                                cv::circle(prev_frame, cv::Point2f(width - x_alarm + 200, height - y_alarm), 7, cv::Scalar(255, 255, 255), 1);

                            // Thread: Drowsiness Alarm
                            std::thread thread_drowsiness(alarmDrowsiness,prev_frame, yawn_total, blinl_total, width, height, x_alarm, y_alarm);

                            // Thread: Drowsiness Alarm
                            thread_drowsiness.join();

                            }
                        i++;
                    }

                    // End Thread 1: Driver Recognition
                    thread_recognition.join();

                    // End Thread 2: Driver Behavior
                    //thread_behavior.join();

                    // Truck Label
                    cv::rectangle(prev_frame, cv::Rect(width - (x + 20), 20, x, y), cv::Scalar(0, 0, 0), -1);
                    cv::rectangle(prev_frame, cv::Rect(width - (x + 20), 20, x, y), cv::Scalar(255, 255, 255), 2);

                    cv::putText(prev_frame, "Truck Information", cv::Point2f(x_truck_i, 40), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 2);
                    if (truck.getEngine())
                        cv::putText(prev_frame, "Engine: ON", cv::Point2f(x_truck_i, 60), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 1.8);
                    else
                        cv::putText(prev_frame, "Engine: OFF", cv::Point2f(x_truck_i, 60), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255), 1.8);
                    cv::putText(prev_frame, cv::format("Speed (Km/h): %3.2f", truck.getSpeed() * 3.6), cv::Point2f(x_truck_i, 80), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1.2);
                    cv::putText(prev_frame, "RPM: " + std::to_string(truck.getRpm()), cv::Point2f(x_truck_i, 100), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1.2);
                    cv::putText(prev_frame, "Gear: " + std::to_string(truck.getGear()), cv::Point2f(x_truck_i, 120), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1.2);
                    if (truck.getTrailer())
                        cv::putText(prev_frame, "Trailer: ON", cv::Point2f(x_truck_i, 140), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 1.2);
                    else
                        cv::putText(prev_frame, "Trailer: OFF", cv::Point2f(x_truck_i, 140), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255), 1.2);

                    // Driver Label
                    cv::putText(prev_frame, "Driver Information", cv::Point2f(x_truck_i, y_driver_i + 20), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 2);

                    if (face_identified)
                    {
                        cv::putText(prev_frame, "Driver: " + driver_name, cv::Point2f(x_truck_i, y_driver_i + 40), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 1);
                    }
                    else if (!face_identified && driver_name == "Unknown")
                    {
                        cv::putText(prev_frame, "Driver: " + driver_name, cv::Point2f(x_truck_i, y_driver_i + 40), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255), 1);
                    }
                }

                // Sample of Results
                cv::imshow("Detection results", prev_frame);
                timer.finish("visualization");
            }

            // End of file (or a single frame file like an image). We just keep last frame displayed to let user check what was shown
            if (isLastFrame)
            {
                timer.finish("total");
                if (!FLAGS_no_wait)
                {
                    std::cout << "No more frames to process. Press any key to exit" << std::endl;
                    cv::waitKey(0);
                }
                break;
            }
            else if (!FLAGS_no_show && -1 != cv::waitKey(1))
            {
                timer.finish("total");
                break;
            }

            prev_frame = frame;
            frame = next_frame;
            next_frame = cv::Mat();
        }

        processing_finished = true;
        beep_thread.join();
        slog::info << "Number of processed frames: " << framesCounter << slog::endl;
        slog::info << "Total image throughput: " << framesCounter * (1000.f / timer["total"].getTotalDuration()) << " fps" << slog::endl;

        // Show performace results.
        if (FLAGS_pc)
        {
            faceDetector.printPerformanceCounts();
            headPoseDetector.printPerformanceCounts();
        }
        // -----------------------------------------------------------------------------------------------------
    }
    catch (const std::exception &error)
    {
        slog::err << error.what() << slog::endl;
        return 1;
    }
    catch (...)
    {
        slog::err << "Unknown/internal exception happened." << slog::endl;
        return 1;
    }

#ifdef SIMULATOR
    truck_data.join();
#endif

    slog::info << "Execution successful" << slog::endl;

    return 0;
}