#!/usr/bin/env python

# external modules
import rospy
import collections
from tf import transformations
from geometry_msgs.msg import Pose

# internal modules
import Simulator as sim
import constants as c
from landmark_localizer.msg import StochasticPoseArray


def rosPoseFromPosAndOrnt(pos, ornt):
    poseMsg = Pose()
    poseMsg.position.x = pos[0]
    poseMsg.position.y = pos[1]
    poseMsg.position.z = pos[2]
    if isinstance(ornt, collections.Iterable):
        quat = ornt
    else:
        quat = transformations.quaternion_from_euler(0, 0, ornt)
    poseMsg.orientation.x = quat[0]
    poseMsg.orientation.y = quat[1]
    poseMsg.orientation.z = quat[2]
    poseMsg.orientation.w = quat[3]
    return poseMsg
    

# publish true pose, estimated poses
def main():
    rospy.init_node('rosSimulator')
    args = sim.getSimulatorArgs()
    simulator = sim.Simulator(args.mapFilename, args.simConfig)
    # TODO make this able to handle 3D
    truePosePub = rospy.Publisher(c.TRUE_POSE_TOPIC, Pose, queue_size=1)
    estPosePub = rospy.Publisher(c.EST_POSE_TOPIC, StochasticPoseArray,
                                 queue_size=1)
    rate = rospy.Rate(2)
    for _ in range(args.n):
        truePos, trueOrnt, estimatedPoses = simulator.runSimulation()
        # convert truePos and trueOrnt into a Pose message
        truePosePub.publish(rosPoseFromPosAndOrnt(truePos, trueOrnt))
        
        # convert estimated pose into a StochasticPose2dArray message
        estPosesMsg = StochasticPoseArray()
        probs, positions, orientations = zip(*estimatedPoses)
        estPosesMsg.ps = probs
        estPosesMsg.poses = map(lambda x: rosPoseFromPosAndOrnt(x[0], x[1]),
                                zip(positions, orientations))
        estPosePub.publish(estPosesMsg)
        rate.sleep()
        # publish the StochasticPose2dArray message


if __name__ == "__main__":
    main()