restore Air Jumps when element is collected midair

Assets/_Game/Scripts/Player/MovementController.cs

@@ -34,6 +34,7 @@ public class MovementController : MonoBehaviour
 	private bool m_Grounded = false; // Whether or not the player is grounded.
 	private bool m_Swimming = false;
 	private bool m_Airborn = true;
+	private bool m_isAir = false;
 	private Rigidbody2D m_Rigidbody2D;
 	private bool m_FacingRight = true; // For determining which way the player is currently facing.
 	private Vector3 m_Velocity = Vector3.zero;
@@ -110,6 +111,7 @@ public class MovementController : MonoBehaviour
 		r_Animator.SetBool("isJumping", m_Jumped);
 		r_Animator.SetBool("isAirborn", m_Airborn);
 		r_Animator.SetBool("isSwimming", m_Swimming);
+		//r_Animator.SetBool("isAirCharged", m_isAir);
 	}
 
 
@@ -155,6 +157,7 @@ public class MovementController : MonoBehaviour
 				m_Grounded = false;
 				m_TempJumpCount--;
 				if (r_State.GetCurrentElement() == InfusedElement.Air) r_State.ChangeCharges(-1);
+				print("AirJump: Charges left = " + m_TempJumpCount);
 				m_Rigidbody2D.velocity = Vector2.up * m_JumpForce;
 				OnJumpEvent.Invoke();
 				StartCoroutine(Wait(m_SecondsWaitToJump));
@@ -180,6 +183,14 @@ public class MovementController : MonoBehaviour
 		else
 		{
 			MovementState state = r_State.GetCurrentMovementState();
+			if (r_State.GetCurrentElement() == InfusedElement.Air)
+			{
+				m_isAir = true;
+			}
+			else
+			{
+				m_isAir = false;
+			}
 			m_JumpForce = state.jumpForce;
 			m_MovementSmoothing = state.smoothing;
 			m_AirControl = state.airControl;
@@ -191,6 +202,12 @@ public class MovementController : MonoBehaviour
 		}
 	}
 
+	public void SetAirJumpCount()
+	{
+		if (r_State.GetCurrentElement() != InfusedElement.Air) return;
+		m_TempJumpCount = r_State.GetCharges();
+	}
+
 	public void SetSwimming(bool value)
 	{
 		m_Swimming = value;