include "mjmodel.pxd"


cdef extern from "mjdata.h" nogil:

    #---------------------------- primitive types (mjt) ------------------------------------
        
    ctypedef enum mjtWarning:            # warning types
        mjWARN_INERTIA      = 0,        # (near) singular inertia matrix
        mjWARN_CONTACTFULL,             # too many contacts in contact list
        mjWARN_CNSTRFULL,               # too many constraints
        mjWARN_VGEOMFULL,               # too many visual geoms
        mjWARN_BADQPOS,                 # bad number in qpos
        mjWARN_BADQVEL,                 # bad number in qvel
        mjWARN_BADQACC,                 # bad number in qacc
        mjWARN_BADCTRL,                 # bad number in ctrl

    enum: mjNWARNING                      # number of warnings


    ctypedef enum mjtTimer:
        # main api
        mjTIMER_STEP        = 0,        # step
        mjTIMER_FORWARD,                # forward
        mjTIMER_INVERSE,                # inverse

        # breakdown of step/forward
        mjTIMER_POSITION,               # fwdPosition
        mjTIMER_VELOCITY,               # fwdVelocity
        mjTIMER_ACTUATION,              # fwdActuation
        mjTIMER_ACCELERATION,           # fwdAcceleration
        mjTIMER_CONSTRAINT,             # fwdConstraint

        # breakdown of fwdPosition
        mjTIMER_POS_KINEMATICS,         # kinematics, com, tendon, transmission
        mjTIMER_POS_INERTIA,            # inertia computations
        mjTIMER_POS_COLLISION,          # collision detection
        mjTIMER_POS_MAKE,               # make constraints
        mjTIMER_POS_PROJECT,            # project constraints

    enum: mjNTIMER                        # number of timers


    #------------------------------ mjContact ----------------------------------------------

    ctypedef struct mjContact:                   # result of collision detection functions
        # contact parameters set by geom-specific collision detector
        mjtNum dist                     # distance between nearest points; neg: penetration
        mjtNum pos[3]                   # position of contact point: midpoint between geoms
        mjtNum frame[9]                 # normal is in [0-2]

        # contact parameters set by mj_collideGeoms
        mjtNum includemargin            # include if dist<includemargin=margin-gap
        mjtNum friction[5]              # tangent1, 2, spin, roll1, 2
        mjtNum solref[mjNREF]           # constraint solver reference
        mjtNum solimp[mjNIMP]           # constraint solver impedance

        # storage used internally by constraint solver
        mjtNum mu                       # friction of regularized cone
        mjtNum H[36]                    # cone Hessian, set by mj_updateConstraint

        # contact descriptors set by mj_collideGeoms
        int dim                         # contact space dimensionality: 1, 3, 4 or 6
        int geom1                       # id of geom 1
        int geom2                       # id of geom 2

        # flag set by mj_fuseContact or mj_instantianteEquality
        int exclude                     # 0: include, 1: in gap, 2: fused, 3: equality

        # address computed by mj_instantiateContact
        int efc_address                 # address in efc; -1: not included, -2-i: distance constraint i ???

    #------------------------------ diagnostics --------------------------------------------

    ctypedef struct mjWarningStat:      # warning statistics
        int lastinfo                    # info from last warning
        int number                      # how many times was warning raised


    ctypedef struct mjTimerStat:        # timer statistics
        mjtNum duration                 # cumulative duration
        int number                      # how many times was timer called


    ctypedef struct mjSolverStat:       # per-iteration solver statistics
        mjtNum improvement              # cost reduction, scaled by 1/trace(M(qpos0))
        mjtNum gradient                 # gradient norm (primal only, scaled)
        mjtNum lineslope                # slope in linesearch
        int nactive                     # number of active constraints
        int nchange                     # number of constraint state changes
        int neval                       # number of cost evaluations in line search
        int nupdate                     # number of Cholesky updates in line search



    #---------------------------------- mjData ---------------------------------------------
    ctypedef struct mjData:
        # constant sizes
        int nstack                      # number of mjtNums that can fit in stack
        int nbuffer                     # size of main buffer in bytes

        # stack pointer
        int pstack                      # first available mjtNum address in stack

        # memory utilization stats
        int maxuse_stack                # maximum stack allocation
        int maxuse_con                  # maximum number of contacts
        int maxuse_efc                  # maximum number of scalar constraints

        # diagnostics
        mjWarningStat warning[mjNWARNING] # warning statistics
        mjTimerStat timer[mjNTIMER]       # timer statistics
        mjSolverStat solver[mjNSOLVER]    # solver statistics per iteration
        int solver_iter                   # number of solver iterations
        int solver_nnz                  # number of non-zeros in Hessian or efc_AR
        mjtNum solver_fwdinv[2]         # forward-inverse comparison: qfrc, efc

        # variable sizes
        int ne                          # number of equality constraints
        int nf                          # number of friction constraints
        int nefc                        # number of constraints
        int ncon                        # number of detected contacts

        # global properties
        mjtNum time                     # simulation time
        mjtNum energy[2]                # potential, kinetic energy

        #-------------------------------- end of info header

        # buffers
        void*     buffer                # main buffer; all pointers point in it    (nbuffer bytes)
        mjtNum*   stack                 # stack buffer                             (nstack mjtNums)

        #-------------------------------- main inputs and outputs of the computation

        # state
        mjtNum*   qpos                  # position                                 (nq x 1)
        mjtNum*   qvel                  # velocity                                 (nv x 1)
        mjtNum*   act                   # actuator activation                      (na x 1)

        # control
        mjtNum*   ctrl                  # control                                  (nu x 1)
        mjtNum*   qfrc_applied          # applied generalized force                (nv x 1)
        mjtNum*   xfrc_applied          # applied Cartesian force/torque           (nbody x 6)

        # dynamics
        mjtNum*   qacc                  # acceleration                             (nv x 1)
        mjtNum*   act_dot               # time-derivative of actuator activation   (na x 1)

        # mocap data
        mjtNum*  mocap_pos              # positions of mocap bodies                (nmocap x 3)
        mjtNum*  mocap_quat             # orientations of mocap bodies             (nmocap x 4)

        # user data
        mjtNum*  userdata               # user data, not touched by engine         (nuserdata x 1)

        # sensors
        mjtNum*  sensordata             # sensor data array                        (nsensordata x 1)

        #-------------------------------- POSITION dependent

        # computed by mj_fwdPosition/mj_kinematics
        mjtNum*   xpos                  # Cartesian position of body frame         (nbody x 3)
        mjtNum*   xquat                 # Cartesian orientation of body frame      (nbody x 4)
        mjtNum*   xmat                  # Cartesian orientation of body frame      (nbody x 9)
        mjtNum*   xipos                 # Cartesian position of body com           (nbody x 3)
        mjtNum*   ximat                 # Cartesian orientation of body inertia    (nbody x 9)
        mjtNum*   xanchor               # Cartesian position of joint anchor       (njnt x 3)
        mjtNum*   xaxis                 # Cartesian joint axis                     (njnt x 3)
        mjtNum*   geom_xpos             # Cartesian geom position                  (ngeom x 3)
        mjtNum*   geom_xmat             # Cartesian geom orientation               (ngeom x 9)
        mjtNum*   site_xpos             # Cartesian site position                  (nsite x 3)
        mjtNum*   site_xmat             # Cartesian site orientation               (nsite x 9)
        mjtNum*   cam_xpos              # Cartesian camera position                (ncam x 3)
        mjtNum*   cam_xmat              # Cartesian camera orientation             (ncam x 9)
        mjtNum*   light_xpos            # Cartesian light position                 (nlight x 3)
        mjtNum*   light_xdir            # Cartesian light direction                (nlight x 3)

        # computed by mj_fwdPosition/mj_comPos
        mjtNum*   subtree_com           # center of mass of each subtree           (nbody x 3)
        mjtNum*   cdof                  # com-based motion axis of each dof        (nv x 6)
        mjtNum*   cinert                # com-based body inertia and mass          (nbody x 10)

        # computed by mj_fwdPosition/mj_tendon
        int*      ten_wrapadr           # start address of tendon's path           (ntendon x 1)
        int*      ten_wrapnum           # number of wrap points in path            (ntendon x 1)
        mjtNum*   ten_length            # tendon lengths                           (ntendon x 1)
        mjtNum*   ten_moment            # tendon moment arms                       (ntendon x nv)
        int*      wrap_obj              # geom id; -1: site; -2: pulley            (nwrap*2 x 1)
        mjtNum*   wrap_xpos             # Cartesian 3D points in all path          (nwrap*2 x 3)

        # computed by mj_fwdPosition/mj_transmission
        mjtNum*   actuator_length       # actuator lengths                         (nu x 1)
        mjtNum*   actuator_moment       # actuator moment arms                     (nu x nv)

        # computed by mj_fwdPosition/mj_crb
        mjtNum*   crb                   # com-based composite inertia and mass     (nbody x 10)
        mjtNum*   qM                    # total inertia                            (nM x 1)

        # computed by mj_fwdPosition/mj_factorM
        mjtNum*   qLD                   # L'*D*L factorization of M                (nM x 1)
        mjtNum*   qLDiagInv             # 1/diag(D)                                (nv x 1)
        mjtNum*   qLDiagSqrtInv         # 1/sqrt(diag(D))                          (nv x 1)

        # computed by mj_fwdPosition/mj_collision
        mjContact* contact              # list of all detected contacts            (nconmax x 1)

        # computed by mj_fwdPosition/mj_makeConstraint
        int*      efc_type              # constraint type (mjtConstraint)          (njmax x 1)
        int*      efc_id                # id of object of specified type           (njmax x 1)
        int*      efc_J_rownnz          # number of non-zeros in Jacobian row      (njmax x 1)
        int*      efc_J_rowadr          # row start address in colind array        (njmax x 1)
        int*      efc_J_colind          # column indices in sparse Jacobian        (njmax x nv)
        int*      efc_JT_rownnz         # number of non-zeros in Jacobian row  T   (nv x 1)
        int*      efc_JT_rowadr         # row start address in colind array    T   (nv x 1)
        int*      efc_JT_colind         # column indices in sparse Jacobian    T   (nv x njmax)
        mjtNum*   efc_solref            # constraint solver reference              (njmax x mjNREF)
        mjtNum*   efc_solimp            # constraint solver impedance              (njmax x mjNIMP)
        mjtNum*   efc_margin            # inclusion margin (contact)               (njmax x 1)
        mjtNum*   efc_frictionloss      # frictionloss (friction)                  (njmax x 1)
        mjtNum*   efc_pos               # constraint position (equality, contact)  (njmax x 1)
        mjtNum*   efc_J                 # constraint Jacobian                      (njmax x nv)
        mjtNum*   efc_JT                # sparse constraint Jacobian transposed    (nv x njmax)
        mjtNum*   efc_diagApprox        # approximation to diagonal of A           (njmax x 1)
        mjtNum*   efc_D                 # constraint mass                          (njmax x 1)
        mjtNum*   efc_R                 # inverse constraint mass                  (njmax x 1)

        # computed by mj_fwdPosition/mj_projectConstraint
        int*      efc_AR_rownnz         # number of non-zeros in AR                (njmax x 1)
        int*      efc_AR_rowadr         # row start address in colind array        (njmax x 1)
        int*      efc_AR_colind         # column indices in sparse AR              (njmax x njmax)
        mjtNum*   efc_AR                # J*inv(M)*J' + R                          (njmax x njmax)

        #-------------------------------- POSITION, VELOCITY dependent

        # computed by mj_fwdVelocity
        mjtNum*   ten_velocity          # tendon velocities                        (ntendon x 1)
        mjtNum*   actuator_velocity     # actuator velocities                      (nu x 1)

        # computed by mj_fwdVelocity/mj_comVel
        mjtNum*   cvel                  # com-based velocity [3D rot; 3D tran]     (nbody x 6)
        mjtNum*   cdof_dot              # time-derivative of cdof                  (nv x 6)

        # computed by mj_fwdVelocity/mj_rne (without acceleration)
        mjtNum*   qfrc_bias             # C(qpos,qvel)                             (nv x 1)

        # computed by mj_fwdVelocity/mj_passive
        mjtNum*   qfrc_passive          # passive force                            (nv x 1)

        # computed by mj_fwdVelocity/mj_referenceConstraint
        mjtNum*   efc_vel               # velocity in constraint space: J*qvel     (njmax x 1)
        mjtNum*   efc_aref              # reference pseudo-acceleration            (njmax x 1)

        # computed by mj_sensorVel
        mjtNum*   subtree_linvel        # linear velocity of subtree com           (nbody x 3)
        mjtNum*   subtree_angmom        # angular momentum about subtree com       (nbody x 3)

        #-------------------------------- POSITION, VELOCITY, CONTROL/ACCELERATION dependent

        # computed by mj_fwdActuation
        mjtNum*   actuator_force        # actuator force in actuation space        (nu x 1)
        mjtNum*   qfrc_actuator         # actuator force                           (nv x 1)

        # computed by mj_fwdAcceleration
        mjtNum*   qfrc_unc              # net unconstrained force                  (nv x 1)
        mjtNum*   qacc_unc              # unconstrained acceleration               (nv x 1)

        # computed by mj_fwdConstraint/mj_inverse
        mjtNum*   efc_b                 # linear cost term: J*qacc_unc - aref      (njmax x 1)
        mjtNum*   efc_force             # constraint force in constraint space     (njmax x 1)
        int*      efc_state             # constraint state (mjtConstraintState)    (njmax x 1)
        mjtNum*   qfrc_constraint       # constraint force                         (nv x 1)
        mjtNum*   qacc_warmstart        # acceleration used for warmstart          (nv x 1)

        # computed by mj_inverse
        mjtNum*   qfrc_inverse          # net external force; should equal:        (nv x 1)
                                        #  qfrc_applied + J'*xfrc_applied + qfrc_actuator

        # computed by mj_sensorAcc/mj_rnePostConstraint; rotation:translation format
        mjtNum*   cacc                  # com-based acceleration                   (nbody x 6)
        mjtNum*   cfrc_int              # com-based interaction force with parent  (nbody x 6)
        mjtNum*   cfrc_ext              # com-based external force on body         (nbody x 6)


# #---------------------------------- callback function types ----------------------------
#
# # generic MuJoCo function
# typedef void (*mjfGeneric)(const mjModel* m, mjData* d)
#
# # sensor simulation
# typedef void (*mjfSensor)(const mjModel* m, mjData* d, int stage)
#
# # timer
# typedef long long int (*mjfTime)(void);
#
# # actuator dynamics, gain, bias
# typedef mjtNum (*mjfAct)(const mjModel* m, const mjData* d, int id);
#
# # solver impedance
# typedef mjtNum (*mjfSolImp)(const mjModel* m, const mjData* d, int id,
#                             mjtNum distance, mjtNum* constimp);
#
# # solver reference
# typedef void (*mjfSolRef)(const mjModel* m, const mjData* d, int id,
#                           mjtNum constimp, mjtNum imp, int dim, mjtNum* ref);
#
# # collision detection
# typedef int (*mjfCollision)(const mjModel* m, const mjData* d,
#                             mjContact* con, int g1, int g2, mjtNum margin);
