diff --git a/AState.hpp b/AState.hpp
index d109163..10ffea3 100644
--- a/AState.hpp
+++ b/AState.hpp
@@ -23,6 +23,11 @@ public:
 	{
 		m_pStateMachine = pStateMachine;
 	}
+
+	virtual void onTransition()
+	{
+		printf("%s\n", __PRETTY_FUNCTION__);
+	}
 	
 protected:
 	void setState(IStateMachine::StateId state)
diff --git a/AStateMachine.hpp b/AStateMachine.hpp
index 5582fd0..7a924b7 100644
--- a/AStateMachine.hpp
+++ b/AStateMachine.hpp
@@ -15,9 +15,8 @@ class AStateMachine : public IStateMachine
 {
 public:
 
-	AStateMachine()
-	: m_stateCurr(NUM_STATES)
-	, m_stateNext(INIT)
+	AStateMachine(StateId stateInit)
+	: m_stateCurr(stateInit)
 	{
 	}
 	
@@ -33,20 +32,17 @@ public:
 	
 	void setState(StateId stateId)
 	{
-		m_stateNext = stateId;
+		bool isTransition = (m_stateCurr != stateId);
+		
+		m_stateCurr = stateId;
+		if (isTransition)
+		{
+			m_states[m_stateCurr]->onTransition();
+		}
 	}
 	
-	bool transition()
-	{
-		bool isTransition = (m_stateNext != m_stateCurr);
-		
-		m_stateCurr = m_stateNext;
-		
-		return isTransition;
-	}
 protected:
 	StateId m_stateCurr;
-	StateId m_stateNext;
 	AState *m_states[NUM_STATES];
 	
 };
diff --git a/main.cpp b/main.cpp
index 6d4cdbd..ffe37ad 100644
--- a/main.cpp
+++ b/main.cpp
@@ -31,7 +31,7 @@ public:
 	int m_count;
 	
 };
-class IState
+class IState : public AState
 {
 public:
 	IState() 
@@ -44,11 +44,6 @@ public:
 		
 	}
 	
-	virtual void onTransition(int x)
-	{
-		printf("%s() x=%d\n", __PRETTY_FUNCTION__, x);
-	}
-
 	virtual void process(int x)
 	{
 		printf("%s() x=%d\n", __PRETTY_FUNCTION__, x);
@@ -60,8 +55,8 @@ class StateMachine : public AStateMachine
 {
 public:
 
-	StateMachine()
-	: AStateMachine()
+	StateMachine(StateId stateInit)
+	: AStateMachine(stateInit)
 	{
 	}
 	
@@ -71,16 +66,12 @@ public:
 	
 	void process(int x)
 	{
-		if(transition())
-		{
-			dynamic_cast<IState*>(m_states[m_stateCurr])->onTransition(x);
-		}
 		dynamic_cast<IState*>(m_states[m_stateCurr])->process(x);
 	}
 	
 };
 
-class StateInit : public AState, public IState
+class StateInit : public IState
 {
 public:
 	StateInit(Context &context)
@@ -101,7 +92,7 @@ private:
 	Context &m_context;
 };
 
-class StateIdle : public AState, public IState
+class StateIdle : public IState
 {
 public:
 	StateIdle(Context &context)
@@ -112,9 +103,9 @@ public:
 	
 	~StateIdle() {}
 	
-	void onTransition(int x)
+	void onTransition()
 	{
-		printf("%s() x=%d\n", __PRETTY_FUNCTION__, x);
+		printf("%s\n", __PRETTY_FUNCTION__);
 		m_context.m_count = 0;
 	}
 	
@@ -131,7 +122,7 @@ private:
 	Context &m_context;
 };
 
-class StateActive0 : public AState, public IState
+class StateActive0 : public IState
 {
 public:
 	StateActive0(Context &context)
@@ -142,6 +133,12 @@ public:
 	
 	~StateActive0() {}
 	
+	void onTransition()
+	{
+		m_context.m_count = 0;
+		printf("%s, Count = %d\n", __PRETTY_FUNCTION__, m_context.m_count);
+	}
+
 	void process(int x)
 	{
 		printf("%s() x=%d\n", __PRETTY_FUNCTION__, x);
@@ -157,7 +154,7 @@ private:
 	Context &m_context;
 };
 
-class StateActive1 : public AState, public IState
+class StateActive1 : public IState
 {
 public:
 	StateActive1(Context &context)
@@ -168,13 +165,19 @@ public:
 	
 	~StateActive1() {}
 	
+	void onTransition()
+	{
+		m_context.m_count = 5;
+		printf("%s, Count = %d\n", __PRETTY_FUNCTION__, m_context.m_count);
+	}
+
 	void process(int x)
 	{
 		printf("%s() x=%d\n", __PRETTY_FUNCTION__, x);
 		m_context.m_count--;
 		if (m_context.m_count == 0)
 		{
-			setState(IStateMachine::IDLE);
+			setState(IStateMachine::ACTIVE_0);
 		}
 	}
 	
@@ -185,7 +188,7 @@ private:
 int main(int argc, char** argv)
 {
 	Context ctx;
-	StateMachine fsm;
+	StateMachine fsm(IStateMachine::INIT);
 	StateInit init(ctx);
 	StateIdle idle(ctx);
 	StateActive0 active0(ctx);