import time
import random
import glob
import argparse
import logging
import sys
import os
from PIL import Image
import uuid
import copy
import numpy as np
import carla_utils
import pickle
import dynamic_weather as weather
import datetime
import csv
try:
    import queue
except ImportError:
    import Queue as queue
# import ipdb
# st = ipdb.set_trace
try:
    sys.path.append(glob.glob('/hdd/carla97/PythonAPI/carla/dist/carla-*%d.%d-%s.egg' % (
        sys.version_info.major,
        sys.version_info.minor,
        'win-amd64' if os.name == 'nt' else 'linux-x86_64'))[0])
except IndexError:
    pass
import carla
import carla_utils
from carla_utils import ClientSideBoundingBoxes
import image_converter
from carla import ColorConverter as cc

from carla_sync_mode import CarlaSyncMode

def save_npy(data, filename):
    np.save(filename, data)

# Built on carla_two_unique.py
# can be used to spawn 2-3 vehicles
'''
camR ranges: 8-14, -5,5, 1-3
'''
class CarlaMultiviewRunner():
    def __init__(self, start_episode, end_episode, mod, save_dir):
        logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.INFO)

        self.base_path = os.path.join(save_dir, "surveil_{}".format(mod))
        if not os.path.exists(self.base_path):
            os.mkdir(self.base_path)
        self.randomize_each_camera = True

        self.start_frame = 0
        self.start_episode = start_episode
        self.end_episode = end_episode
        self.frames_per_episode = 1
        self.num_vehicles = 2
        self.today = str(datetime.date.today())
        self.episode_path_format = '{:s}_episode_{:0>4d}_vehicles_{:0>3d}'
        self.out_filename_format = '{:s}_episode_{:0>4d}_vehicles_{:0>3d}/{:s}/{:0>6d}'

        self.nearest_x = 2.0
        self.dis_to_center = 15.0#5.0
        self.height = 0.65
        self.d_height = 1.0

        # Camera specific params
        self.image_size_x = 768
        self.image_size_y = 256
        self.fov = 90
        self.focal = self.image_size_x/(2.0*np.tan(self.fov*np.pi/360.0))
        self.fps = 10
        self.half_sqrt2 = np.sqrt(2)/2

        self.host = "127.0.0.1"
        self.port = 2000
        self.filterv = "vehicle.*"
        self.sensor_types = ['sensor.camera.rgb', 'sensor.camera.depth', 'sensor.camera.semantic_segmentation']
        self.sensor_names_format = ['CameraRGB{}', 'CameraDepth{}', 'CameraSemantic{}']

        self.category = {'vehicle': 0}

        self.calculate_camera_locations()

        self.visualize = True
        self.visual_dir = 'visual'
        if self.visualize and not os.path.exists(self.visual_dir):
            os.mkdir(self.visual_dir)

    def calculate_camera_locations(self):
        self.center = np.array([self.nearest_x+self.dis_to_center, 0.0, self.height])
        self.positions = []
        self.rotations = []
        #hori_rs = [self.dis_to_center, self.dis_to_center, self.dis_to_center]
        verti_hs = [self.height+self.d_height, self.height, self.height+2*self.d_height]
        pitches = [-0.08/np.pi*180 for i in range(3)]
        #for i in [0]: #range(len(hori_rs)):
        #hori_r = hori_rs[0]
        h = 1.65 # verti_hs[1]
        p_angle = pitches[0]
        #theta = 15.0
        # for angle in np.arange(5.0, 355.0, 20.0):#np.arange(0.0, 359.0, 30.0):
        #     a_rad = angle/180.0*np.pi
        #     scale = 1.
        #     trans = np.zeros(3)
        #     #pos = self.center + scale*np.array([-hori_r*np.cos(a_rad), hori_r*np.sin(a_rad), h]) + trans
        #     pos = scale*np.array([-hori_r*np.cos(a_rad), hori_r*np.sin(a_rad), h]) + trans
        #     self.positions.append(pos)
        #     self.rotations.append([p_angle, -angle, 0])

        # specify angles to rotate around
        bin_angle_size = 20.0
        num_views_per_bin = 2
        angles = np.arange(0, 360.0, bin_angle_size)
        angles = np.tile(angles, num_views_per_bin)

        #hori_rs = [10.0, 15.0, 20.0]

        # pick radii for each angle
        hori_rs = np.random.uniform(low=3.0, high=15.0, size=angles.shape[0])

        ind = 0
        for angle in angles: #np.arange(0.0, 359.0, 30.0):
            a_rad = angle/180.0*np.pi
            scale = 1.
            trans = np.zeros(3)
            hori_r = hori_rs[ind]
            #pos = self.center + scale*np.array([-hori_r*np.cos(a_rad), hori_r*np.sin(a_rad), h]) + trans
            pos = scale*np.array([-hori_r*np.cos(a_rad), hori_r*np.sin(a_rad), h]) + trans
            self.positions.append(pos)
            self.rotations.append([p_angle, -angle, 0])
            ind += 1

        self.K = np.identity(3)
        self.K[0, 2] = self.image_size_x / 2.0
        self.K[1, 2] = self.image_size_y / 2.0
        # use -focal to be the same as the previous version
        self.K[0, 0] = self.K[1, 1] = self.image_size_x / \
            (2.0 * np.tan(self.fov * np.pi / 360.0))

    def destroy_actors(self):
        print('Destroying actors.')
        for actor in self.vehicles_list:
            if actor.is_alive:
                actor.destroy()
        for actor in self.sensors_list:
            if actor.is_alive:
                actor.destroy()

        print("Destroyed all actors")

    def run_carla_client(self):
        # print(self.positions)
        # print(self.rotations)
        # print(self.start_episode, self.end_episode)
        client = carla.Client(self.host, self.port)
        client.set_timeout(20.0)

        try:
            self.available_maps = client.get_available_maps()
            print('loaded available maps')

            logging.info('listening to server %s:%s', self.host, self.port)

            num_batches = 2

            for episode in range(self.start_episode, self.end_episode):

                print("Starting episode number %d" % episode)

                uuid_run = str(uuid.uuid1())

                cur_map = random.choice(self.available_maps)
                print("About to load the map %s" %cur_map)
                map_name = cur_map.split('/')[-1]

                episode_path = os.path.join(self.base_path, self.episode_path_format.format(map_name, episode, self.num_vehicles))
                if not os.path.exists(episode_path):
                    os.mkdir(episode_path)

                self.world = client.load_world(cur_map)

                self.world.tick()

                # Initialize the actor lists
                self.vehicles_list = []
                # self.sensors_list = []
                # self.sensors_name = []
                self.vehicle_dir_paths = []
                self.actors = []
                # self.position_list = []
                # self.rotation_list = []
                self.camR_sensor_indices = []
                self.vehicle_bbox_list = []
                self.vehicle_extrinsics = []
                self.vehicle_names = []

                world_T_agents_f = open(os.path.join(episode_path, 'world_T_agent.txt'), 'w')
                world_T_agents_f.write('roll,pitch,yaw\n')
                fcsv = open(os.path.join(episode_path, 'bboxes.csv'), 'w')
                csv_writer = csv.writer(fcsv)
                csv_writer.writerow(['episode', 'frame', 'obj_id', 'category','x','y','z','l','w','h','r','x1','y1','x2','y2','depth', 'occluded'])

                # Get all the blueprints
                blueprints = self.world.get_blueprint_library().filter(self.filterv)

                # remove bikes
                cars = [x for x in blueprints if int(x.get_attribute('number_of_wheels')) == 4]
                blueprints = cars
                #for blueprint in blueprints:
                    #print(blueprint)

                # for blueprint in blueprints:
                #     print(blueprint.id)
                #     for attr in blueprint:
                #         print('  - {}'.format(attr))

                weather = carla.WeatherParameters(
                        cloudiness=np.random.randint(0, 70),
                        precipitation=np.random.randint(0, 75),
                        sun_altitude_angle=np.random.randint(30, 90))
                self.world.set_weather(weather)
                self.world.tick()

                # Get all the spawn points
                spawn_points = self.world.get_map().get_spawn_points()
                random.shuffle(spawn_points)

                self.spawn_vehicles(blueprints, episode_path, spawn_points)

                batches = list(np.floor(np.linspace(0, len(self.positions), num_batches)).astype(int))
                print("BATCH INDICES: ", batches)
                self.idx = 0 # for naming
                self.positions_ = [] # actual rotations of agents
                self.rotations_ = [] # actual rotations of agents
                self.first_batch = True
                for batch in range(len(batches)-1):

                    print("BATCH #", batches[batch], " to ", batches[batch+1])

                    self.sensors_list = []
                    self.sensors_name = []
                    self.position_list = []
                    self.rotation_list = []

                    # if batch == len(batches)-2:
                    #     positions_batch = self.positions[batches[batch]:]
                    #     rotations_batch = self.rotations[batches[batch]:]
                    # else:
                    positions_batch = self.positions[batches[batch]:batches[batch+1]]
                    rotations_batch = self.rotations[batches[batch]:batches[batch+1]]

                    print(positions_batch)
                    print(rotations_batch)
                    # print("SPAWN POINT: ", spawn_points[0])

                    self.spawn_sensors(positions_batch, rotations_batch, spawn_points[0], episode_path) # set the position relative to the first spawn point

                    print("Done with actor creation")
                    print("Total number of sensors are: ", len(self.sensors_list))

                    self.world.tick()

                    last_saved_frame = 0
                    # Create a synchronous mode context.
                    with CarlaSyncMode(self.world, self.sensors_list, fps=self.fps) as sync_mode:
                        cnt = 0
                        for v in self.vehicles_list:
                            # print("Bounding box for this vehicle is: ", v.bounding_box.location, v.bounding_box.extent)
                            bbox_loc, bbox_ext = v.bounding_box.location, v.bounding_box.extent
                            bbox = [bbox_loc.x - bbox_ext.x, bbox_loc.y - bbox_ext.y, bbox_loc.z - bbox_ext.z, bbox_loc.x + bbox_ext.x, bbox_loc.y + bbox_ext.y, bbox_loc.z + bbox_ext.z]
                            self.vehicle_bbox_list.append(bbox)
                            print("bbox coords are: ", bbox)
                        #     v.set_autopilot(False)
                            v.set_simulate_physics(False)

                        print("VEHICLES: ", self.vehicles_list)

                        # print("All vehicles put to autopilot")
                        # set weather

                        while True:
                            print(cnt)

                            if cnt == self.frames_per_episode + self.start_frame:
                                print("Done with episode %d." %episode)
                                time.sleep(5)
                                break

                            # self.world.tick()
                            # print("Getting the data")
                            # Advance the simulation and wait for the data.
                            data, frame = sync_mode.tick(timeout=12.0)
                            # print('Got the data :))')
                            if self.first_batch:
                                frame_a = frame
                            #frame += batches[batch]

                            print("FRAME ", frame_a)

                            if cnt >= self.start_frame:
                                data = data[1:] # Remove the world tick datapoint
                                valid = self.save_data(data, episode, frame_a, episode_path, world_T_agents_f, csv_writer)
                                if not valid:
                                    print("throwing out view ", cnt, " due to depth")
                            cnt += 1

                    for actor in self.sensors_list:
                        if actor.is_alive:
                            # print("ALIVE! ", actor)
                            actor.destroy()

                    self.first_batch = False

                # these are the translation relative to the first spawn point
                # print("ALL POSITIONS", self.positions_)
                # print("LENGTH ", len(self.positions_))
                # print("ALL ROTATIONS", self.rotations_)
                # print("LENGTH ", len(self.rotations_))
                save_npy(np.array(self.positions_), os.path.join(episode_path, 'all_cam_positions'))
                save_npy(np.array(self.rotations_), os.path.join(episode_path, 'all_cam_rotations'))

        except Exception as e: print(e)

        finally:
            #self.world.destroy()
            # for actor in self.world.get_actors().filter("vehicle*"):
            #     if actor.is_alive:
            #         #print("2", actor)
            #         #sys.stdout.write('2')
            #         actor.destroy()
            # for actor in self.sensors_list:
            #     if actor.is_alive:
            #         #print("5", actor)
            #         #sys.stdout.write('5')
            #         actor.destroy()
            client.apply_batch([carla.command.DestroyActor(x) for x in self.world.get_actors().filter("vehicle*")])
            client.apply_batch([carla.command.DestroyActor(x) for x in self.sensors_list])
        #     self.destroy_actors()

    def save_data(self, data, episode, framenum, episode_path, world_T_agents_f, csv_writer):
        # Check if too many occlusions (e.g. in a wall)
        # for idx, name in enumerate(self.sensors_name):
        #     if "Depth" in name:
        #         data_instance = data[idx]
        #         processed_data = image_converter.depth_in_meters(data_instance)
        #         valid = self.check_depth(processed_data)
        #         if valid == 0:
        #             return False

        # Save data
        #print("SENSOR LENGTH:", len(self.sensors_name))
        #print("SENSOR NAMES:", self.sensors_name)
        for idx, name in enumerate(self.sensors_name):
            camera_path = os.path.join(episode_path, name)
            if not os.path.exists(camera_path):
                os.mkdir(camera_path)
            if not os.path.exists(os.path.join(episode_path, 'boxes')):
                os.mkdir(os.path.join(episode_path, 'boxes'))
            data_instance = data[idx]
            if "RGB" in name:
                processed_data = image_converter.to_rgb_array(data_instance)
            elif "Depth" in name:
                processed_data = image_converter.depth_in_meters(data_instance)
                #valid = self.get_valid_view(processed_data)
                # print("DEPTH")
                # print(processed_data.shape)
                #print(processed_data)
            elif "Semantic" in name:
                # TODO: handle R channel properly here.
                print("SEMANTIC")
                # print(data_instance.shape)
                # print(data_instance)
                #print(np.unique(data_instance, axis=0))
                processed_data = image_converter.to_rgb_array(data_instance)
                print(processed_data.shape)
                print(np.unique(processed_data[:,:,0]))
                # data_instance.convert(cc.CityScapesPalette)
                # #data_instance.save_to_disk(os.path.join(self.visual_dir, 'img_e{:0>4d}_{:s}_f{:0>4d}.png'.format(0, name, framenum)),carla.ColorConverter.CityScapesPalette)
                # #processed_data = image_converter.to_rgb_array(data_instance)
                # print(processed_data.shape)
                # print(np.unique(processed_data[:,:,0]))

                # get segmentation
                sensor = self.sensors_list[idx]
                corners_cam, corners_cam_unproject, w2c, v2w = ClientSideBoundingBoxes.get_bounding_boxes(self.vehicles_list, sensor, self.K)
                c2w = np.linalg.inv(w2c)
                v2w = np.concatenate(v2w, axis=0)

                masks = []
                for obj_ind in range(len(corners_cam)):

                    # get 2D bounding box
                    xmin = int(np.min(corners_cam[obj_ind][:,0]))
                    xmax = int(np.max(corners_cam[obj_ind][:,0]))
                    ymin = int(np.min(corners_cam[obj_ind][:,1]))
                    ymax = int(np.max(corners_cam[obj_ind][:,1]))
                    if ymin < 0:
                        ymin = 0
                    elif ymax > self.image_size_y:
                        ymax = self.image_size_y
                    if xmin < 0:
                        xmin = 0
                    elif xmax > self.image_size_x:
                        xmax = self.image_size_x

                    veh_id = 10 # vehicle segmentation ID for semantic sensor

                    # get segmentation mask
                    img = processed_data.copy()
                    img2 = processed_data.copy()
                    img2 = img2[ymin:ymax, xmin:xmax, :]
                    #print(np.unique(img2[:,:,0]))
                    img2[img2!=veh_id] = 0.0
                    #print(np.sum(img2[:,:,:]==10))
                    img2[img2[:,:,0]==veh_id, :] = 1.0
                    #print(np.sum(img2==255))
                    img[:,:,:] = 0.0
                    img[ymin:ymax, xmin:xmax, :] = img2

                    if np.sum(img) != 0:
                        masks.append(img)

                    if self.visualize:
                        img = img * 255 # make mask white
                        img = Image.fromarray(img)
                        #img = ClientSideBoundingBoxes.draw_bounding_boxes(img, corners_cam)
                        #data_instance.save_to_disk(os.path.join(self.visual_dir, 'img_e{:0>4d}_{:s}_f{:0>4d}.png'.format(0, name, framenum)),carla.ColorConverter.CityScapesPalette)
                        img.save(os.path.join(self.visual_dir, 'img_e{:0>4d}_{:s}_f{:0>4d}_obj{}.png'.format(0, name, framenum, obj_ind)))

                masks = np.array(masks)

                processed_data = masks

                print("MASKS", masks.shape)

            else:
                print("Invalid sensor type %s. Quitting" %j)
                exit(1)
            save_npy(processed_data, os.path.join(camera_path, '{:0>6d}'.format(framenum)))

            if 'RGB' in name:#frame % 1 == 0:
                #output non-player agents
                img = processed_data.copy()
                sensor = self.sensors_list[idx]
                corners_cam, corners_cam_unproject, w2c, v2w = ClientSideBoundingBoxes.get_bounding_boxes(self.vehicles_list, sensor, self.K)
                c2w = np.linalg.inv(w2c)
                v2w = np.concatenate(v2w, axis=0)

                sensor_to_world = ClientSideBoundingBoxes.get_matrix(sensor.get_transform())
                save_npy(c2w, os.path.join(camera_path, '{:0>6d}_c2w'.format(framenum)))

                if self.visualize:
                    img = Image.fromarray(processed_data)
                    img = ClientSideBoundingBoxes.draw_bounding_boxes(img, corners_cam)
                    img.save(os.path.join(self.visual_dir, 'img_e{:0>4d}_{:s}_f{:0>4d}.png'.format(0, name, framenum)))
                if 'RGB0' in name:
                    # boxes = np.array(corners_cam_unproject, dtype=np.float32)
                    save_npy(v2w, os.path.join(episode_path, 'boxes', '{:0>6d}'.format(framenum)))
            # if 'Semantic' in name:#frame % 1 == 0:

            #     if self.visualize:
            #         # print("1", corners_cam)
            #         # print(corners_cam[0].shape)
            #         # print("2", corners_cam_unproject)
            #         # print("3", w2c)
            #         # print("4", v2w)
            #         xmin = int(np.min(corners_cam[0][:,0]))
            #         xmax = int(np.max(corners_cam[0][:,0]))
            #         ymin = int(np.min(corners_cam[0][:,1]))
            #         ymax = int(np.max(corners_cam[0][:,1]))
            #         if ymin < 0:
            #             ymin = 0
            #         elif ymax > self.image_size_y:
            #             ymax = self.image_size_y
            #         if xmin < 0:
            #             xmin = 0
            #         elif xmax > self.image_size_x:
            #             xmax = self.image_size_x
            #         print(xmin)
            #         print(xmax)
            #         print(ymin)
            #         print(ymax)
            #         print(processed_data.flags)
            #         img = processed_data.copy()
            #         img2 = processed_data.copy()
            #         img2 = img2[ymin:ymax, xmin:xmax, :]
            #         print(np.unique(img2[:,:,0]))
            #         img2[img2!=10] = 0.0
            #         print(np.sum(img2[:,:,:]==10))
            #         img2[img2[:,:,0]==10, :] = 255
            #         print(np.sum(img2==255))
            #         #processed_data = image_converter.to_rgb_array(data_instance)
            #         #img = np.zeros(img2.shape)
            #         img[:,:,:] = 0.0
            #         img[ymin:ymax, xmin:xmax, :] = img2
            #         img = Image.fromarray(img)
            #         img.save(os.path.join(self.visual_dir, 'img_e{:0>4d}_{:s}_f{:0>4d}.png'.format(0, name, framenum)))
        return True

    def get_valid_view(self, data):
        for idx, name in enumerate(self.sensors_name):
            if "Depth" in name:
                data_instance = data[idx]
                processed_data = image_converter.depth_in_meters(data_instance)
                valid = self.check_depth(processed_data)
                if valid == 0:
                    return False
        return True

    def check_depth(self, depth):
        #all_ok = True
        # num_views = len(depths)

        # # determine how many objects present in each view
        # obj_in_views = np.zeros(num_views)
        # for n in list(range(MAX_OBJS)):
        #         # print('ok, working on object %d' % n)
        #     score = scorelist_s[0][n]
        #     tid = tidlist_s[0][n]
        #     lrt = lrtlist_camRs[0][n]
        #     if score == 1.0:
        #         vis_multiview = get_target_multiview_vis(tid, lrt, origin_T_camXs, origin_T_camRs[0], xyz_camXs)
        #         obj_in_views = np.add(vis_multiview, obj_in_views)
        #         # print(vis_multiview)
        #         #print(obj_in_views)
        #         # print(S)

        # # print(obj_in_views)
        # # print(obj_in_views.shape)

        #valid_views = []
        # num_bad = 0
        # for view in range(num_views):
        # mindepth = np.min(depth)
        #depth = depths[view]
        not_ok = depth < 2.0
        yes_ok = depth > 2.0
        A = np.sum(yes_ok)
        B = np.sum(not_ok)

        #print(scorelist_s[view])
        if A < B*3: # or (obj_in_views[view]==0): # view does not have space in front of it or no objects throw out
        #if (obj_in_views[view]==0): # view does not have space in front of it or no objects throw out
            print("NOT VALID")
            print('yes = %d; no = %d;' % (A, B))
            #print(obj_in_views[view])
            print("min depth: ", np.min(depth))
            #valid_views.append([0])
            #num_bad += 1
            valid_view = 0

            #if obj_in_views[view]==0:
                #print("No objects found in view ", view)

        else: # ok view
            #valid_views.append([1])
            print("VALID")
            # print("min depth: ", np.min(depth))
            print('yes = %d; no = %d;' % (A, B))
        #print("Number of invalid views: ", num_bad)
            valid_view = 1
        #valid_views = np.array(valid_views).astype(np.float32)
        #valid_views = valid_views.reshape((len(valid_views)))
        return valid_view

    def pts_dist(self, locsi, locsj):
        dist = (locsi.x - locsj.x)**2 + (locsi.y - locsj.y)**2 + (locsi.z - locsj.z)**2
        return np.sqrt(dist)
    '''
    Spawns the vehicles and attaches sensors to it.
    '''
    def spawn_vehicles(self, blueprints, episode_path, spawn_points):
        #print(spawn_points)
        # print(spawn_points[:min(self.num_vehicles, len(spawn_points))])
        # print(spawn_points[0])
        #for spawn_point in spawn_points: #spawn_points[:min(self.num_vehicles, len(spawn_points))]:
            # setup blueprint

        # spawn first vehicle at center of rotation
        spawn_point = spawn_points[0]
        blueprint = random.choice(blueprints)
        if blueprint.has_attribute('driver_id'):
            driver_id = random.choice(blueprint.get_attribute('driver_id').recommended_values)
            blueprint.set_attribute('driver_id', driver_id)
        color = random.choice(blueprint.get_attribute('color').recommended_values)
        blueprint.set_attribute('color', color)
        blueprint.set_attribute('role_name', 'autopilot')

        # position = carla.Location(x=0, y=0, z=0) + spawn_point.location
        # transform = carla.Transform(position, carla.Rotation(pitch=0, yaw=0, roll=0))
        vehicle = self.world.try_spawn_actor(blueprint, spawn_point)
        #vehicle = self.world.try_spawn_actor(blueprint, transform)
        self.vehicle_names.append(blueprint.id)
        self.actors.append(vehicle)
        self.vehicles_list.append(vehicle)

        # Spawn second vehicle closest to first vehicle
        #dist_best = 100000000
        spts_closest = []
        center_loc = spawn_point.location
        for spts in spawn_points:
            loc = spts.location
            #print(loc)
            #print(loc.shape)
            #print(loc.x)
            #print(loc.y)
            dist = np.sqrt((loc.x - center_loc.x)**2 + (loc.y - center_loc.y)**2 + (loc.z - center_loc.z)**2)
            spts_closest.append(dist)
            # print(dist)
            # if dist < dist_best:
            #     dist_best = dist
            #     spts_closest = spts_closest
            #     print(dist_best)
            #     print(spts.location)
            #     print(spawn_point)
        spts_closest = np.array(spts_closest)
        spts_closest_inds = np.argsort(spts_closest)
        spts_closest_sort = np.sort(spts_closest)
        spts_closest_inds = spts_closest_inds[1:] # remove first object
        spts_closest_sort = spts_closest_sort[1:] # remove first object

        dist_threshold = 5 # must be atleast this away from first object
        while True:
            if spts_closest_sort[0] < dist_threshold:
                spts_closest_sort = spts_closest_sort[1:]
                spts_closest_inds = spts_closest_inds[1:]
            else:
                break

        # print(spts_closest)
        print(spts_closest_inds)
        print(spts_closest_sort)

        if self.num_vehicles==1:
            return

        for i in range(self.num_vehicles-1):
            blueprint = random.choice(blueprints)
            if blueprint.has_attribute('driver_id'):
                driver_id = random.choice(blueprint.get_attribute('driver_id').recommended_values)
                blueprint.set_attribute('driver_id', driver_id)
            color = random.choice(blueprint.get_attribute('color').recommended_values)
            blueprint.set_attribute('color', color)
            blueprint.set_attribute('role_name', 'autopilot')
            # offset = np.array([self.nearest_x+self.dis_to_center, 0.0, 0.0]) # -1.0+10*np.random.rand(*(np.array(positions).shape))
            # position = carla.Location(x=offset[0], y=offset[1], z=offset[2]) + spawn_point.location
            # trans = -1.0+2*np.random.rand(3)
            # transform = carla.Transform(position, carla.Rotation(pitch=trans[0], yaw=trans[1], roll=trans[2]))
            spawn_point_close = spawn_points[spts_closest_inds[i]]
            vehicle = self.world.try_spawn_actor(blueprint, spawn_point_close)
            self.vehicle_names.append(blueprint.id)
            self.actors.append(vehicle)
            self.vehicles_list.append(vehicle)

        print('finish spawning {} vehicles'.format(len(self.vehicles_list)))

    def spawn_sensors(self, positions, rotations, spawn_point, episode_path):

        #camera_indices = [0] + camera_indices[0:3] + [len(positions) - 1]
        if self.first_batch:
            camera_indices = list(range(1, len(positions) - 1))
            random.shuffle(camera_indices)
            camera_indices = [0] + camera_indices
        else:
            camera_indices = list(range(0, len(positions) - 1))
            random.shuffle(camera_indices)
        # camera_indices = [0] # for traj data, we only need one camera

        #trans = -1.0+2*np.random.rand(*(np.array(positions).shape))
        positions_ = positions #+ trans
        self.positions_.extend(positions_) # for saving out positions
        #trans = -1.0+2*np.random.rand(*(np.array(positions).shape))
        rotations_ = rotations #+ trans
        self.rotations_.extend(rotations_) # for saving out rotations

        # # these are the translation relative to the first spawn point
        # save_npy(np.array(positions_), os.path.join(episode_path, 'all_cam_positions'))
        # save_npy(np.array(rotations_), os.path.join(episode_path, 'all_cam_rotations'))

        nsp = 0
        for idx in camera_indices:
            position, rotation = positions_[idx], rotations_[idx]
            for sensor_type, sensor_name in zip(self.sensor_types, self.sensor_names_format):
                print("DOING ", sensor_name.format(self.idx))
                sensor = self.get_sensor(spawn_point, position, rotation, sensor_type, sensor_name.format(idx))
                self.sensors_list.append(sensor)
                self.sensors_name.append(sensor_name.format(self.idx))
                #time.sleep(1)
                print("FINISHED ", nsp)
                nsp += 1
            self.idx += 1

        print('finish spawning {} sensors'.format(len(self.sensors_list)))

    def get_sensor(self, relative_to, position, rotation, sensor_type, sensor_id):
        blueprint = self.world.get_blueprint_library().find(sensor_type)

        blueprint.set_attribute('image_size_x', str(self.image_size_x))
        blueprint.set_attribute('image_size_y', str(self.image_size_y))
        blueprint.set_attribute('fov', str(self.fov))
        blueprint.set_attribute('role_name', sensor_id)

        # Set the time in seconds between sensor captures
        blueprint.set_attribute('sensor_tick', '0.0')
        # print("RELATIVE TO: ", relative_to.location)
        # print("POSITON: ", position)

        # Provide the position of the sensor relative to the vehicle.
        position = carla.Location(x=position[0], y=position[1], z=position[2]) + relative_to.location
        transform = carla.Transform(position, carla.Rotation(pitch=rotation[0], yaw=rotation[1], roll=rotation[2]))

        self.position_list.append(position)
        self.rotation_list.append(rotation)
        # use the commented line to attach cameras to the vehicle actor.
        # sensor = self.world.spawn_actor(blueprint, transform, attach_to=vehicle)
        # fixed cameras
        sensor = self.world.spawn_actor(blueprint, transform)

        return sensor

    def initiate(self):
        if not os.path.exists(self.base_path):
            os.mkdir(self.base_path)

        while True:
            # try:
            self.run_carla_client()
            return

            # except Exception as error:
            #     print("Error occured while interacting with carla.")
            #     print(error)
            #     logging.error(error)
            #     time.sleep(1)

if __name__ == '__main__':
    episode = int(sys.argv[1])
    mod = sys.argv[2]
    save_dir = sys.argv[3]
    print(mod)
    print(save_dir)
    cmr = CarlaMultiviewRunner(episode, episode + 1, mod, save_dir)
    cmr.initiate()


import logging; log = logging.getLogger(__name__)
from importlib import resources
import struct
import math
import sys
import gl
import numpy as np
from .SfaProgram import SfaProgram
from . import shaders

class LineShader(SfaProgram):
    separable       = True
    vertex_shader   = (shaders, 'line.vert')
    geometry_shader = (shaders, 'line.geom')

class LineRenderer(gl.Pipeline):
    """Subprogram for rendering thick lines."""
    MAX_LINES = 1024
    vtxBufferFmt = "6f" # (x,y,z) head, (x,y,z) tail
    colorBufferFmt = "8f" # (r,g,b,a) per line (head, tail)

    def __init__(self, parent):
        self.parent = parent
        self.shader = LineShader(self.parent.ctx)
        self.lineWidth = 8
        self._nLines = 0

        super().__init__(parent.ctx,
            vertex_shader   = self.shader,
            geometry_shader = self.shader,
            fragment_shader = parent.fragShader,
        )

        self._vtxSize   = struct.calcsize(self.vtxBufferFmt)
        self._colorSize = struct.calcsize(self.colorBufferFmt)
        self.bufVtxs    = self.ctx.Buffer(self._vtxSize   * self.MAX_LINES, True)
        self.bufColors  = self.ctx.Buffer(self._colorSize * self.MAX_LINES, True)
        self.useDepthTest = True

    def addLine(self, pointA, pointB, colorA, colorB):
        """Add a line from pointA to pointB."""
        if self._nLines >= self.MAX_LINES: return

        x1, y1, z1 = pointA
        x2, y2, z2 = pointB
        r1, g1, b1, a1 = colorA
        r2, g2, b2, a2 = colorB

        vtxData = struct.pack(self.vtxBufferFmt, x1, y1, z1, x2, y2, z2)
        colData = struct.pack(self.colorBufferFmt, r1, g1, b1, a1, r2, g2, b2, a2)
        vtxOffs = self._nLines * self._vtxSize
        colOffs = self._nLines * self._colorSize
        self.bufVtxs.write(vtxData, vtxOffs)
        self.bufColors.write(colData, colOffs)
        self._nLines += 1

    def reset(self):
        self._nLines = 0

    def run(self):
        if self._nLines == 0: return
        self.ctx.glBlendFunc(self.ctx.GL_SRC_ALPHA, self.ctx.GL_ONE_MINUS_SRC_ALPHA)
        self.ctx.glEnable(self.ctx.GL_BLEND)

        if self.useDepthTest: self.ctx.glEnable(self.ctx.GL_DEPTH_TEST)
        else: self.ctx.glDisable(self.ctx.GL_DEPTH_TEST)
        self.ctx.glDisable(self.ctx.GL_CULL_FACE)
        self.ctx.glLineWidth(self.lineWidth)

        with self:
            self._bindBuffers()
            self.shader.vao.render(self.ctx.GL_POINTS, count=self._nLines)

    def setMtxs(self, projection, modelView):
        self.shader.setMtxs(projection, modelView)

    def _bindBuffers(self):
        with self.shader as shader:
            # in vec3 head
            shader.attribs['head'].bindBuffer(
                self.bufVtxs, self.ctx.GL_FLOAT, 3,
                stride=6 * gl.Util.SIZEOF_FLOAT,
                offset=0 * gl.Util.SIZEOF_FLOAT)

            # in vec3 tail
            shader.attribs['tail'].bindBuffer(
                self.bufVtxs, self.ctx.GL_FLOAT, 3,
                stride=6 * gl.Util.SIZEOF_FLOAT,
                offset=3 * gl.Util.SIZEOF_FLOAT)

            # in vec4 colHead
            shader.attribs['colHead'].bindBuffer(
                self.bufColors, self.ctx.GL_FLOAT, 4,
                stride=8 * gl.Util.SIZEOF_FLOAT,
                offset=0 * gl.Util.SIZEOF_FLOAT)

            # in vec4 colTail
            shader.attribs['colTail'].bindBuffer(
                self.bufColors, self.ctx.GL_FLOAT, 4,
                stride=8 * gl.Util.SIZEOF_FLOAT,
                offset=4 * gl.Util.SIZEOF_FLOAT)


import pickle
import logging
import numpy as np
import pandas as pd
from collections import Counter
from typing import Union, Sequence, Optional, List, Any
from sklearn.base import BaseEstimator
from indra.statements import Evidence, Statement, get_all_descendants
from indra.belief import BeliefScorer, check_extra_evidence, \
                         get_stmt_evidence, SimpleScorer

logger = logging.getLogger(__name__)

class SklearnScorer(BeliefScorer):
    """Use a pre-trained Sklearn classifier to predict belief scores.

    An implementing instance of this base class has two personalities: as a
    subclass of BeliefScorer, it implements the functions required by the
    BeliefEngine, `score_statements` and `check_prior_probs`. It also behaves
    like an sklearn model by composition, implementing methods `fit`,
    `predict`, `predict_proba`, and `predict_log_proba`, which are passed
    through to an internal sklearn model.

    A key role of this wrapper class is to implement the preprocessing of
    statement properties into a feature matrix in a standard way, so that
    a classifier trained on one corpus of statement data will still work when
    used on another corpus.

    Implementing subclasses must implement at least one of the methods for
    building the feature matrix, `stmts_to_matrix` or `df_to_matrix`.

    Parameters
    ----------
    model :
        Any instance of a classifier object supporting the methods `fit`,
        `predict_proba`, `predict`, and `predict_log_proba`.
    """
    def __init__(
        self,
        model: BaseEstimator,
    ):
        self.model = model

    def check_prior_probs(
        self,
        statements: Sequence[Statement],
    ) -> None:
        """Empty implementation for now."""
        pass

    def score_statements(
        self,
        statements: Sequence[Statement],
        extra_evidence: Optional[List[List[Evidence]]] = None,
    ) -> Sequence[float]:
        return self.predict_proba(statements, extra_evidence)[:, 1]

    def stmts_to_matrix(
        self,
        stmts: Sequence[Statement],
        extra_evidence: Optional[List[List[Evidence]]] = None,
    ) -> np.ndarray:
        """Convert a list of Statements to a feature matrix."""
        raise NotImplementedError('Need to implement the stmts_to_matrix '
                                  'method')

    def df_to_matrix(
        self,
        df: pd.DataFrame,
    ) -> np.ndarray:
        """Convert a statement DataFrame to a feature matrix."""
        raise NotImplementedError('Need to implement the df_to_matrix '
                                   'method')

    def to_matrix(self,
        stmt_data: Union[np.ndarray, Sequence[Statement], pd.DataFrame],
        extra_evidence: Optional[List[List[Evidence]]] = None,
    ) -> np.ndarray:
        """Get stmt feature matrix by calling appropriate method.

        If `stmt_data` is already a matrix (e.g., obtained after performing a
        train/test split on a matrix generated for a full statement corpus), it
        is returned directly; if a DataFrame of Statement metadata,
        `self.df_to_matrix` is called; if a list of Statements,
        `self.stmts_to_matrix` is called.

        Parameters
        ----------
        stmt_data :
            Statement content to be used to generate a feature matrix.
        extra_evidence :
            A list corresponding to the given list of statements, where
            each entry is a list of Evidence objects providing additional
            support for the corresponding statement (i.e., Evidences that
            aren't already included in the Statement's own evidence list).

        Returns
        -------
        :
            Feature matrix for the statement data.
        """
        # If we got a Numpy array, just use it!
        if isinstance(stmt_data, np.ndarray):
            stmt_arr = stmt_data
        # Otherwise check if we have a dataframe or a list of statements
        # and call the appropriate *_to_matrix method
        elif isinstance(stmt_data, pd.DataFrame):
            if extra_evidence is not None:
                raise NotImplementedError(
                   'extra_evidence cannot be used with a statement DataFrame.')
            stmt_arr = self.df_to_matrix(stmt_data)
        # Check if stmt_data is a list/tuple (i.e., of Statements):
        elif isinstance(stmt_data, (list, tuple)):
            # Check that the first entry is a Statement
            if not isinstance(stmt_data[0], Statement):
                raise ValueError('stmt_data must contain Statements.')
            stmt_arr = self.stmts_to_matrix(stmt_data, extra_evidence)
        # If it's something else, error
        else:
            raise TypeError(f'stmt_data is type {type(stmt_data)}: '
                            'must be a numpy array, DataFrame, or '
                            'list/tuple of Statements')
        return stmt_arr

    def fit(self,
        stmt_data: Union[np.ndarray, Sequence[Statement], pd.DataFrame],
        y_arr: Sequence[float],
        extra_evidence: Optional[List[List[Evidence]]] = None,
        *args,
        **kwargs,
    ):
        """Preprocess stmt data and run sklearn model `fit` method.

        Additional `args` and `kwargs` are passed to the `fit` method of the
        wrapped sklearn model.

        Parameters
        ----------
        stmt_data :
            Statement content to be used to generate a feature matrix.
        y_arr :
            Class values for the statements (e.g., a vector of 0s and 1s
            indicating correct or incorrect).
        extra_evidence :
            A list corresponding to the given list of statements, where
            each entry is a list of Evidence objects providing additional
            support for the corresponding statement (i.e., Evidences that
            aren't already included in the Statement's own evidence list).
        """
        # Check dimensions of stmts (x) and y_arr
        if len(stmt_data) != len(y_arr):
            raise ValueError("Number of stmts/rows must match length of y_arr.")
        # Get the data matrix based on the stmt list or stmt DataFrame
        stmt_arr = self.to_matrix(stmt_data, extra_evidence)
        # Call the fit method of the internal sklearn model
        self.model.fit(stmt_arr, y_arr, *args, **kwargs)

    def predict_proba(
        self,
        stmt_data: Union[np.ndarray, Sequence[Statement], pd.DataFrame],
        extra_evidence: Optional[List[List[Evidence]]] = None,
        *args,
        **kwargs,
    ) -> np.ndarray:
        """Preprocess stmt data and run sklearn model `predict_proba` method.

        Additional `args` and `kwargs` are passed to the `predict_proba` method
        of the wrapped sklearn model.

        Parameters
        ----------
        stmt_data :
            Statement content to be used to generate a feature matrix.
        extra_evidence :
            A list corresponding to the given list of statements, where
            each entry is a list of Evidence objects providing additional
            support for the corresponding statement (i.e., Evidences that
            aren't already included in the Statement's own evidence list).
        """
        # Call the prediction method of the internal sklearn model
        stmt_arr = self.to_matrix(stmt_data, extra_evidence)
        return self.model.predict_proba(stmt_arr, *args, **kwargs)

    def predict(
        self,
        stmt_data: Union[np.ndarray, Sequence[Statement], pd.DataFrame],
        extra_evidence: Optional[List[List[Evidence]]] = None,
        *args,
        **kwargs,
    ) -> np.ndarray:
        """Preprocess stmt data and run sklearn model `predict` method.

        Additional `args` and `kwargs` are passed to the `predict` method of
        the wrapped sklearn model.

        Parameters
        ----------
        stmt_data :
            Statement content to be used to generate a feature matrix.
        extra_evidence :
            A list corresponding to the given list of statements, where
            each entry is a list of Evidence objects providing additional
            support for the corresponding statement (i.e., Evidences that
            aren't already included in the Statement's own evidence list).
        """
        stmt_arr = self.to_matrix(stmt_data, extra_evidence)
        return self.model.predict(stmt_arr, *args, **kwargs)

    def predict_log_proba(
        self,
        stmt_data: Union[np.ndarray, Sequence[Statement], pd.DataFrame],
        extra_evidence: Optional[List[List[Evidence]]] = None,
        *args,
        **kwargs,
    ) -> np.ndarray:
        """Preprocess stmt data and run sklearn model `predict_log_proba`.

        Additional `args` and `kwargs` are passed to the `predict` method of
        the wrapped sklearn model.

        Parameters
        ----------
        stmt_data :
            Statement content to be used to generate a feature matrix.
        extra_evidence :
            A list corresponding to the given list of statements, where
            each entry is a list of Evidence objects providing additional
            support for the corresponding statement (i.e., Evidences that
            aren't already included in the Statement's own evidence list).
        """
        stmt_arr = self.to_matrix(stmt_data, extra_evidence)
        return self.model.predict_log_proba(stmt_arr, *args, **kwargs)

class CountsScorer(SklearnScorer):
    """Belief model learned from evidence counts and other stmt properties.

    If using a DataFrame for Statement data, it should have the following
    columns:

    * `stmt_type`
    * `source_counts`

    Alternatively, if the DataFrame doesn't have a `source_counts` column, it
    should have columns with names matching the sources in `self.source_list`.

    Parameters
    ----------
    model :
        Any instance of a classifier object supporting the methods `fit`,
        `predict_proba`, `predict`, and `predict_log_proba`.
    source_list :
        List of strings denoting the evidence sources (evidence.source_api
        values) to be used for prediction.
    include_more_specific :
        If True, will add extra columns to the statement data matrix for the
        source counts drawn from more specific evidences; if use_num_pmids is
        True, will also add an additional column for the number of PMIDs from
        more specific evidences. If False, these columns will not be included
        even if the `extra_evidence` argument is passed to the
        `stmts_to_matrix` method. This is to ensure that the featurization of
        statements is consistent between training and prediction.
    use_stmt_type :
        Whether to include statement type as a feature.
    use_num_members :
        Whether to include a feature denoting the number of members of the
        statement. Primarily for stratifying belief predictions about Complex
        statements with more than two members. Cannot be used for statement
        data passed in as a DataFrame.
    use_num_pmids :
        Whether to include a feature for the total number of unique PMIDs
        supporting each statement. Cannot be used for statement passed in as a
        DataFrame.
    use_promoter :
        Whether to include a feature giving the fraction of evidence (0 to 1)
        containing the (case-insensitive) word "promoter". Tends to improve
        misclassification of Complex statements that actually refer to
        protein-DNA binding.
    use_avg_evidence_len :
        Whether to include a feature giving the average evidence sentence
        length (in space-separated tokens).

    Example
    -------
    .. code-block:: python

        from sklearn.linear_model import LogisticRegression
        clf = LogisticRegression()
        all_stmt_sources = CountsScorer.get_all_sources(stmts)
        scorer = CountsScorer(clf, all_stmt_sources, use_stmt_type=True,
                              use_num_pmids=True)
        scorer.fit(stmts, y_arr)
        be = BeliefEngine(scorer)
        be.set_hierarchy_probs(stmts)
    """
    def __init__(
        self,
        model: BaseEstimator,
        source_list: List[str],
        include_more_specific: bool = False,
        use_stmt_type: bool = False,
        use_num_members: bool = False,
        use_num_pmids: bool = False,
        use_promoter: bool = False,
        use_avg_evidence_len: bool = False,
    ):
        # Call superclass constructor to store the model
        super(CountsScorer, self).__init__(model)
        self.source_list = source_list
        self.include_more_specific = include_more_specific
        self.use_stmt_type = use_stmt_type
        self.use_num_members = use_num_members
        self.use_num_pmids = use_num_pmids
        self.use_promoter = use_promoter
        self.use_avg_evidence_len = use_avg_evidence_len
        # Build dictionary mapping INDRA Statement types to integers
        if use_stmt_type:
            all_stmt_types = get_all_descendants(Statement)
            self.stmt_type_map = {t.__name__: ix
                                  for ix, t in enumerate(all_stmt_types)}

    @staticmethod
    def get_all_sources(
        stmts: Sequence[Statement],
        include_more_specific: bool = True,
        include_less_specific: bool = True,
    ) -> List[str]:
        """Get a list of all the source_apis supporting the given statements.

        Useful for determining the set of sources to be used for fitting
        and prediction.

        Parameters
        ----------
        stmts :
            A list of INDRA Statements to collect source APIs for.
        include_more_specific :
            If True (default), then includes the source APIs for the more
            specific statements in the `supports` attribute of each statement.
        include_less_specific :
            If True (default), then includes the source APIs for the less
            specific statements in the `supported_by` attribute of each
            statement.

        Returns
        -------
        :
            A list of (unique) source_apis found in the set of statements.
        """
        stmt_sources = set([ev.source_api for s in stmts for ev in s.evidence])
        if include_more_specific:
            stmt_sources.update([ev.source_api
                                 for stmt in stmts
                                 for supp_stmt in stmt.supports
                                 for ev in supp_stmt.evidence])
        if include_less_specific:
            stmt_sources.update([ev.source_api
                                 for stmt in stmts
                                 for supp_by_stmt in stmt.supported_by
                                 for ev in supp_by_stmt.evidence])
        return list(stmt_sources)

    def stmts_to_matrix(
        self,
        stmts: Sequence[Statement],
        extra_evidence: Optional[List[List[Evidence]]] = None,
    ) -> np.ndarray:
        """Convert a list of Statements to a feature matrix.

        Features are encoded as follows:

        * One column for every source listed in `self.source_list`, containing
          the number of statement evidences from that source. If
          `self.include_more_specific` is True and `extra_evidence` is
          provided, these are used in combination with the Statement's own
          evidence in determining source counts.
        * If `self.use_stmt_type` is set, statement type is included via
          one-hot encoding, with one column for each statement type.
        * If `self.use_num_members` is set, a column is added for the number
          of agents in the Statement.
        * If `self.use_num_pmids` is set, a column is added with the total
          total number of unique PMIDs supporting the Statement.  If
          `extra_evidence` is provided, these are used in combination with the
          Statement's own evidence in determining the number of PMIDs.

        Parameters
        ----------
        stmts :
            A list or tuple of INDRA Statements to be used to generate a
            feature matrix.
        extra_evidence :
            A list corresponding to the given list of statements, where
            each entry is a list of Evidence objects providing additional
            support for the corresponding statement (i.e., Evidences that
            aren't already included in the Statement's own evidence list).

        Returns
        -------
        :
            Feature matrix for the statement data.
        """
        # Check arguments for including more specific evidences
        if self.include_more_specific and extra_evidence is None:
            logger.info("CountScorer is set to include_more_specific "
                        "evidences but no extra_evidence was included.")
            extra_evidence = [[] for stmt in stmts]
        elif not self.include_more_specific and extra_evidence is not None:
            logger.warning("extra_evidence was included but CountScorer "
                           "instance is not set to include_more_specific "
                           "evidences so extra_evidence will be ignored.")
        # Check our list of extra evidences
        check_extra_evidence(extra_evidence, len(stmts))
        # Add categorical features and collect source_apis
        cat_features = []
        stmt_sources = set()
        for ix, stmt in enumerate(stmts):
            # Collect all source_apis from stmt evidences
            dir_pmids = set()
            promoter_ct = 0
            evidence_lens = []
            for ev in stmt.evidence:
                stmt_sources.add(ev.source_api)
                dir_pmids.add(ev.pmid)
                if ev.text is not None:
                    evidence_lens.append(len(ev.text.split()))
                    if 'promoter' in ev.text.lower():
                        promoter_ct += 1

            indir_pmids = set()
            if self.include_more_specific and extra_evidence:
                for ev in extra_evidence[ix]:
                    stmt_sources.add(ev.source_api)
                    indir_pmids.add(ev.pmid)
            # Collect non-source count features (e.g. type) from stmts
            feature_row: List[Any] = [] # Appease the Type Hint Gods
            # One-hot encoding of stmt type
            if self.use_stmt_type:
                stmt_type_ix = self.stmt_type_map[type(stmt).__name__]
                type_features = [1 if ix == stmt_type_ix else 0
                                 for ix in range(len(self.stmt_type_map))]
                feature_row.extend(type_features)
            # Add field for number of members
            if self.use_num_members:
                feature_row.append(len(stmt.agent_list()))
            # Add field with number of unique PMIDs
            if self.use_num_pmids:
                feature_row.append(len(dir_pmids))
                if self.include_more_specific and extra_evidence:
                    feature_row.append(len(indir_pmids))
            # Add a field specifying the percentage of evidences containing
            # the word "promoter":
            if self.use_promoter:
                promoter_pct = promoter_ct / len(stmt.evidence) \
                                        if len(stmt.evidence) > 0 else 0
                feature_row.append(promoter_pct)
            # Add a field giving length of the sentence in words
            if self.use_avg_evidence_len:
                avg_evidence_len = np.mean(evidence_lens) if evidence_lens \
                                                          else 0
                feature_row.append(avg_evidence_len)
            # Only add a feature row if we're using some of the features.
            if feature_row:
                cat_features.append(feature_row)

        # Before proceeding, check whether all source_apis are in
        # source_list
        if stmt_sources.difference(set(self.source_list)):
            logger.info("source_list does not include all source_apis "
                             "in the statement data.")
        # Get source count features
        # If we have extra_evidence, we double the source count features
        if self.include_more_specific:
            num_cols = len(self.source_list) * 2
        else:
            num_cols = len(self.source_list)
        num_rows = len(stmts)
        x_arr = np.zeros((num_rows, num_cols))
        for stmt_ix, stmt in enumerate(stmts):
            # Source from the stmt itself
            direct_sources = [ev.source_api for ev in stmt.evidence]
            dsrc_ctr = Counter(direct_sources)
            for src_ix, src in enumerate(self.source_list):
                x_arr[stmt_ix, src_ix] = dsrc_ctr.get(src, 0)
            # Get indirect evidences
            if self.include_more_specific and extra_evidence:
                indirect_sources = [ev.source_api
                                    for ev in extra_evidence[stmt_ix]]
                idsrc_ctr = Counter(indirect_sources)
                for src_ix, src in enumerate(self.source_list):
                    x_arr[stmt_ix, src_ix + len(self.source_list)] = \
                                            idsrc_ctr.get(src, 0)
        # If we have any categorical features, turn them into an array and
        # add them to matrix
        if cat_features:
            cat_arr = np.array(cat_features)
            x_arr = np.hstack((x_arr, cat_arr))
        return x_arr

    def df_to_matrix(
        self,
        df: pd.DataFrame,
    ) -> np.ndarray:
        """Convert a DataFrame of statement data to a feature matrix.

        Based on information available in a DataFrame of statement data, this
        implementation uses only source counts and statement type in building a
        feature matrix, and will raise a ValueError if either
        `self.use_num_members` or `self.use_num_pmids` is set.

        Features are encoded as follows:

        * One column for every source listed in `self.source_list`, containing
          the number of statement evidences from that source. If
          `extra_evidence` is provided, these are used in combination with the
          Statement's own evidence in determining source counts.
        * If `self.use_stmt_type` is set, statement type is included via
          one-hot encoding, with one column for each statement type.

        Parameters
        ----------
        df :
            A pandas DataFrame with statement metadata. It should have columns
            `stmt_type` and `source_counts`; alternatively, if it doesn't have
            a `source_counts` column, it should have columns with names
            matching the sources in `self.source_list`.

        Returns
        -------
        :
            Feature matrix for the statement data.
        """
        required_cols = {'stmt_type'}
        # Currently, statement DataFrames are not expected to contain
        # number of members or num_pmids as a data column, hence we raise a
        # ValueError if either of these are set
        if self.use_num_members:
            raise ValueError('use_num_members not supported for statement '
                             'DataFrames.')
        if self.use_num_pmids:
            raise ValueError('use_num_pmids not supported for statement '
                             'DataFrames.')
        # Make sure that the dataframe contains at least all of the above
        # columns
        if not required_cols.issubset(set(df.columns)):
            raise ValueError('Statement DataFrame is missing required '
                             'columns.')
        # Check for the source_counts column. If it's there, we're good
        if 'source_counts' in df.columns:
            has_sc_col = True
        # If it's not, make sure that we have columns named for sources in
        # self.source_list:
        else:
            has_sc_col = False
            for source in self.source_list:
                if source not in df.columns:
                    raise ValueError(f'Expected column "{source}" not in the '
                                      'given statement DataFrame')

        # Add categorical features and collect source_apis
        cat_features = []
        stmt_sources = set()
        # For every statement entry in the dataframe...
        for rowtup in df.itertuples():
            # Collect statement sources
            # ...if there's a source_counts col with dicts
            if has_sc_col:
                stmt_sources |= set(rowtup.source_counts.keys())
            # Collect non-source count features (e.g. type) from stmts
            feature_row = []
            # One-hot encoding of stmt type
            if self.use_stmt_type:
                stmt_type_ix = self.stmt_type_map[rowtup.stmt_type]
                type_features = [1 if ix == stmt_type_ix else 0
                                 for ix in range(len(self.stmt_type_map))]
                feature_row.extend(type_features)
            # Only add a feature row if we're using some of the features.
            if feature_row:
                cat_features.append(feature_row)

        # Before proceeding, check whether all source_apis are in
        # source_list. If we don't have a source_counts dict, we don't look
        # for columns beyond the sources in the source list, and we are
        # guaranteed to have all of them because of the check performed above
        source_diff = stmt_sources.difference(set(self.source_list))
        if has_sc_col and source_diff:
            logger.warning("source_list does not include all source_apis "
                           f"in the statement data: {str(source_diff)}")

        # Get source count features
        num_cols = len(self.source_list)
        num_rows = len(df)
        x_arr = np.zeros((num_rows, num_cols))
        for stmt_ix, rowtup in enumerate(df.itertuples()):
            for src_ix, src in enumerate(self.source_list):
                # Get counts from the source_count dictionary
                if has_sc_col:
                    x_arr[stmt_ix, src_ix] = rowtup.source_counts.get(src, 0)
                # ...or get counts from named source column
                else:
                    x_arr[stmt_ix, src_ix] = rowtup._asdict()[src]

        # If we have any categorical features, turn them into an array and
        # add them to matrix
        if cat_features:
            cat_arr = np.array(cat_features)
            x_arr = np.hstack((x_arr, cat_arr))
        return x_arr

class HybridScorer(BeliefScorer):
    """Use CountsScorer for known sources, SimpleScorer priors for any others.

    Allows the use of a CountsScorer to make belief predictions based on
    sources seen in training data, while falling back to SimpleScorer
    priors for any sources not accounted for by the CountsScorer.
    Like the SimpleScorer, uses an independence assumption to combine
    beliefs from the two scorers (i.e., `hybrid_bel = 1 - (1 - cs_bel) *
    (1 - ss_bel)`).

    Parameters
    ----------
    counts_scorer :
        Instance of CountsScorer.
    simple_scorer :
        Instance of SimpleScorer.
    """
    def __init__(
        self,
        counts_scorer: CountsScorer,
        simple_scorer: SimpleScorer,
    ):
        self.counts_scorer = counts_scorer
        self.simple_scorer = simple_scorer

    def check_prior_probs(
        self,
        statements: Sequence[Statement],
    ) -> None:
        """Check that sources in the set of statements are accounted for."""
        # Get all sources for the set of statements
        sources = CountsScorer.get_all_sources(statements,
                                               include_more_specific=True)
        non_cs_sources = set(sources).difference(
                                        set(self.counts_scorer.source_list))
        return self.simple_scorer._check_sources(non_cs_sources)

    def score_statements(
        self,
        statements: Sequence[Statement],
        extra_evidence: Optional[List[List[Evidence]]] = None,
    ) -> Sequence[float]:
        """
        Parameters
        ----------
        statements :
            INDRA Statements whose belief scores are to be calculated.
        extra_evidence :
            A list corresponding to the given list of statements, where
            each entry is a list of Evidence objects providing additional
            support for the corresponding statement (i.e., Evidences that
            aren't already included in the Statement's own evidence list).

        Returns
        -------
        :
            The computed probabilities for each statement.
        """
        # Get beliefs from the sklearn model, using the sources in the
        # CountScorer source_list as features
        skl_beliefs = self.counts_scorer.predict_proba(statements,
                                                       extra_evidence)[:, 1]
        skl_sources = self.counts_scorer.source_list
        hybrid_beliefs = []
        # Iterate over the statements...
        for ix, stmt in enumerate(statements):
            # ...get both the statement's own evidence and the more-specific
            # (extra) evidences
            all_evidence = get_stmt_evidence(stmt, ix, extra_evidence)
            # Next, filter out any evidences that have sources in the skl
            # model source list, leaving behind the rest. At the same time,
            # record whether we've found any sources the skl model source list.
            filt_evidence = []
            has_skl_source = False
            for ev in all_evidence:
                if ev.source_api in skl_sources:
                    has_skl_source = True
                else:
                    filt_evidence.append(ev)
            # Get the simple belief
            simple_bel = self.simple_scorer.score_evidence_list(filt_evidence)
            # Calculate hybrid belief: the probability that all sources, both
            # those evaluated by the sklearn model and the simplescorer, are
            # not jointly incorrect. If there are no sources from the skl
            # model list, we set the skl belief to 0 so the probability comes
            # only from the simple scorer
            skl_bel = skl_beliefs[ix] if has_skl_source else 0
            hybrid_bel = 1 - (1 - skl_bel) * (1 - simple_bel)
            hybrid_beliefs.append(hybrid_bel)
        return hybrid_beliefs


import numpy as np
import os
import time

from rlgpu.utils.torch_jit_utils import *
from rlgpu.tasks.base.base_task import BaseTask
from isaacgym import gymtorch
from isaacgym import gymapi

import torch
#from torch.tensor import Tensor
from typing import Tuple, Dict

class Foos(BaseTask):

    def __init__(self, cfg, sim_params, physics_engine, device_type, device_id, headless):

        self.cfg = cfg
        self.sim_params = sim_params
        self.physics_engine = physics_engine

        #load parameters
        self.rew_scales = {}

        #load randomization
        # (domain and initial condition)

        # plane params (probs not needed)

        # base init state

        #load default joint positions

        # other
        self.dt = sim_params.dt
        self.max_episode_length_s = self.cfg["env"]["learn"]["episodeLength_s"]
        self.max_episode_length = int(self.max_episode_length_s / self.dt + 0.5)
        self.Kp = self.cfg["env"]["control"]["stiffness"]
        self.Kd = self.cfg["env"]["control"]["damping"]

        for key in self.rew_scales.keys():
            self.rew_scales[key] *= self.dt

        self.cfg["env"]["numObservations"] = 48
        self.cfg["env"]["numActions"] = 12

        self.cfg["device_type"] = device_type
        self.cfg["device_id"] = device_id
        self.cfg["headless"] = headless

        super().__init__(cfg=self.cfg)

        if self.viewer != None:
            p = self.cfg["env"]["viewer"]["pos"]
            lookat = self.cfg["env"]["viewer"]["lookat"]
            cam_pos = gymapi.Vec3(p[0], p[1], p[2])
            cam_target = gymapi.Vec3(lookat[0], lookat[1], lookat[2])
            self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target)

        # get gym state tensors ( CHECK THIS WITH FOOS ASSET)
        actor_root_state = self.gym.acquire_actor_root_state_tensor(self.sim)
        dof_state_tensor = self.gym.acquire_dof_state_tensor(self.sim)
        net_contact_forces = self.gym.acquire_net_contact_force_tensor(self.sim)
        torques = self.gym.acquire_dof_force_tensor(self.sim)

        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)
        self.gym.refresh_dof_force_tensor(self.sim)

        # create some wrapper tensors for different slices

        '''
        self.root_states = gymtorch.wrap_tensor(actor_root_state)
        self.dof_state = gymtorch.wrap_tensor(dof_state_tensor)
        self.dof_pos = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 0]
        self.dof_vel = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 1]
        self.contact_forces = gymtorch.wrap_tensor(net_contact_forces).view(self.num_envs, -1, 3)  # shape: num_envs, num_bodies, xyz axis
        self.torques = gymtorch.wrap_tensor(torques).view(self.num_envs, self.num_dof)
        self.last_dof_vel = torch.zeros_like(self.dof_vel)
        self.last_actions = torch.zeros((self.num_envs, 12), device= self.device, dtype=torch.float, requires_grad=False)
        self.last_base_lin_vel = torch.zeros((self.num_envs, 3), device= self.device, dtype=torch.float, requires_grad=False)
        self.last_base_ang_vel = torch.zeros((self.num_envs, 3), device= self.device, dtype=torch.float, requires_grad=False)
        self.foot_in_contact = torch.zeros((self.num_envs, 4), device= self.device, dtype=torch.float, requires_grad=False) == 1
        self.foot_first_contact = torch.zeros((self.num_envs, 4), device= self.device, dtype=torch.float, requires_grad=False) == 0
        self.foot_air_time = torch.zeros((self.num_envs, 4), device= self.device, dtype=torch.float, requires_grad=False)

        self.commands = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device, requires_grad=False)
        self.commands_y = self.commands.view(self.num_envs, 3)[..., 1]
        self.commands_x = self.commands.view(self.num_envs, 3)[..., 0]
        self.commands_yaw = self.commands.view(self.num_envs, 3)[..., 2]
        self.default_dof_pos = torch.zeros_like(self.dof_pos, dtype=torch.float, device=self.device, requires_grad=False)

        for i in range(self.cfg["env"]["numActions"]):
            name = self.dof_names[i]
            angle = self.named_default_joint_angles[name]
            self.default_dof_pos[:, i] = angle
        '''
        # initialize some data used later on
        self.extras = {}
        torch_zeros = lambda : torch.zeros((self.num_envs, ), device=self.device, dtype= torch.float, requires_grad=False)
        self.reward_terms = {'rew1': torch_zeros(),
                             'rew2': torch_zeros(),
                             }

        #self.initial_root_states = self.root_states.clone()
        #self.initial_root_states[:] = to_torch(self.base_init_state, device=self.device, requires_grad=False)
        #self.gravity_vec = to_torch(get_axis_params(-1., self.up_axis_idx), device=self.device).repeat((self.num_envs, 1))
        #self.actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False)
        #self.time_out_buf = torch.zeros_like(self.reset_buf)

        self.reset(torch.arange(self.num_envs, device=self.device))

    def create_sim(self):
        self.up_axis_idx = self.set_sim_params_up_axis(self.sim_params, 'z')
        self.sim = super().create_sim(self.device_id, self.graphics_device_id, self.physics_engine, self.sim_params)
        self._create_ground_plane()
        self._create_envs(self.num_envs, self.cfg["env"]['envSpacing'], int(np.sqrt(self.num_envs)))

    def _create_ground_plane(self):
        plane_params = gymapi.PlaneParams()
        plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0)
        plane_params.static_friction = self.plane_static_friction
        plane_params.dynamic_friction = self.plane_dynamic_friction
        self.gym.add_ground(self.sim, plane_params)

    def _create_envs(self, num_envs, spacing, num_per_row):
        asset_root = os.getcwd()+"/assets"
        asset_file = "urdf/yobo_model/yobotics_description/urdf/yobotics.urdf"
        asset_path = os.path.join(asset_root, asset_file)
        asset_root = os.path.dirname(asset_path)
        asset_file = os.path.basename(asset_path)

        asset_options = gymapi.AssetOptions()
        asset_options.default_dof_drive_mode = gymapi.DOF_MODE_NONE
        asset_options.collapse_fixed_joints = True
        asset_options.replace_cylinder_with_capsule = False
        asset_options.flip_visual_attachments = False
        asset_options.fix_base_link = self.cfg["env"]["urdfAsset"]["fixBaseLink"]
        asset_options.density = 0.001
        asset_options.angular_damping = 0.0
        asset_options.linear_damping = 0.0
        asset_options.armature = 0.0
        asset_options.thickness = 0.01
        asset_options.disable_gravity = False

        cheetah_asset = self.gym.load_asset(self.sim, asset_root, asset_file, asset_options)
        self.num_dof = self.gym.get_asset_dof_count(cheetah_asset)
        self.num_bodies = self.gym.get_asset_rigid_body_count(cheetah_asset)

        start_pose = gymapi.Transform()
        start_pose.p = gymapi.Vec3(*self.base_init_state[:3])

        body_names = self.gym.get_asset_rigid_body_names(cheetah_asset)
        self.dof_names = self.gym.get_asset_dof_names(cheetah_asset)
        hip_names = [s for s in body_names if 'hip' in s]
        extremity_name = "calf" if asset_options.collapse_fixed_joints else "foot"
        feet_names = [s for s in body_names if extremity_name in s]
        self.feet_indices = torch.zeros(len(feet_names), dtype=torch.long, device=self.device, requires_grad=False)
        self.hip_indices = torch.zeros(len(feet_names), dtype=torch.long, device=self.device, requires_grad=False)
        knee_names = [s for s in body_names if "THIGH" in s]
        self.knee_indices = torch.zeros(len(knee_names), dtype=torch.long, device=self.device, requires_grad=False)
        self.base_index = 0

        dof_props = self.gym.get_asset_dof_properties(cheetah_asset)
        for i in range(self.num_dof):
            dof_props['driveMode'][i] = gymapi.DOF_MODE_POS
            dof_props['stiffness'][i] = self.Kp
            dof_props['damping'][i] = self.Kd

        env_lower = gymapi.Vec3(-spacing, -spacing, 0.0)
        env_upper = gymapi.Vec3(spacing, spacing, spacing)
        self.cheetah_handles = []
        self.envs = []

        for i in range(self.num_envs):
            # create env instance
            env_ptr = self.gym.create_env(self.sim, env_lower, env_upper, num_per_row)
            cheetah_handle = self.gym.create_actor(env_ptr, cheetah_asset, start_pose, "cheetah", i, 1, 0)
            self.gym.set_actor_dof_properties(env_ptr, cheetah_handle, dof_props)
            self.gym.enable_actor_dof_force_sensors(env_ptr, cheetah_handle)
            self.envs.append(env_ptr)
            self.cheetah_handles.append(cheetah_handle)

        self.base_index = self.gym.find_actor_rigid_body_handle(self.envs[0], self.cheetah_handles[0], "name_base")

    def pre_physics_step(self, actions):
        self.actions = actions.clone().to(self.device)
        targets = self.action_scale * self.actions + self.default_dof_pos
        self.gym.set_dof_position_target_tensor(self.sim, gymtorch.unwrap_tensor(targets))

    def post_physics_step(self):
        self.progress_buf += 1

        env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
        if self.push_cheetah:
            self.push_robots()
        if len(env_ids) > 0:
            self.reset(env_ids)

        self.compute_observations()
        self.compute_reward(self.actions)
        self.last_dof_vel[:] = self.dof_vel[:]
        self.last_actions[:] = self.actions[:]
        self.last_base_lin_vel[:] = self.base_lin_vel
        self.last_base_ang_vel[:] = self.base_ang_vel

    def compute_reward(self, actions):
        rew_1 = 0
        rew_2 = 0
        self.rew_buf = torch.clip(self.rew_buf, 0., None)

        # reset agents
        reset = False# RESET CONDITION
        time_out = self.progress_buf > self.max_episode_length  # no terminal reward for time-outs
        self.reset_buf = reset | time_out

        self.reward_terms['rew1'] += rew_1
        self.reward_terms['rew2'] += rew_2

    def compute_observations(self):
        self.gym.refresh_dof_state_tensor(self.sim)  # done in step
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)
        self.gym.refresh_dof_force_tensor(self.sim)

        self.obs_buf[:] = compute_obs(self.root_states)

    def reset(self, env_ids):
        # Randomization can happen only at reset time, since it can reset actor positions on GPU
        '''
        if self.randomize:
            self.apply_randomizations(self.randomization_params)

        positions_offset = torch_rand_float(0.5, 1.5, (len(env_ids), self.num_dof), device=self.device)
        velocities = torch_rand_float(-0.1, 0.1, (len(env_ids), self.num_dof), device=self.device)

        self.dof_pos[env_ids] = self.default_dof_pos[env_ids] * positions_offset
        self.dof_vel[env_ids] = velocities

        env_ids_int32 = env_ids.to(dtype=torch.int32)

        self.gym.set_actor_root_state_tensor_indexed(self.sim,
                                                     gymtorch.unwrap_tensor(self.initial_root_states),
                                                     gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32))

        self.gym.set_dof_state_tensor_indexed(self.sim,
                                              gymtorch.unwrap_tensor(self.dof_state),
                                              gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32))

        self.commands_x[env_ids] = torch_rand_float(self.command_x_range[0], self.command_x_range[1], (len(env_ids), 1), device=self.device).squeeze()
        self.commands_y[env_ids] = torch_rand_float(self.command_y_range[0], self.command_y_range[1], (len(env_ids), 1), device=self.device).squeeze()
        self.commands_yaw[env_ids] = torch_rand_float(self.command_yaw_range[0], self.command_yaw_range[1], (len(env_ids), 1), device=self.device).squeeze()

        norm_small = torch.norm(self.commands[env_ids, :2], dim =1 ) < 0.2

        idx = torch.where(norm_small)[0]
        self.commands[idx, :2] = 0.
        norm_small = torch.abs(self.commands[env_ids, 2]) < 0.2
        idx = torch.where(norm_small)[0]
        self.commands[idx, 2] = 0.

        self.progress_buf[env_ids] = 0
        self.reset_buf[env_ids] = 1
        self.last_dof_vel[env_ids] = 0.
        self.last_actions[env_ids] = 0.
        self.last_base_ang_vel[env_ids] = 0.
        self.last_base_lin_vel[env_ids] = 0.
        self.foot_air_time[env_ids, :] = 0.
        self.foot_first_contact[env_ids, :] = True

        self.extras = {}
        for key in self.reward_terms.keys():
            self.extras['reward_'+key] = (torch.mean(self.reward_terms[key][env_ids])/self.max_episode_length_s).view(-1,1)
            self.reward_terms[key][env_ids] = 0
        '''

#####################################################################
###=========================jit functions=========================###
#####################################################################

@torch.jit.script
def compute_obs(root_states: Tensor) -> Tensor:
    pass


# -*- coding: utf-8 -*-

"""This module contains transistor characterization and optimization related classes.
"""

import os

# import pyoptsparse
import numpy as np

from .core import CharDB

class MosCharDB(CharDB):
    """The mosfet characterization database.

    This class holds transistor characterization data and provides useful query methods.

    Parameters
    ----------
    root_dir : str
        path to the root characterization data directory.
    mos_type : str
        the transistor type.  Either 'pch' or 'nch'.
    discrete_params : list[str]
        a list of parameters that should take on discrete values instead of being interpolated.
        Usually intent, length, or transistor width (for finfets).
    env_list : list[str]
        list of simulation environments to consider.
    update : bool
        True to update post-processed data from raw simulation data.
    intent : str or None
        the threshold flavor name.
    l : float or None
        the channel length, in meters.
    w : int or float or None
        the transistor width, in meters or number of fins.
    vgs : float or None
        the Vgs voltage.
    vds : float or None
        the Vds voltage.
    vbs : float or None
        the Vbs voltage.
    **kwargs :
        additional characterization database parameters.  See documentation for CharDB.
    """

    _raw_data_names = ['ids', 'y11', 'y12', 'y13', 'y21', 'y22', 'y23', 'y31', 'y32', 'y33']

    def __init__(self, root_dir, mos_type, discrete_params, env_list,
                 intent=None, l=None, w=None, vgs=None, vds=None, vbs=None,
                 **kwargs):
        constants = dict(mos_type=mos_type)
        init_params = dict(intent=intent, l=l, w=w, vgs=vgs, vds=vds, vbs=vbs)
        CharDB.__init__(self, root_dir, constants, discrete_params, init_params, env_list, **kwargs)

    @classmethod
    def get_sim_file(cls, root_dir, constants):
        """Returns the simulation data file name.

        Parameters
        ----------
        root_dir : str
            absolute path to the root characterization data directory.
        constants : dict[string, any]
            constants dictionary.

        Returns
        -------
        fname : str
            the simulation data file name.
        """
        return os.path.join(root_dir, '%s.hdf5' % constants['mos_type'])

    @classmethod
    def get_cache_file(cls, root_dir, constants):
        """Returns the post-processed characterization data file name.

        Parameters
        ----------
        root_dir : str
            absolute path to the root characterization data directory.
        constants : dict[string, any]
            constants dictionary.

        Returns
        -------
        fname : str
            the post-processed characterization data file name.
        """
        return os.path.join(root_dir, '%s__%s.hdf5' % (constants['mos_type'], cls.__name__))

    @classmethod
    def post_process_data(cls, sim_data, sweep_params, sweep_values, constants):
        """Postprocess simulation data.

        Parameters
        ----------
        sim_data : dict[string, np.array]
            the simulation data as a dictionary from output name to numpy array.
        sweep_params : list[str]
            list of parameter name for each dimension of numpy array.
        sweep_values : list[numpy.array]
            list of parameter values for each dimension.
        constants : dict[string, any]
            the constants dictionary.

        Returns
        -------
        data : dict[str, np.array]
            a dictionary of post-processed data.
        """
        # compute small signal parameters
        w = 2 * np.pi * constants['char_freq']
        fg = constants['fg']

        ids = sim_data['ids']
        gm = (sim_data['y21'].real - sim_data['y31'].real) / 2.0
        gds = (sim_data['y22'].real - sim_data['y32'].real) / 2.0
        gb = (sim_data['y33'].real - sim_data['y23'].real) / 2.0 - gm - gds

        cgd = -0.5 / w * (sim_data['y12'].imag + sim_data['y21'].imag)
        cgs = -0.5 / w * (sim_data['y13'].imag + sim_data['y31'].imag)
        cds = -0.5 / w * (sim_data['y23'].imag + sim_data['y32'].imag)
        cgb = sim_data['y11'].imag / w - cgd - cgs
        cdb = sim_data['y22'].imag / w - cds - cgd
        csb = sim_data['y33'].imag / w - cgs - cds

        ss_data = dict(
            ids=ids / fg,
            gm=gm / fg,
            gds=gds / fg,
            gb=gb / fg,
            cgd=cgd / fg,
            cgs=cgs / fg,
            cds=cds / fg,
            cgb=cgb / fg,
            cdb=cdb / fg,
            csb=csb / fg,
        )

        return ss_data

    @classmethod
    def derived_parameters(cls):
        """Returns a list of derived parameters."""
        return ['cgg', 'cdd', 'css', 'cbb', 'vstar', 'gain', 'ft']

    @classmethod
    def compute_derived_parameters(cls, fdict):
        """Compute derived parameter functions.

        Parameters
        ----------
        fdict : dict[string, bag.math.dfun.DiffFunction]
            a dictionary from core parameter name to the corresponding function.

        Returns
        -------
        deriv_dict : dict[str, bag.math.dfun.DiffFunction]
            a dictionary from derived parameter name to the corresponding function.
        """
        cgg = fdict['cgd'] + fdict['cgs'] + fdict['cgb']
        return dict(
            cgg=cgg,
            cdd=fdict['cgd'] + fdict['cds'] + fdict['cdb'],
            css=fdict['cgs'] + fdict['cds'] + fdict['csb'],
            cbb=fdict['cgb'] + fdict['cdb'] + fdict['csb'],
            vstar=2.0 * (fdict['ids'] / fdict['gm']),
            gain=fdict['gm'] / fdict['gds'],
            ft=fdict['gm'] / (2.0 * np.pi * cgg),
        )

class MosCharGDDB(CharDB):
    """The mosfet characterization database.

    This class holds transistor characterization data and provides useful query methods.

    Parameters
    ----------
    root_dir : str
        path to the root characterization data directory.
    mos_type : str
        the transistor type.  Either 'pch' or 'nch'.
    discrete_params : list[str]
        a list of parameters that should take on discrete values instead of being interpolated.
        Usually intent, length, or transistor width (for finfets).
    env_list : list[str]
        list of simulation environments to consider.
    update : bool
        True to update post-processed data from raw simulation data.
    intent : str or None
        the threshold flavor name.
    l : float or None
        the channel length, in meters.
    w : int or float or None
        the transistor width, in meters or number of fins.
    vgs : float or None
        the Vgs voltage.
    vds : float or None
        the Vds voltage.
    vbs : float or None
        the Vbs voltage.
    **kwargs :
        additional characterization database parameters.  See documentation for CharDB.
    """

    _raw_data_names = ['ids', 'y11', 'y12', 'y21', 'y22']

    def __init__(self, root_dir, mos_type, discrete_params, env_list,
                 intent=None, l=None, w=None, vgs=None, vds=None,
                 **kwargs):
        constants = dict(mos_type=mos_type)
        init_params = dict(intent=intent, l=l, w=w, vgs=vgs, vds=vds)
        CharDB.__init__(self, root_dir, constants, discrete_params, init_params, env_list, **kwargs)

    @classmethod
    def get_sim_file(cls, root_dir, constants):
        """Returns the simulation data file name.

        Parameters
        ----------
        root_dir : str
            absolute path to the root characterization data directory.
        constants : dict[string, any]
            constants dictionary.

        Returns
        -------
        fname : str
            the simulation data file name.
        """
        return os.path.join(root_dir, '%s.hdf5' % constants['mos_type'])

    @classmethod
    def get_cache_file(cls, root_dir, constants):
        """Returns the post-processed characterization data file name.

        Parameters
        ----------
        root_dir : str
            absolute path to the root characterization data directory.
        constants : dict[string, any]
            constants dictionary.

        Returns
        -------
        fname : str
            the post-processed characterization data file name.
        """
        return os.path.join(root_dir, '%s__%s.hdf5' % (constants['mos_type'], cls.__name__))

    @classmethod
    def post_process_data(cls, sim_data, sweep_params, sweep_values, constants):
        """Postprocess simulation data.

        Parameters
        ----------
        sim_data : dict[string, np.array]
            the simulation data as a dictionary from output name to numpy array.
        sweep_params : list[str]
            list of parameter name for each dimension of numpy array.
        sweep_values : list[numpy.array]
            list of parameter values for each dimension.
        constants : dict[string, any]
            the constants dictionary.

        Returns
        -------
        data : dict[str, np.array]
            a dictionary of post-processed data.
        """
        # compute small signal parameters
        w = 2 * np.pi * constants['char_freq']
        fg = constants['fg']

        ids = sim_data['ids']
        gm = sim_data['y21'].real
        gds = sim_data['y22'].real

        cgd = -0.5 / w * (sim_data['y12'].imag + sim_data['y21'].imag)
        cgs = sim_data['y11'].imag / w - cgd
        cds = sim_data['y22'].imag / w - cgd

        ss_data = dict(
            ids=ids / fg,
            gm=gm / fg,
            gds=gds / fg,
            cgd=cgd / fg,
            cgs=cgs / fg,
            cds=cds / fg,
        )

        return ss_data

    @classmethod
    def derived_parameters(cls):
        """Returns a list of derived parameters."""
        return ['cgg', 'cdd', 'vstar', 'gain', 'ft']

    @classmethod
    def compute_derived_parameters(cls, fdict):
        """Compute derived parameter functions.

        Parameters
        ----------
        fdict : dict[string, bag.math.dfun.DiffFunction]
            a dictionary from core parameter name to the corresponding function.

        Returns
        -------
        deriv_dict : dict[str, bag.math.dfun.DiffFunction]
            a dictionary from derived parameter name to the corresponding function.
        """
        cgg = fdict['cgd'] + fdict['cgs']
        return dict(
            cgg=cgg,
            cdd=fdict['cgd'] + fdict['cds'],
            vstar=2.0 * (fdict['ids'] / fdict['gm']),
            gain=fdict['gm'] / fdict['gds'],
            ft=fdict['gm'] / (2.0 * np.pi * cgg),
        )