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)

Problem:

SEAT_POS = (-0.4, 0.4) # Seat position

Simulation works.

SEAT_POS = (None, 0.4) # Seat position

Simulation not works.