"""Extended G1 arm IK solver with elbow position constraints.

Subclasses G1_29_ArmIK to add an elbow position cost term that resolves
the elbow ambiguity inherent in wrist-only IK. Uses l_lower/r_lower
positions from OpenXR body tracking as elbow targets.
"""

import numpy as np
import casadi
import pinocchio as pin
from pinocchio import casadi as cpin

from robot_arm_ik import G1_29_ArmIK


class G1_29_ArmIK_Elbow(G1_29_ArmIK):
    """IK solver that constrains both wrist AND elbow positions.

    Cost function:
        50 * wrist_position² + rotation² + elbow_weight * elbow_position²
        + 0.02 * regularization² + 0.1 * smoothing²
    """

    def __init__(self, elbow_weight=30.0, **kwargs):
        super().__init__(**kwargs)

        # Find elbow frame IDs in the reduced model
        self.L_elbow_id = self.reduced_robot.model.getFrameId("left_elbow_joint")
        self.R_elbow_id = self.reduced_robot.model.getFrameId("right_elbow_joint")

        # Symbolic elbow target parameters
        self.cElbow_l = casadi.SX.sym("elbow_l", 3, 1)
        self.cElbow_r = casadi.SX.sym("elbow_r", 3, 1)

        # Elbow position error: FK(q) elbow pos - tracked elbow pos
        self.elbow_error = casadi.Function(
            "elbow_error",
            [self.cq, self.cElbow_l, self.cElbow_r],
            [casadi.vertcat(
                self.cdata.oMf[self.L_elbow_id].translation - self.cElbow_l,
                self.cdata.oMf[self.R_elbow_id].translation - self.cElbow_r,
            )],
        )

        # Rebuild the optimization problem with elbow cost
        self.opti = casadi.Opti()
        self.var_q = self.opti.variable(self.reduced_robot.model.nq)
        self.var_q_last = self.opti.parameter(self.reduced_robot.model.nq)
        self.param_tf_l = self.opti.parameter(4, 4)
        self.param_tf_r = self.opti.parameter(4, 4)
        self.param_elbow_l = self.opti.parameter(3, 1)
        self.param_elbow_r = self.opti.parameter(3, 1)

        # Re-evaluate cost functions with new opti variables
        self.translational_cost = casadi.sumsqr(
            self.translational_error(self.var_q, self.param_tf_l, self.param_tf_r))
        self.rotation_cost = casadi.sumsqr(
            self.rotational_error(self.var_q, self.param_tf_l, self.param_tf_r))
        self.regularization_cost = casadi.sumsqr(self.var_q)
        self.smooth_cost = casadi.sumsqr(self.var_q - self.var_q_last)
        self.elbow_cost = casadi.sumsqr(
            self.elbow_error(self.var_q, self.param_elbow_l, self.param_elbow_r))

        # Joint limits
        self.opti.subject_to(self.opti.bounded(
            self.reduced_robot.model.lowerPositionLimit,
            self.var_q,
            self.reduced_robot.model.upperPositionLimit)
        )

        # Combined objective
        self.opti.minimize(
            50 * self.translational_cost +
            self.rotation_cost +
            0.02 * self.regularization_cost +
            0.1 * self.smooth_cost +
            elbow_weight * self.elbow_cost
        )

        opts = {
            'expand': True,
            'detect_simple_bounds': True,
            'calc_lam_p': False,
            'print_time': False,
            'ipopt.sb': 'yes',
            'ipopt.print_level': 0,
            'ipopt.max_iter': 30,
            'ipopt.tol': 1e-4,
            'ipopt.acceptable_tol': 5e-4,
            'ipopt.acceptable_iter': 5,
            'ipopt.warm_start_init_point': 'yes',
            'ipopt.derivative_test': 'none',
            'ipopt.jacobian_approximation': 'exact',
        }
        self.opti.solver("ipopt", opts)

    def solve_ik(self, left_wrist, right_wrist,
                 current_lr_arm_motor_q=None, current_lr_arm_motor_dq=None,
                 left_elbow_pos=None, right_elbow_pos=None):
        """Solve IK with wrist + elbow position targets.

        Args:
            left_wrist: 4x4 SE(3) left wrist target
            right_wrist: 4x4 SE(3) right wrist target
            current_lr_arm_motor_q: current 14-DOF arm angles (warm start)
            current_lr_arm_motor_dq: current velocities (API compat, unused)
            left_elbow_pos: 3-vector left elbow position (waist frame)
            right_elbow_pos: 3-vector right elbow position (waist frame)

        Returns:
            (sol_q, sol_tauff) — 14-DOF joint angles and feedforward torques
        """
        if current_lr_arm_motor_q is not None:
            self.init_data = current_lr_arm_motor_q
        self.opti.set_initial(self.var_q, self.init_data)

        self.opti.set_value(self.param_tf_l, left_wrist)
        self.opti.set_value(self.param_tf_r, right_wrist)
        self.opti.set_value(self.var_q_last, self.init_data)

        self.opti.set_value(self.param_elbow_l,
                            left_elbow_pos if left_elbow_pos is not None else np.zeros(3))
        self.opti.set_value(self.param_elbow_r,
                            right_elbow_pos if right_elbow_pos is not None else np.zeros(3))

        try:
            sol = self.opti.solve()
            sol_q = self.opti.value(self.var_q)
        except Exception as e:
            sol_q = self.opti.debug.value(self.var_q)

        self.smooth_filter.add_data(sol_q)
        sol_q = self.smooth_filter.filtered_data

        v = (current_lr_arm_motor_dq * 0.0 if current_lr_arm_motor_dq is not None
             else (sol_q - self.init_data) * 0.0)

        self.init_data = sol_q
        sol_tauff = pin.rnea(
            self.reduced_robot.model, self.reduced_robot.data,
            sol_q, v, np.zeros(self.reduced_robot.model.nv))

        return sol_q, sol_tauff