import stormpy
import stormpy.simulator

newline = "\n"

def printFrame(streetLength, lanes, carPos, pedPos):
    sidewalk = "#" * (streetLength+1)
    frame = list(sidewalk + (lanes-2)*("_"*(streetLength+1)) + sidewalk)

    frame[(int(carPos[0])-1)*(streetLength+1)+int(carPos[1])+1] = "C"
    frame[(int(pedPos[0])-1)*(streetLength+1)+int(pedPos[1])+1] = "P"
    for i in range(lanes):
        frame[i*(streetLength+1)] = newline
    return "".join(frame)

def mdp():
    options = stormpy.BuilderOptions([])
    options.set_build_state_valuations(True)
    options.set_build_choice_labels(True)
    options.set_build_all_labels()
    prism_program = stormpy.parse_prism_program("/media/car_pedestrian_mdp.prism")

    streetLength = int(str(prism_program.get_constant("streetLength").definition))
    lanes = int(str(prism_program.get_constant("lanes").definition))
    formula_str = "Pmax=? [G !\"crashed\" ];"
    #formula_str = "Pmax=? [F \"crashed\" ];"

    print(stormpy.build_model(prism_program))
    formulas = stormpy.parse_properties_for_prism_program(formula_str, prism_program)
    model = stormpy.build_sparse_model_with_options(prism_program, options)
    initial_state = model.initial_states[0]
    assert initial_state == 0
    result = stormpy.model_checking(model, formulas[0], extract_scheduler=True)
    assert result.has_scheduler
    scheduler = result.scheduler
    assert scheduler.memoryless
    assert scheduler.deterministic

    dtmc = model.apply_scheduler(scheduler)
    print(dtmc)
    simulator = stormpy.simulator.create_simulator(dtmc, seed=42)
    simulator.set_observation_mode(stormpy.simulator.SimulatorObservationMode.PROGRAM_LEVEL)

    i=0
    while True:
        observation, reward, labels = simulator.step()
        print(observation, labels)
        frame = printFrame(int(streetLength), int(lanes), (observation["car_lane_pos"], observation["car_street_pos"]), (observation["ped_lane_pos"], observation["ped_street_pos"]))
        print(frame)
        if simulator.is_done() or "pedCrossed" in labels:
            print("Pedestrian arrived on the other side!")
            break
        input("")
        with open(f"/media/frame{i:03}.txt", "w") as text_file:
            text_file.write(frame)
        i+=1
        if "crashed" in labels:
            print("Pedestrian did not arrive on the other side!")
            break



if __name__ == '__main__':
    mdp()