#ifndef WALKMODULEFAVORZ_H
#define WALKMODULEFAVORZ_H

#include "WalkModule.h"
#include "WalkStateHeading.h"

#include "../../../../geo/Heading.h"
#include "../../../../math/distribution/Normal.h"

#include "../../../../Assertions.h"

/** state-parameter needed for WalkModuleFavorZ */
struct WalkStateFavorZ {

	/** nested struct to prevent name clashes */
	struct {

		/**
		 * 0 =	up / down / stay is legal
		 * > 0	force states to walk upwards
		 * < 0	force states to walk downwards
		 *
		 * shifted towards 0 after every taken edge
		 * so: we force states to walk into the same z-direction for some time
		 */
		int zTendence = 0;

	} favorZ;

};


/** favor z-transitions */
template <typename Node, typename WalkState> class WalkModuleFavorZ : public WalkModule<Node, WalkState> {

private:

	// force states to walk into the same z-direction for 30 edges
	const int keepForXEdges = 8;


public:

	/** ctor */
	WalkModuleFavorZ() {

		// ensure the template WalkState inherits from 'WalkStateFavorZ'
		StaticAssert::AinheritsB<WalkState, WalkStateFavorZ>();

	}

	virtual void updateBefore(WalkState& state, const Node& startNode) override {
		(void) state;
		(void) startNode;
	}

	virtual void updateAfter(WalkState& state, const Node& startNode, const Node& endNode) override {
		(void) state;
		(void) startNode;
		(void) endNode;

	}

	virtual void step(WalkState& state, const Node& curNode, const Node& nextNode) override {

		// currently no walk-tendence configured
		if (state.favorZ.zTendence == 0) {

			// does the taken edge indicate a z-change?
			const int diff = nextNode.z_cm - curNode.z_cm;

			// if so, keep this z-direction for the next few edges to come!
			if (diff != 0) {
				state.favorZ.zTendence = (diff > 0) ? (+keepForXEdges) : (-keepForXEdges);
			}

		// currently there IS a walk-tendence configured
		} else {

			// update the tendence (shift towards 0)
			if		(state.favorZ.zTendence < 0) {++state.favorZ.zTendence;}
			else if	(state.favorZ.zTendence > 0) {--state.favorZ.zTendence;}

		}

	}

	double getProbability(const WalkState& state, const Node& startNode, const Node& curNode, const Node& potentialNode) const override {

		(void) state;
		(void) startNode;

		const int tendence = state.favorZ.zTendence;
		const int diff = potentialNode.z_cm - curNode.z_cm;

		// tendence available + tendence match? -> high score!
		if (tendence > 0 && diff >= 0) {return 0.90;}
		if (tendence < 0 && diff <= 0) {return 0.90;}

		// tendence available + tendence mismatch? -> very low score!
		if (tendence > 0 && diff < 0) {return 0.10;}
		if (tendence < 0 && diff > 0) {return 0.10;}

		// no tendence available -> just favor z-transitions over non-z-transitions
		return (diff != 0) ? (0.75) : (0.25);

	}


};

#endif // WALKMODULEFAVORZ_H