sim = require('sim')
simROS2 = require('simROS2')
THETA0 = -math.pi/2
YAW_SIGN = -1     
function sysCall_init()
    wheel_joints = {
        fl = sim.getObject('/Omnirob/FLwheel_motor'),
        fr = sim.getObject('/Omnirob/FRwheel_motor'),
        rl = sim.getObject('/Omnirob/RLwheel_motor'),
        rr = sim.getObject('/Omnirob/RRwheel_motor')
    }
    
    -- ? Set wheel radius and correct L and W (from center to wheel)
    wheel_radius = 0.108   -- meters
    L = 0.694    
    W = 0.4765 
    kinematic_factor = (L/2 + W/2)
    -- Sim base ? ROS base_link frame correction:

    laserFrameObj = sim.getObject('/Omnirob/laser_frame')
    baseLinkObj   = sim.getObject('/Omnirob')
        
    -- ROS2 setup
    cmdSub   = simROS2.createSubscription('/Omnirob/cmd_vel', 'geometry_msgs/msg/Twist', 'cmdVelCallback')
    tf_static_Pub = simROS2.createPublisher('/tf_static','tf2_msgs/msg/TFMessage')
    tfPub    = simROS2.createPublisher('/tf', 'tf2_msgs/msg/TFMessage')
    odomPub  = simROS2.createPublisher('/odom', 'nav_msgs/msg/Odometry')
    clockPub = simROS2.createPublisher('/clock', 'rosgraph_msgs/msg/Clock')
    laserPub = simROS2.createPublisher('/scan', 'sensor_msgs/msg/LaserScan')
    
    sensors = {
        {vs=sim.getObject('/Omnirob/SickS300/sensor1'), joint=sim.getObject('/Omnirob/SickS300/joint1')},
        {vs=sim.getObject('/Omnirob/SickS300/sensor2'), joint=sim.getObject('/Omnirob/SickS300/joint2')},
        {vs=sim.getObject('/Omnirob/SickS3001/sensor1'), joint=sim.getObject('/Omnirob/SickS3001/joint1')},
        {vs=sim.getObject('/Omnirob/SickS3001/sensor2'), joint=sim.getObject('/Omnirob/SickS3001/joint2')}
    }

    maxScanDistance = 5.0
    scanningAngle1 = math.rad(260)
    scanningAngle2 = math.rad(135)

    for i = 1, 2 do
        sim.setObjectFloatParam(sensors[i].vs, sim.visionfloatparam_far_clipping, maxScanDistance)
        sim.setObjectFloatParam(sensors[i].vs, sim.visionfloatparam_perspective_angle, scanningAngle1 / 2)
    end
    for i = 3, 4 do
        sim.setObjectFloatParam(sensors[i].vs, sim.visionfloatparam_far_clipping, maxScanDistance)
        sim.setObjectFloatParam(sensors[i].vs, sim.visionfloatparam_perspective_angle, scanningAngle2 / 2)
    end
    sim.setJointPosition(sensors[1].joint, -scanningAngle1 / 4)
    sim.setJointPosition(sensors[2].joint,  scanningAngle1 / 4)
    sim.setJointPosition(sensors[3].joint, -scanningAngle2 / 4)
    sim.setJointPosition(sensors[4].joint,  scanningAngle2 / 4)
    


    -- Internal state
    odom_x, odom_y, odom_yaw = 0.0, 0.0, 0.0
    last_time = sim.getSimulationTime()
    last_vx_ros, last_vy_ros = 0.0, 0.0
    vx_cmd, vy_cmd, omega_cmd = 0.0, 0.0, 0.0
    -- Laser scan
    angle_min = -math.pi
    angle_max = math.pi
    num_readings = 720
    angle_increment = (angle_max - angle_min) / (num_readings - 1)
    maxScanDistance_ = maxScanDistance * 0.9999
    
    sigma_x   = 0.03   -- m
    sigma_y   = 0.03   -- m
    sigma_yaw = 0.08   -- rad
    bigVar    = 1e3    -- very uncertain on z/roll/pitch (2D)

    pose_cov = {
        sigma_x^2, 0,         0,      0,      0,      0,
        0,         sigma_y^2, 0,      0,      0,      0,
        0,         0,         bigVar, 0,      0,      0,
        0,         0,         0,      bigVar, 0,      0,
        0,         0,         0,      0,      bigVar, 0,
        0,         0,         0,      0,      0,      sigma_yaw^2
    }

    -- Odometry twist (vx, vy, vz, wx, wy, wz)
    sigma_v    = 0.04  -- m/s
    sigma_omega= 0.10  -- rad/s
    twist_cov = {
        sigma_v^2, 0,          0,         0,         0,         0,
        0,         sigma_v^2,  0,         0,         0,         0,
        0,         0,          bigVar,    0,         0,         0,
        0,         0,          0,         bigVar,    0,         0,
        0,         0,          0,         0,         bigVar,    0,
        0,         0,          0,         0,         0,         sigma_omega^2
    }

end

function cmdVelCallback(msg)
    vx_cmd = msg.linear.x
    vy_cmd = msg.linear.y
    omega_cmd = msg.angular.z
     if math.abs(omega_cmd) > 1e-3 then
        sim.addLog(sim.verbosity_scriptinfos, string.format("cmd ?=%.3f", omega_cmd))
    end
end

function sysCall_actuation()
    local now = sim.getSimulationTime()
    local dt = now - last_time
    if dt <= 0 then return end
    local base = sim.getObject('/Omnirob')
    local fwd = sim.multiplyVector(sim.getObjectMatrix(base, -1), {1,0,0})
    local left = sim.multiplyVector(sim.getObjectMatrix(base, -1), {0,1,0})
    --print(string.format("FORWARD points to: (%.2f, %.2f, %.2f)", fwd[1], fwd[2], fwd[3]))
    --print(string.format("LEFT points to: (%.2f, %.2f, %.2f)", left[1], left[2], left[3]))

    -- ? Drive robot (Mecanum wheel inverse kinematics)
    local s_fl, s_fr, s_rl, s_rr = -1, 1, 1, -1  
    local scale = 1 / wheel_radius
    -- desired *math* wheel speeds (rad/s) per standard mecanum IK (x fwd, y left):
    local w_fl_math =  (vy_cmd - vx_cmd - kinematic_factor * omega_cmd) * scale
    local w_fr_math = (-vy_cmd - vx_cmd + kinematic_factor * omega_cmd) * scale
    local w_rl_math =  (vy_cmd + vx_cmd + kinematic_factor * omega_cmd) * scale
    local w_rr_math = (-vy_cmd + vx_cmd - kinematic_factor * omega_cmd) * scale

    -- convert to joint targets using the SAME signs:
    local w_fl_joint = s_fl * w_fl_math
    local w_fr_joint = s_fr * w_fr_math
    local w_rl_joint = s_rl * w_rl_math
    local w_rr_joint = s_rr * w_rr_math

    sim.setJointTargetVelocity(wheel_joints.fl, w_fl_joint)
    sim.setJointTargetVelocity(wheel_joints.fr, w_fr_joint)
    sim.setJointTargetVelocity(wheel_joints.rl, w_rl_joint)
    sim.setJointTargetVelocity(wheel_joints.rr, w_rr_joint)

    -- ------------- FK (joint -> body twist) read joints, map back to *math* wheels:
    local v_fl_math = s_fl * sim.getJointVelocity(wheel_joints.fl)
    local v_fr_math = s_fr * sim.getJointVelocity(wheel_joints.fr)
    local v_rl_math = s_rl * sim.getJointVelocity(wheel_joints.rl)
    local v_rr_math = s_rr * sim.getJointVelocity(wheel_joints.rr)

    -- standard mecanum FK (x fwd, y left):
    local vx = ( v_fl_math - v_fr_math + v_rl_math - v_rr_math) * wheel_radius / 4.0
    local vy = (-v_fl_math - v_fr_math + v_rl_math + v_rr_math) * wheel_radius / 4.0
    local omega = (-v_fl_math + v_fr_math + v_rl_math - v_rr_math) * wheel_radius / (4.0 * kinematic_factor)
    
    -- rotate twist from SIM base to ROS base_link
    local c0, s0 = math.cos(THETA0), math.sin(THETA0)
    local vx_ros    =  c0*vx - s0*vy
    local vy_ros    =  s0*vx + c0*vy
    local omega_ros =  YAW_SIGN * omega

    -- integrate odom *in the ROS base frame*
    local cy = math.cos(odom_yaw)
    local sy = math.sin(odom_yaw)
    local dx =  vx_ros*cy - vy_ros*sy
    local dy =  -(vx_ros*sy + vy_ros*cy)
    odom_x   = odom_x + dx*dt
    odom_y   = odom_y + dy*dt
    odom_yaw = odom_yaw + omega_ros*dt

    -- DEBUG: show wheel pattern during pure rotation tests
    --if math.abs(omega_cmd) > 1e-6 and math.abs(vx_cmd) < 1e-6 and math.abs(vy_cmd) < 1e-6 then
    --    print(string.format("w_math [fl, fr, rl, rr] = [%.2f, %.2f, %.2f, %.2f]",
    --        w_fl_math, w_fr_math, w_rl_math, w_rr_math))
    --end

    --print(string.format("vx=%.3f, vy=%.3f, omega=%.3f", vx_ros, vy_ros, omega_ros))


    if odom_yaw > math.pi then odom_yaw = odom_yaw - 2 * math.pi end
    if odom_yaw < -math.pi then odom_yaw = odom_yaw + 2 * math.pi end

    local ax = (vx_ros - last_vx_ros) / dt
    local ay = (vy_ros - last_vy_ros) / dt
    last_vx_ros, last_vy_ros = vx_ros, vy_ros

    -- ? ROS2 time stamp
    local sec = math.floor(now)
    local nsec = math.floor((now - sec) * 1e9)
    simROS2.publish(clockPub, { clock = { sec = sec, nanosec = nsec } })

    -- ? TF: odom ? base_link and base_link ? laser_frame
    
    -- Get relative transform from base_link to laser_frame
    local pos_sim = sim.getObjectPosition(laserFrameObj, baseLinkObj)
    local quat_sim = sim.getObjectQuaternion(laserFrameObj, baseLinkObj)
    local euler = sim.getObjectOrientation(baseLinkObj, -1)
    -- helper: rotate a vector in XY by angle ang
    local function rot2d(p, ang)
      local c, s = math.cos(ang), math.sin(ang)
      return { c*p[1] - s*p[2], s*p[1] + c*p[2], p[3] }
    end

    -- helper: compose yaw-only quaternion Rz(ang) * q
    local function yawQuat(ang)
      return {0,0, math.sin(ang/2), math.cos(ang/2)}
    end
    local function quatMul(a,b) -- (ax,ay,az,aw) * (bx,by,bz,bw)
      return {
        a[4]*b[1] + a[1]*b[4] + a[2]*b[3] - a[3]*b[2],
        a[4]*b[2] - a[1]*b[3] + a[2]*b[4] + a[3]*b[1],
        a[4]*b[3] + a[1]*b[2] - a[2]*b[1] + a[3]*b[4],
        a[4]*b[4] - a[1]*b[1] - a[2]*b[2] - a[3]*b[3]
      }
    end

    -- convert sim-base-relative transform to ROS base_link
    local pos_ros  = rot2d(pos_sim, -THETA0)
    local q_off    = yawQuat(-THETA0)
    local quat_ros = quatMul(q_off, {quat_sim[1], quat_sim[2], quat_sim[3], quat_sim[4]})

    local tfMsgs = {{
      header = { stamp = { sec = sec, nanosec = nsec }, frame_id = 'odom' },
      child_frame_id = 'base_link',
      transform = {
        translation = { x = odom_x, y = odom_y, z = 0.0 },
        rotation    = { x = 0.0, y = 0.0, z = math.sin(odom_yaw/2), w = math.cos(odom_yaw/2) }
      }
    }}
    simROS2.publish(tfPub, { transforms = tfMsgs })

    local pos = sim.getObjectPosition(baseLinkObj, -1)
    local ori = sim.getObjectOrientation(baseLinkObj, -1)
    
    --print(string.format("Odom: (%.3f, %.3f, %.3f)", odom_x, odom_y, odom_yaw))
    --print(string.format("True: (%.3f, %.3f, %.3f)", pos[1], pos[2], ori[3]))
    --print(string.format("Yaw in deg: %.1f", math.deg(euler[3])))
    -- ? Odometry message
    simROS2.publish(odomPub, {
        header = { stamp = { sec = sec, nanosec = nsec }, frame_id = 'odom' },
        child_frame_id = 'base_link',
        pose = {
          pose = {
            position = { x = odom_x, y = odom_y, z = 0.0 },
            orientation = {
              x = 0.0, y = 0.0,
              z = math.sin(odom_yaw/2), w = math.cos(odom_yaw/2)
            }
          },
          covariance = pose_cov
        },

        twist = {
            twist = {
              linear  = { x = vx_ros,  y = vy_ros,  z = 0.0 },
              angular = { x = 0.0, y = 0.0, z = omega_ros }
            },
            covariance = twist_cov
        }
    })



    -- ? Laser scan
    local scan_points = {}
    for _, s in ipairs(sensors) do
        local _, _, u = sim.readVisionSensor(s.vs)
        local relMatrix = sim.getObjectMatrix(s.vs, laserFrameObj)
        collectPoints(u, relMatrix, scan_points)
    end

    table.sort(scan_points, function(a, b) return a[1] < b[1] end)

    local ranges, bin_min = {}, {}
    for _, p in ipairs(scan_points) do
        local angle, d = p[1], p[2]
        if angle >= angle_min and angle <= angle_max then
            local idx = math.floor((angle - angle_min) / angle_increment) + 1
            if idx >= 1 and idx <= num_readings then
                if bin_min[idx] == nil or d < bin_min[idx] then
                    bin_min[idx] = d
                end
            end
        end
    end
    for i = 1, num_readings do
        ranges[i] = bin_min[i] or maxScanDistance
    end

    local scan_dt = dt
    simROS2.publish(laserPub, {
      header = { stamp = { sec = sec, nanosec = nsec }, frame_id = 'laser_frame' },
      angle_min = angle_min,
      angle_max = angle_max,
      angle_increment = angle_increment,
      time_increment = 0,
      scan_time = 0,
      range_min = 0.1,
      range_max = maxScanDistance,
      ranges = ranges,
      intensities = {}
    })


    last_time = now
end

function collectPoints(unitData, matrix, out)
    if not unitData then return end
    for j = 0, unitData[2] - 1 do
        for i = 0, unitData[1] - 1 do
            local w = 2 + 4 * (j * unitData[1] + i)
            local x, y, z, d = unitData[w+1], unitData[w+2], unitData[w+3], unitData[w+4]
            if d < maxScanDistance_ then
                local p = sim.multiplyVector(matrix, {x, y, z})
                local dist = math.sqrt(p[1]^2 + p[2]^2)
                local ang  = math.atan2(p[2], p[1])
                table.insert(out, {ang, dist})
                --print(d)
            end
        end
    end
end