import os
import lgsvl
import time
import psutil
import atexit
import math
from svl_script.object_types import static_types, pedestrian_types, vehicle_types
from customized_utils import emptyobject
import numpy as np
from svl_script.apollo_configs import get_modules_for_id

accident_happen = False

# temporary, can be imported from customized_utils
def exit_handler():
    PROC_NAME = "mainboard"
    for proc in psutil.process_iter():
        # check whether the process to kill name matches
        if proc.name() == PROC_NAME:
            proc.kill()
##################################################

def norm_2d(loc_1, loc_2):
    return np.sqrt((loc_1.x - loc_2.x) ** 2 + (loc_1.z - loc_2.z) ** 2)

def on_waypoint(agent, index):
    print("Waypoint {} reached".format(index))

def initialize_simulator(map, sim_specific_arguments):
    SIMULATOR_HOST = os.environ.get("SIMULATOR_HOST", "127.0.0.1")
    SIMULATOR_PORT = int(os.environ.get("SIMULATOR_PORT", 8181))
    BRIDGE_HOST = os.environ.get("BRIDGE_HOST", "127.0.0.1")
    BRIDGE_PORT = int(os.environ.get("BRIDGE_PORT", 9090))

    if not sim_specific_arguments.sim:
        sim = lgsvl.Simulator(SIMULATOR_HOST, SIMULATOR_PORT)
        sim_specific_arguments.sim = sim
    else:
        sim = sim_specific_arguments.sim
    # hack: TBD: make map name consistent for Apollo and carla

    if sim.current_scene == map:
        sim.reset()
    else:
        # seed make sure the weather and NPC behvaiors deterministic
        sim.load(map, seed=0)

    return sim, BRIDGE_HOST, BRIDGE_PORT

def initialize_dv_and_ego(sim, map, model_id, start, destination, BRIDGE_HOST, BRIDGE_PORT, events_path):

    global accident_happen
    accident_happen = False

    def on_collision(agent1, agent2, contact):
        global accident_happen
        accident_happen = True
        name1 = "STATIC OBSTACLE" if agent1 is None else agent1.name
        name2 = "STATIC OBSTACLE" if agent2 is None else agent2.name
        print("{} collided with {} at {}".format(name1, name2, contact))

        agent1_info = [agent1.transform.position.x, agent1.transform.position.y, agent1.transform.position.z, agent1.state.velocity.x, agent1.state.velocity.y, agent1.state.velocity.z]

        if not agent2:
            agent2_name = 'static'
            agent2_uid = 0
            agent2_info = [0]*6
        else:
            agent2_name = agent2.name
            agent2_uid = agent2.uid
            agent2_info = [
            agent2.transform.position.x, agent2.transform.position.y, agent2.transform.position.z, agent2.state.velocity.x, agent2.state.velocity.y, agent2.state.velocity.z]

        contact_loc = [0, 0, 0]
        if contact:
            print('contact', contact.__dict__)
            print('type(contact)', type(contact))
            print('contact.x', contact.x)
            contact_loc = [contact.x, contact.y, contact.z]

        data_row = ['collision', agent2_name, agent2_uid] + agent1_info + agent2_info + contact_loc

        data_row = ','.join([str(data) for data in data_row])
        with open(events_path, 'a') as f_out:
            f_out.write(data_row+'\n')
        time.sleep(2)
        sim.stop()

    times = 0
    success = False
    while times < 3:
        try:
            state = lgsvl.AgentState()
            state.transform = start

            ego = sim.add_agent(model_id, lgsvl.AgentType.EGO, state)
            ego.connect_bridge(BRIDGE_HOST, BRIDGE_PORT)
            ego.on_collision(on_collision)

            # Dreamview setup
            dv = lgsvl.dreamview.Connection(sim, ego, BRIDGE_HOST)
            dv.set_hd_map(map)
            dv.set_vehicle('Lincoln2017MKZ_LGSVL')

            modules = get_modules_for_id(model_id)

            start = lgsvl.Transform(position=ego.transform.position, rotation=ego.transform.rotation)

            print('start', start)
            print('destination', destination)
            dv.setup_apollo(destination.position.x, destination.position.z, modules, default_timeout=60)
            print('finish setup_apollo')
            success = True
            break
        except Exception as e:
            print(e)
            print('fail to spin up apollo, try again!')
            times += 1
    if not success:
        raise Exception('fail to spin up apollo')

    return ego, dv

#######################################################################
# for step-wise simulation
def save_camera(ego, main_camera_folder, counter, i):
    import os
    for sensor in ego.get_sensors():
        if sensor.name == "Main Camera":
            rel_path = 'run_results_svl'+'/'+"main_camera_"+str(counter)+'_'+str(i)+".png"
            try:
                sensor.save(rel_path, compression=9)
            except:
                print('exception happens when saving camera image')

def save_measurement(ego, measurements_path):
    state = ego.state
    pos = state.position
    rot = state.rotation
    speed = state.speed * 3.6
    with open(measurements_path, 'a') as f_out:
        f_out.write(','.join([str(v) for v in [speed, pos.x, pos.y, pos.z, rot.x, rot.y, rot.z]])+'\n')

def get_bbox(agent):
    x_min = agent.bounding_box.min.x
    x_max = agent.bounding_box.max.x
    z_min = agent.bounding_box.min.z
    z_max = agent.bounding_box.max.z
    bbox = [
        agent.transform.position+lgsvl.Vector(x_min, 0, z_min),
        agent.transform.position+lgsvl.Vector(x_min, 0, z_max),
        agent.transform.position+lgsvl.Vector(x_max, 0, z_min),
        agent.transform.position+lgsvl.Vector(x_max, 0, z_max)
    ]

    return bbox

def gather_info(ego, other_agents, cur_values, deviations_path):
    ego_bbox = get_bbox(ego)
    # TBD: only using the front two vertices
    # ego_front_bbox = ego_bbox[:2]

    # hack: smaller than default values to make sure the files are created to have at least one line
    min_d = 9999
    d_angle_norm = 0.99
    for i, other_agent in enumerate(other_agents):

        # d_angle_norm_i = angle_from_center_view_fov(other_agent, ego, fov=90)
        # d_angle_norm = np.min([d_angle_norm, d_angle_norm_i])
        # if d_angle_norm_i == 0:
        # TBD: temporarily disable d_angle_norm since Apollo uses LiDAR
        d_angle_norm = 0
        other_bbox = get_bbox(other_agent)
        for other_b in other_bbox:
            for ego_b in ego_bbox:
                d = norm_2d(other_b, ego_b)
                min_d = np.min([min_d, d])

    if min_d < cur_values.min_d:
        cur_values.min_d = min_d
        with open(deviations_path, 'a') as f_out:
            f_out.write('min_d,'+str(cur_values.min_d)+'\n')

    if d_angle_norm < cur_values.d_angle_norm:
        cur_values.d_angle_norm = d_angle_norm
        with open(deviations_path, 'a') as f_out:
            f_out.write('d_angle_norm,'+str(cur_values.d_angle_norm)+'\n')

    # TBD: out-of-road violation related data

#######################################################################
# for continuous simulation
def bind_socket(socket, port_num):
    import subprocess
    while True:
        try:
            socket.bind("tcp://*:"+str(port_num))
            break
        except:
            subprocess.run("kill -9 $(lsof -t -i:" + str(port_num) + " -sTCP:LISTEN)", shell=True)

def receive_zmq(q, path_list, record_every_n_step):
    import zmq
    print('receive zmq')

    def get_d(x1, y1, z1, x2, y2, z2):
        return np.sqrt((x1-x2)**2+(y1-y2)**2)

    odometry_path, perception_obstacles_path, perception_obstacles_apollo_path, main_camera_folder, deviations_path = path_list

    print('start binding sockets')
    context = zmq.Context()
    socket_odometry = context.socket(zmq.PAIR)
    socket_perception_obstacles = context.socket(zmq.PAIR)
    socket_front_camera = context.socket(zmq.PAIR)
    socket_perception_obstacles_apollo = context.socket(zmq.PAIR)

    bind_socket(socket_odometry, 5561)
    bind_socket(socket_perception_obstacles, 5562)
    bind_socket(socket_front_camera, 5563)
    bind_socket(socket_perception_obstacles_apollo, 5564)
    print('finish binding sockets')

    min_ego_i_d = 100
    odometry_dict = {}
    perception_obstacles_dict = {}
    perception_obstacles_apollo_dict = {}
    with open(odometry_path, 'a') as f_out_odometry:
        with open(perception_obstacles_path, 'a') as f_out_perception_obstacles:
            with open(perception_obstacles_apollo_path, 'a') as f_out_perception_obstacles_apollo:
                while True:
                    data_str_odometry = None
                    data_str_perception_obstacles = None

                    try:
                        cmd = q.get(timeout=0.0001)
                        if cmd == 'end':
                            with open(deviations_path, 'a') as f_out:
                                for k in perception_obstacles_dict:
                                    if k in odometry_dict:
                                        ego_x, ego_y, ego_z = odometry_dict[k]
                                        npc_num = len(perception_obstacles_dict[k])
                                        for i in range(npc_num):
                                            i_x, i_y, i_z = perception_obstacles_dict[k][i]

                                            ego_i_d = get_d(ego_x, ego_y, ego_z, i_x, i_y, i_z)
                                            if ego_i_d < min_ego_i_d:
                                                min_ego_i_d = ego_i_d
                                                f_out.write('min_d,'+str(min_ego_i_d)+'\n')
                                                # print('time step', k, 'min_d', min_ego_i_d)
                            socket_odometry.close()
                            socket_perception_obstacles.close()
                            socket_front_camera.close()
                            socket_perception_obstacles_apollo.close()
                            context.term()
                            print('free sockets and context')
                            return
                    except:
                        pass

                    try:
                        data_str_odometry = socket_odometry.recv_string(flags=zmq.NOBLOCK)
                        f_out_odometry.write(data_str_odometry+'\n')

                        odometry_tokens = data_str_odometry.split(':')[1].split(',')
                        ego_x, ego_y, ego_z = [float(x) for x in odometry_tokens[2:]]
                        time_step = int(perception_obstacles_tokens[1])
                        odometry_dict[time_step] = (ego_x, ego_y, ego_z)
                    except Exception:
                        pass

                    try:
                        data_str_perception_obstacles = socket_perception_obstacles.recv_string(flags=zmq.NOBLOCK)
                        f_out_perception_obstacles.write(data_str_perception_obstacles+'\n')

                        perception_obstacles_tokens = data_str_perception_obstacles.split(':')[1].split(',')
                        npc_num = int(perception_obstacles_tokens[2])
                        if npc_num > 0:
                            time_step = int(perception_obstacles_tokens[1])
                            for i in range(npc_num):
                                i_x, i_y, i_z = [float(x) for x in perception_obstacles_tokens[3*(1+i):3*(2+i)]]
                                if time_step not in perception_obstacles_dict:
                                    perception_obstacles_dict[time_step] = [(i_x, i_y, i_z)]
                                else:
                                    perception_obstacles_dict[time_step].append((i_x, i_y, i_z))
                    except Exception:
                        pass

                    try:
                        data_str_front_camera = socket_front_camera.recv(flags=zmq.NOBLOCK)
                        timestamp_sec, sequence_num, front_image = data_str_front_camera.split(b':data_delimiter:')

                        # record image after warm-up stage
                        if int(sequence_num) > 150 and int(sequence_num) % record_every_n_step == 0:
                            img_path = os.path.join(main_camera_folder, sequence_num.decode()+'_'+timestamp_sec.decode()+'.jpg')
                            with open(img_path, 'wb') as f_out_front_camera:
                                f_out_front_camera.write(front_image)
                    except Exception:
                        pass

                    try:
                        data_str_perception_obstacles_apollo = socket_perception_obstacles_apollo.recv_string(flags=zmq.NOBLOCK)
                        f_out_perception_obstacles_apollo.write(data_str_perception_obstacles_apollo+'\n')

                        perception_obstacles_apollo_tokens = data_str_perception_obstacles_apollo.split(':')[1].split(',')
                        npc_num = int(perception_obstacles_apollo_tokens[2])
                        if npc_num > 0:
                            time_step = int(perception_obstacles_apollo_tokens[1])
                            for i in range(npc_num):
                                i_x, i_y, i_z = [float(x) for x in perception_obstacles_apollo_tokens[3*(1+i):3*(2+i)]]
                                if time_step not in perception_obstacles_apollo_dict:
                                    perception_obstacles_apollo_dict[time_step] = [(i_x, i_y, i_z)]
                                else:
                                    perception_obstacles_apollo_dict[time_step].append((i_x, i_y, i_z))
                    except Exception as e:
                        # print('repr(e)', repr(e))
                        pass

#######################################################################

def rotate(x, y, rot_rad):
    x_rot = x * np.cos(rot_rad) - y * np.sin(rot_rad)
    y_rot = x * np.sin(rot_rad) + y * np.cos(rot_rad)
    return x_rot, y_rot

def initialize_sim(map, sim_specific_arguments, arguments, customized_data, model_id, events_path, npc_events_path):

    sim, BRIDGE_HOST, BRIDGE_PORT = initialize_simulator(map, sim_specific_arguments)

    def npc_on_collision(agent1, agent2, contact):
        name1 = "STATIC OBSTACLE" if agent1 is None else agent1.name
        name2 = "STATIC OBSTACLE" if agent2 is None else agent2.name
        print("{} collided with {} at {}".format(name1, name2, contact))

        loc = agent1.transform.position
        ego_speed = np.linalg.norm([agent1.state.velocity.x, agent1.state.velocity.y, agent1.state.velocity.z])

        collision_type = 'npc collision'
        if name2 == model_id:
            collision_type = 'ego collision'

        data_row = [collision_type, ego_speed, name2, loc.x, loc.y]
        data_row = ','.join([str(data) for data in data_row])
        with open(npc_events_path, 'a') as f_out:
            f_out.write(data_row+'\n')

    if len(arguments.route_info["location_list"]) == 0:
        spawns = sim.get_spawn()
        start = spawns[0]
        destination = spawns[0].destinations[0]
    else:
        start, destination = arguments.route_info["location_list"]

        start = lgsvl.Transform(position=lgsvl.Vector(start[0], start[1], start[2]), rotation=lgsvl.Vector(start[3], start[4], start[5]))
        destination = lgsvl.Transform(position=lgsvl.Vector(destination[0], destination[1], destination[2]), rotation=lgsvl.Vector(destination[3], destination[4], destination[5]))

    try:
        sim.weather = lgsvl.WeatherState(rain=customized_data['rain'], fog=customized_data['fog'], wetness=customized_data['wetness'], cloudiness=customized_data['cloudiness'], damage=customized_data['damage'])

        from datetime import datetime
        dt = datetime(
              year=2020,
              month=12,
              day=25,
              hour=int(customized_data['hour']),
              minute = 0,
              second = 0
            )
        sim.set_date_time(dt, fixed=True)
    except:
        import traceback
        traceback.print_exc()

    ego, dv = initialize_dv_and_ego(sim, map, model_id, start, destination, BRIDGE_HOST, BRIDGE_PORT, events_path)

    # extra destination request to avoid previous request lost?
    dv.set_destination(destination.position.x, destination.position.z)
    # hack: delay several seconds to make sure the routing request succeeds
    delay_t = 5
    # print('before', delay_t)
    time.sleep(delay_t)
    # print('after', delay_t)

    middle_point = lgsvl.Transform(position=(destination.position + start.position) * 0.5, rotation=start.rotation)

    for k, v in customized_data['customized_center_transforms'].items():
        if v[0] == "absolute_location":
            middle_point_i = lgsvl.Transform(position=lgsvl.Vector(v[1], v[2], v[3]), rotation=lgsvl.Vector(v[4], v[5], v[6]))
            customized_data[k] = middle_point_i

    other_agents = []
    for static in customized_data['static_list']:
        state = lgsvl.ObjectState()
        state.transform.position = lgsvl.Vector(static.x,0,static.y)
        state.transform.rotation = lgsvl.Vector(0,0,0)
        state.velocity = lgsvl.Vector(0,0,0)
        state.angular_velocity = lgsvl.Vector(0,0,0)

        static_object = sim.controllable_add(static_types[static.model], state)

    # for i, vehicle in enumerate(customized_data['static_vehicles_list']):
    #     center_key_i = "static_vehicle_center_transform_"+str(i)
    #     if center_key_i in customized_data:
    #         middle_point_i = customized_data[center_key_i]
    #     else:
    #         middle_point_i = middle_point
    #     rot_rad = np.deg2rad(360 - middle_point_i.rotation.y)
    #
    #     vehicle_x, vehicle_y = rotate(vehicle.x, vehicle.y, rot_rad)
    #     vehicle_position_offset = lgsvl.Vector(vehicle_x, 0, vehicle_y)
    #     vehicle_rotation_offset = lgsvl.Vector(0, 0, 0)
    #
    #     vehicle_point = lgsvl.Transform(position=middle_point_i.position+vehicle_position_offset, rotation=middle_point_i.rotation+vehicle_rotation_offset)
    #
    #
    #
    #     state = lgsvl.AgentState()
    #     state.transform = vehicle_point
    #     p = sim.add_agent(vehicle_types[vehicle.model], lgsvl.AgentType.NPC, state)

    for i, ped in enumerate(customized_data['pedestrians_list']):
        center_key_i = "pedestrian_center_transform_"+str(i)
        if center_key_i in customized_data:
            middle_point_i = customized_data[center_key_i]
        else:
            middle_point_i = middle_point
        rot_rad = np.deg2rad(360 - middle_point_i.rotation.y)

        ped_x, ped_y = rotate(ped.x, ped.y, rot_rad)

        ped_position_offset = lgsvl.Vector(ped_x, 0, ped_y)
        ped_rotation_offset = lgsvl.Vector(0, 0, 0)

        ped_point = lgsvl.Transform(position=middle_point_i.position+ped_position_offset, rotation=middle_point_i.rotation+ped_rotation_offset)

        wps = [lgsvl.WalkWaypoint(position=ped_point.position, idle=ped.waypoints[0].idle, trigger_distance=ped.waypoints[0].trigger_distance, speed=ped.speed)]

        if ped.travel_distance > 0:
            print('ped.travel_distance', ped.travel_distance, 'ped.yaw', ped.yaw)
            wp_x = ped.travel_distance * np.sin(np.deg2rad(ped.yaw))
            wp_y = ped.travel_distance * np.cos(np.deg2rad(ped.yaw))
            # print('before', 'wp_x', wp_x, 'wp_y', wp_y)
            # do not consider the rotation of the middle point in this case for easiter visualization
            wp_x, wp_y = rotate(wp_x, wp_y, rot_rad)
            loc = middle_point_i.position+lgsvl.Vector(wp_x, 0, wp_y)

            j_next = np.min([1, len(ped.waypoints)-1])
            wp_next = ped.waypoints[j_next]

            wps.append(lgsvl.WalkWaypoint(position=loc, idle=wp_next.idle, trigger_distance=wp_next.trigger_distance, speed=ped.speed))
            # print('after', 'wp_x', wp_x, 'wp_y', wp_y)
            print('wps', wps[0].position, wps[0].idle, wps[0].trigger_distance, wps[0].speed, wps[1].position, wps[1].idle, wps[1].trigger_distance, wps[1].speed)
        else:
            for j, wp in enumerate(ped.waypoints):
                center_key_i_j = "pedestrian_"+str(i)+"_center_transform_"+str(j)
                if center_key_i_j in customized_data:
                    # print(center_key_i_j)
                    middle_point_i = customized_data[center_key_i_j]
                    rot_rad = np.deg2rad(360 - middle_point_i.rotation.y)

                j_next = np.min([j+1, len(ped.waypoints)-1])
                wp_next = ped.waypoints[j_next]

                wp_x, wp_y = rotate(wp.x, wp.y, rot_rad)
                loc = middle_point_i.position+lgsvl.Vector(wp_x, 0, wp_y)

                wps.append(lgsvl.WalkWaypoint(position=loc, idle=wp_next.idle, trigger_distance=wp_next.trigger_distance, speed=ped.speed))

        state = lgsvl.AgentState()
        state.transform = ped_point
        # print('\n'*3, 'ped.model', ped.model, '\n'*3)
        p = sim.add_agent(pedestrian_types[ped.model], lgsvl.AgentType.PEDESTRIAN, state)
        p.on_collision(npc_on_collision)

        if ped.travel_distance > 0:
            p.follow(wps, False)
        else:
            p.follow(wps, False, waypoints_path_type='BezierSpline')
        other_agents.append(p)

    for i, vehicle in enumerate(customized_data['vehicles_list']):
        center_key_i = "vehicle_center_transform_"+str(i)
        if center_key_i in customized_data:
            middle_point_i = customized_data[center_key_i]
        else:
            middle_point_i = middle_point
        rot_rad = np.deg2rad(360 - middle_point_i.rotation.y)

        vehicle_x, vehicle_y = rotate(vehicle.x, vehicle.y, rot_rad)
        vehicle_position_offset = lgsvl.Vector(vehicle_x, 0, vehicle_y)
        vehicle_rotation_offset = lgsvl.Vector(0, 0, 0)

        vehicle_point = lgsvl.Transform(position=middle_point_i.position+vehicle_position_offset, rotation=middle_point_i.rotation+vehicle_rotation_offset)

        wps = [lgsvl.DriveWaypoint(position=vehicle_point.position, speed=vehicle.speed, acceleration=0, angle=vehicle_point.rotation, idle=vehicle.waypoints[0].idle, deactivate=False, trigger_distance=vehicle.waypoints[0].trigger_distance)]

        for j, wp in enumerate(vehicle.waypoints):
            center_key_i_j = "vehicle_"+str(i)+"_center_transform_"+str(j)
            if center_key_i_j in customized_data:
                # print(center_key_i_j)
                middle_point_i = customized_data[center_key_i_j]
                rot_rad = np.deg2rad(360 - middle_point_i.rotation.y)

            j_next = np.min([j+1, len(vehicle.waypoints)-1])
            wp_next = vehicle.waypoints[j_next]

            wp_x, wp_y = rotate(wp.x, wp.y, rot_rad)
            pos = middle_point_i.position + lgsvl.Vector(wp_x, 0, wp_y)

            # to avoid vehicle going underground
            # pos.y += 0.3

            wps.append(lgsvl.DriveWaypoint(position=pos, speed=vehicle.speed, acceleration=0, angle=middle_point_i.rotation, idle=wp_next.idle, deactivate=False, trigger_distance=wp_next.trigger_distance))

        state = lgsvl.AgentState()
        state.transform = vehicle_point
        # print('\n'*3, 'vehicle.model', vehicle.model, '\n'*3)
        p = sim.add_agent(vehicle_types[vehicle.model], lgsvl.AgentType.NPC, state)
        p.on_collision(npc_on_collision)
        p.follow(wps, False, waypoints_path_type='BezierSpline')
        other_agents.append(p)

    controllables = sim.get_controllables()
    for i in range(len(controllables)):
        signal = controllables[i]
        if signal.type == "signal":
            control_policy = signal.control_policy
            control_policy = "trigger=500;green=20;yellow=2;red=5;loop"
            signal.control(control_policy)

    with open(npc_events_path, 'a') as f_out:
        f_out.write('npc_agents_uids,'+','.join(list(sim.agents.keys()))+'\n')

    return sim, ego, destination

def run_sim_with_initialization(q, duration, time_scale, map, sim_specific_arguments, arguments, customized_data, model_id):

    events_path = os.path.join(arguments.deviations_folder, "events.txt")
    npc_events_path = os.path.join(arguments.deviations_folder, "npc_events.txt")

    sim, ego, destination = initialize_sim(map, sim_specific_arguments, arguments, customized_data, model_id, events_path, npc_events_path)
    print('start run sim')
    sim.run(time_limit=duration, time_scale=time_scale)

    d_to_dest = norm_2d(ego.transform.position, destination.position)
    if d_to_dest > 10:
        with open(events_path, 'a') as f_out:
            f_out.write('fail_to_finish,'+str(d_to_dest))
    print('ego car final transform:', ego.transform, 'destination', destination, 'd_to_dest', d_to_dest)
    time.sleep(1)
    q.put('end')
    return

def start_simulation(customized_data, arguments, sim_specific_arguments, launch_server, episode_max_time):

    deviations_path = os.path.join(arguments.deviations_folder, 'deviations.txt')
    main_camera_folder = os.path.join(arguments.deviations_folder, 'main_camera_data')

    if not os.path.exists(main_camera_folder):
        os.mkdir(main_camera_folder)

    model_id = arguments.model_id
    map = arguments.route_info["town_name"]
    counter = arguments.counter

    duration = episode_max_time
    time_scale = 1
    continuous = True
    if continuous == True:
        odometry_path = os.path.join(arguments.deviations_folder, 'odometry.txt')
        perception_obstacles_path = os.path.join(arguments.deviations_folder, 'perception_obstacles.txt')
        perception_obstacles_apollo_path = os.path.join(arguments.deviations_folder, 'perception_obstacles_apollo.txt')
        path_list = [odometry_path, perception_obstacles_path, perception_obstacles_apollo_path, main_camera_folder, deviations_path]

        from multiprocessing import Process, Queue
        q = Queue()
        p = Process(target=run_sim_with_initialization, args=(q, duration, time_scale, map, sim_specific_arguments, arguments, customized_data, model_id))
        p.daemon = True
        p.start()
        receive_zmq(q, path_list, arguments.record_every_n_step)
        p.terminate()

    # else:
    #     measurements_path = os.path.join(arguments.deviations_folder, 'measurements.txt')
    #
    #     sim = initialize_sim()
    #     step_time = 30
    #     step_rate = 1.0 / step_time
    #     steps = int(duration * step_rate)
    #
    #     cur_values = emptyobject(min_d=10000, d_angle_norm=1)
    #
    #     for i in range(steps):
    #         sim.run(time_limit=step_time, time_scale=1)
    #         if i % arguments.record_every_n_step == 0:
    #             save_camera(ego, main_camera_folder, counter, i)
    #
    #         save_measurement(ego, measurements_path)
    #         gather_info(ego, other_agents, cur_values, deviations_path)
    #
    #         d_to_dest = norm_2d(ego.transform.position, destination.position)
    #         # print('d_to_dest', d_to_dest)
    #         if d_to_dest < 5:
    #             print('ego car reachs destination successfully')
    #             break
    #
    #         if accident_happen:
    #             break

if __name__ == '__main__':
    atexit.register(exit_handler)
    map = "Borregas Ave"
    config = [4, 4, 2, 3, 10, 50]
    run_svl_simulation(map, config)


import cv2
import numpy as np
import copy
import torch
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
import time
import autopy
import mediapipe as mp
import math
class Cnn(nn.Module):
    def __init__(self):
        super(Cnn, self).__init__()
        # convolutional layer
        self.conv1 = nn.Conv2d(1, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 3 * 3, 256)
        self.fc2 = nn.Linear(256, 26)
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        # Adding sequence of convolutional and max pooling layers
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 64 * 3 * 3)
        x = self.dropout(x)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)

        return x

def get_histogram(frame):
    roi1 = frame[rect1_tl[1]:rect1_tl[1] + width, rect1_tl[0]:rect1_tl[0] + height]
    roi2 = frame[rect2_tl[1]:rect2_tl[1] + width, rect2_tl[0]:rect2_tl[0] + height]
    roi3 = frame[rect3_tl[1]:rect3_tl[1] + width, rect3_tl[0]:rect3_tl[0] + height]
    roi4 = frame[rect4_tl[1]:rect4_tl[1] + width, rect4_tl[0]:rect4_tl[0] + height]
    roi5 = frame[rect5_tl[1]:rect5_tl[1] + width, rect5_tl[0]:rect5_tl[0] + height]
    roi = np.concatenate((roi1, roi2, roi3, roi4, roi5), axis=0)
    roi_hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)

    return cv2.calcHist([roi_hsv], [0, 1], None, [180, 256], [0, 180, 0, 256])

def get_ROI(canvas):
    gray = cv2.bitwise_not(canvas)
    ret, thresh = cv2.threshold(gray, 90, 255, cv2.THRESH_BINARY_INV)
    ctrs, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    areas = []
    for i in range(len(ctrs)):
        x, y, w, h = cv2.boundingRect(ctrs[i])
        areas.append((w * h, i))

    def sort_second(val):
        return val[0]

    areas.sort(key=sort_second, reverse=True)
    x, y, w, h = cv2.boundingRect(ctrs[areas[1][1]])
    cv2.rectangle(canvas, (x, y), (x + w, y + h), (255, 255, 0), 1)
    roi = gray[y:y + h, x:x + w]
    return roi

def character_prediction(roi, model):
    """Predicts character written with image processing"""
    img = cv2.resize(roi, (28, 28))
    img = cv2.GaussianBlur(img, (3, 3), 0)
    img = Image.fromarray(img)

    normalize = transforms.Normalize(
        mean=[0.5, 0.5, 0.5],
        std=[0.5, 0.5, 0.5]
    )
    preprocess = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))
    ])

    p_img = preprocess(img)

    model.eval()
    p_img = p_img.reshape([1, 1, 28, 28]).float()
    output = model(torch.transpose(p_img, 2, 3))
    _, preds_tensor = torch.max(output, 1)
    preds = np.squeeze(preds_tensor.numpy())
    return preds
class handDetector():
    def __init__(self, mode=False, maxHands=1, detectionCon=0.8, trackCon=0.8):
        self.mode = mode
        self.maxHands = maxHands
        self.detectionCon = detectionCon
        self.trackCon = trackCon

        self.mpHands = mp.solutions.hands
        self.hands = self.mpHands.Hands(self.mode, self.maxHands,
                                        self.detectionCon, self.trackCon)
        self.mpDraw = mp.solutions.drawing_utils
        self.tipIds = [4, 8, 12, 16, 20]

    def findHands(self, img, draw=True):
        imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        self.results = self.hands.process(imgRGB)
        # print(results.multi_hand_landmarks)

        if self.results.multi_hand_landmarks:
            for handLms in self.results.multi_hand_landmarks:
                if draw:
                    self.mpDraw.draw_landmarks(img, handLms,
                                               self.mpHands.HAND_CONNECTIONS)

        return img

    def findPosition(self, img, handNo=0, draw=True):
        xList = []
        yList = []
        bbox = []
        self.lmList = []
        if self.results.multi_hand_landmarks:
            myHand = self.results.multi_hand_landmarks[handNo]
            for id, lm in enumerate(myHand.landmark):
                # print(id, lm)
                h, w, c = img.shape
                cx, cy = int(lm.x * w), int(lm.y * h)
                xList.append(cx)
                yList.append(cy)
                # print(id, cx, cy)
                self.lmList.append([id, cx, cy])
                if draw:
                    cv2.circle(img, (cx, cy), 5, (255, 0, 255), cv2.FILLED)

            xmin, xmax = min(xList), max(xList)
            ymin, ymax = min(yList), max(yList)
            bbox = xmin, ymin, xmax, ymax

            if draw:
                cv2.rectangle(img, (xmin - 20, ymin - 20), (xmax + 20, ymax + 20),
                              (0, 255, 0), 2)

        return self.lmList, bbox

    def fingersUp(self):
        fingers = []
        # Thumb
        if self.lmList[self.tipIds[0]][1] > self.lmList[self.tipIds[0] - 1][1]:
            fingers.append(1)
        else:
            fingers.append(0)

        # Fingers
        for id in range(1, 5):

            if self.lmList[self.tipIds[id]][2] < self.lmList[self.tipIds[id] - 2][2]:
                fingers.append(1)
            else:
                fingers.append(0)

        # totalFingers = fingers.count(1)

        return fingers

    def findDistance(self, p1, p2, img, draw=True,r=15, t=3):
        x1, y1 = self.lmList[p1][1:]
        x2, y2 = self.lmList[p2][1:]
        cx, cy = (x1 + x2) // 2, (y1 + y2) // 2

        if draw:
            cv2.line(img, (x1, y1), (x2, y2), (255, 0, 255), t)
            cv2.circle(img, (x1, y1), r, (255, 0, 255), cv2.FILLED)
            cv2.circle(img, (x2, y2), r, (255, 0, 255), cv2.FILLED)
            cv2.circle(img, (cx, cy), r, (0, 0, 255), cv2.FILLED)
        length = math.hypot(x2 - x1, y2 - y1)

        return length, img, [x1, y1, x2, y2, cx, cy]

def main():
    pTime = 0
    cTime = 0
    cap = cv2.VideoCapture(0)
    detector = handDetector()
    canvas = np.zeros((720, 1280), np.uint8)
    far_points = []
    pressed = False
    is_drawing = False
    made_prediction = False
    # Creating the model
    model = Cnn()
    model.load_state_dict(torch.load('model_emnist.pt', map_location='cpu'))
    while True:
        _, frame = cap.read()
        original_frame = copy.deepcopy(frame)
        success, img = cap.read()
        img = detector.findHands(img)
        lmList, bbox = detector.findPosition(img)
        if len(lmList) != 0:
            print(lmList[8][1:])
            #far = lmList[8][1:]
            far = (lmList[8][1],lmList[8][2])

        cTime = time.time()
        fps = 1 / (cTime - pTime)
        pTime = cTime

        cv2.putText(img, str(int(fps)), (10, 70), cv2.FONT_HERSHEY_PLAIN, 3,
                    (255, 0, 255), 3)

        canvas[:, :] = 255

        key = cv2.waitKey(1)

        # To start drawing
        if key & 0xFF == ord('d'):
            is_drawing = True

        # To clear drawing
        if key & 0xFF == ord('c'):
            canvas[:, :] = 255
            is_drawing = False
            far_points.clear()
            made_prediction = False

        if is_drawing:
            if len(far_points) > 100:
                far_points.pop(0)
            far_points.append(far)
            #print(far)
            for i in range(len(far_points) - 1):
                cv2.line(original_frame, far_points[i], far_points[i + 1], (255, 5, 255), 20)
                cv2.line(canvas, far_points[i], far_points[i + 1], (0, 0, 0), 20)

        # To predict the character drawn
        if key & 0xFF == ord('p'):
            is_drawing = False
            roi = get_ROI(canvas)
            print(roi)
            prediction = character_prediction(roi, model)
            print(prediction)
            made_prediction = True
            name = str(prediction) + '.jpg'
            cv2.imwrite(name, roi)

        if pressed:
            mask = get_mask(frame, histogram)
            max_contour = get_max_contour(mask)
            hull = cv2.convexHull(max_contour, returnPoints=False)
            draw_defects(original_frame, max_contour, hull)
            defects = cv2.convexityDefects(max_contour, hull)
            far = get_farthest_point(defects, max_contour, get_centroid(max_contour))
            cv2.circle(original_frame, far, 10, [0, 200, 255], -1)

        if made_prediction:
            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(original_frame,"Character Written: " + chr(prediction + 65), (8, 250), font, 1,
                        (255, 255, 255), 2, cv2.LINE_AA)
            canvas[:, :] = 255
            is_drawing = False
            far_points.clear()

        # To quit the drawing
        if key & 0xFF == ord('q'):
            break

        cv2.imshow("frame", original_frame)
        cv2.waitKey(1)

    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()


import re
from os.path import exists, join as join_path

import numpy as np

from constants.app_constants import DATA_SCP_FILE, MFCC_DIR, VAD_DIR, FEATS_SCP_FILE, UTT2NUM_FRAMES_FILE, TMP_DIR, \
    VAD_SCP_FILE
from services.common import load_array, run_parallel, run_command
from services.kaldi import Kaldi, spaced_file_to_dict

class MFCC:
    def __init__(self, fs=8000, fl=100, fh=4000, frame_len_ms=25, n_jobs=20, n_ceps=20, save_loc='../save'):
        mfcc_loc = join_path(save_loc, MFCC_DIR)
        params_file = join_path(mfcc_loc, 'mfcc.params')
        config_file = join_path(mfcc_loc, 'mfcc.conf')

        with open(params_file, 'w') as f:
            f.write('nj={}\n'.format(n_jobs))
            f.write('compress={}\n'.format('true'))
            f.write('mfcc_loc={}\n'.format(mfcc_loc))
            f.write('mfcc_config={}\n'.format(config_file))

        with open(config_file, 'w') as f:
            f.write('--sample-frequency={}\n'.format(fs))
            f.write('--low-freq={}\n'.format(fl))
            f.write('--high-freq={}\n'.format(fh))
            f.write('--frame-length={}\n'.format(frame_len_ms))
            f.write('--num-ceps={}\n'.format(n_ceps))
            f.write('--snip-edges={}\n'.format('false'))

        self.mfcc_loc = mfcc_loc
        self.save_loc = save_loc
        self.params_file = params_file
        self.n_ceps = n_ceps
        self.n_jobs = n_jobs

    def extract(self, data_loc, split):
        return Kaldi().run_command('sh ./kaldi/make_mfcc.sh {} {} {}'.format(data_loc, split, self.params_file), print_error=True)

    def extract_with_vad_and_normalization(self, data_loc, split, threshold=5.5, mean_scale=0.5, cmvn_window=300,
                                           var_norm=False):
        vad_loc = join_path(self.save_loc, VAD_DIR)
        tmp_loc = join_path(self.save_loc, TMP_DIR)

        feats_scp = join_path(data_loc, FEATS_SCP_FILE)
        vad_scp = join_path(data_loc, VAD_SCP_FILE)

        print('MFCC: Extracting features...')
        self.extract(data_loc, split)

        print('MFCC: Computing VAD...')
        vad = VAD(threshold, mean_scale, n_jobs=self.n_jobs, save_loc=self.save_loc)
        vad.compute(data_loc, split)

        print('MFCC: Normalizing features and selecting voiced frames..')
        feats_scp_dict = spaced_file_to_dict(feats_scp)
        vad_scp_dict = spaced_file_to_dict(vad_scp)

        splits = np.array_split(list(feats_scp_dict.keys()), self.n_jobs)
        for i in range(self.n_jobs):
            split_feat_scp = open(join_path(tmp_loc, 'feats.{}.scp'.format(i + 1)), 'w')
            split_vad_scp = open(join_path(tmp_loc, 'vad.{}.scp'.format(i + 1)), 'w')
            for key in splits[i]:
                split_feat_scp.write('{} {}\n'.format(key, feats_scp_dict[key]))
                split_vad_scp.write('{} {}\n'.format(key, vad_scp_dict[key]))
            split_feat_scp.close()
            split_vad_scp.close()

        Kaldi().queue('JOB=1:{nj} {mfcc_loc}/log/voiced_feats.JOB.log '
                      'apply-cmvn-sliding --norm-vars={var_norm} --center=true --cmn-window={window} scp:{tmp_loc}/feats.JOB.scp ark:- \| '
                      'select-voiced-frames ark:- scp,ns,cs:{tmp_loc}/vad.JOB.scp ark:- \| '
                      'copy-feats --compress=true --write-num-frames=ark,t:{mfcc_loc}/log/utt2num_frames.{name}.JOB ark:- '
                      'ark,scp:{mfcc_loc}/voiced_feats.{name}.JOB.ark,{mfcc_loc}/voiced_feats.{name}.JOB.scp || exit 1;'
                      .format(mfcc_loc=self.mfcc_loc, tmp_loc=tmp_loc, vad_loc=vad_loc,
                              var_norm='true' if var_norm else 'false',
                              nj=self.n_jobs, window=cmvn_window, name=split))

        run_command('for n in $(seq {nj}); do \n'
                    '   cat {mfcc_loc}/voiced_feats.{name}.$n.scp || exit 1;\n'
                    'done > {data_loc}/voiced_feats.scp || exit 1'.format(mfcc_loc=self.mfcc_loc, data_loc=data_loc,
                                                                          nj=self.n_jobs, name=split))

        run_command('for n in $(seq {nj}); do \n'
                    '   cat {mfcc_loc}/log/utt2num_frames.{name}.$n || exit 1;\n'
                    'done > {data_loc}/utt2num_frames || exit 1'.format(mfcc_loc=self.mfcc_loc, data_loc=data_loc,
                                                                        nj=self.n_jobs, name=split))

class VAD:
    def __init__(self, threshold=5.5, mean_scale=0.5, n_jobs=20, save_loc='../save'):
        vad_loc = join_path(save_loc, VAD_DIR)
        params_file = join_path(vad_loc, 'vad.params')
        config_file = join_path(vad_loc, 'vad.conf')

        with open(params_file, 'w') as f:
            f.write('nj={}\n'.format(n_jobs))
            f.write('vad_loc={}\n'.format(vad_loc))
            f.write('vad_config={}\n'.format(config_file))

        with open(config_file, 'w') as f:
            f.write('--vad-energy-threshold={}\n'.format(threshold))
            f.write('--vad-energy-mean-scale={}\n'.format(mean_scale))

        self.params_file = params_file

    def compute(self, data_loc, split):
        return Kaldi().run_command('sh ./kaldi/compute_vad.sh {} {} {}'.format(data_loc, split, self.params_file))

def add_frames_to_args(args_list, frame_dict):
    frames = []
    for key in args_list[:, 0]:
        frames.append(frame_dict[key])
    return np.vstack([args_list.T, frames]).T

def generate_data_scp(save_loc, args_list, append=False):
    data_scp_file = join_path(save_loc, DATA_SCP_FILE)
    with open(data_scp_file, 'a' if append else 'w') as f:
        for args in args_list:
            f.write('{} {} |\n'.format(args[0], args[4]))

def get_frame(file_loc):
    return load_array(file_loc).shape[1]

def get_mfcc_frames(save_loc, args):
    utt2num_frames = join_path(save_loc, UTT2NUM_FRAMES_FILE)
    utt2num_frames_dict = spaced_file_to_dict(utt2num_frames)
    frames = []
    count = 0
    for a in args:
        try:
            frames.append(utt2num_frames_dict[a])
        except KeyError:
            count = count + 1
    if count > 0:
        print('MFCC FRAMES: Can not find {} utterances in utt2num_frames'.format(count))
    return np.array(frames).reshape([-1, 1])

def load_feature(file_name):
    return load_array(file_name)

def remove_bad_files(args_list, save_loc='../save'):
    feats_scp = join_path(save_loc, FEATS_SCP_FILE)
    feats_scp_dict = spaced_file_to_dict(feats_scp)

    bad_files = []
    for i, key in enumerate(args_list[:, 0]):
        try:
            _ = feats_scp_dict[key]
        except KeyError:
            bad_files.append(i)
    return np.delete(args_list, bad_files, axis=0)

def remove_present_from_scp(save_loc, n_jobs=10):
    data_scp_file = join_path(save_loc, DATA_SCP_FILE)
    mfcc_loc = join_path(save_loc, MFCC_DIR)
    file_list = []
    index_list = []
    scp_list = []
    with open(data_scp_file, 'r') as f:
        for line in f.readlines():
            tokens = re.split('[\s]+', line.strip())
            file_list.append('{}/{}.npy'.format(mfcc_loc, tokens[0]))
            index_list.append(tokens[0])
            scp_list.append(line)
    absent = np.invert(run_parallel(exists, file_list, n_jobs, p_bar=False), dtype=bool)
    scp_list = np.array(scp_list)[absent]
    with open(data_scp_file, 'w') as f:
        f.writelines(scp_list)
    return sum(absent)


# -*- coding: utf-8 -*-
"""
Created on Wed Dec  9 17:21:24 2020

@author: dasssvis
"""

import numpy as np
from tensorflow.keras.preprocessing.sequence import pad_sequences

import pickle

tokenizer=pickle.load(open('tokenizer.pkl','rb+'))
from tensorflow.keras.models import load_model

model=load_model('model.h5')

def inp(name):
  q=tokenizer.texts_to_sequences(name)
  q=[i[0] for i in q]
  q1=[]
  q1.append(q)
  s=pad_sequences(q1,maxlen=15,padding='post')
  return s

def decoder(ans):
    decode=np.argmax(ans,axis=1)[0]
    return decode

def gender(name):
    sequences=inp(name)
    predict=model.predict(sequences)
    result=decoder(predict)
    if result==1:
        print(name," is a male")
    if result==0:
        print(name," is a female")
    return result


"""
An Example of Diversity Is All You Need.
https://arxiv.org/abs/1802.06070

"""

import os
import gym
import torch
import argparse
import numpy as np
import pybullet_envs

import torch
import torch.nn as nn
import torch.nn.functional as F

from machina import logger
from machina.traj import Traj
from machina.envs import SkillEnv
from machina.utils import measure
from machina.algos import diayn_sac, diayn
from machina.pols import GaussianPol
from machina.samplers import EpiSampler
from machina.utils import set_device, measure
from machina.traj import epi_functional as ef
from machina.vfuncs import DeterministicSAVfunc, DeterministicSVfunc
from simple_net import PolNet, QNet, DiaynDiscrimNet

parser = argparse.ArgumentParser()
parser.add_argument('--log', type=str, default='garbage',
                    help='Directory name of log.')
parser.add_argument('--env_name', type=str,
                    default='ReacherBulletEnv-v0', help='Name of environment.')
parser.add_argument('--cuda', type=int, default=-1, help='cuda device number.')
parser.add_argument('--seed', type=int, default=256)
parser.add_argument('--max_episodes', type=int, default=1e8,
                    help='Number of episodes to run.')
parser.add_argument('--num_parallel', type=int, default=4,
                    help='Number of processes to sample.')
parser.add_argument('--max_steps_per_iter', type=int, default=2000,
                    help='Number of steps to use in an iteration.')
parser.add_argument('--batch_size', type=int, default=256)
parser.add_argument('--discrim_batch_size', type=int, default=32)
parser.add_argument('--discrim_h_size', type=int,
                    default=100, help='Hidden size of Discriminator.')
parser.add_argument('--pol_lr', type=float, default=1e-4,
                    help='Policy learning rate.')
parser.add_argument('--qf_lr', type=float, default=3e-4,
                    help='Q function learning rate.')
parser.add_argument('--discrim_lr', type=float,
                    default=0.001, help='Discriminator learning rate.')
parser.add_argument('--discrim_momentum', type=float,
                    default=0.9, help='Discriminator momentum.')
parser.add_argument('--epoch_per_iter', type=int,
                    default=100, help='Number of epoch in an iteration.')
parser.add_argument('--gamma', type=float,
                    default=0.99, help='Discount factor.')
parser.add_argument('--tau', type=float, default=5e-3,
                    help='Coefficient of target function.')
parser.add_argument('--sampling', type=int, default=1,
                    help='Number of sampling in calculation of expectation.')
parser.add_argument('--no_reparam', action='store_true', default=False)
parser.add_argument('--steps_per_save_models', type=float,
                    default=2e5, help='Frequency of saving model per steps.')
parser.add_argument('--num_skill', type=int, default=4,
                    help='Number of skills.')
args = parser.parse_args()

'''
Feature extractor of the discriminator.
please see the paper 4.4.2 question 7.
e.g.:
def discrim_f(x): return x
f_dim = 9  # dimension of feature space
'''

def discrim_f(x): return x[:, 0:2]+x[:, 2:4]

f_dim = 2

env = gym.make(args.env_name)
env = SkillEnv(env, num_skill=args.num_skill)
obs = env.reset()
observation_space = env.real_observation_space
skill_space = env.skill_space
ob_skill_space = env.observation_space
action_space = env.action_space
ob_dim = ob_skill_space.shape[0] - args.num_skill
device_name = 'cpu' if args.cuda < 0 else "cuda:{}".format(args.cuda)
device = torch.device(device_name)
set_device(device)

# policy
pol_net = PolNet(ob_skill_space, action_space)
pol = GaussianPol(ob_skill_space, action_space, pol_net)

# q-function
qf_net1 = QNet(ob_skill_space, action_space)
qf1 = DeterministicSAVfunc(ob_skill_space, action_space, qf_net1)
targ_qf_net1 = QNet(ob_skill_space, action_space)
targ_qf_net1.load_state_dict(qf_net1.state_dict())
targ_qf1 = DeterministicSAVfunc(ob_skill_space, action_space, targ_qf_net1)
qf_net2 = QNet(ob_skill_space, action_space)
qf2 = DeterministicSAVfunc(ob_skill_space, action_space, qf_net2)
targ_qf_net2 = QNet(ob_skill_space, action_space)
targ_qf_net2.load_state_dict(qf_net2.state_dict())
targ_qf2 = DeterministicSAVfunc(ob_skill_space, action_space, targ_qf_net2)
qfs = [qf1, qf2]
targ_qfs = [targ_qf1, targ_qf2]

log_alpha = nn.Parameter(torch.ones((), device=device))

high = np.array([np.finfo(np.float32).max]*f_dim)
f_space = gym.spaces.Box(-high, high, dtype=np.float32)
discrim_net = DiaynDiscrimNet(
    f_space, skill_space, h_size=args.discrim_h_size, discrim_f=discrim_f).to(device)

discrim = DeterministicSVfunc(f_space, discrim_net, rnn=False)

# set optimizer to both models
optim_pol = torch.optim.Adam(pol_net.parameters(), args.pol_lr)
optim_qf1 = torch.optim.Adam(qf_net1.parameters(), args.qf_lr)
optim_qf2 = torch.optim.Adam(qf_net2.parameters(), args.qf_lr)
optim_qfs = [optim_qf1, optim_qf2]
optim_alpha = torch.optim.Adam([log_alpha], args.pol_lr)
optim_discrim = torch.optim.SGD(discrim.parameters(
), lr=args.discrim_lr, momentum=args.discrim_momentum)

off_traj = Traj()
sampler = EpiSampler(
    env, pol, num_parallel=args.num_parallel, seed=args.seed)

if not os.path.exists(args.log):
    os.makedirs(args.log)
    os.makedirs(args.log+'/models')
score_file = os.path.join(args.log, 'progress.csv')
logger.add_tabular_output(score_file)
logger.add_tensorboard_output(args.log)
np.random.seed(args.seed)
torch.manual_seed(args.seed)

# counter and record for loop
total_epi = 0
total_step = 0
mean_rew = 0

discrim_loss = []
discrim_total_step = []

# train loop
while args.max_episodes > total_epi:
    print('totalepi:', total_epi, 'mean_rew', mean_rew)
    # sample trajectories
    with measure('sample'):
        epis = sampler.sample(pol, max_steps=args.max_steps_per_iter)

    with measure('train'):
        on_traj = Traj()
        on_traj.add_epis(epis)

        on_traj = ef.add_next_obs(on_traj)
        on_traj = ef.compute_diayn_rews(
            on_traj, lambda x: diayn_sac.calc_rewards(x, args.num_skill, discrim))
        on_traj.register_epis()
        off_traj.add_traj(on_traj)

        total_epi += on_traj.num_epi
        step = on_traj.num_step
        total_step += step
        log_alpha = nn.Parameter(
            np.log(0.1)*torch.ones((), device=device))  # fix alpha

        result_dict = diayn_sac.train(
            off_traj,
            pol, qfs, targ_qfs, log_alpha,
            optim_pol, optim_qfs, optim_alpha,
            step, args.batch_size,
            args.tau, args.gamma, args.sampling,
            discrim, args.num_skill,
            not args.no_reparam)
        discrim_losses = diayn.train(discrim, optim_discrim, on_traj,
                                     args.discrim_batch_size, args.epoch_per_iter,
                                     args.num_skill)
    # update counter and record
    rewards = [np.sum(epi['rews']) for epi in epis]
    result_dict['discrimloss'] = discrim_losses
    mean_rew = np.mean(rewards)
    logger.record_results(args.log, result_dict, score_file,
                          total_epi, step, total_step,
                          rewards,
                          plot_title=args.env_name)

    # save models regular intervals
    steps_as = str(int(
        int(total_step / args.steps_per_save_models + 1) * args.steps_per_save_models))
    if 'prev_as' in locals():
        if not prev_as == steps_as:
            torch.save(pol.state_dict(), os.path.join(
                args.log, 'models', 'pol_'+steps_as+'.pkl'))
            torch.save(qf1.state_dict(), os.path.join(
                args.log, 'models', 'qf1_'+steps_as+'.pkl'))
            torch.save(qf2.state_dict(), os.path.join(
                args.log, 'models', 'qf2_'+steps_as+'.pkl'))
            torch.save(discrim.state_dict(), os.path.join(
                args.log, 'models', 'discrim_'+steps_as+'.pkl'))
            torch.save(optim_pol.state_dict(), os.path.join(
                args.log, 'models', 'optim_pol_'+steps_as+'.pkl'))
            torch.save(optim_qf1.state_dict(), os.path.join(
                args.log, 'models', 'optim_qf1_'+steps_as+'.pkl'))
            torch.save(optim_qf2.state_dict(), os.path.join(
                args.log, 'models', 'optim_qf2_'+steps_as+'.pkl'))
            torch.save(optim_discrim.state_dict(), os.path.join(
                args.log, 'models', 'optim_discrim_'+steps_as+'.pkl'))
    prev_as = str(int(
        int(total_step / args.steps_per_save_models + 1) * args.steps_per_save_models))
    del on_traj
del sampler


import sys
import numpy as np
import cv2
from scipy.spatial.transform import Rotation
import imutils
import itertools

import torchvision.models as models
import torch
from torchvision import transforms
from torchvision.models.detection.anchor_utils import DefaultBoxGenerator
from rrc_example_package import rearrange_dice_env
import trifinger_simulation.tasks.rearrange_dice as task
from trifinger_object_tracking.py_lightblue_segmenter import segment_image
import torch.nn.functional as F
import warnings
from torchvision.models.detection import _utils as det_utils
from torchvision.models.detection.backbone_utils import _validate_trainable_layers
from torchvision.models.detection.transform import GeneralizedRCNNTransform
from torchvision.models import vgg
from torchvision.ops import boxes as box_ops
from collections import OrderedDict
from torch import nn, Tensor
from typing import Any, Dict, List, Optional, Tuple
from torchvision.models.utils import load_state_dict_from_url

__all__ = ['SSD', 'ssd300_vgg16']

model_urls = {
    'ssd300_vgg16_coco': 'https://download.pytorch.org/models/ssd300_vgg16_coco-b556d3b4.pth',
}

backbone_urls = {
    # We port the features of a VGG16 backbone trained by amdegroot because unlike the one on TorchVision, it uses the
    # same input standardization method as the paper. Ref: https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth
    'vgg16_features': 'https://download.pytorch.org/models/vgg16_features-amdegroot.pth'
}

def _xavier_init(conv: nn.Module):
    for layer in conv.modules():
        if isinstance(layer, nn.Conv2d):
            torch.nn.init.xavier_uniform_(layer.weight)
            if layer.bias is not None:
                torch.nn.init.constant_(layer.bias, 0.0)

class SSDHead(nn.Module):
    def __init__(self, in_channels: List[int], num_anchors: List[int], num_classes: int):
        super().__init__()
        self.classification_head = SSDClassificationHead(in_channels, num_anchors, num_classes)
        self.regression_head = SSDRegressionHead(in_channels, num_anchors)

    def forward(self, x: List[Tensor]) -> Dict[str, Tensor]:
        return {
            'bbox_regression': self.regression_head(x),
            'cls_logits': self.classification_head(x),
        }

class SSDScoringHead(nn.Module):
    def __init__(self, module_list: nn.ModuleList, num_columns: int):
        super().__init__()
        self.module_list = module_list
        self.num_columns = num_columns

    def _get_result_from_module_list(self, x: Tensor, idx: int) -> Tensor:
        """
        This is equivalent to self.module_list[idx](x),
        but torchscript doesn't support this yet
        """
        num_blocks = len(self.module_list)
        if idx < 0:
            idx += num_blocks
        i = 0
        out = x
        for module in self.module_list:
            if i == idx:
                out = module(x)
            i += 1
        return out

    def forward(self, x: List[Tensor]) -> Tensor:
        all_results = []

        for i, features in enumerate(x):
            results = self._get_result_from_module_list(features, i)

            # Permute output from (N, A * K, H, W) to (N, HWA, K).
            N, _, H, W = results.shape
            results = results.view(N, -1, self.num_columns, H, W)
            results = results.permute(0, 3, 4, 1, 2)
            results = results.reshape(N, -1, self.num_columns)  # Size=(N, HWA, K)

            all_results.append(results)

        return torch.cat(all_results, dim=1)

class SSDClassificationHead(SSDScoringHead):
    def __init__(self, in_channels: List[int], num_anchors: List[int], num_classes: int):
        cls_logits = nn.ModuleList()
        for channels, anchors in zip(in_channels, num_anchors):
            cls_logits.append(nn.Conv2d(channels, num_classes * anchors, kernel_size=3, padding=1))
        _xavier_init(cls_logits)
        super().__init__(cls_logits, num_classes)

class SSDRegressionHead(SSDScoringHead):
    def __init__(self, in_channels: List[int], num_anchors: List[int]):
        bbox_reg = nn.ModuleList()
        for channels, anchors in zip(in_channels, num_anchors):
            bbox_reg.append(nn.Conv2d(channels, 4 * anchors, kernel_size=3, padding=1))
        _xavier_init(bbox_reg)
        super().__init__(bbox_reg, 4)

class SSD(nn.Module):
    """
    Implements SSD architecture from `"SSD: Single Shot MultiBox Detector" <https://arxiv.org/abs/1512.02325>`_.

    The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each
    image, and should be in 0-1 range. Different images can have different sizes but they will be resized
    to a fixed size before passing it to the backbone.

    The behavior of the model changes depending if it is in training or evaluation mode.

    During training, the model expects both the input tensors, as well as a targets (list of dictionary),
    containing:
        - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with
          ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.
        - labels (Int64Tensor[N]): the class label for each ground-truth box

    The model returns a Dict[Tensor] during training, containing the classification and regression
    losses.

    During inference, the model requires only the input tensors, and returns the post-processed
    predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as
    follows, where ``N`` is the number of detections:

        - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with
          ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.
        - labels (Int64Tensor[N]): the predicted labels for each detection
        - scores (Tensor[N]): the scores for each detection

    Args:
        backbone (nn.Module): the network used to compute the features for the model.
            It should contain an out_channels attribute with the list of the output channels of
            each feature map. The backbone should return a single Tensor or an OrderedDict[Tensor].
        anchor_generator (DefaultBoxGenerator): module that generates the default boxes for a
            set of feature maps.
        size (Tuple[int, int]): the width and height to which images will be rescaled before feeding them
            to the backbone.
        num_classes (int): number of output classes of the model (excluding the background).
        image_mean (Tuple[float, float, float]): mean values used for input normalization.
            They are generally the mean values of the dataset on which the backbone has been trained
            on
        image_std (Tuple[float, float, float]): std values used for input normalization.
            They are generally the std values of the dataset on which the backbone has been trained on
        head (nn.Module, optional): Module run on top of the backbone features. Defaults to a module containing
            a classification and regression module.
        score_thresh (float): Score threshold used for postprocessing the detections.
        nms_thresh (float): NMS threshold used for postprocessing the detections.
        detections_per_img (int): Number of best detections to keep after NMS.
        iou_thresh (float): minimum IoU between the anchor and the GT box so that they can be
            considered as positive during training.
        topk_candidates (int): Number of best detections to keep before NMS.
        positive_fraction (float): a number between 0 and 1 which indicates the proportion of positive
            proposals used during the training of the classification head. It is used to estimate the negative to
            positive ratio.
    """
    __annotations__ = {
        'box_coder': det_utils.BoxCoder,
        'proposal_matcher': det_utils.Matcher,
    }

    def __init__(self, backbone: nn.Module, anchor_generator: DefaultBoxGenerator,
                 size: Tuple[int, int], num_classes: int,
                 image_mean: Optional[List[float]] = None, image_std: Optional[List[float]] = None,
                 head: Optional[nn.Module] = None,
                 score_thresh: float = 0.00,
                 nms_thresh: float = 0.45,
                 detections_per_img: int = 25,
                 iou_thresh: float = 0.5,
                 topk_candidates: int = 400,
                 positive_fraction: float = 0.25):
        super().__init__()

        self.backbone = backbone

        self.anchor_generator = anchor_generator

        self.box_coder = det_utils.BoxCoder(weights=(10., 10., 5., 5.))

        if head is None:
            if hasattr(backbone, 'out_channels'):
                out_channels = backbone.out_channels
            else:
                out_channels = det_utils.retrieve_out_channels(backbone, size)

            assert len(out_channels) == len(anchor_generator.aspect_ratios)

            num_anchors = self.anchor_generator.num_anchors_per_location()
            head = SSDHead(out_channels, num_anchors, num_classes)
        self.head = head

        self.proposal_matcher = det_utils.SSDMatcher(iou_thresh)

        if image_mean is None:
            image_mean = [0.485, 0.456, 0.406]
        if image_std is None:
            image_std = [0.229, 0.224, 0.225]
        self.transform = GeneralizedRCNNTransform(min(size), max(size), image_mean, image_std,
                                                  size_divisible=1, fixed_size=size)

        self.score_thresh = score_thresh
        self.nms_thresh = nms_thresh
        self.detections_per_img = detections_per_img
        self.topk_candidates = topk_candidates
        self.neg_to_pos_ratio = (1.0 - positive_fraction) / positive_fraction

        # used only on torchscript mode
        self._has_warned = False

    def eager_outputs(self, losses: Dict[str, Tensor],
                      detections: List[Dict[str, Tensor]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]:
        if self.training:
            return losses

        return detections

    def compute_loss(self, targets: List[Dict[str, Tensor]], head_outputs: Dict[str, Tensor], anchors: List[Tensor],
                     matched_idxs: List[Tensor]) -> Dict[str, Tensor]:
        bbox_regression = head_outputs['bbox_regression']
        cls_logits = head_outputs['cls_logits']

        # Match original targets with default boxes
        num_foreground = 0
        bbox_loss = []
        cls_targets = []
        for (targets_per_image, bbox_regression_per_image, cls_logits_per_image, anchors_per_image,
             matched_idxs_per_image) in zip(targets, bbox_regression, cls_logits, anchors, matched_idxs):
            # produce the matching between boxes and targets
            foreground_idxs_per_image = torch.where(matched_idxs_per_image >= 0)[0]
            foreground_matched_idxs_per_image = matched_idxs_per_image[foreground_idxs_per_image]
            num_foreground += foreground_matched_idxs_per_image.numel()

            # Calculate regression loss
            matched_gt_boxes_per_image = targets_per_image['boxes'][foreground_matched_idxs_per_image]
            bbox_regression_per_image = bbox_regression_per_image[foreground_idxs_per_image, :]
            anchors_per_image = anchors_per_image[foreground_idxs_per_image, :]
            target_regression = self.box_coder.encode_single(matched_gt_boxes_per_image, anchors_per_image)
            bbox_loss.append(torch.nn.functional.smooth_l1_loss(
                bbox_regression_per_image,
                target_regression,
                reduction='sum'
            ))

        bbox_loss = torch.stack(bbox_loss)
        N = max(1, num_foreground)
        out_loss = bbox_loss.sum() / N
        print("Loss: {}".format(out_loss))
        return out_loss

    def forward(self, images: List[Tensor],
                targets: Optional[List[Dict[str, Tensor]]] = None) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]:
        if self.training and targets is None:
            raise ValueError("In training mode, targets should be passed")

        if self.training:
            assert targets is not None
            for target in targets:
                boxes = target["boxes"]
                if isinstance(boxes, torch.Tensor):
                    if len(boxes.shape) != 2 or boxes.shape[-1] != 4:
                        raise ValueError("Expected target boxes to be a tensor"
                                         "of shape [N, 4], got {:}.".format(
                                             boxes.shape))
                else:
                    raise ValueError("Expected target boxes to be of type "
                                     "Tensor, got {:}.".format(type(boxes)))

        # get the original image sizes
        original_image_sizes: List[Tuple[int, int]] = []
        for img in images:
            val = img.shape[-2:]
            assert len(val) == 2
            original_image_sizes.append((val[0], val[1]))

        # transform the input
        images, targets = self.transform(images, targets)

        # Check for degenerate boxes
        if targets is not None:
            for target_idx, target in enumerate(targets):
                boxes = target["boxes"]
                degenerate_boxes = boxes[:, 2:] <= boxes[:, :2]
                if degenerate_boxes.any():
                    bb_idx = torch.where(degenerate_boxes.any(dim=1))[0][0]
                    degen_bb: List[float] = boxes[bb_idx].tolist()
                    raise ValueError("All bounding boxes should have positive height and width."
                                     " Found invalid box {} for target at index {}."
                                     .format(degen_bb, target_idx))

        # get the features from the backbone
        features = self.backbone(images.tensors)
        if isinstance(features, torch.Tensor):
            features = OrderedDict([('0', features)])

        features = list(features.values())

        # compute the ssd heads outputs using the features
        head_outputs = self.head(features)

        # create the set of anchors
        anchors = self.anchor_generator(images, features)

        losses = {}
        detections: List[Dict[str, Tensor]] = []
        if self.training:
            assert targets is not None

            matched_idxs = []
            for anchors_per_image, targets_per_image in zip(anchors, targets):
                if targets_per_image['boxes'].numel() == 0:
                    matched_idxs.append(torch.full((anchors_per_image.size(0),), -1, dtype=torch.int64,
                                                   device=anchors_per_image.device))
                    continue

                match_quality_matrix = box_ops.box_iou(targets_per_image['boxes'], anchors_per_image)
                matched_idxs.append(self.proposal_matcher(match_quality_matrix))

            losses = self.compute_loss(targets, head_outputs, anchors, matched_idxs)
        else:
            detections = self.postprocess_detections(head_outputs, anchors, images.image_sizes)
            detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)

        if torch.jit.is_scripting():
            if not self._has_warned:
                warnings.warn("SSD always returns a (Losses, Detections) tuple in scripting")
                self._has_warned = True
            return losses, detections
        return self.eager_outputs(losses, detections)

    def postprocess_detections(self, head_outputs: Dict[str, Tensor], image_anchors: List[Tensor],
                               image_shapes: List[Tuple[int, int]]) -> List[Dict[str, Tensor]]:
        bbox_regression = head_outputs['bbox_regression']
        pred_scores = F.softmax(head_outputs['cls_logits'], dim=-1)

        num_classes = pred_scores.size(-1)
        device = pred_scores.device

        detections: List[Dict[str, Tensor]] = []

        for boxes, scores, anchors, image_shape in zip(bbox_regression, pred_scores, image_anchors, image_shapes):
            print(scores)
            boxes = self.box_coder.decode_single(boxes, anchors)
            boxes = box_ops.clip_boxes_to_image(boxes, image_shape)
            print(boxes.shape)
            print(boxes)

            image_boxes = []
            image_scores = []
            image_labels = []
            for label in range(0, num_classes):
                score = np.stack(([1],)*len(boxes))

                #keep_idxs = score > self.score_thresh
                #score = score[keep_idxs]
                #box = boxes[keep_idxs]

                # keep only topk scoring predictions
                #num_topk = min(self.topk_candidates, 25)
                #score, idxs = score.topk(num_topk)
                #box = boxes[idxs]

                image_boxes.append(boxes)
                image_scores.append(score)
                image_labels.append(torch.full_like(torch.from_numpy(score), fill_value=label, dtype=torch.int64, device=device))

            image_boxes = torch.cat(image_boxes, dim=0)
            image_scores = torch.from_numpy(np.array(image_scores))
            image_labels = torch.cat(image_labels, dim=0)

            # non-maximum suppression
            keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh)
            keep = keep[:self.detections_per_img]

            detections.append({
                'boxes': image_boxes[keep],
                'scores': image_scores[keep],
                'labels': image_labels[keep],
            })
        return detections

class SSDFeatureExtractorVGG(nn.Module):
    def __init__(self, backbone: nn.Module, highres: bool):
        super().__init__()

        _, _, maxpool3_pos, maxpool4_pos, _ = (i for i, layer in enumerate(backbone) if isinstance(layer, nn.MaxPool2d))

        # Patch ceil_mode for maxpool3 to get the same WxH output sizes as the paper
        backbone[maxpool3_pos].ceil_mode = True

        # parameters used for L2 regularization + rescaling
        self.scale_weight = nn.Parameter(torch.ones(512) * 20)

        # Multiple Feature maps - page 4, Fig 2 of SSD paper
        self.features = nn.Sequential(
            *backbone[:maxpool4_pos]  # until conv4_3
        )

        # SSD300 case - page 4, Fig 2 of SSD paper
        extra = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(1024, 256, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(256, 512, kernel_size=3, padding=1, stride=2),  # conv8_2
                nn.ReLU(inplace=True),
            ),
            nn.Sequential(
                nn.Conv2d(512, 128, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(128, 256, kernel_size=3, padding=1, stride=2),  # conv9_2
                nn.ReLU(inplace=True),
            ),
            nn.Sequential(
                nn.Conv2d(256, 128, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(128, 256, kernel_size=3),  # conv10_2
                nn.ReLU(inplace=True),
            ),
            nn.Sequential(
                nn.Conv2d(256, 128, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(128, 256, kernel_size=3),  # conv11_2
                nn.ReLU(inplace=True),
            )
        ])
        if highres:
            # Additional layers for the SSD512 case. See page 11, footernote 5.
            extra.append(nn.Sequential(
                nn.Conv2d(256, 128, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(128, 256, kernel_size=4),  # conv12_2
                nn.ReLU(inplace=True),
            ))
        _xavier_init(extra)

        fc = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1, ceil_mode=False),  # add modified maxpool5
            nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3, padding=6, dilation=6),  # FC6 with atrous
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=1024, out_channels=1024, kernel_size=1),  # FC7
            nn.ReLU(inplace=True)
        )
        _xavier_init(fc)
        extra.insert(0, nn.Sequential(
            *backbone[maxpool4_pos:-1],  # until conv5_3, skip maxpool5
            fc,
        ))
        self.extra = extra

    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        # L2 regularization + Rescaling of 1st block's feature map
        x = self.features(x)
        rescaled = self.scale_weight.view(1, -1, 1, 1) * F.normalize(x)
        output = [rescaled]

        # Calculating Feature maps for the rest blocks
        for block in self.extra:
            x = block(x)
            output.append(x)

        return OrderedDict([(str(i), v) for i, v in enumerate(output)])

def _vgg_extractor(backbone_name: str, highres: bool, progress: bool, pretrained: bool, trainable_layers: int):
    if backbone_name in backbone_urls:
        # Use custom backbones more appropriate for SSD
        arch = backbone_name.split('_')[0]
        backbone = vgg.__dict__[arch](pretrained=False, progress=progress).features
        if pretrained:
            state_dict = load_state_dict_from_url(backbone_urls[backbone_name], progress=progress)
            backbone.load_state_dict(state_dict)
    else:
        # Use standard backbones from TorchVision
        backbone = vgg.__dict__[backbone_name](pretrained=pretrained, progress=progress).features

    # Gather the indices of maxpools. These are the locations of output blocks.
    stage_indices = [i for i, b in enumerate(backbone) if isinstance(b, nn.MaxPool2d)]
    num_stages = len(stage_indices)

    # find the index of the layer from which we wont freeze
    assert 0 <= trainable_layers <= num_stages
    freeze_before = len(backbone) if trainable_layers == 0 else stage_indices[num_stages - trainable_layers]

    for b in backbone[:freeze_before]:
        for parameter in b.parameters():
            parameter.requires_grad_(False)

    return SSDFeatureExtractorVGG(backbone, highres)

def my_ssd300_vgg16(pretrained= False, progress= True, num_classes= 91,
                 pretrained_backbone = True, trainable_backbone_layers = None, **kwargs):
    """Constructs an SSD model with input size 300x300 and a VGG16 backbone.

    Reference: `"SSD: Single Shot MultiBox Detector" <https://arxiv.org/abs/1512.02325>`_.

    The input to the model is expected to be a list of tensors, each of shape [C, H, W], one for each
    image, and should be in 0-1 range. Different images can have different sizes but they will be resized
    to a fixed size before passing it to the backbone.

    The behavior of the model changes depending if it is in training or evaluation mode.

    During training, the model expects both the input tensors, as well as a targets (list of dictionary),
    containing:

        - boxes (``FloatTensor[N, 4]``): the ground-truth boxes in ``[x1, y1, x2, y2]`` format, with
          ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.
        - labels (Int64Tensor[N]): the class label for each ground-truth box

    The model returns a Dict[Tensor] during training, containing the classification and regression
    losses.

    During inference, the model requires only the input tensors, and returns the post-processed
    predictions as a List[Dict[Tensor]], one for each input image. The fields of the Dict are as
    follows, where ``N`` is the number of detections:

        - boxes (``FloatTensor[N, 4]``): the predicted boxes in ``[x1, y1, x2, y2]`` format, with
          ``0 <= x1 < x2 <= W`` and ``0 <= y1 < y2 <= H``.
        - labels (Int64Tensor[N]): the predicted labels for each detection
        - scores (Tensor[N]): the scores for each detection

    Example:

        >>> model = torchvision.models.detection.ssd300_vgg16(pretrained=True)
        >>> model.eval()
        >>> x = [torch.rand(3, 300, 300), torch.rand(3, 500, 400)]
        >>> predictions = model(x)

    Args:
        pretrained (bool): If True, returns a model pre-trained on COCO train2017
        progress (bool): If True, displays a progress bar of the download to stderr
        num_classes (int): number of output classes of the model (including the background)
        pretrained_backbone (bool): If True, returns a model with backbone pre-trained on Imagenet
        trainable_backbone_layers (int): number of trainable (not frozen) resnet layers starting from final block.
            Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
    """
    if "size" in kwargs:
        warnings.warn("The size of the model is already fixed; ignoring the argument.")

    trainable_backbone_layers = models.detection.ssd._validate_trainable_layers(
        pretrained or pretrained_backbone, trainable_backbone_layers, 5, 5)

    if pretrained:
        # no need to download the backbone if pretrained is set
        pretrained_backbone = False

    backbone = _vgg_extractor("vgg16_features", False, progress, pretrained_backbone, trainable_backbone_layers)
    anchor_generator = DefaultBoxGenerator([[2], [2, 3], [2, 3], [2, 3], [2], [2]],
                                           min_ratio=0.07, max_ratio=0.15)

    defaults = {
        # Rescale the input in a way compatible to the backbone
        "image_mean": [0.48235, 0.45882, 0.40784],
        "image_std": [1.0 / 255.0, 1.0 / 255.0, 1.0 / 255.0],  # undo the 0-1 scaling of toTensor
    }
    kwargs = {**defaults, **kwargs}
    model = SSD(backbone, anchor_generator, (270, 270), num_classes, **kwargs)
    return model

FACE_CORNERS = (
    (0, 1, 2, 3),
    (4, 5, 1, 0),
    (5, 6, 2, 1),
    (7, 6, 2, 3),
    (4, 7, 3, 0),
    (4, 5, 6, 7),
)

def get_cell_corners_3d(pos):
    """Get 3d positions of the corners of the cell at the given position."""
    d = 0.022 / 2
    nppos = np.asarray(pos)

    # order of the corners is the same as in the cube model of the
    # trifinger_object_tracking package
    # people.tue.mpg.de/mpi-is-software/robotfingers/docs/trifinger_object_tracking/doc/cube_model.html
    return np.array(
        (
            nppos + (d, -d, d),
            nppos + (d, d, d),
            nppos + (-d, d, d),
            nppos + (-d, -d, d),
            nppos + (d, -d, -d),
            nppos + (d, d, -d),
            nppos + (-d, d, -d),
            nppos + (-d, -d, -d),
        )
    )

def generate_goal_mask(camera_parameters, goal):
    """Generate goal masks that can be used with :func:`evaluate_state`.
    A goal mask is a single-channel image where the areas at which dice are
    supposed to be placed are white and everything else is black.  So it
    corresponds more or less to a segmentation mask where all dice are at the
    goal positions.
    For rendering the mask, :data:`TARGET_WIDTH` is used for the die width to
    add some tolerance.
    Args:
        camera_parameters: List of camera parameters, one per camera.
        goal: The goal die positions.
    Returns:
        List of masks.  The number and order of masks corresponds to the input
        ``camera_parameters``.
    """
    #masks = []
    #for cam in camera_parameters:
    #mask = np.zeros((camera_parameters.image_height, camera_parameters.image_width), dtype=np.uint8)
    mask = np.zeros((270, 270), dtype=np.uint8)

    # get camera position and orientation separately
    tvec = camera_parameters.tf_world_to_camera[:3, 3]
    rmat = camera_parameters.tf_world_to_camera[:3, :3]
    rvec = Rotation.from_matrix(rmat).as_rotvec()

    for pos in goal:
        corners = get_cell_corners_3d(pos)

        # project corner points into the image
        projected_corners, _ = cv2.projectPoints(
            corners,
            rvec,
            tvec,
            camera_parameters.camera_matrix,
            camera_parameters.distortion_coefficients,
        )

        # draw faces in mask
        for face_corner_idx in FACE_CORNERS:
            points = np.array(
                [projected_corners[i] for i in face_corner_idx],
                dtype=np.int32,
            )
            mask = cv2.fillConvexPoly(mask, points, 255)

        #masks.append(mask)

    return mask

def bbox_generator(camera_params, goal, i):
  #print(goal)
  mask = generate_goal_mask(camera_params, [goal])
  #cv2.imwrite('mask{}.png'.format(i), mask)
  contour = cv2.findContours(mask.astype(np.uint8).copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
  #contour = contour[0] if len(contour) == 2 else contour[1]
  x, y, w, h = cv2.boundingRect(contour[0])
  return np.array([x, y, x + w, y + h])

def generate_batch(env, batch_size):
  batch = np.ones((batch_size, 3, 270, 270))
  detections = []
  for i in range(batch_size):
    g_ = task.sample_goal()
    #print(g_)
    g_mask = generate_goal_mask(env.camera_params[0], g_)
    batch[i] = np.stack((g_mask,)*3, axis=0)
    bboxes=[]
    for idx, g in enumerate(g_):
        bboxes.append(bbox_generator(env.camera_params[0], g, idx))
    detections.append({'boxes': torch.from_numpy(np.array(bboxes)).float(),
		   'labels': torch.from_numpy(np.array([1])).float()})
    #detections.append(bboxes)
  return batch, detections

"""def generate_batch(env, batch_size):
  batch = np.ones((batch_size, 3, 270, 270))
	goals = np.ones((batch_size, 25 * 3))
	#goals = np.ones((batch_size, 25 * 2))
	for i in range(batch_size):
		seg_mask = np.ones((3, 270, 270))
		g = task.sample_goal()
		goal = list(itertools.chain(*g))
		#goal = [g for i, g in enumerate(goal) if ((i+1) % 3) !=0]
		goals[i] = np.array(goal)
		for idx, c in enumerate(env.camera_params):
			seg_mask[idx,:,:] = generate_goal_mask(c, g)
		#segmentation_masks = np.array([segment_image(cv2.cvtColor(c.image, cv2.COLOR_RGB2BGR)) for c in obs.cameras])
		batch[i] = seg_mask
	return batch, goals"""

train = False
model = my_ssd300_vgg16(num_classes=1)
env = rearrange_dice_env.RealRobotRearrangeDiceEnv(rearrange_dice_env.ActionType.POSITION,goal= None,step_size=1,)
env.reset()
#mask, bboxes = generate_batch(env, 16)
#print(bboxes)
if train:
	optim = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.0001)
	min_cost = 20000
	while True:
		  optim.zero_grad()
		  mask, bboxes = generate_batch(env, 4)
		  loss = model.forward(torch.from_numpy(mask).float(), bboxes)
		  loss.backward()
		  optim.step()
		  if loss < min_cost:
		  	min_cost = loss
		  	print('Saving model')
		  	torch.save(model.state_dict(), './ssd_test_newAnchors.pth')
else:
	mask, bboxes = generate_batch(env, 1)
	model.load_state_dict(torch.load('./ssd_test_newAnchors.pth', map_location = torch.device('cpu')))
	model.eval()
	preds = model.forward(torch.from_numpy(mask).float())
	print(preds)


import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np
import logging
from torch.nn.parameter import Parameter
from torch.autograd import Variable
from milvus import Milvus, IndexType, MetricType, Status

client = Milvus(host='192.168.1.85', port='19666')

logger = logging.getLogger()
logger.setLevel(logging.INFO)

class EMA():
    def __init__(self, decay=0.999):
        self.decay = decay
        self.shadow = {}

    def register(self, name, val):
        self.shadow[name] = val.cpu().detach()

    def get(self, name):
        return self.shadow[name]

    def update(self, name, x):
        assert name in self.shadow
        new_average = (1.0 - self.decay) * x.cpu().detach() + self.decay * self.shadow[name]
        self.shadow[name] = new_average.clone()

def create_collection(gallery):
    param = {'collection_name':'test01', 'dimension':512, 'index_file_size':1024, 'metric_type':MetricType.IP}
    status = client.create_collection(param)
    ivf_param = {'nlist': 2048}
    status = client.create_index('test01', IndexType.IVF_FLAT, ivf_param)
    status, inserted_vector_ids = client.insert(collection_name='test01', records=gallery)
    #print(len(inserted_vector_ids))

def pairwise_distance(A, B):
    """
    Compute distance between points in A and points in B
    :param A:  (m,n) -m points, each of n dimension. Every row vector is a point, denoted as A(i).
    :param B:  (k,n) -k points, each of n dimension. Every row vector is a point, denoted as B(j).
    :return:  Matrix with (m, k). And the ele in (i,j) is the distance between A(i) and B(j)
    """
    A_square = torch.sum(A * A, dim=1, keepdim=True)
    B_square = torch.sum(B * B, dim=1, keepdim=True)

    distance = A_square + B_square.t() - 2 * torch.matmul(A, B.t())

    return distance

def one_hot_coding(index, k):
    if type(index) is torch.Tensor:
        length = len(index)
    else:
        length = 1
    out = torch.zeros((length, k), dtype=torch.int64).cuda()
    index = index.reshape((len(index), 1))
    out.scatter_(1, index, 1)
    return out

# deprecated due to the large memory usage
def constraints_old(features, labels):
    distance = pairwise_distance(features, features)
    labels_reshape = torch.reshape(labels, (features.shape[0], 1))
    labels_dist = labels_reshape - labels_reshape.t()
    labels_mask = (labels_dist == 0).float()

    # Average loss with each matching pair
    num = torch.sum(labels_mask) - features.shape[0]
    if num == 0:
        con_loss = 0.0
    else:
        con_loss = torch.sum(distance * labels_mask) / num

    return con_loss

def constraints(features, labels):
    labels = torch.reshape(labels, (labels.shape[0],1))
    con_loss = AverageMeter()
    index_dict = {k.item() for k in labels}
    for index in index_dict:
        labels_mask = (labels == index)
        feas = torch.masked_select(features, labels_mask)
        feas = feas.view(-1, features.shape[1])
        distance = pairwise_distance(feas, feas)
        #torch.sqrt_(distance)
        num = feas.shape[0] * (feas.shape[0] - 1)
        loss = torch.sum(distance) / num
        con_loss.update(loss, n = num / 2)
    return con_loss.avg

def constraints_loss(data_loader, network, args):
    network.eval()
    max_size = args.batch_size * len(data_loader)
    images_bank = torch.zeros((max_size, args.feature_size)).cuda()
    text_bank = torch.zeros((max_size,args.feature_size)).cuda()
    labels_bank = torch.zeros(max_size).cuda()
    index = 0
    con_images = 0.0
    con_text = 0.0
    with torch.no_grad():
        for images, captions, labels, captions_length in data_loader:
            images = images.cuda()
            captions = captions.cuda()
            interval = images.shape[0]
            image_embeddings, text_embeddings = network(images, captions, captions_length)
            images_bank[index: index + interval] = image_embeddings
            text_bank[index: index + interval] = text_embeddings
            labels_bank[index: index + interval] = labels
            index = index + interval
        images_bank = images_bank[:index]
        text_bank = text_bank[:index]
        labels_bank = labels_bank[:index]

    if args.constraints_text:
        con_text = constraints(text_bank, labels_bank)
    if args.constraints_images:
        con_images = constraints(images_bank, labels_bank)

    return con_images, con_text

class Loss(nn.Module):
    def __init__(self, args):
        super(Loss, self).__init__()
        self.CMPM = args.CMPM
        self.CMPC = args.CMPC
        self.epsilon = args.epsilon
        self.num_classes = args.num_classes
        if args.resume:
            checkpoint = torch.load(args.model_path)
            self.W = Parameter(checkpoint['W'])
            print('=========> Loading in parameter W from pretrained models')
        else:
            self.W = Parameter(torch.randn(args.feature_size, args.num_classes))
            self.init_weight()

    def init_weight(self):
        nn.init.xavier_uniform_(self.W.data, gain=1)

    def compute_cmpc_loss(self, image_embeddings, text_embeddings, labels):
        """
        Cross-Modal Projection Classfication loss(CMPC)
        :param image_embeddings: Tensor with dtype torch.float32
        :param text_embeddings: Tensor with dtype torch.float32
        :param labels: Tensor with dtype torch.int32
        :return:
        """
        criterion = nn.CrossEntropyLoss(reduction='mean')
        self.W_norm = self.W / self.W.norm(dim=0)
        #labels_onehot = one_hot_coding(labels, self.num_classes).float()
        image_norm = image_embeddings / image_embeddings.norm(dim=1, keepdim=True)
        text_norm = text_embeddings / text_embeddings.norm(dim=1, keepdim=True)

        image_proj_text = torch.sum(image_embeddings * text_norm, dim=1, keepdim=True) * text_norm
        text_proj_image = torch.sum(text_embeddings * image_norm, dim=1, keepdim=True) * image_norm

        image_logits = torch.matmul(image_proj_text, self.W_norm)
        text_logits = torch.matmul(text_proj_image, self.W_norm)

        #labels_one_hot = one_hot_coding(labels, num_classes)
        '''
        ipt_loss = criterion(input=image_logits, target=labels)
        tpi_loss = criterion(input=text_logits, target=labels)
        cmpc_loss = ipt_loss + tpi_loss
        '''
        cmpc_loss = criterion(image_logits, labels) + criterion(text_logits, labels)
        #cmpc_loss = - (F.log_softmax(image_logits, dim=1) + F.log_softmax(text_logits, dim=1)) * labels_onehot
        #cmpc_loss = torch.mean(torch.sum(cmpc_loss, dim=1))
        # classification accuracy for observation
        image_pred = torch.argmax(image_logits, dim=1)
        text_pred = torch.argmax(text_logits, dim=1)

        image_precision = torch.mean((image_pred == labels).float())
        text_precision = torch.mean((text_pred == labels).float())

        return cmpc_loss, image_precision, text_precision

    def compute_cmpm_loss(self, image_embeddings, text_embeddings, labels):
        """
        Cross-Modal Projection Matching Loss(CMPM)
        :param image_embeddings: Tensor with dtype torch.float32
        :param text_embeddings: Tensor with dtype torch.float32
        :param labels: Tensor with dtype torch.int32
        :return:
            i2t_loss: cmpm loss for image projected to text
            t2i_loss: cmpm loss for text projected to image
            pos_avg_sim: average cosine-similarity for positive pairs
            neg_avg_sim: averate cosine-similarity for negative pairs
        """

        batch_size = image_embeddings.shape[0]
        labels_reshape = torch.reshape(labels, (batch_size, 1))
        labels_dist = labels_reshape - labels_reshape.t()
        labels_mask = (labels_dist == 0)

        image_norm = image_embeddings / image_embeddings.norm(dim=1, keepdim=True)
        text_norm = text_embeddings / text_embeddings.norm(dim=1, keepdim=True)
        image_proj_text = torch.matmul(image_embeddings, text_norm.t())
        text_proj_image = torch.matmul(text_embeddings, image_norm.t())

        # normalize the true matching distribution
        labels_mask_norm = labels_mask.float() / labels_mask.float().norm(dim=1)

        i2t_pred = F.softmax(image_proj_text, dim=1)
        #i2t_loss = i2t_pred * torch.log((i2t_pred + self.epsilon)/ (labels_mask_norm + self.epsilon))
        i2t_loss = i2t_pred * (F.log_softmax(image_proj_text, dim=1) - torch.log(labels_mask_norm + self.epsilon))

        t2i_pred = F.softmax(text_proj_image, dim=1)
        #t2i_loss = t2i_pred * torch.log((t2i_pred + self.epsilon)/ (labels_mask_norm + self.epsilon))
        t2i_loss = t2i_pred * (F.log_softmax(text_proj_image, dim=1) - torch.log(labels_mask_norm + self.epsilon))

        cmpm_loss = torch.mean(torch.sum(i2t_loss, dim=1)) + torch.mean(torch.sum(t2i_loss, dim=1))

        sim_cos = torch.matmul(image_norm, text_norm.t())

        pos_avg_sim = torch.mean(torch.masked_select(sim_cos, labels_mask))
        neg_avg_sim = torch.mean(torch.masked_select(sim_cos, labels_mask == 0))

        return cmpm_loss, pos_avg_sim, neg_avg_sim

    def forward(self, image_embeddings, text_embeddings, labels):
        cmpm_loss = 0.0
        cmpc_loss = 0.0
        image_precision = 0.0
        text_precision = 0.0
        neg_avg_sim = 0.0
        pos_avg_sim =0.0
        if self.CMPM:
            cmpm_loss, pos_avg_sim, neg_avg_sim = self.compute_cmpm_loss(image_embeddings, text_embeddings, labels)
        if self.CMPC:
            cmpc_loss, image_precision, text_precision = self.compute_cmpc_loss(image_embeddings, text_embeddings, labels)

        loss = cmpm_loss + cmpc_loss

        return cmpm_loss, cmpc_loss, loss, image_precision, text_precision, pos_avg_sim, neg_avg_sim

class AverageMeter(object):
    """
    Computes and stores the averate and current value
    Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py #L247-262
    """
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += n * val
        self.count += n
        self.avg = self.sum / self.count

def compute_topk(query, gallery, target_query, target_gallery, k=[1,10], reverse=False):
    result = []
    query = query / query.norm(dim=1,keepdim=True)
    gallery = gallery / gallery.norm(dim=1,keepdim=True)
    #print("query:", query, "size:", query.size())
    #print("gallery:", gallery, "size:", gallery.size())
    sim_cosine = torch.matmul(query, gallery.t())
    #create_collection(gallery.tolist())
    search_param = {'nprobe': 16}
    status, results = client.search(collection_name='test01', query_records=query.tolist()[:2], top_k=5, params=search_param)
    result.extend(topk(sim_cosine, target_gallery, target_query, k=[1,10]))
    if reverse:
        result.extend(topk(sim_cosine, target_query, target_gallery, k=[1,10], dim=0))
    return result

def topk(sim, target_gallery, target_query, k=[1,10], dim=1):
    result = []
    maxk = max(k)
    size_total = len(target_gallery)
    _, pred_index = sim.topk(maxk, dim, True, True)
    pred_labels = target_gallery[pred_index]
    print('pred_labels:', pred_labels)
    if dim == 1:
        pred_labels = pred_labels.t()
    print('pred_labels:', pred_labels)
    correct = pred_labels.eq(target_query.view(1,-1).expand_as(pred_labels))
    print('correct:', correct)

    for topk in k:
        #correct_k = torch.sum(correct[:topk]).float()
        correct_k = torch.sum(correct[:topk], dim=0)
        correct_k = torch.sum(correct_k > 0).float()
        result.append(correct_k * 100 / size_total)
    return result