import casadi as ca
import numpy as np
import matplotlib.pyplot as plt
from triple_pend import TriplePendulum, Link
from hybrid_trajectory_optimizer import HybridTrajectoryOptimizer, HybridSystemParams, Mode, Transition
from two_link_ik import solve_ik_optimization

F_EPS_FRICTION = 1E-1  # Small value to avoid singularity in friction constraints

class STS(TriplePendulum):
    def __init__(self, link1: Link, link2: Link, link3: Link,
                 seat_pos: tuple | None = None,
                 foot_x_lims: tuple | None = None,
                 **kwargs):
        super().__init__(link1, link2, link3, **kwargs)
        if seat_pos is None:
            self.seat_pos = (-link2.l, link1.l)
        else:
            self.seat_pos = seat_pos

        self.foot_x_lims = foot_x_lims

    def _solve_constrained_sit_system(self, x, u):
        """
        Core solver for the constrained sit phase. 
        Calculates both accelerations and contact forces.
        """
        q = x[0:3]
        q_dot = x[3:6]
        q1 = q[0]
        q2 = q[1]
        q1_dot = q_dot[0]
        q2_dot = q_dot[1]

        # Get the standard manipulator coefficients
        M, C, G = self.manipulator_coefficients_fn(x)

        # Get the constraint Jacobian
        l1 = self.link1.l
        l2 = self.link2.l
        Jh = ca.vertcat(
            ca.horzcat(-l1 * ca.cos(q1), -l2 * ca.cos(q2), 0),
            ca.horzcat(-l1 * ca.sin(q1), -l2 * ca.sin(q2), 0),
        )

        # Calculate the J_dot * q_dot term
        Jh_dot_q_dot = ca.vertcat(
            -l1 * ca.cos(q1) * q1_dot**2 - l2 * ca.cos(q2) * q2_dot**2,
            -l1 * ca.sin(q1) * q1_dot**2 - l2 * ca.sin(q2) * q2_dot**2
        )

        # Build the large matrix for the constrained system
        # A = [[M, Jh.T],
        #      [Jh, 0]]
        A_top = ca.horzcat(M, Jh.T)
        if type(x) is ca.MX or type(x) is ca.SX:
            A_bot = ca.horzcat(Jh, ca.MX.zeros((Jh.shape[0], Jh.shape[0])))
        else:
            A_bot = ca.horzcat(Jh, ca.DM.zeros((Jh.shape[0], Jh.shape[0])))
        A = ca.vertcat(A_top, A_bot)

        # Build the right-hand side vector
        # b = [u - C - G,
        #      -Jh_dot * q_dot]
        b_top = u - C - G
        b_bot = -Jh_dot_q_dot
        b = ca.vertcat(b_top, b_bot)

        # Solve the linear system
        solution = ca.solve(A, b, 'symbolicqr')
        q_ddot = solution[0:3]  # Joint accelerations
        contact_force = -solution[3:]  # Contact forces
        
        return q_ddot, contact_force

    def sit_dynamics_fn(self, x, u):
        """
        Calculate the dynamics of the sit system.
        x: state vector (6D)
        u: control input (3D)
        """
        q_dot = x[3:6]
        q_ddot, _ = self._solve_constrained_sit_system(x, u)
        return ca.vertcat(q_dot, q_ddot)
    
    def sit_contact_force_fn(self, x, u):
        """
        Calculate the contact force during the sit phase.
        x: state vector (6D)
        u: control input (3D)
        """
        _, contact_force = self._solve_constrained_sit_system(x, u)
        return contact_force
    
    def stand_dynamics_fn(self, x, u):
        """
        Calculate the dynamics of the stand system.
        x: state vector (6D)
        u: control input (3D)
        """
        return self.dynamics_fn(x, u)
    
    def j2_pos(self, x):
        """
        Calculate the position of joint 2.
        x: state vector (6D)
        """
        q1 = x[0]
        q2 = x[1]
        # x position of joint 2
        x_j2 = -self.link1.l * ca.sin(q1) - self.link2.l * ca.sin(q2)
        # y position of joint 2
        y_j2 = self.link1.l * ca.cos(q1) + self.link2.l * ca.cos(q2)
        return ca.vertcat(x_j2, y_j2)
    

class STSTrajectoryOptimizer:
    def __init__(self, sts: STS,
                 state_weights: tuple = (1, 1, 1, 1, 1, 1),
                 control_weights: tuple = (1, 1, 1)):
        self.sts = sts
        self.Q = ca.diag(state_weights)
        self.R = ca.diag(control_weights)

    def sit_stand_map_fn(self, x, u):
        """
        Map function to transform the state from sit to stand phase.
        x: state vector (6D)
        u: control input (3D)
        """
        # In this case, we don't need to change the state
        return x
    
    def cost_function(self, X_vars: list, U_vars: list, H_vars: list):
        total_cost = 0
        # total_cost += 1E7 * self.cost_function_cop(X_vars, U_vars, H_vars)
        total_cost += self.cost_function_control_effort(X_vars, U_vars, H_vars)
        # total_cost += self.cost_function_time(X_vars, U_vars, H_vars)
        # total_cost += 1E5 * self.cost_function_positive_energy(X_vars, U_vars, H_vars)
        # total_cost += 1E5 * self.cost_function_max_torque(X_vars, U_vars, H_vars)
        total_cost += 1E-3 * self.cost_function_dynamic_effort(X_vars, U_vars, H_vars)
        return total_cost
    
    def cost_function_max_torque(self, X_vars: list, U_vars: list, H_vars: list):
        """
        Cost function to minimize the maximum torque.
        X_vars: state variables
        U_vars: control variables
        H_vars: hybrid variables
        """
        U = ca.hcat(U_vars)
        # Calculate the maximum torque
        max_torque = ca.norm_inf(U)
        return max_torque
    
    def cost_function_cop(self, X_vars: list, U_vars: list, H_vars: list):
        """
        Cost function to minimize the center of pressure in the standing phase.
        X_vars: state variables
        U_vars: control variables
        H_vars: hybrid variables
        """
        total_cost = 0
        X = X_vars[1] # Use the state variables from the stand phase
        U = U_vars[1]
        H = H_vars[1]
        for n in range(H.shape[1]):
            x = X[:, n]
            u = U[:, n]
            h = H[0, n]
            # Calculate the CoP
            cop_x = self.sts.cop_x(x, u)
            # Cost function: minimize the CoP
            total_cost += cop_x**2 * h
        return total_cost

    def cost_function_control_effort(self, X_vars: list, U_vars: list, H_vars: list):
        """
        Cost function to minimize the control effort.
        X_vars: state variables
        U_vars: control variables
        H_vars: hybrid variables
        """
        total_cost = 0
        for u, h in zip(U_vars, H_vars):
            for n in range(h.shape[1]):
                # Cost function: minimize the control effort
                uk = u[:, n]
                uk_next = u[:, n + 1]
                hk = h[0, n]
                total_cost += hk * (uk.T @ self.R @ uk + uk_next.T @ self.R @ uk_next)
        return total_cost
    
    def cost_function_positive_energy(self, X_vars: list, U_vars: list, H_vars: list):
        """
        Cost function to minimize the positive energy.
        X_vars: state variables
        U_vars: control variables
        H_vars: hybrid variables
        """
        total_cost = 0
        for X, U, H in zip(X_vars, U_vars, H_vars):
            for n in range(H.shape[1]):
                # Cost function: minimize the positive energy
                q_dot = X[3:, n]
                v = ca.vertcat(q_dot[0],
                               q_dot[1] - q_dot[0],
                               q_dot[2] - q_dot[1])
                u = U[:, n]
                h = H[0, n]
                # Calculate the kinetic energy
                T = u.T @ v
                # Calculate the positive energy
                T_pos = ca.fmax(0, T)
                total_cost += h * T_pos
        return total_cost
    
    def cost_function_time(self, X_vars: list, U_vars: list, H_vars: list):
        """
        Cost function to minimize the time.
        X_vars: state variables
        U_vars: control variables
        H_vars: hybrid variables
        """
        total_cost = 0
        for h in H_vars:
            total_cost += ca.sum(h)
        return total_cost
    
    def cost_function_dynamic_effort(self, X_vars: list, U_vars: list, H_vars: list):
        """
        Cost function to minimize the dynamic effort.
        X_vars: state variables
        U_vars: control variables
        H_vars: hybrid variables
        """
        total_cost = 0
        for u, h in zip(U_vars, H_vars):
            for n in range(h.shape[1]):
                # Cost function: minimize changing control effort
                uk = u[:, n]
                uk_next = u[:, n + 1]
                hk = h[0, n]
                u_dot = (uk_next - uk) / hk
                total_cost += hk * (u_dot.T @ u_dot)
        return total_cost
                

    def j2_y_pos_higher_than_seat_constraint_fn(self, opti, Xk, Uk, Hk):
        """
        Constraint function to ensure that the y position of joint 2 is higher than the seat in the stand phase.
        """
        for n in range(Xk.shape[1]):
            x = Xk[:, n]
            y_j2 = self.sts.j2_pos(x)[1]
            opti.subject_to(y_j2 >= self.sts.seat_pos[1])

    def base_contact_no_slip_global_constraint_fn(self, opti, X_vars, U_vars, H_vars):
        """
        Constraint function to ensure that the base does not slip.
        Applies to all modes.
        """
        # Important:
        # In the sit phase, the base contact force calculated by the base_force_fn needs
        # to be corrected to account for the fact that the seat is fixed.
        # Thisis done by subtracting the seat contact force from the base contact force.
        # M*a_com = F_base + F_seat + W
        # F_base = M*a_com - F_seat - W = base_force_fn(x, u) - sit_contact_force_fn(x, u)
        X_seat = X_vars[0]  # State variables in the sit phase
        U_seat = U_vars[0]  # Control variables in the sit phase
        X_stand = X_vars[1]  # State variables in the stand phase
        U_stand = U_vars[1]  # Control variables in the stand phase
        for n in range(X_seat.shape[1]):
            x_seat = X_seat[:, n]
            u_seat = U_seat[:, n]
            # Base contact force in the sit phase
            base_force_seat = self.sts.base_force_fn(x_seat, u_seat) - self.sts.sit_contact_force_fn(x_seat, u_seat)
            # Base contact force in the y direction in the sit phase
            base_force_x_seat = base_force_seat[0]
            base_force_y_seat = base_force_seat[1]
            # Constraint: base contact force in the y direction should be greater than 0
            opti.subject_to(base_force_y_seat >= 0)
            if 'mu' in self.sts.kwargs:
                # Constraint: base contact force in the x direction should be less than the friction force
                mu = self.sts.kwargs['mu']
                opti.subject_to(ca.fabs(base_force_x_seat) <= mu * base_force_y_seat + F_EPS_FRICTION)
            else:
                print("Warning: Friction coefficient not provided. No friction constraint applied.")
        # Now apply the same constraints to the stand phase
        for n in range(X_stand.shape[1]):
            x_stand = X_stand[:, n]
            u_stand = U_stand[:, n]
            # Base contact force in the stand phase
            base_force_stand = self.sts.base_force_fn(x_stand, u_stand)
            # Base contact force in the y direction in the stand phase
            base_force_y_stand = base_force_stand[1]
            base_force_x_stand = base_force_stand[0]
            # Constraint: base contact force in the y direction should be greater than 0
            opti.subject_to(base_force_y_stand >= 0)
            if 'mu' in self.sts.kwargs:
                # Constraint: base contact force in the x direction should be less than the friction force
                mu = self.sts.kwargs['mu']
                opti.subject_to(ca.fabs(base_force_x_stand) <= mu * base_force_y_stand + F_EPS_FRICTION)
            else:
                print("Warning: Friction coefficient not provided. No friction constraint applied.")

    def seat_contact_no_slip_constraint_fn(self, opti, Xk, Uk, Hk):
        """
        Constraint function to ensure that the seat does not slip.
        """
        # Exclude the last knot point
        # The reason is that in the last knot point, the normal force is forced to be 0
        # which makes it very difficult (if not impossible) to satisfy the no-slip constraint.
        # This can also be avoided using F_EPS_FRICTION.
        for n in range(Xk.shape[1]):
            x = Xk[:, n]
            u = Uk[:, n]
            # Seat contact force
            seat_force = self.sts.sit_contact_force_fn(x, u)
            # Seat contact force in the y direction
            seat_force_y = seat_force[1]
            seat_force_x = seat_force[0]
            # Constraint: seat contact force in the y direction should be greater than 0
            # opti.subject_to(seat_force_y >= 0) # This is already ensured by the guard constraint
            if 'mu' in self.sts.kwargs:
                # Constraint: seat contact force in the x direction should be less than the friction force
                mu = self.sts.kwargs['mu']
                opti.subject_to(ca.fabs(seat_force_x) <= mu * seat_force_y + F_EPS_FRICTION)
            else:
                print("Warning: Friction coefficient not provided. No friction constraint applied.")

    def equal_time_segment_global_constraint_fn(self, opti, X_vars, U_vars, H_vars):
        """
        Constraint function to ensure that the time segments are equal.
        """
        for Hk in H_vars:
            for n in range(Hk.shape[1] - 1):
                hk = Hk[0, n]
                hk_next = Hk[0, n + 1]
                opti.subject_to(hk == hk_next)

    def knee_joint_constraint_fn(self, opti, Xk, Uk, Hk):
        """
        Constraint function to ensure that the knee joint cannot bend backwards.
        """
        for n in range(Xk.shape[1]):
            x = Xk[:, n]
            q1 = x[0]
            q2 = x[1]
            knee_joint_angle = q2 - q1
            # Constraint: knee joint angle should be greater than 0
            opti.subject_to(knee_joint_angle >= 0)

    def sit_path_constraint_fn(self, opti, Xk, Uk, Hk):
        """
        Constraint to ensure that for the entire sit phase, the hip joint (j2) remains fixed on the seat.
        """
        # Initial position of j2
        j2_initial_pos = self.sts.j2_pos(Xk[:, 0])
        # Constraint: j2 position should be equal to the seat position
        if self.sts.seat_pos[0] is not None:
            opti.subject_to(j2_initial_pos[0] == self.sts.seat_pos[0])
        if self.sts.seat_pos[1] is not None:
            opti.subject_to(j2_initial_pos[1] == self.sts.seat_pos[1])

    def COP_global_constraint_fn(self, opti, X_vars, U_vars, H_vars):
        """
        Constraint function to ensure that the COP remains inside the support polygon.
        """
        if sts.foot_x_lims is None:
            print("Warning: Foot x limits not provided. No COP constraint applied.")
            return
        for n in range(X_vars[0].shape[1]):
            x_seat = X_vars[0][:, n]
            u_seat = U_vars[0][:, n]
            # Calculate the CoP
            cop_x = self.sts.cop_x(x_seat, u_seat)
            # Constraint: CoP should be inside the foot limits
            # opti.subject_to(cop_x >= self.sts.foot_x_lims[0])
            opti.subject_to(cop_x <= self.sts.foot_x_lims[1])
        for n in range(X_vars[1].shape[1]):
            x_stand = X_vars[1][:, n]
            u_stand = U_vars[1][:, n]
            # Calculate the CoP
            cop_x = self.sts.cop_x(x_stand, u_stand)
            # Constraint: CoP should be inside the foot limits
            opti.subject_to(cop_x >= self.sts.foot_x_lims[0])
            opti.subject_to(cop_x <= self.sts.foot_x_lims[1])

    def control_effort_vanishes_global_constraint_fn(self, opti, X_vars, U_vars, H_vars):
        """
        Constraint function to ensure that the control effort stops out.
        """
        U_last = U_vars[-1]
        u_last = U_last[:, -1]
        # Constraint: last control input should be equal to 0
        opti.subject_to(u_last == 0)
            

###
MU = 0.5 # Coefficient of friction
T = 1.5 # Total time
T_UB = 2.0 # Upper bound on total time
N_knots = 25 # Number of knot points per segment
M = 74 # Mass of the pendulum (kg)
L = 1.74 # Length of the pendulum (m)
lleg = Link(m=0.093 * M, l=0.246 * L)  # Lower leg link
uleg = Link(m=0.200 * M, l=0.245 * L)  # Upper leg link
torso = Link(m=0.678 * M, l=0.47 * L)  # Torso link
# lleg = Link(m=1, l=1)  # Link 1
# uleg = Link(m=1, l=1)  # Link 2
# torso = Link(m=1, l=1)  # Link 3
SEAT_POS = (-0.4, 0.4) # Seat position
FOOT_X_LIMS = (-0.05, 0.2) # Foot x limits
R_weights_init = (1, 1, 1) # Control weights
R_weights = (1, 3, 1)

U_UB = (None,) * 3 # Upper bound on control inputs
U_LB = (None,) * 3 # Lower bound on control inputs

# Define the system
sts = STS(link1=lleg, link2=uleg, link3=torso,
          mu=MU, seat_pos=SEAT_POS, foot_x_lims=FOOT_X_LIMS)
sts_optimizer = STSTrajectoryOptimizer(sts,
                                       control_weights=R_weights_init,
)

x0 = (None, None, -0.01, 0, 0, 0) # Initial state
xf = (0, 0, 0, 0, 0, 0)

j1, j2 = solve_ik_optimization(*SEAT_POS, lleg.l, uleg.l)
# Adjust the joint angles to match STS convention
j2 = j1 + j2 - np.pi/2
j1 = j1 - np.pi/2
print(j1, j2)
sit_mode_initial_guess_x0 = np.array([j1, j2, -0.01, 0, 0, 0]).reshape(-1, 1) 
sit_mode_initial_guess_xf = np.array([j1, j2, -np.pi/6, 0, 0, -0.1]).reshape(-1, 1) 
sit_mode_initial_guess = np.hstack((sit_mode_initial_guess_x0, sit_mode_initial_guess_xf))

mode_sit = Mode(
    name='sit',
    dynamics_fn=sts.sit_dynamics_fn,
    n_x=6,
    n_u=3,
    constraints_fns=[
                sts_optimizer.seat_contact_no_slip_constraint_fn,
                sts_optimizer.sit_path_constraint_fn,
                     ],
    initial_x_guess=sit_mode_initial_guess,
    u_lb=U_LB,
    u_ub=U_UB,
)

mode_stand = Mode(
    name='stand',
    dynamics_fn=sts.stand_dynamics_fn,
    n_x=6,
    n_u=3,
    constraints_fns=[   
                # sts_optimizer.COP_constraint_fn,
                sts_optimizer.knee_joint_constraint_fn,
                    ],
    u_lb=U_LB,
    u_ub=U_UB,
)

transition_sit_stand = Transition(
    guard_fn=lambda x, u: sts.sit_contact_force_fn(x, u)[1], # y component of the contact force
    reset_map_fn=sts_optimizer.sit_stand_map_fn,
)

hybrid_params = HybridSystemParams(
    mode_sequence=[mode_sit, mode_stand],
    transitions=[transition_sit_stand],
    x0=x0, # Initial state
    xf=xf, # Final state
    # T = T, # Total time
    T_ub = T_UB, # Upper bound on total time
    min_time_step = 1E-6,
    num_knot_points_per_segment = N_knots, # Number of knot points per segment
)

hybrid_optimizer = HybridTrajectoryOptimizer(
    params=hybrid_params,
    cost_fn=sts_optimizer.cost_function,
    global_constraints_fns=[sts_optimizer.base_contact_no_slip_global_constraint_fn,
                            sts_optimizer.COP_global_constraint_fn,
                            sts_optimizer.equal_time_segment_global_constraint_fn,
                            sts_optimizer.control_effort_vanishes_global_constraint_fn,
                            ],
)

# Solve the optimization problem
solver_name = 'ipopt'
solver_options = {
    'ipopt.print_level': 5,
    'ipopt.max_iter': 3000,
    'ipopt.fixed_variable_treatment': 'relax_bounds',
    'ipopt.tol': 1E-3,
}

solution = hybrid_optimizer.solve(solver_name=solver_name, opts=solver_options)

## Now use the solution to warm start the simulation ##
# DIRTY JOB #
# sts_optimizer.R = ca.diag(R_weights)  # Update the control weights
# hybrid_optimizer._setup_objective()  # Re-setup the objective function with the new weights
# solution = hybrid_optimizer.solve(solver_name, solver_options, warm_start=True)

SEAT_POS_NEW = (-0.4, 0.2)  # New seat position
# Update the initial guess for the sit mode
mode_sit.initial_x_guess = solution['states'][0]
mode_stand.initial_x_guess = solution['states'][1]
sts_new = STS(link1=lleg, link2=uleg, link3=torso,
              mu=MU, seat_pos=SEAT_POS_NEW, foot_x_lims=FOOT_X_LIMS)
sts_optimizer_new = STSTrajectoryOptimizer(sts_new,
                                        control_weights=R_weights_init,
)
hybrid_params_new = HybridSystemParams(
    mode_sequence=[mode_sit, mode_stand],
    transitions=[transition_sit_stand],
    x0=x0,  # Initial state
    xf=xf,  # Final state
    # T = T,  # Total time
    T_ub=T_UB,  # Upper bound on total time
    min_time_step=1E-6,
    num_knot_points_per_segment=N_knots,  # Number of knot points per segment
)
# Recreate the hybrid optimizer with the new system
hybrid_optimizer_new = HybridTrajectoryOptimizer(
    params=hybrid_params_new,
    cost_fn=sts_optimizer_new.cost_function,
    global_constraints_fns=[sts_optimizer_new.base_contact_no_slip_global_constraint_fn,
                            sts_optimizer_new.COP_global_constraint_fn,
                            sts_optimizer_new.equal_time_segment_global_constraint_fn,
                            sts_optimizer_new.control_effort_vanishes_global_constraint_fn,
                            ],
)
# Solve the optimization problem again with the new system
solution = hybrid_optimizer_new.solve(solver_name=solver_name, opts=solver_options)
# End of dirty job #
##

times = solution['times']
transition_time = times[0][-1] # Time at the transition from sit to stand
states = solution['states']
controls = solution['controls']
t_seat = times[0]  # Times in the sit phase
t_stand = times[1]  # Times in the stand phase
x_seat = states[0]  # States in the sit phase
x_stand = states[1]  # States in the stand phase
u_seat = controls[0]  # Controls in the sit phase
u_stand = controls[1]  # Controls in the stand phase
t_all = np.hstack(times)
x_all = np.hstack(states)
u_all = np.hstack(controls)
base_force_seat = np.zeros((2, len(t_seat)))  # Base contact force (T, N)
base_force_stand = np.zeros((2, len(t_stand)))  # Base contact force (T, N)
seat_force = np.zeros((2, len(t_seat)))  # Seat contact force (T, N)
for i in range(len(t_seat)):
    all_contact_force = sts.base_force_fn(x_seat[:, i], u_seat[:, i])
    seat_contact_force = sts.sit_contact_force_fn(x_seat[:, i], u_seat[:, i])
    base_force_seat[:, i] = np.array(all_contact_force - seat_contact_force).flatten()
    seat_force[:, i] = np.array(seat_contact_force).flatten()
for i in range(len(t_stand)):
    base_force_stand[:, i] = np.array(sts.base_force_fn(x_stand[:, i], u_stand[:, i])).flatten()
cop = np.zeros((1, len(t_all)))  # Center of pressure
for i in range(len(t_all)):
    cop[0, i] = sts.cop_x(x_all[:, i], u_all[:, i])
# Combine the base forces from both phases
base_force = np.hstack((base_force_seat, base_force_stand))

t_linear = np.linspace(0, t_all[-1], 100)
dt = t_linear[1] - t_linear[0]
x_sampled = np.zeros((x_all.shape[0], len(t_linear)))
u_sampled = np.zeros((u_all.shape[0], len(t_linear)))
base_force_sampled = np.zeros((base_force.shape[0], len(t_linear)))
cop_sampled = np.zeros((1, len(t_linear)))
for i in range(x_all.shape[0]):
    # Interpolate the state variables
    x_sampled[i, :] = np.interp(t_linear, t_all, x_all[i, :])
for i in range(u_all.shape[0]):
    # Interpolate the control variables
    u_sampled[i, :] = np.interp(t_linear, t_all, u_all[i, :])
# base_force_sampled[0, :] = np.interp(t_linear, t_all, base_force[0, :])
# base_force_sampled[1, :] = np.interp(t_linear, t_all, base_force[1, :])
# cop_sampled[0, :] = np.interp(t_linear, t_all, cop[0, :])

ankle_angle = x_sampled[0, :]
knee_angle = x_sampled[1, :] - x_sampled[0, :]
hip_angle = x_sampled[2, :] - x_sampled[1, :]

ankle_velocity = x_sampled[3, :]
knee_velocity = x_sampled[4, :] - x_sampled[3, :]
hip_velocity = x_sampled[5, :] - x_sampled[4, :]


# Plot the results
fig, ax = plt.subplots(5, 1, figsize=(10, 8))
ax[0].plot(t_linear, ankle_angle, label='ankle')
ax[0].plot(t_linear, -knee_angle, label='knee')
ax[0].plot(t_linear, hip_angle, label='hip')
ax[0].set_title('Joint Angles')
ax[0].set_ylabel('Angle (rad)')

ax[1].plot(t_linear, ankle_velocity, label='ankle')
ax[1].plot(t_linear, -knee_velocity, label='knee')
ax[1].plot(t_linear, hip_velocity, label='hip')
ax[1].set_title('Joint Velocities')
ax[1].set_ylabel('Velocity (rad/s)')

ax[2].plot(t_all, u_all[0, :] / 2 / M, label='ankle')
ax[2].plot(t_all, -u_all[1, :] / 2 / M, label='knee')
ax[2].plot(t_all, u_all[2, :] / 2 / M, label='hip') 
ax[2].set_title('Control Inputs')
ax[2].set_ylabel('Normalized Torque (N m/kg)')

ax[3].plot(t_all, base_force.T, label=['T', 'N'])
ax[3].plot(t_seat, seat_force.T, label=['Seat T', 'Seat N'], linestyle='--')
ax[3].set_title('Base Contact Force')
ax[3].set_ylabel('Force (N)')

ax[4].plot(t_all, cop.T, label='CoP')
ax[4].set_title('Center of Pressure')
ax[4].set_ylabel('CoP (m)')

ax[-1].set_xlabel('Time (s)')

for a in ax:
    a.axvline(x=transition_time, color='r', linestyle='--', linewidth=2)
    a.legend()
    a.grid()
    a.set_xlim([0, t_linear[-1]])
plt.tight_layout()

# Animate the pendulum
ani = sts.animate(x_sampled, dt=dt)

plt.show()


Help me debug this code. The problem is that I get the same result after the dirty job run.