import argparse
import json
from google.protobuf.json_format import MessageToDict
from data.chum import chumpy
from vehicle.common.proto.decision_pb2 import Decision


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("--chum-uri", "-i", required=True, type=str,
            help="Chum URI to get data from")
    parser.add_argument("--outfile", "-o", required=True, type=str,
            help="Path to store data to")
    return parser.parse_args()

def get_stv_data(uri):
    "Prints out s position, velocity, and timing data"
    chum_store, chum_range = chumpy.parseChumUri(uri)
    chum_range.topics = {'/planner/decision'}

    reader = chumpy.Reader.create(chum_store, chum_range)

    for msg in reader:
        ts = msg.message_publish_time_ns
        decision = Decision()
        decision.ParseFromString(msg.getData())

        yield {
            'ts': ts,
            'state': decision.state,
            's': decision.route_stz.s,
            'ey': decision.route_stz.t,
            'vx': decision.hero_state.body_velocity.v_x,
            'vy': decision.hero_state.body_velocity.v_y,
            'z': decision.route_stz.z,
            'decision_traj': MessageToDict(decision.trajectory),
        }

def main():
    args = parse_args()
    with open(args.outfile, "wb") as f:
        for record in get_stv_data(args.chum_uri):
            f.write(json.dumps(record))
            f.write("\n")


if __name__ == "__main__":
    main()
