import mmap
import random
import sys
from time import sleep
import numpy as np
from beamngpy import BeamNGpy, Scenario, Vehicle, setup_logging, ScenarioObject
from beamngpy.sensors import Lidar


def generate_script():
    # Create the script
    script = list()
    points = list()
    point_colors = list()
    spheres = list()
    sphere_colors = list()

    # Hardcoded traj
    traj = [[-499.55,238.29,102.01], [-504.67,232.95,102.06], [-510.00,227.36,102.09], [-515.57,221.53,102.12], [-522.3,214.45,102.17], [-526.53,207.74,102.18], [-529.22,207.19,102.24], [-531.59,208.04,102.28], [-533.72,209.82,102.31], [-535.93, 211.59, 102.31], [-537.26,212.84, 102.4]]

    t = 1
    for p in traj:
        node = {'x': p[0], 'y':p[1], 'z': p[2], 't': t}
        t += 1
        script.append(node)
        points.append([node['x'], node['y'], node['z']])
        point_colors.append([0, 1, 0, 0.1])
        spheres.append([node['x'], node['y'], node['z'], 0.25])
        sphere_colors.append([1, 0, 0, 0.8])

    # Return data
    return script, points, point_colors, spheres, sphere_colors

def load_locations(filename):
    # Open the file
    with open(filename) as f:
        # Read the line
        contents = f.read()
        lines = contents.split()
        # Holds the results
        result = []
        # Holds the individual location
        individual_location = []
        # For each line
        line_count = 0
        for line in lines:
            # Ignore the count, only worry about x,y,z
            if line_count % 4 != 0:
                # Process the file to get the x,y and z co-oridinates
                line = line[2:]
                line = line.replace(',',"")
                individual_location.append(float(line))
            else:
                # Save result and reset
                if len(individual_location) > 0:
                    result.append(individual_location)
                individual_location = []
            # Increment line count
            line_count += 1
    # Return results
    return result

def load_vehicle_types(filename):
    # Open the file
    with open(filename) as f:
        # Read the line
        contents = f.read()
        lines = contents.split()
        # Holds the results
        result = []
        # For each line
        for line in lines:
            result.append(line)
    # Return results
    return result

def create_lidar(viz="physcov"):
    lidar = Lidar()
    if viz == "physcov":
        lidar.__init__(offset=(2.4, 1, 1), direction=(-0.707, -0.707, 0), vres=2,
                       vangle=0.1, rps=10000, hz=60, angle=180, max_dist=200,
                       visualized=True)

    elif viz == "physcovthick":
        lidar.__init__(offset=(2.4, 1, 1), direction=(-0.707, -0.707, 0), vres=2,
                       vangle=0.1, rps=2200000, hz=60, angle=180, max_dist=200,
                       visualized=True)
    elif viz == "roof":
        lidar.__init__(offset=(0, 0, 1.7), direction=(-0.707, -0.707, 0), vres=32,
                       vangle=0.01, rps=2200000, hz=20, angle=360, max_dist=200,
                       visualized=True)
    elif viz == "default":
        lidar.__init__(offset=(0, 0, 1.7), direction=(0, -1, 0), vres=32,
                 vangle=26.9, rps=2200000, hz=20, angle=360, max_dist=200,
                 visualized=True)
    else: # no viz
        lidar = Lidar(offset=(0, 0, 1.6), visualized=False)
    return lidar

def main():
    number_parked_vehicles = 5

    # random.seed(1703)

    setup_logging()

    beamng = BeamNGpy('localhost', 64256, home='H:/BeamNG.research.v1.7.0.1clean')
    bng = beamng.open(launch=True)

    # Create a scenario in west_coast_usa
    scenario = Scenario("west_coast_usa", 'tech_test', description='Random driving for research')

    # https://documentation.beamng.com/vehicle_system/vehicles/index.html
    vehicle = Vehicle('ego_vehicle', model='hatch', licence='Big C', color='Red')

    # Create sensors
    lidar = create_lidar("physcovthick") # physcov, physcovthick, default, no viz
    vehicle.attach_sensor('lidar', lidar)

    # Add the vehicle
    scenario.add_vehicle(vehicle, pos=(-496.16, 241.5, 102), rot=None, rot_quat=(0, 0, 0.258819, 0.9659258))

    # Compile the scenario and place it in BeamNG's map folder
    scenario.make(bng)

    # Get the parking locations and vehicle types
    parking_locations = load_locations("all_parking_locations.txt")
    vehicle_types = load_vehicle_types("car_types.txt")
    orientations = [[0.423, -0.906], [0.940, 0.342]]

    # Add the objects
    if number_parked_vehicles > -1:
        parking_locations = random.sample(parking_locations, number_parked_vehicles)
    for i in range(len(parking_locations)):
        # Randomly select a car type
        vehicle_type = random.choice(vehicle_types)
        # Get an orientation
        orient = random.choice(orientations)
        # Get the car position
        x, y, z = parking_locations[i]
        z += 0.1
        # Create the object
        scenario_obj = ScenarioObject(oid='car' + str(i), name=vehicle_type, otype='BeamNGVehicle', pos=(x, y, z), rot=None, rot_quat=(0, 0, orient[0], orient[1]), scale=(1, 1, 1), JBeam = vehicle_type, datablock="default_vehicle")
        # Add an object
        scenario.add_object(scenario_obj)

    scenario.make(beamng)

    # Start BeamNG and enter the main loop
    try:
        bng.hide_hud()
        bng.set_deterministic()  # Set simulator to be deterministic
        bng.set_steps_per_second(60)  # With 60hz temporal resolution

        # Load and start the scenario
        bng.load_scenario(scenario)
        bng.start_scenario()
        # file1 = open("parking_locations.txt","w")

        # keep track of the number of parking spots checked
        parking_counter = 1
        was_pressed = False

        # Generate a vehicle script
        script, points, point_colors, spheres, sphere_colors = generate_script()
        # bng.add_debug_line(points, point_colors, spheres=spheres, sphere_colors=sphere_colors, cling=True, offset=0.1)
        vehicle.ai_set_script(script)
        vehicle.ai_set_speed(100, mode='limit')
        vehicle.ai_set_aggression(0.2)

        while True:

            # # Retrieve sensor data and show the camera data.
            # vehicle.poll_sensors()
            # sensors = vehicle.sensors
            #
            # # Get the current vehicle position
            # current_x, current_y, current_z = np.round(sensors["state"].data["pos"], 4)
            #
            # # print("-----------------------------")
            # # current_x, current_y, current_z = np.round(sensors["state"].data["pos"], 4)
            # # print("Position {}, {}".format(current_x, current_y))
            # # file1.write("x:{}, y:{}, z:{}\n".format(current_x, current_y, current_z))
            # # print("-----------------------------")
            vehicle.update_vehicle()
            sensors = bng.poll_sensors(vehicle)
            state = vehicle.state
            print("-----------------------------")
            print("SENSORS:\n{}".format(sensors))
            print("VEHICLE STATE:\n{}".format(state))
            print("-----------------------------")

    finally:
        bng.close()
        # file1.close()

if __name__ == '__main__':
    main()