import numpy as np
import tensorflow as tf
import time

import opt_solver
from make_animation import make_animation

class Quadrotor(opt_solver.DiscreteSystemModel):
    """
    Models a simple 2d quadrotor system with the state space
    [x_pos, y_pos, x_vel, y_vel, theta, theta_dot]
    where x and y are the location of quadrotor with respect to world frame.
    theta=0 means the quadrotor is horizontal (parallel with the ground).
    As theta increases, the quadrotor rotates counterclockwise.
    Theta is expressed in radians.    

    The control input is [u1, u2], where u1 is the force generated
    by the left propeller and u2 is the force generated by the right.
    """
    def __init__(self):
        super().__init__(state_dim=6, control_dim=2)
        self.mass = 1.0
        self.length = 1.0
        self.dt = 0.1
        self.g = 9.81

    def step(self, x, u):
        x_deriv = x[:, 2]
        y_deriv = x[:, 3]
        vx_deriv = -(u[:, 0] + u[:, 1]) * tf.sin(x[:, 4])
        vy_deriv = (u[:, 0] + u[:, 1]) * tf.cos(x[:, 4]) - self.g
        theta_deriv = x[:, 5]
        theta_dot_deriv = (u[:, 1] - u[:, 0]) / (self.mass * self.length)

        x_dot = tf.stack([x_deriv, y_deriv, 
                          vx_deriv, vy_deriv,
                          theta_deriv, theta_dot_deriv], axis=1)
        x = x + self.dt * x_dot
        return x

class QuadrotorProblem(opt_solver.ScipyMPCOptProblem):
    """
    An MPC problem that tries to get the quadrotor
    to move to the origin and stay there.
    You can use set_initial_state() to give the problem
    different initial states to start the quadrotor at.
    """
    def __init__(self, T=50):
        super().__init__(
            T=T,
            sys_model=Quadrotor())

        self.problem_kwargs = {
            "control_max": tf.convert_to_tensor(10.),
            "control_min": tf.convert_to_tensor(0.),
            "initial_loc": 
                tf.convert_to_tensor(np.array([4., 10., 0, 0, 0, 0], dtype=np.float32)),
            "target_loc": tf.convert_to_tensor(np.array([0, 0.], dtype=np.float32))
        }

    def initial_guess(self):
        state = np.zeros((self.T, self.state_dim))
        control = np.zeros((self.T - 1, self.control_dim)) 
        
        state[:, 0] = np.linspace(
            self.problem_kwargs["initial_loc"][0], 
            self.problem_kwargs["target_loc"][0], 
            self.T)
        state[:, 1] = np.linspace(
            self.problem_kwargs["initial_loc"][1], 
            self.problem_kwargs["target_loc"][1], 
            self.T)

        state = tf.convert_to_tensor(state)
        control = tf.convert_to_tensor(control)
        return (state, control)

    def eval_obj(self, state, control, **kwargs):
        diff = state[:, 0:2] - kwargs["target_loc"]
        obj = tf.reduce_mean(tf.square(diff))
        return obj

    def eval_ineq_constraints(self, state, control, **kwargs):
        u_max = control - kwargs['control_max']
        u_min = kwargs['control_min'] - control
        constraints = tf.reshape(tf.stack([u_min, u_max]), (-1, 1))
        return constraints

    def eval_eq_constraints(self, state, control, **kwargs):
        constraints = state[0, :] - kwargs["initial_loc"]
        constraints = tf.reshape(constraints, (-1, 1))
        return constraints

    def get_initial_state(self):
        return self.problem_kwargs["initial_loc"]

    def set_initial_state(self, initial_loc):
        initial_loc = np.array(initial_loc, dtype=np.float32)
        initial_loc = np.reshape(initial_loc, (6,))
        self.problem_kwargs["initial_loc"] = tf.convert_to_tensor(initial_loc)



if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Run Quadrotor optimization")
    parser.add_argument(
        "-T",
        action="store",
        type=int,
        default=50,
        help="Time horizon for problem.")
    parser.add_argument(
        "-x",
        action="store",
        type=float,
        default=4.0, 
        help="Starting x location.")
    parser.add_argument(
        "-y",
        action="store",
        type=float,
        default=10.0,
        help="Starting y location.")
    parser.add_argument(
        "--theta",
        action="store",
        type=float,
        default=-np.pi/2,
        help="Starting orientation (in radians).")
    subparser = parser.add_subparsers(
        help="Run the optimizer once and plot results, or run mpc.",
        dest="cmd")
    once_subparser = subparser.add_parser("once")
    mpc_subparser = subparser.add_parser("mpc")
    mpc_subparser.add_argument(
        "--niter",
        action="store",
        type=int,
        default=20,
        help="Number of times to rerun optimizer for MPC.")
    args = parser.parse_args()


    problem = QuadrotorProblem(T=args.T)
    problem.set_initial_state(np.array([args.x, args.y, 0, 0, args.theta, 0]))


    if args.cmd == "mpc":
        num_steps = 10
        curr_state = problem.get_initial_state()
        # initial solve to do all the tensorflow tracing
        problem.solve(maxiter=200, ftol=1e-3)
        x_list = []
        u_list = []
        start_time = time.clock()
        for t in range(args.niter):
            problem.set_initial_state(curr_state)
            sol, x, u = problem.solve()
            for i in range(num_steps):
                x_list.append(x[i, :])
                u_list.append(u[i, :])
            curr_state = x_list[-1]
        x_list.append(x[-1, :])
        end_time = time.clock()
        x_list = np.array(x_list)
        u_list = np.array(u_list)
        print("Average time / solve: {} seconds".format((end_time - start_time) / args.niter))
    else:
        sol, x_list, u_list = problem.solve()

    make_animation(x_list, u_list)



