/**
 * @file methods/reinforcement_learning/environment/double_pole_cart.hpp
 * @author Rahul Ganesh Prabhu
 *
 * This file is an implementation of Double Pole Cart Balancing Task.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */

#ifndef MLPACK_METHODS_RL_ENVIRONMENT_DOUBLE_POLE_CART_HPP
#define MLPACK_METHODS_RL_ENVIRONMENT_DOUBLE_POLE_CART_HPP

#include <mlpack/prereqs.hpp>

namespace mlpack {
namespace rl {

/**
 * Implementation of Double Pole Cart Balancing task. This is an extension of
 * the existing CartPole environment. The environment comprises of a cart of a
 * cart with two poles of different lengths and masses. The agent is meant to
 * balance the poles by applying force on the cart.
 */
class DoublePoleCart
{
 public:
  /**
   * Implementation of the state of Double Pole Cart. The state is expressed as
   * a vector (position, velocity, angle, angular velocity, angle, angular
   * velocity)
   */
  class State
  {
   public:
    /**
     * Construct a state instance.
     */
    State() : data(dimension)
    { /* Nothing to do here. */ }

    /**
     * Construct a state instance from given data.
     *
     * @param data Data for the position, velocity, angle and angular velocity.
     */
    State(const arma::colvec& data) : data(data)
    { /* Nothing to do here */ }

    //! Get the internal representation of the state
    arma::colvec Data() const { return data; }
    //! Modify the internal representation of the state.
    arma::colvec& Data() { return data; }

    //! Get the position of the cart.
    double Position() const { return data[0]; }
    //! Modify the position of the cart.
    double& Position() { return data[0]; }

    //! Get the velocity of the cart.
    double Velocity() const { return data[1]; }
    //! Modify the velocity of the cart.
    double& Velocity() { return data[1]; }

    //! Get the angle of the $i^{th}$ pole.
    double Angle(const size_t i) const { return data[2 * i]; }
    //! Modify the angle of the $i^{th}$ pole.
    double& Angle(const size_t i) { return data[2 * i]; }

    //! Get the angular velocity of the $i^{th}$ pole.
    double AngularVelocity(const size_t i) const { return data[2 * i + 1]; }
    //! Modify the angular velocity of the $i^{th}$ pole.
    double& AngularVelocity(const size_t i) { return data[2 * i + 1]; }

    //! Encode the state to a vector..
    const arma::colvec& Encode() const { return data; }

    //! Dimension of the encoded state.
    static constexpr size_t dimension = 6;

   private:
    //! Locally-stored state data.
    arma::colvec data;
  };

  /**
   * Implementation of action of Double Pole Cart.
   */
  class Action
  {
   public:
    enum actions
    {
      backward,
      forward
    };
    // To store the action.
    Action::actions action;

    // Track the size of the action space.
    static const size_t size = 2;
  };

  /**
   * Construct a Double Pole Cart instance using the given constants.
   *
   * @param maxSteps The number of steps after which the episode
   *    terminates. If the value is 0, there is no limit.
   * @param m1 The mass of the first pole.
   * @param m2 The mass of the second pole.
   * @param l1 The length of the first pole.
   * @param l2 The length of the second pole.
   * @param gravity The gravity constant.
   * @param massCart The mass of the cart.
   * @param forceMag The magnitude of the applied force.
   * @param tau The time interval.
   * @param thetaThresholdRadians The maximum angle.
   * @param xThreshold The maximum position.
   * @param doneReward Reward recieved by agent on success.
   */
  DoublePoleCart(const size_t maxSteps = 0,
                 const double m1 = 0.1,
                 const double m2 = 0.01,
                 const double l1 = 0.5,
                 const double l2 = 0.05,
                 const double gravity = 9.8,
                 const double massCart = 1.0,
                 const double forceMag = 10.0,
                 const double tau = 0.02,
                 const double thetaThresholdRadians = 36 * 2 * 3.1416 / 360,
                 const double xThreshold = 2.4,
                 const double doneReward = 0.0) :
      maxSteps(maxSteps),
      m1(m1),
      m2(m2),
      l1(l1),
      l2(l2),
      gravity(gravity),
      massCart(massCart),
      forceMag(forceMag),
      tau(tau),
      thetaThresholdRadians(thetaThresholdRadians),
      xThreshold(xThreshold),
      doneReward(doneReward),
      stepsPerformed(0)
  { /* Nothing to do here */ }

  /**
   * Dynamics of Double Pole Cart instance. Get reward and next state based on current
   * state and current action.
   *
   * @param state The current state.
   * @param action The current action.
   * @param nextState The next state.
   * @return reward, it's always 1.0.
   */
  double Sample(const State& state,
                const Action& action,
                State& nextState)
  {
    // Update the number of steps performed.
    stepsPerformed++;

    arma::vec dydx(6, arma::fill::zeros);
    dydx[0] = state.Velocity();
    dydx[2] = state.AngularVelocity(1);
    dydx[4] = state.AngularVelocity(2);
    Dsdt(state, action, dydx);
    RK4(state, action, dydx, nextState);

    // Check if the episode has terminated.
    bool done = IsTerminal(nextState);

    // Do not reward agent if it failed.
    if (done && maxSteps != 0 && stepsPerformed >= maxSteps)
      return doneReward;
    else if (done)
      return 0;

    /**
     * When done is false, it means that the cartpole has fallen down.
     * For this case the reward is 1.0.
     */
    return 1.0;
  }

  /**
   * This is the ordinary differential equations required for estimation of
   * next state through RK4 method.
   *
   * @param state The current state.
   * @param action The action taken.
   * @param dydx The differential.
   */
  void Dsdt(const State& state,
            const Action& action,
            arma::vec& dydx)
  {
    double totalForce = action.action ? forceMag : -forceMag;
    double totalMass = massCart;
    double omega1 = state.AngularVelocity(1);
    double omega2 = state.AngularVelocity(2);
    double sinTheta1 = std::sin(state.Angle(1));
    double sinTheta2 = std::sin(state.Angle(2));
    double cosTheta1 = std::cos(state.Angle(1));
    double cosTheta2 = std::cos(state.Angle(2));

    // Calculate total effective force.
    totalForce += m1 * l1 * omega1 * omega1 * sinTheta1 + 0.375 * m1 * gravity *
        std::sin(2 * state.Angle(1));
    totalForce += m2 * l2 * omega2 * omega2 * sinTheta1 + 0.375 * m2 * gravity *
        std::sin(2 * state.Angle(2));

    // Calculate total effective mass.
    totalMass += m1 * (0.25 + 0.75 * sinTheta1 * sinTheta1);
    totalMass += m2 * (0.25 + 0.75 * sinTheta2 * sinTheta2);

    // Calculate acceleration.
    double xAcc = totalForce / totalMass;
    dydx[1] = xAcc;

    // Calculate angular acceleration.
    dydx[3] = -0.75 * (xAcc * cosTheta1 + gravity * sinTheta1) / l1;
    dydx[5] = -0.75 * (xAcc * cosTheta2 + gravity * sinTheta2) / l2;
  }

  /**
   * This function calls the RK4 iterative method to estimate the next state
   * based on given ordinary differential equation.
   *
   * @param state The current state.
   * @param action The action to be applied.
   * @param dydx The differential.
   * @param nextState The next state.
   */
  void RK4(const State& state,
           const Action& action,
           arma::vec& dydx,
           State& nextState)
  {
    const double hh = tau * 0.5;
    const double h6 = tau / 6;
    arma::vec yt(6);
    arma::vec dyt(6);
    arma::vec dym(6);

    yt = state.Data() + (hh * dydx);
    Dsdt(State(yt), action, dyt);
    dyt[0] = yt[1];
    dyt[2] = yt[3];
    dyt[4] = yt[5];
    yt = state.Data() + (hh * dyt);

    Dsdt(State(yt), action, dym);
    dym[0] = yt[1];
    dym[2] = yt[3];
    dym[4] = yt[5];
    yt = state.Data() + (tau * dym);
    dym += dyt;

    Dsdt(State(yt), action, dyt);
    dyt[0] = yt[1];
    dyt[2] = yt[3];
    dyt[4] = yt[5];
    nextState.Data() = state.Data() + h6 * (dydx + dyt + 2 * dym);
  }

  /**
   * Dynamics of Double Pole Cart. Get reward based on current state and current
   * action.
   *
   * @param state The current state.
   * @param action The current action.
   * @return reward, it's always 1.0.
   */
  double Sample(const State& state, const Action& action)
  {
    State nextState;
    return Sample(state, action, nextState);
  }

  /**
   * Initial state representation is randomly generated within [-0.05, 0.05].
   *
   * @return Initial state for each episode.
   */
  State InitialSample()
  {
    stepsPerformed = 0;
    return State((arma::randu<arma::vec>(6) - 0.5) / 10.0);
  }

  /**
   * This function checks if the car has reached the terminal state.
   *
   * @param state The desired state.
   * @return true if state is a terminal state, otherwise false.
   */
  bool IsTerminal(const State& state) const
  {
    if (maxSteps != 0 && stepsPerformed >= maxSteps)
    {
      Log::Info << "Episode terminated due to the maximum number of steps"
          "being taken.";
      return true;
    }
    if (std::abs(state.Position()) > xThreshold)
    {
      Log::Info << "Episode terminated due to cart crossing threshold";
      return true;
    }
    if (std::abs(state.Angle(1)) > thetaThresholdRadians ||
        std::abs(state.Angle(2)) > thetaThresholdRadians)
    {
      Log::Info << "Episode terminated due to pole falling";
      return true;
    }
    return false;
  }

  //! Get the number of steps performed.
  size_t StepsPerformed() const { return stepsPerformed; }

  //! Get the maximum number of steps allowed.
  size_t MaxSteps() const { return maxSteps; }
  //! Set the maximum number of steps allowed.
  size_t& MaxSteps() { return maxSteps; }

 private:
  //! Locally-stored maximum number of steps.
  size_t maxSteps;

  //! Locally-stored mass of the first pole.
  double m1;

  //! Locally-stored mass of the second pole.
  double m2;

  //! Locally-stored length of the first pole.
  double l1;

  //! Locally-stored length of
  double l2;

  //! Locally-stored gravity.
  double gravity;

  //! Locally-stored mass of the cart.
  double massCart;

  //! Locally-stored magnitude of the applied force.
  double forceMag;

  //! Locally-stored time interval.
  double tau;

  //! Locally-stored maximum angle.
  double thetaThresholdRadians;

  //! Locally-stored maximum position.
  double xThreshold;

  //! Locally-stored done reward.
  double doneReward;

  //! Locally-stored number of steps performed.
  size_t stepsPerformed;
};

} // namespace rl
} // namespace mlpack

#endif
