import numpy as np
from pydrake.multibody import inverse_kinematics
from pydrake.trajectories import PiecewisePolynomial

from pydrake.math import RollPitchYaw, RotationMatrix, RigidTransform
from pydrake.examples.manipulation_station import ManipulationStation
import pydrake.solvers.mathematicalprogram as mp
from manip_station_sim.plan_utils import ConnectPointsWithCubicPolynomial

# starting position of point Q in world frame.
p_WQ_home = np.array([0.5, 0, 0.4])

# position of point Q in EE frame.  Point Q is fixed in the EE frame.
p_L7Q = np.array([0., 0., 0.21])

# robot home position
q_home = np.array([0, 0, 0, -1.75, 0, 1.0, 0])


def GenerateIiwaTrajectoriesAndGripperSetPoints(X_WO, is_printing=True):
    """
    :param X_WO: the pose of the foam brick in world frame.
    :param is_printing: set to True to print IK solution results.
    :return: 1. a list of joint space trajectories.
        2. a list of gripper set points.
        The two lists must have the same length.
    """
    station = ManipulationStation()
    station.SetupManipulationClassStation()
    station.Finalize()
    plant = station.get_controller_plant()

    # Go to p_WQ_home
    ik_iiwa = inverse_kinematics.InverseKinematics(plant)
    world_frame = plant.world_frame()
    l7_frame = plant.GetFrameByName("iiwa_link_7")

    theta_bound = 0.005 * np.pi
    R_WL7 = RollPitchYaw(0, np.pi * 5 / 6, 0).ToRotationMatrix()

    ik_iiwa.AddOrientationConstraint(
        frameAbar=world_frame, R_AbarA=R_WL7,
        frameBbar=l7_frame, R_BbarB=RotationMatrix(),
        theta_bound=theta_bound)

    ik_iiwa.AddPositionConstraint(
        frameB=l7_frame, p_BQ=p_L7Q,
        frameA=world_frame,
        p_AQ_lower=p_WQ_home - 0.001,
        p_AQ_upper=p_WQ_home + 0.001)

    prog = ik_iiwa.prog()
    prog.SetInitialGuess(ik_iiwa.q(), np.zeros(7))
    result = mp.Solve(prog)
    if is_printing:
        print(result.get_solution_result())
    q_val_0 = result.GetSolution(ik_iiwa.q())

    qtraj_leave_home = ConnectPointsWithCubicPolynomial(q_home, q_val_0, 4.0)

    # open fingers
    q_knots = np.zeros((2, 7))
    q_knots[0] = qtraj_leave_home.value(qtraj_leave_home.end_time()).squeeze()
    q_knots[1] = q_knots[0]
    qtraj_open_gripper = PiecewisePolynomial.ZeroOrderHold([0, 1], q_knots.T)

    ik_iiwa = inverse_kinematics.InverseKinematics(plant)
    ik_iiwa.AddOrientationConstraint(
        frameAbar=world_frame, R_AbarA=R_WL7,
        frameBbar=l7_frame, R_BbarB=RotationMatrix(),
        theta_bound=theta_bound)

    ik_iiwa.AddPositionConstraint(
        frameB=l7_frame, p_BQ=p_L7Q,
        frameA=world_frame,
        p_AQ_lower=X_WO.translation() + np.array([0.0, -0.01, 0.0]) - 0.0001,
        p_AQ_upper=X_WO.translation() + np.array([0.0, -0.01, 0.0]) + 0.0001)

    prog = ik_iiwa.prog()
    prog.SetInitialGuess(ik_iiwa.q(), np.zeros(7))
    result = mp.Solve(prog)
    if is_printing:
        print(result.get_solution_result())
    q_val_1 = result.GetSolution(ik_iiwa.q())

    qtraj_reach_foam = ConnectPointsWithCubicPolynomial(q_val_0, q_val_1, 6.0)

    # close fingers
    q_knots = np.zeros((2, 7))
    q_knots[0] = qtraj_reach_foam.value(qtraj_reach_foam.end_time()).squeeze()
    q_knots[1] = q_knots[0]
    qtraj_close_gripper = PiecewisePolynomial.ZeroOrderHold([0, 1], q_knots.T)

    qtraj_return_home = ConnectPointsWithCubicPolynomial(q_val_1, q_val_0, 6.0)

    R_WL8 = RollPitchYaw(0, np.pi * 3 / 6, 0).ToRotationMatrix()

    ik_iiwa = inverse_kinematics.InverseKinematics(plant)
    ik_iiwa.AddOrientationConstraint(
        frameAbar=world_frame, R_AbarA=R_WL8,
        frameBbar=l7_frame, R_BbarB=RotationMatrix(),
        theta_bound=theta_bound)

    ik_iiwa.AddPositionConstraint(
        frameB=l7_frame, p_BQ=p_L7Q,
        frameA=world_frame,
        p_AQ_lower=p_WQ_home + np.array([0.4, 0.0, 0.0]) - 0.0001,
        p_AQ_upper=p_WQ_home + np.array([0.4, 0.0, 0.0]) + 0.0001)

    prog = ik_iiwa.prog()
    prog.SetInitialGuess(ik_iiwa.q(), np.zeros(7))
    result = mp.Solve(prog)
    if is_printing:
        print(result.get_solution_result())
    q_val_2 = result.GetSolution(ik_iiwa.q())

    qtraj_reach_shelf = ConnectPointsWithCubicPolynomial(q_val_0, q_val_2, 6.0)

    # open fingers
    q_knots = np.zeros((2, 7))
    q_knots[0] = qtraj_reach_shelf.value(qtraj_reach_shelf.end_time()).squeeze()
    q_knots[1] = q_knots[0]
    qtraj_open_gripper2 = PiecewisePolynomial.ZeroOrderHold([0, 1], q_knots.T)

    qtraj_return_home2 = ConnectPointsWithCubicPolynomial(q_val_2, q_val_0, 6.0)
    qtraj_return_home3 = ConnectPointsWithCubicPolynomial(q_val_0, q_home, 4.0)

    # Complete your pick and place trajectories.

    q_traj_list = [qtraj_leave_home,
                   qtraj_open_gripper,
                   qtraj_reach_foam,
                   qtraj_close_gripper,
                   qtraj_return_home,
                   qtraj_reach_shelf,
                   qtraj_open_gripper2,
                   qtraj_return_home2,
                   qtraj_return_home3,
                  ]
    gripper_setpoint_list = [0.01, 0.1, 0.1, 0.01, 0.01, 0.01, 0.1, 0.1, 0.1]

    return q_traj_list, gripper_setpoint_list