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

"""

@author: bartulem

Read in sync events from all desired data streams.

Neuropixel recordings are saved into a *1D binary vector*, from a 2D array organised as 385 rows (channels) x n columns (samples).
The data are written to file from this matrix in column-major (F) order, i.e., the first sample in the recording was written to file
for every channel, then the second sample was written for every channel, etc. Neuropixel raw data is provided as an int16 binary.
Neuropixel ADCs are 10 bits, with a range of -0.6 to 0.6 V, and acquisition was at 500x gain, yielding a resolution of 2.34 µV/bit.
To obtain readings in µV, you should multiply the int16 values by 2.34.

This script has the purpose of extracting LEDon sync events from four independent data streams: (1) NPX recording files (may be one or
two in a given session, depending on the number of probes), (2) the exported Motive tracking file, and (3) the IMU sensor file,
(4) the sync IPI generator, all of which keep track of the LEDon occurrences. Unlike v2.0.0, future versions will assume the LED
pulses are generated continuously and randomly (by default with a fixed 250ms duration, and 250-1500ms IPIs).

The code first goes through all of those files and picks up the LED events, (taking into account that the LED signal jitters,
which can be probe independent). It then selects a portion of the starting and ending LED pulses in the tracking data stream
(10 by default) and tries to match these templates in the other data streams, aligning the start and end of tracking to different
recordings. All the corresponding LED events in all data streams are stored in a DataFrame, which is saved to a separate
binary .pkl file. This code might consume a lot of RAM (64Gb should be enough in most cases).

"""

import os
import sys
import numpy as np
import gc
import time
import csv
import pandas as pd
import pickle
from collections import Counter
from tqdm.notebook import tqdm
import operator
import warnings
warnings.simplefilter('ignore')
from astropy.convolution import convolve
from astropy.convolution import Gaussian1DKernel

class EventReader:

    # initializer / instance attributes
    def __init__(self, npx_files, sync_df):
        self.npx_files = npx_files
        self.sync_df = sync_df

    def read_se(self, **kwargs):

        """
        Inputs
        ----------
        **kwargs: dictionary
        nchan : int/float
            Total number of channels on the NPX probe, for probe 3B2 should be 385; defaults to 385.
        sync_chan : int/float
            Sync events channel number, for probe 3B2 should be 385; defaults to 385.
        track_file : str/boolean
            The absolute path to the tracking data .csv file; defaults to 0.
        imu_file : str/boolean
            The absolute path to the IMU data .txt file; defaults to 0.
        sync_led_duration : int/float
            The duration of ON time for sync LEDs; defaults to 250 (ms).
        sync_led_error : int/float
            The possible error in the duration of ON time for sync LEDs; defaults to 50 (ms).
        ground_probe : int
            In a dual probe setting, the probe the other is synced to; defaults to 0.
        frame_rate : int/float
            The tracking camera frame rate for that session; defaults to 120.
        npx_sampling_rate : int/float
            The sampling rate of the NPX system; defaults to 3e4.
        sync_sequence : int/float
            The length of the sequence the LED events should be matched across data streams; defaults to 10.
        sample_error : int/float
            The time the presumed IMEC/IMU LEDs could be allowed to err around; defaults to 30 (ms).
        which_imu_time : int/float
            The IMU time to be used in the analyses, loop.starttime (0) or sample.time (1); defaults to 1.
        sync_ipi_file : str/boolean
            The absolute path to the sync IPI data .txt file; defaults to 0.
        ----------

        Outputs
        ----------
        track_file : csv object
            The corrected and LED-shortened tracking data; saved as a separate .csv file.
        imu_file : pd.DataFrame
            The IMU data; saved as .pkl file.
        sync_df : pd.DataFrame
            The sync data; saved as .pkl file.
        ----------
        """

        nchan = int(kwargs['nchan'] if 'nchan' in kwargs.keys() and (type(kwargs['nchan']) == int or type(kwargs['nchan']) == float) else 385)
        sync_chan = int(kwargs['sync_chan'] if 'sync_chan' in kwargs.keys() and (type(kwargs['sync_chan']) == int or type(kwargs['sync_chan']) == float) else 385)
        track_file = kwargs['track_file'] if 'track_file' in kwargs.keys() and type(kwargs['track_file']) == str else 'eldiablomuerte'
        imu_file = kwargs['imu_file'] if 'imu_file' in kwargs.keys() and type(kwargs['imu_file']) == str else 'eldiablomuerte'
        sync_led_duration = int(kwargs['sync_led_duration'] if 'sync_led_duration' in kwargs.keys() and (type(kwargs['sync_led_duration']) == int or type(kwargs['sync_led_duration']) == float) else 250)
        sync_led_error = int(kwargs['sync_led_error'] if 'sync_led_error' in kwargs.keys() and (type(kwargs['sync_led_error']) == int or type(kwargs['sync_led_error']) == float) else 50)
        ground_probe = int(kwargs['ground_probe'] if 'ground_probe' in kwargs.keys() else 0)
        frame_rate = float(kwargs['frame_rate'] if 'frame_rate' in kwargs.keys() else 120.)
        npx_sampling_rate = int(kwargs['npx_sampling_rate'] if 'npx_sampling_rate' in kwargs.keys() else 3e4)
        sync_sequence = int(kwargs['sync_sequence'] if 'sync_sequence' in kwargs.keys() else 10)
        sample_error = int(kwargs['sample_error'] if 'sample_error' in kwargs.keys() and (type(kwargs['sample_error']) == int or type(kwargs['sample_error']) == float) else 30)
        which_imu_time = int(kwargs['which_imu_time'] if 'which_imu_time' in kwargs.keys() and (type(kwargs['which_imu_time']) == int or type(kwargs['which_imu_time']) == float) else 1)
        sync_ipi_file = kwargs['sync_ipi_file'] if 'sync_ipi_file' in kwargs.keys() and type(kwargs['sync_ipi_file']) == str else 'eldiablomuerte'

        # check that the NPX files are there
        for anpxfile in self.npx_files:
            if not os.path.exists(anpxfile):
                print('Could not find file {}, try again.'.format(anpxfile))
                sys.exit()

        start_time = time.time()
        print('Extracting sync data from NPX file(s), please be patient - this could take >5 minutes.\n')

        # initialize dictionary to store the data in
        sync_dict = {}

        # extract LED events from the NPX file(s)
        for anpxfile in self.npx_files:

            # memory maps are used for accessing small segments of large files on disk, without reading the entire file into memory.
            npx_recording = np.memmap(anpxfile, mode='r', dtype=np.int16, order='C')

            # integer divide the length of the recording by channel count to get number of samples
            npx_samples = npx_recording.shape[0] // nchan

            # reshape the array such that channels are rows and samples are columns
            npx_recording = npx_recording.reshape((nchan, npx_samples), order='F')

            # find sync events and collect them in a dict / this warrants a more detailed explanation
            # the sync channel consists of 16 1-bit channels, and when its values are read as np.int16 they can assume
            # different values such as 0, 64, 128, 132, etc. (to be converted back to binary, use np.binary_repr(num, 16)).
            # The sync LED event should be a sequence of 64s that is 250 ms in duration (this is the default value that
            # can be set as an input parameter), a LED off period should be an, at least equal, sequence of zeros.
            # The other values occur mainly because of noise. So, in theory the LED events should last 7500 samples
            # (250ms * 30) but for different reasons may last less and might include jitter values (128, 132, etc.),
            # so in screening for LED events that is taken into account.

            counter = sorted(dict(Counter(npx_recording[sync_chan - 1, :])).items(), key=operator.itemgetter(1), reverse=True)
            most_freq_two_values = [counter[0][0], counter[1][0]]

            session_proportion = (counter[0][1] + counter[1][1]) / npx_samples
            if session_proportion < .99:
                print('The two most dominant values, {} and {} appear together only {:.3f} of the total session, '
                      'so something is wrong. Check it out!'.format(counter[0][0], counter[1][0], session_proportion))
                sys.exit()

            high_val = np.nanmax(most_freq_two_values)
            low_val = np.nanmin(most_freq_two_values)

            # how many samples to screen LED events for
            imec_check_samples = 30 * (sync_led_duration - sync_led_error)

            probe_sync = []
            counter_on = 0
            probe_sync.append(0)
            for inxSync, itemSync in enumerate(tqdm(npx_recording[sync_chan - 1, :])):
                if 0 < inxSync < npx_samples - imec_check_samples \
                        and itemSync == high_val and npx_recording[sync_chan - 1, inxSync - 1] == low_val \
                        and (npx_recording[sync_chan - 1, max(0, inxSync - imec_check_samples):inxSync] != high_val).all() \
                        and (npx_recording[sync_chan - 1, (inxSync + 1):(inxSync + imec_check_samples + 1)] != low_val).all():
                    probe_sync.append(inxSync)
                    counter_on += 1
            probe_sync.append(npx_samples)

            # delete the map object from memory
            del npx_recording
            gc.collect()

            # save data to sync_dict where imec number is key
            sync_dict['imec{}'.format(anpxfile[anpxfile.find('imec') + len('imec')])] = probe_sync
            print('There are {} total LED events in the imec{} file.'.format(counter_on, anpxfile[anpxfile.find('imec') + len('imec')]))

        # extract LED events from the tracking file
        if os.path.exists(track_file):

            print('Extracting sync data from the tracking file, please remain patient.')

            # read in the 4th line of the tracking .csv and find the column where LED appears for the first time
            with open(track_file, 'r') as thecsv:
                for rowindx, row in enumerate(csv.reader(thecsv, delimiter=',')):
                    if rowindx == 3:
                        columnofint = next(indx for indx, obj in enumerate(row) if 'LED' in obj)
                        break

            # correct fully empty frames / this warrants a more detailed explanation
            # tracking files, for some reason, sometimes get frames where all markers drop out. This happens
            # infrequently and when it does it is only a couple of frames that are surrounded by frames that have
            # all markers present. This can be a problem for LED detection because one empty frame would separate
            # something that was a unitary LED event into two different LED events. There's also the problem
            # where only the LED rigid body markers get lost for a couple of frames. To prevent this staying in
            # the data, the next snippet of code goes through the tracking file and interpolates those empty
            # frames by smoothing over a window of the LED duration (250 ms or 30 rows above and below the frame).

            # the duration of sync LEDs in tracking frames (rounded to the nearest decade)
            tracking_check_samples = int(round(sync_led_duration / (1 / frame_rate * 1000), -1))

            # the half of tracking check samples
            track_smooth_window = int(round(tracking_check_samples / 2))

            # define led columns
            led_cols = list(range(columnofint, columnofint + 9))

            # load original tracking data
            original_tracking_data = pd.read_csv(track_file, sep=',', skiprows=6)
            total_frame_number = original_tracking_data.shape[0]

            print('Correcting empty frames.\n')

            corrected_frames = {}
            for i in tqdm(range(total_frame_number - 1)):

                # The following conditions are remedied by this snippet of code >

                # (1) values are missing for all animal points outside of a LED event
                statement_1 = (original_tracking_data.iloc[i, 2:].isnull().values.all() and
                               original_tracking_data.iloc[max(0, i - track_smooth_window):i, columnofint:(columnofint + 9)].isnull().values.all() and
                               original_tracking_data.iloc[(i + 1):min(i + track_smooth_window + 1, total_frame_number - 1), columnofint:(columnofint + 9)].isnull().values.all())

                # (2) values are missing for all points during a LED event
                statement_2 = (original_tracking_data.iloc[i, 2:].isnull().values.all() and
                               original_tracking_data.iloc[max(0, i - track_smooth_window):i, columnofint:(columnofint + 9)].notnull().values.any() and
                               original_tracking_data.iloc[(i + 1):min(i + track_smooth_window + 1, total_frame_number - 1), columnofint:(columnofint + 9)].notnull().values.any())

                # (3) values are missing for just the LED points (during a LED event)
                statement_3 = (original_tracking_data.iloc[i, 2:].notnull().values.any() and
                               original_tracking_data.iloc[i, columnofint:(columnofint + 9)].isnull().values.all() and
                               original_tracking_data.iloc[max(0, i - track_smooth_window):i, columnofint:(columnofint + 9)].notnull().values.any() and
                               original_tracking_data.iloc[(i + 1):min(i + track_smooth_window + 1, total_frame_number - 1), columnofint:(columnofint + 9)].notnull().values.any())

                if statement_1 or statement_2 or statement_3:
                    # print('Frame {} is being corrected.'.format(original_tracking_data.iloc[i, 0]))
                    changed_array = np.zeros(original_tracking_data.shape[1])
                    for inx in range(len(changed_array)):
                        if inx < 2:
                            if inx < 1:
                                changed_array[inx] = original_tracking_data.iloc[i, inx]
                            else:
                                changed_array[inx] = round(original_tracking_data.iloc[i, inx], 6)
                        elif 1 < inx and inx not in led_cols:
                            if statement_1 or statement_2:
                                changed_array[inx] = round(convolve(original_tracking_data.iloc[max(0, i - track_smooth_window):min(i + track_smooth_window + 1, total_frame_number - 1), inx],
                                                                    kernel=Gaussian1DKernel(stddev=2), nan_treatment='interpolate', preserve_nan=False)[track_smooth_window], 6)
                            else:
                                changed_array[inx] = round(original_tracking_data.iloc[i, inx], 6)
                        else:
                            if statement_1:
                                changed_array[inx] = np.nan
                            else:
                                changed_array[inx] = round(convolve(original_tracking_data.iloc[max(0, i - track_smooth_window):min(i + track_smooth_window + 1, total_frame_number - 1), inx],
                                                                    kernel=Gaussian1DKernel(stddev=2), nan_treatment='interpolate', preserve_nan=False)[track_smooth_window], 6)

                    # format strings accordingly
                    changed_list = []
                    for idx, item in enumerate(changed_array):
                        if idx == 0:
                            changed_list.append(str(int(item)))
                        else:
                            if not np.isnan(item):
                                changed_list.append('{:.6f}'.format(item))
                            else:
                                changed_list.append('')
                    corrected_frames[i + 7] = changed_list

            # save the interpolated tracking .csv
            with open(track_file, 'r') as rtf, open('{}_interpolated.csv'.format(track_file[:-4]), 'w', newline='') as wtf:
                writer = csv.writer(wtf)
                reader = csv.reader(rtf)
                for rowindx, row in enumerate(reader):
                    if rowindx not in corrected_frames.keys():
                        writer.writerow(row)
                    else:
                        writer.writerow(corrected_frames[rowindx])

            # load the df and search for LEDons/offs
            tracking_data = pd.read_csv('{}_interpolated.csv'.format(track_file[:-4]), sep=',', skiprows=6)

            tracking_sync = {}
            tracking_on = 1
            all_led_frames = []

            print('Looking for LED events in the corrected tracking file.\n')

            for row in tqdm(range(tracking_data.shape[0])):
                if row > 0 and not tracking_data.iloc[row, columnofint:(columnofint + 9)].isnull().values.all() \
                        and tracking_data.iloc[max(row - tracking_check_samples, 0):row, columnofint:(columnofint + 9)].isnull().values.all() \
                        and tracking_data.iloc[row:min(row + tracking_check_samples, tracking_data.shape[0]), columnofint:(columnofint + 9)].isnull().all(axis=1).sum() == 0:
                    tracking_sync['{}LEDon'.format(tracking_on)] = tracking_data.loc[row, 'Frame']
                    all_led_frames.append(tracking_data.loc[row, 'Frame'])
                    tracking_on += 1

            # save data to sync_dict
            sync_dict['tracking'] = tracking_sync
            print('Completed! There are {} total LED events in the tracking file.'.format(len(tracking_sync.keys())))

            # save the tracking .csv with everything before the first LEDon and after the last LEDon removed
            with open('{}_interpolated.csv'.format(track_file[:-4]), 'r') as inp, \
                    open('{}final.csv'.format('{}_interpolated.csv'.format(track_file[:-4])[:-16]), 'w', newline='') as out:
                writer = csv.writer(out)
                for rowindx, row in enumerate(csv.reader(inp)):
                    if rowindx < 7 or (all_led_frames[0] + 7) <= rowindx <= (all_led_frames[-1] + 7):
                        writer.writerow(row)

            # remove interpolated file from the tracking directory
            os.remove('{}_interpolated.csv'.format(track_file[:-4]))

        else:
            print('Tracking file not given or found.')

        # extract LED events from the IMU file
        if os.path.exists(imu_file):

            print('Extracting sync data from the IMU file, please remain patient.\n')

            # load the files as df
            imu_df = pd.read_csv(imu_file, sep=',', header=None)

            # get rid of the date in the first column (loop.starttime)
            imu_df.iloc[:, 0] = [np.int64(x.split('\t')[1]) for x in imu_df.iloc[:, 0]]

            # add the imu header
            imu_df.columns = ['loop.starttime (ms)', 'sample.time (ms)', 'acc.x', 'acc.y', 'acc.z',
                              'linacc.x', 'linacc.y', 'linacc.z', 'gyr.x', 'gyr.y', 'gyr.z',
                              'mag.x', 'mag.y', 'mag.z', 'euler.x', 'euler.y', 'euler.z',
                              'LED', 'sound', 'sys.cal', 'gyr.cal', 'acc.cal', 'mag.cal']

            # get LED times
            sample_array = imu_df.iloc[:, which_imu_time].tolist()
            led_array = imu_df['LED'].tolist()
            teensy_time = []
            imu_led = []
            imu_on = 0
            for idx, item in enumerate(tqdm(led_array)):
                if idx > 0 and item != 0 and led_array[idx - 1] == 0:
                    imu_led.append(idx)
                    teensy_time.append(sample_array[idx])
                    imu_on += 1

            # save data to sync_dict
            sync_dict['imu_frame_number'] = imu_led
            sync_dict['teensy_sample_time'] = teensy_time
            print('Completed! here are {} total LED events in the IMU file.'.format(imu_on))

            # save IMU data to file
            with open('{}.pkl'.format(imu_file[:-4]), 'wb') as imu_data:
                pickle.dump(imu_df, imu_data)

        else:
            print('IMU file not given or found.')

        # align LED events from the tracking data stream to other data streams
        if os.path.exists(track_file):

            # find the matching sequence of IPIs in the IPI .txt file
            if os.path.exists(sync_ipi_file):

                print('Template matching IPIs to tracking data.')

                # load IPI file and clean up the first column
                ipi_df = pd.read_csv(sync_ipi_file, sep=',', header=None)
                ipi_df.iloc[:, 0] = [np.int64(x.split('\t')[1]) for x in ipi_df.iloc[:, 0]]

                # add sync_led_duration to every IPI
                ipis = ipi_df.to_numpy().ravel() + sync_led_duration

                # find sequence in tracking data
                first_diffs = (np.diff(list(sync_dict['tracking'].values())[:sync_sequence]) + 1) * (1e3 / frame_rate)
                last_diffs = (np.diff(list(sync_dict['tracking'].values())[-sync_sequence:]) + 1) * (1e3 / frame_rate)

                start_end_ipis = []

                for idx, item in enumerate(ipis):
                    if idx < len(ipis) - sync_sequence:
                        temp_diffs = ipis[idx:idx + sync_sequence - 1]
                        if (np.absolute(temp_diffs - first_diffs) <= sample_error).all():
                            print('Found first matching IPI as {} ms.'.format(item))
                            start_end_ipis.append(idx)
                        elif (np.absolute(temp_diffs - last_diffs) <= sample_error).all():
                            print('Found last matching IPI as {} ms.'.format(ipis[idx + sync_sequence - 2]))
                            start_end_ipis.append(idx + sync_sequence - 2)
                            break

                if len(start_end_ipis) == 0 or start_end_ipis[0] == start_end_ipis[-1]:
                    print('At least one IPI not matched, check it out!')
                    sys.exit()

                # convert IPIs to time elapsed since first pulse
                tracking_ipis = ipis[start_end_ipis[0]:(start_end_ipis[-1] + 1)]

                ipis_time = np.zeros(len(tracking_ipis) + 1)
                for idx in range(len(tracking_ipis)):
                    ipis_time[idx + 1] = np.sum(tracking_ipis[:(idx + 1)])

                # save data to sync_dict
                ipis_dict = {}
                counter_on = 1
                for idx, led_time in enumerate(ipis_time):
                    ipis_dict['{}LEDon'.format(counter_on)] = led_time
                    counter_on += 1

                sync_dict['time (ms)'] = ipis_dict

                print('After template matching, there are now {} clock times based on IPIs.'.format(len(ipis_time)))

            for data_key in sync_dict.keys():
                if 'imec' in data_key:
                    print('Template matching tracking to {} LED events.'.format(data_key))

                    if os.path.exists(sync_ipi_file):
                        first_diffs = tracking_ipis[:(sync_sequence - 1)] * 30
                        last_diffs = tracking_ipis[(-sync_sequence + 1):] * 30
                    else:
                        first_diffs = (np.diff(list(sync_dict['tracking'].values())[:sync_sequence]) + 1) * (npx_sampling_rate / frame_rate)
                        last_diffs = (np.diff(list(sync_dict['tracking'].values())[-sync_sequence:]) + 1) * (npx_sampling_rate / frame_rate)

                    imec_leds = sync_dict[data_key][1:-1]
                    important_led_positions = []

                    for idx, item in enumerate(imec_leds):
                        if idx < len(imec_leds) - sync_sequence:
                            temp_diffs = (np.diff(imec_leds[idx:(idx + sync_sequence)]) + 1)
                            if (np.absolute(temp_diffs - first_diffs) <= 30 * sample_error).all():
                                print('Found first matching LED at sample number {}.'.format(item))
                                important_led_positions.append(idx)
                            elif (np.absolute(temp_diffs - last_diffs) <= 30 * sample_error).all():
                                print('Found last matching LED at sample number {}.'.format(imec_leds[idx + sync_sequence - 1]))
                                important_led_positions.append(idx + sync_sequence - 1)
                                break

                    if len(important_led_positions) == 0 or important_led_positions[0] == important_led_positions[-1]:
                        print('At least one matching LED not found, check it out!')
                        sys.exit()

                    imec_leds = imec_leds[important_led_positions[0]:important_led_positions[-1] + 1]
                    print('After template matching, there are now {} LED events in the {} file.'.format(len(imec_leds), data_key))

                    imec_dict = {}
                    counter_on = 1
                    imec_dict['Session start'] = 0
                    for led_event in imec_leds:
                        imec_dict['{}LEDon'.format(counter_on)] = led_event
                        counter_on += 1
                    imec_dict['Session stop'] = sync_dict[data_key][-1]
                    sync_dict[data_key] = imec_dict

                elif 'teensy' in data_key:
                    print('Template matching tracking to {} LED events.'.format(data_key))

                    if os.path.exists(sync_ipi_file):
                        first_diffs = tracking_ipis[:(sync_sequence - 1)]
                        last_diffs = tracking_ipis[(-sync_sequence + 1):]
                    else:
                        first_diffs = (np.diff(list(sync_dict['tracking'].values())[:sync_sequence]) + 1) * (1e3 / frame_rate)
                        last_diffs = (np.diff(list(sync_dict['tracking'].values())[-sync_sequence:]) + 1) * (1e3 / frame_rate)

                    imu_leds = sync_dict[data_key]
                    important_led_positions = []

                    for idx, item in enumerate(imu_leds):
                        if idx < len(imu_leds) - sync_sequence:
                            temp_diffs = (np.diff(imu_leds[idx:idx + sync_sequence]) + 1)
                            if (np.absolute(temp_diffs - first_diffs) <= sample_error).all():
                                print('Found first matching LED at sample time {}.'.format(item))
                                important_led_positions.append(idx)
                            elif (np.absolute(temp_diffs - last_diffs) <= sample_error).all():
                                print('Found last matching LED at sample time {}.'.format(imu_leds[idx + sync_sequence - 1]))
                                important_led_positions.append(idx + sync_sequence - 1)
                                break

                    if len(important_led_positions) == 0 or important_led_positions[0] == important_led_positions[-1]:
                        print('At least one matching LED not found, check it out!')
                        sys.exit()

                    imu_leds = imu_leds[important_led_positions[0]:important_led_positions[-1] + 1]
                    imu_frame_numbers = sync_dict['imu_frame_number'][important_led_positions[0]:important_led_positions[-1] + 1]
                    print('After template matching, there are now {} LED events in the {} file.'.format(len(imu_leds), data_key))

                    teensy_dict = {}
                    imu_frame_dict = {}
                    counter_on = 1
                    for idx, led_event in enumerate(imu_leds):
                        teensy_dict['{}LEDon'.format(counter_on)] = led_event
                        imu_frame_dict['{}LEDon'.format(counter_on)] = imu_frame_numbers[idx]
                        counter_on += 1
                    sync_dict['imu_frame_number'] = imu_frame_dict
                    sync_dict['teensy_sample_time'] = teensy_dict

            # pack sync_dict into a df and save it to a .pkl file
            choose_imec = 'imec{}'.format(ground_probe)
            export_sync_df = pd.DataFrame(index=sync_dict[choose_imec].keys(), columns=sync_dict.keys())
            for data_key in sync_dict.keys():
                for event_key in sync_dict[choose_imec].keys():
                    if event_key in sync_dict[data_key].keys():
                        export_sync_df.loc[event_key, data_key] = sync_dict[data_key][event_key]

            # check if DataFrame contains NANs
            if export_sync_df.iloc[1:-1, :].isnull().values.any():
                print('There is at least one NAN value as a LED event in the saved sync_df! Check it out!')

            with open('{}'.format(self.sync_df), 'wb') as df:
                pickle.dump(export_sync_df, df)

        print('Extraction complete! It took {:.2f} minutes.\n'.format((time.time() - start_time) / 60))


import os
import time
import datetime

import numpy as np
from mpi4py import MPI

from common import logger
from common.misc_util import dbg_tf_init
import common.tf_util as tf_util
from ddpg.pretrain import pretrain_demo, load_history

def train(env, eval_env, agent, render=False, render_eval=False, sanity_run=False,
          nb_epochs=500, nb_epoch_cycles=20, nb_rollout_steps=100, nb_train_steps=50,
          param_noise_adaption_interval=50, hist_files=None, start_ckpt=None, demo_files=None):

    rank = MPI.COMM_WORLD.Get_rank()
    mpi_size = MPI.COMM_WORLD.Get_size()
    if rank == 0:
        logdir = logger.get_dir()
    else:
        logdir = None

    memory = agent.memory
    batch_size = agent.batch_size

    with tf_util.single_threaded_session() as sess:
        # Prepare everything.
        agent.initialize(sess, start_ckpt=start_ckpt)
        sess.graph.finalize()
        agent.reset()
        dbg_tf_init(sess, agent.dbg_vars)

        total_nb_train = 0
        total_nb_rollout = 0
        total_nb_eval = 0

        # pre-train demo and critic_step
        # train_params: (nb_steps, lr_scale)
        total_nb_train = pretrain_demo(agent, env, demo_files, total_nb_train,
                                       train_params=[(100, 1.0)], start_ckpt=start_ckpt)
        load_history(agent, env, hist_files)

        # main training
        obs = env.reset()
        reset = False
        episode_step = 0
        last_episode_step = 0

        for i_epoch in range(nb_epochs):
            t_epoch_start = time.time()
            logger.info('\n%s epoch %d starts:' % (datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), i_epoch))
            for i_cycle in range(nb_epoch_cycles):
                logger.info('\n%s cycles_%d of epoch_%d' % (datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f'),
                                                            i_cycle, i_epoch))

                # rollout
                rcd_obs, rcd_action, rcd_r, rcd_new_obs, rcd_done = [], [], [], [], []
                if not sanity_run and mpi_size == 1 and last_episode_step != 0:
                    # todo: use mpi_max(last_episode_step)
                    # dynamically set nb_rollout_steps
                    nb_rollout_steps = max(last_episode_step * 4, batch_size)
                logger.info('[%d, %d] rollout for %d steps.' % (total_nb_rollout, memory.nb_entries, nb_rollout_steps))
                t_rollout_start = time.time()

                for i_rollout in range(nb_rollout_steps):
                    rollout_log = i_cycle == 0
                    # 50% param_noise, 40% action_noise
                    action, q = agent.pi(obs, total_nb_rollout, compute_Q=True, rollout_log=rollout_log,
                                         apply_param_noise=i_rollout % 10 < 5, apply_action_noise=i_rollout % 10 > 5)
                    assert action.shape == env.action_space.shape
                    new_obs, r, done, reset, info = env.step(action)

                    if rank == 0 and render:
                        env.render()

                    episode_step += 1
                    total_nb_rollout += 1

                    if rollout_log:
                        summary_list = [('rollout/%s' % tp, info[tp]) for tp in ['rwd_walk', 'rwd_total']]
                        tp = 'rwd_agent'
                        summary_list += [('rollout/%s_x%d' % (tp, info['rf_agent']), info[tp] * info['rf_agent'])]
                        summary_list += [('rollout/q', q)]
                        if r != 0:
                            summary_list += [('rollout/q_div_r', q / r)]
                        agent.add_list_summary(summary_list, total_nb_rollout)

                    # store at the end of cycle to speed up MPI rollout
                    # agent.store_transition(obs, action, r, new_obs, done)
                    rcd_obs.append(obs)
                    rcd_action.append(action)
                    rcd_r.append(r)
                    rcd_new_obs.append(new_obs)
                    rcd_done.append(done)

                    obs = new_obs
                    if reset:
                        # Episode done.
                        last_episode_step = episode_step
                        episode_step = 0

                        agent.reset()
                        obs = env.reset()

                agent.store_multrans(memory, rcd_obs, rcd_action, rcd_r, rcd_new_obs, rcd_done)

                t_train_start = time.time()
                steps_per_second = float(nb_rollout_steps) / (t_train_start - t_rollout_start)
                agent.add_list_summary([('rollout/steps_per_second', steps_per_second)], total_nb_rollout)

                # Train.
                if not sanity_run:
                    # dynamically set nb_train_steps
                    if memory.nb_entries > batch_size * 20:
                        # using 1% of data for training every step?
                        nb_train_steps = max(int(memory.nb_entries * 0.01 / batch_size), 1)
                    else:
                        nb_train_steps = 0
                logger.info('[%d] training for %d steps.' % (total_nb_train, nb_train_steps))
                for _ in range(nb_train_steps):
                    # Adapt param noise, if necessary.
                    if memory.nb_entries >= batch_size and total_nb_train % param_noise_adaption_interval == 0:
                        agent.adapt_param_noise(total_nb_train)

                    agent.train_main(total_nb_train)
                    agent.update_target_net()
                    total_nb_train += 1

                if i_epoch == 0 and i_cycle < 5:
                    rollout_duration = t_train_start - t_rollout_start
                    train_duration = time.time() - t_train_start
                    logger.info('rollout_time(%d) = %.3fs, train_time(%d) = %.3fs' % (
                        nb_rollout_steps, rollout_duration, nb_train_steps, train_duration))
                    logger.info('rollout_speed=%.3fs/step, train_speed = %.3fs/step' % (
                        np.divide(rollout_duration, nb_rollout_steps), np.divide(train_duration, nb_train_steps)))

            logger.info('')
            mpi_size = MPI.COMM_WORLD.Get_size()
            # Log stats.
            stats = agent.get_stats(memory)
            combined_stats = stats.copy()

            def as_scalar(x):
                if isinstance(x, np.ndarray):
                    assert x.size == 1
                    return x[0]
                elif np.isscalar(x):
                    return x
                else:
                    raise ValueError('expected scalar, got %s' % x)
            combined_stats_sums = MPI.COMM_WORLD.allreduce(np.array([as_scalar(x) for x in combined_stats.values()]))
            combined_stats = {k: v / mpi_size for (k, v) in zip(combined_stats.keys(), combined_stats_sums)}

            # exclude logging zobs_dbg_%d, zobs_dbg_%d_normalized
            summary_list = [(key, combined_stats[key]) for key, v in combined_stats.items() if 'dbg' not in key]
            agent.add_list_summary(summary_list, i_epoch)

            # only print out train stats for epoch_0 for sanity check
            if i_epoch > 0:
                combined_stats = {}

            # Evaluation and statistics.
            if eval_env is not None:
                logger.info('[%d, %d] run evaluation' % (i_epoch, total_nb_eval))
                total_nb_eval = eval_episode(eval_env, render_eval, agent, combined_stats, total_nb_eval)

            logger.info('epoch %d duration: %.2f mins' % (i_epoch, (time.time() - t_epoch_start) / 60))
            for key in sorted(combined_stats.keys()):
                logger.record_tabular(key, combined_stats[key])
            logger.dump_tabular()
            logger.info('')

            if rank == 0:
                agent.store_ckpt(os.path.join(logdir, '%s.ckpt' % 'ddpg'), i_epoch)

def eval_episode(eval_env, render_eval, agent, combined_stats, total_nb_eval=0):
    steps = 0
    submit_episode_reward = 0.
    reset = False
    obs = eval_env.reset()

    while not reset:
        action, q = agent.pi(obs, total_nb_eval, apply_param_noise=False, apply_action_noise=False, compute_Q=True)
        obs, r, _, reset, info = eval_env.step(action)

        steps += 1
        total_nb_eval += 1
        if render_eval:
            eval_env.render()
        submit_episode_reward += r

        summary_list = [('the_action/%d_eval' % i, a) for i, a in enumerate(action)]
        summary_list += [('eval/q', q)]
        if r != 0:
            summary_list += [('eval/q_div_r', q / r)]

        summary_list += [('eval/%s' % tp, info[tp]) for tp in ['rwd_walk', 'rwd_total',
                                                               'rwd_blnc', 'rwd_dist', 'rwd_enrg']]
        tp = 'rwd_agent'
        summary_list += [('eval/%s_x%d' % (tp, info['rf_agent']), info[tp] * info['rf_agent'])]
        summary_list += [('eval_obs/%d' % i, o) for i, o in enumerate(obs[:agent.nb_demo_kine])]
        agent.add_list_summary(summary_list, total_nb_eval)

    agent.add_list_summary([('eval/episode_steps', steps)], total_nb_eval)
    combined_stats['eval/episode_return'] = submit_episode_reward
    combined_stats['eval/episode_steps'] = steps
    return total_nb_eval

# MPI not used in testing mode
def test(eval_env, agent, render_eval=True, nb_epochs=1, start_ckpt=None, **kwargs):
    logger.info('Start testing:', start_ckpt, '\n')
    with tf_util.single_threaded_session() as sess:
        agent.initialize(sess, start_ckpt=start_ckpt)
        sess.graph.finalize()

        for _ in range(nb_epochs):
            combined_stats = {}
            eval_episode(eval_env, render_eval, agent, combined_stats)

            for key in sorted(combined_stats.keys()):
                logger.record_tabular(key, combined_stats[key])
            logger.dump_tabular()
            logger.info('')

def config(group):
    group.add_argument('--nb_epochs', type=int, default=1)
    group.add_argument('--nb_epoch_cycles', type=int, default=2)
    group.add_argument('--nb_rollout_steps', type=int, default=25)
    group.add_argument('--nb_train_steps', type=int, default=10)
    group.add_argument('--start_ckpt', type=str, default='')
    group.add_argument('--hist_files', type=str, default='')
    group.add_argument('--demo_files', type=str,
                       default='demo/train_pos.csv,demo/train_neg.csv,demo/test0.csv,demo/test1.csv,demo/test2.csv')


# --------------------------------------------------------
# dota_evaluation_task1
# Licensed under The MIT License [see LICENSE for details]
# Written by Jian Ding, based on code from Bharath Hariharan
# --------------------------------------------------------

"""
    To use the code, users should to config detpath, annopath and imagesetfile
    detpath is the path for 15 result files, for the format, you can refer to "http://captain.whu.edu.cn/DOTAweb/tasks.html"
    search for PATH_TO_BE_CONFIGURED to config the paths
    Note, the evaluation is on the large scale images
"""
import xml.etree.ElementTree as ET
import os
#import cPickle
import numpy as np
import cv2
from mmdetection.polyiou import polyiou
from functools import partial
import glob

def parse_gt(filename):
    """
    :param filename: ground truth file to parse
    :return: all instances in a picture
    """
    objects = []
    with  open(filename, 'r') as f:
        while True:
            line = f.readline()
            if line:
                splitlines = line.strip().split(' ')
                object_struct = {}
                if (len(splitlines) < 9):
                    continue
                object_struct['name'] = splitlines[8]

                if (len(splitlines) == 9):
                    object_struct['difficult'] = 0
                elif (len(splitlines) == 10):
                    object_struct['difficult'] = int(splitlines[9])
                object_struct['bbox'] = [float(splitlines[0]),
                                         float(splitlines[1]),
                                         float(splitlines[2]),
                                         float(splitlines[3]),
                                         float(splitlines[4]),
                                         float(splitlines[5]),
                                         float(splitlines[6]),
                                         float(splitlines[7])]
                objects.append(object_struct)
            else:
                break
    return objects

# parse_gt('/content/EfficientDet/DOTA_split_600_150/val/labelTxt/P0003_0000_0000.txt')

def voc_ap(rec, prec, use_07_metric=False):
    """ ap = voc_ap(rec, prec, [use_07_metric])
    Compute VOC AP given precision and recall.
    If use_07_metric is true, uses the
    VOC 07 11 point method (default:False).
    """
    if use_07_metric:
        # 11 point metric
        ap = 0.
        for t in np.arange(0., 1.1, 0.1):
            if np.sum(rec >= t) == 0:
                p = 0
            else:
                p = np.max(prec[rec >= t])
            ap = ap + p / 11.
    else:
        # correct AP calculation
        # first append sentinel values at the end
        mrec = np.concatenate(([0.], rec, [1.]))
        mpre = np.concatenate(([0.], prec, [0.]))

        # compute the precision envelope
        for i in range(mpre.size - 1, 0, -1):
            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

        # to calculate area under PR curve, look for points
        # where X axis (recall) changes value
        i = np.where(mrec[1:] != mrec[:-1])[0]

        # and sum (\Delta recall) * prec
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap

def voc_eval_dota(detpath,
             annopath,
             imagesetfile,
             classname,
            # cachedir,
             ovthresh=0.5,
             use_07_metric=False):
    """rec, prec, ap = voc_eval(detpath,
                                annopath,
                                imagesetfile,
                                classname,
                                [ovthresh],
                                [use_07_metric])
    Top level function that does the PASCAL VOC evaluation.
    detpath: Path to detections
        detpath.format(classname) should produce the detection results file.
    annopath: Path to annotations
        annopath.format(imagename) should be the xml annotations file.
    imagesetfile: Text file containing the list of images, one image per line.
    classname: Category name (duh)
    cachedir: Directory for caching the annotations
    [ovthresh]: Overlap threshold (default = 0.5)
    [use_07_metric]: Whether to use VOC07's 11 point AP computation
        (default False)
    """
    # assumes detections are in detpath.format(classname)
    # assumes annotations are in annopath.format(imagename)
    # assumes imagesetfile is a text file with each line an image name
    # cachedir caches the annotations in a pickle file

    # first load gt
    #if not os.path.isdir(cachedir):
     #   os.mkdir(cachedir)
    #cachefile = os.path.join(cachedir, 'annots.pkl')
    # read list of images
    with open(imagesetfile, 'r') as f:
        lines = f.readlines()
    imagenames = [x.strip() for x in lines]
    #print('imagenames: ', imagenames)
    #if not os.path.isfile(cachefile):
        # load annots
    recs = {}
    for i, imagename in enumerate(imagenames):
        # print('parse_files name: ', annopath.format(imagename))
        recs[imagename] = parse_gt(annopath.format(imagename))
        # recs[imagename] = parse_gt(annopath.format(osp.basename(imagename).split('.png')[0]))
    # extract gt objects for this class
    class_recs = {}
    npos = 0
    for imagename in imagenames:
        R = [obj for obj in recs[imagename] if obj['name'] == classname]
        bbox = np.array([x['bbox'] for x in R])
        difficult = np.array([x['difficult'] for x in R]).astype(np.bool)
        det = [False] * len(R)
        npos = npos + sum(~difficult)
        class_recs[imagename] = {'bbox': bbox,
                                 'difficult': difficult,
                                 'det': det}

    # read dets from Task1* files
    detfile = detpath.format(classname)
    with open(detfile, 'r') as f:
        lines = f.readlines()

    splitlines = [x.strip().split(' ') for x in lines]
    image_ids = [x[0] for x in splitlines]
    confidence = np.array([float(x[1]) for x in splitlines])

    #print('check confidence: ', confidence)

    BB = np.array([[float(z) for z in x[2:]] for x in splitlines])

    # sort by confidence
    sorted_ind = np.argsort(-confidence)
    sorted_scores = np.sort(-confidence)

    #print('check sorted_scores: ', sorted_scores)
    #print('check sorted_ind: ', sorted_ind)

    ## note the usage only in numpy not for list
    BB = BB[sorted_ind, :]
    image_ids = [image_ids[x] for x in sorted_ind]
    #print('check imge_ids: ', image_ids)
    #print('imge_ids len:', len(image_ids))
    # go down dets and mark TPs and FPs
    nd = len(image_ids)
    tp = np.zeros(nd)
    fp = np.zeros(nd)
    for d in range(nd):
        R = class_recs[image_ids[d]]
        bb = BB[d, :].astype(float)
        ovmax = -np.inf
        BBGT = R['bbox'].astype(float)

        ## compute det bb with each BBGT

        if BBGT.size > 0:
            # compute overlaps
            # intersection

            # 1. calculate the overlaps between hbbs, if the iou between hbbs are 0, the iou between obbs are 0, too.
            # pdb.set_trace()
            BBGT_xmin =  np.min(BBGT[:, 0::2], axis=1)
            BBGT_ymin = np.min(BBGT[:, 1::2], axis=1)
            BBGT_xmax = np.max(BBGT[:, 0::2], axis=1)
            BBGT_ymax = np.max(BBGT[:, 1::2], axis=1)
            bb_xmin = np.min(bb[0::2])
            bb_ymin = np.min(bb[1::2])
            bb_xmax = np.max(bb[0::2])
            bb_ymax = np.max(bb[1::2])

            ixmin = np.maximum(BBGT_xmin, bb_xmin)
            iymin = np.maximum(BBGT_ymin, bb_ymin)
            ixmax = np.minimum(BBGT_xmax, bb_xmax)
            iymax = np.minimum(BBGT_ymax, bb_ymax)
            iw = np.maximum(ixmax - ixmin + 1., 0.)
            ih = np.maximum(iymax - iymin + 1., 0.)
            inters = iw * ih

            # union
            uni = ((bb_xmax - bb_xmin + 1.) * (bb_ymax - bb_ymin + 1.) +
                   (BBGT_xmax - BBGT_xmin + 1.) *
                   (BBGT_ymax - BBGT_ymin + 1.) - inters)

            overlaps = inters / uni

            BBGT_keep_mask = overlaps > 0
            BBGT_keep = BBGT[BBGT_keep_mask, :]
            BBGT_keep_index = np.where(overlaps > 0)[0]
            # pdb.set_trace()
            def calcoverlaps(BBGT_keep, bb):
                overlaps = []
                for index, GT in enumerate(BBGT_keep):

                    overlap = polyiou.iou_poly(polyiou.VectorDouble(BBGT_keep[index]), polyiou.VectorDouble(bb))
                    overlaps.append(overlap)
                return overlaps
            if len(BBGT_keep) > 0:
                overlaps = calcoverlaps(BBGT_keep, bb)

                ovmax = np.max(overlaps)
                jmax = np.argmax(overlaps)
                # pdb.set_trace()
                jmax = BBGT_keep_index[jmax]
        if ovmax > ovthresh:
            if not R['difficult'][jmax]:
                if not R['det'][jmax]:
                    tp[d] = 1.
                    R['det'][jmax] = 1
                else:
                    fp[d] = 1.
        else:
            fp[d] = 1.

    # compute precision recall

    # print('check fp:', fp)
    # print('check tp', tp)

    # print('npos num:', npos)
    fp = np.cumsum(fp)
    tp = np.cumsum(tp)

    rec = tp / float(npos)
    # avoid divide by zero in case the first detection matches a difficult
    # ground truth
    prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
    ap = voc_ap(rec, prec, use_07_metric)

    return rec, prec, ap
# from help_utils import tools
wordname_15 = ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle', 'ship', 'tennis-court',
                'basketball-court', 'storage-tank',  'soccer-ball-field', 'roundabout', 'harbor', 'swimming-pool', 'helicopter']
def write_voc_results_file(all_boxes, test_imgid_list, det_save_dir):
  '''
  :param all_boxes: is a list. each item reprensent the detections of a img.
  the detections is a array. shape is [-1, 7]. [category, score, x, y, w, h, theta]
  Note that: if none detections in this img. that the detetions is : []
  :param test_imgid_list:
  :param det_save_path:
  :return:
  '''
  for cls_id,cls  in enumerate(wordname_15):
    if cls == 'back_ground':
      continue
    print("Writing {} VOC resutls file".format(cls))

    # tools.mkdir(det_save_dir)
    det_save_path = os.path.join(det_save_dir, "Task1_"+cls+".txt")
    with open(det_save_path, 'wt') as f:
      for index, img_name in enumerate(test_imgid_list):
        this_img_detections = all_boxes[index]

        this_cls_detections = this_img_detections[this_img_detections[:, 0] == cls_id]
        if this_cls_detections.shape[0] == 0:
          continue # this cls has none detections in this img
        for a_det in this_cls_detections:
          f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.
                  format( os. path. splitext(os.path.basename(img_name))[0], a_det[1],
                         a_det[2], a_det[3],
                         a_det[4], a_det[5],
                         a_det[6],a_det[7],
                         a_det[8],a_det[9]))  # that is [img_name, score, x, y, w, h, theta]

def poly_nms(dets,scores, iou_threshold):
      # scores = dets[:, 8]
      polys = []
      areas = []
      for i in range(len(dets)):
          tm_polygon = polyiou.VectorDouble([dets[i][0], dets[i][1],
                                              dets[i][2], dets[i][3],
                                              dets[i][4], dets[i][5],
                                              dets[i][6], dets[i][7]])
          polys.append(tm_polygon)
      order = scores.argsort()[::-1]

      keep = []
      while order.size > 0:
          ovr = []
          i = order[0]
          keep.append(i)
          for j in range(order.size - 1):
              iou = polyiou.iou_poly(polys[i], polys[order[j + 1]])
              ovr.append(iou)
          ovr = np.array(ovr)
          inds = np.where(ovr <= iou_threshold)[0]
          order = order[inds + 1]
      return keep

def py_cpu_nms_poly_fast(dets,scores, thresh):
    try:
        obbs = dets[:, 0:-1]
    except:
        print('fail index')
        pdb.set_trace()
    x1 = np.min(obbs[:, 0::2], axis=1)
    y1 = np.min(obbs[:, 1::2], axis=1)
    x2 = np.max(obbs[:, 0::2], axis=1)
    y2 = np.max(obbs[:, 1::2], axis=1)
    # scores = dets[:, 8]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    polys = []
    for i in range(len(dets)):
        tm_polygon = polyiou.VectorDouble([dets[i][0], dets[i][1],
                                            dets[i][2], dets[i][3],
                                            dets[i][4], dets[i][5],
                                            dets[i][6], dets[i][7]])
        polys.append(tm_polygon)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        ovr = []
        i = order[0]
        keep.append(i)
        # if order.size == 0:
        #     break
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        # w = np.maximum(0.0, xx2 - xx1 + 1)
        # h = np.maximum(0.0, yy2 - yy1 + 1)
        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)
        hbb_inter = w * h
        hbb_ovr = hbb_inter / (areas[i] + areas[order[1:]] - hbb_inter)
        # h_keep_inds = np.where(hbb_ovr == 0)[0]
        h_inds = np.where(hbb_ovr > 0)[0]
        tmp_order = order[h_inds + 1]
        for j in range(tmp_order.size):
            iou = polyiou.iou_poly(polys[i], polys[tmp_order[j]])
            hbb_ovr[h_inds[j]] = iou
            # ovr.append(iou)
            # ovr_index.append(tmp_order[j])

        # ovr = np.array(ovr)
        # ovr_index = np.array(ovr_index)
        # print('ovr: ', ovr)
        # print('thresh: ', thresh)
        try:
            if math.isnan(ovr[0]):
                pdb.set_trace()
        except:
            pass
        inds = np.where(hbb_ovr <= thresh)[0]

        # order_obb = ovr_index[inds]
        # print('inds: ', inds)
        # order_hbb = order[h_keep_inds + 1]
        order = order[inds + 1]
        # pdb.set_trace()
        # order = np.concatenate((order_obb, order_hbb), axis=0).astype(np.int)
    return keep

def poly_nms_np(quads,scores,lables):
    quadrangles = np.reshape(quads,(-1,8))
    idx = py_cpu_nms_poly_fast(quadrangles.astype(np.float64),scores,.1)
    # print(idx)
    quads = quads[idx]
    scores = scores[idx]
    lables = lables[idx]
    return quads, scores, lables

def rotate(ct,angle):
      # rotation_amount_rad = angle * np.pi / 180.0
      # c1=(ct[:,0].mean(),ct[:,1].mean())
      # a=np.cos(rotation_amount_rad)
      # b=np.sin(rotation_amount_rad)
    #   M=np.array([[  a , b,  (1-a)*c1[0] -  b*c1[1]    ] ,
    #               [  -b , a,   b*c1[0]    + (1-a)*c1[1] ]])
      M = cv2.getRotationMatrix2D(tuple(ct.mean(axis=0)),angle, 1.0)
      # print(M)
      ones = np.ones(shape=(len(ct), 1))

      points_ones = np.hstack([ct, ones])

      # transform points
      return  (M.dot(points_ones.T).T).astype(np.float32)

def result_to_dets(result):
    d = [ np.concatenate((r, np.repeat(i,len(r))[..., np.newaxis] ),axis=1)
                            for i,r in enumerate( result ) if len(r)>0]
    # print(len(d))
    if len(d):
      dets = np.concatenate(d)
      angles = dets[...,0]
      quads_off = dets[...,1:5]
      scores = dets[...,5]
      lables = dets[...,6]

      quads = np.zeros((quads_off.shape[0] , 8))
      quads[:, 0] = quads_off[:, 0]
      quads[:, 1] = quads_off[:, 1]
      quads[:, 2] = quads_off[:, 2]
      quads[:, 3] = quads_off[:, 1]
      quads[:, 4] = quads_off[:, 2]
      quads[:, 5] = quads_off[:, 3]
      quads[:, 6] = quads_off[:, 0]
      quads[:, 7] = quads_off[:, 3]
      quads = np.reshape(quads,(-1,4,2))
      quads = np.array([rotate(q,-r)  for q,r in zip(quads,angles)])
      return quads, scores, lables
    else:
      quads = np.zeros((1,8))
      scores = np.zeros((1))
      lables = np.zeros((1))
      return quads, scores, lables


import os, inspect
currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
print("current_dir=" + currentdir)
os.sys.path.insert(0, currentdir)

from copy import deepcopy
import itertools
import numpy as np
from torch.optim import Adam
import gym
import time
from pytorch_shared import *
import torch
import torch.nn as nn
from common import *
from tensorboardX import SummaryWriter

import pointMass
from gym import wrappers
# @title SAC Model{ display-mode: "form" }
class SAC_model():

    def __init__(self, act_limit, obs_dim, act_dim, hidden_sizes, lr=1e-3, gamma=0.99, alpha=0.2, polyak=0.995,
                 load=False, exp_name='Exp1', replay_buffer = None, path='saved_models/', curiosity=False, curiosity_weighting = 1000,):
        self.act_limit = act_limit
        self.gamma = gamma
        self.alpha = alpha
        self.polyak = polyak
        self.load = load
        self.exp_name = exp_name
        self.path = path
        self.lr = lr
        self.curiosity = curiosity
        self.curiosity_weighting = curiosity_weighting
        self.create_networks(obs_dim, act_dim, hidden_sizes)
        self.replay_buffer = replay_buffer
        self.CQL  = True
        self.AWAC = False

        torch.set_num_threads(torch.get_num_threads())

    def create_networks(self, obs_dim, act_dim, hidden_sizes=[32, 32], batch_size=100, activation=nn.ReLU):

        # Get Env dimensions
        # Action limit for clamping: critically, assumes all dimensions share the same bound!
        # Create actor-critic module and target networks
        self.ac = MLPActorCritic(self.act_limit, obs_dim, act_dim, hidden_sizes=hidden_sizes, activation=activation)
        print(self.load,'-------------------')
        self.actor  = self.ac # TODO COMPATABILITY
        if self.load:
            self.load_weights()
        self.ac_targ = deepcopy(self.ac)

        # Freeze target networks with respect to optimizers (only update via polyak averaging)
        for p in self.ac_targ.parameters():
            p.requires_grad = False

        # List of parameters for both Q-networks (save this for convenience)
        self.q_params = itertools.chain(self.ac.q1.parameters(), self.ac.q2.parameters())

        # Count variables (protip: try to get a feel for how different size networks behave!)
        var_counts = tuple(count_vars(module) for module in [self.ac.pi, self.ac.q1, self.ac.q2])
        print('\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d\n' % var_counts)

        # Set up optimizers for policy and q-function
        self.pi_optimizer = Adam(self.ac.pi.parameters(), lr=self.lr)
        self.q_optimizer = Adam(self.q_params, lr=self.lr)

        if self.curiosity:
            self.forward_model = mlp([obs_dim + act_dim] + list(hidden_sizes) + [obs_dim]).cuda()
            self.forward_model_optimizer = Adam(self.forward_model.parameters(), lr=self.lr)

    # Set up function for computing SAC Q-losses
    def compute_loss_q(self,data):
        o, a, r, o2, d = data['obs'], data['act'], data['rew'], data['obs2'], data['done']
        q1 = self.ac.q1(o, a)
        q2 = self.ac.q2(o, a)

        if self.curiosity:
            # update our prediction model
            o_next = self.forward_model(torch.cat([data['obs'], data['act']], axis=1))
            prediction_loss = ((o_next - data['obs2']) ** 2).mean()
            prediction_loss.backward()
            self.forward_model_optimizer.step()

        # Bellman backup for Q functions
        with torch.no_grad():
            # Target actions come from *current* policy
            a2, logp_a2, _ = self.ac.pi(o2)

            # Target Q-values
            q1_pi_targ = self.ac_targ.q1(o2, a2)
            q2_pi_targ = self.ac_targ.q2(o2, a2)
            q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ)
            backup = r + self.gamma * (1 - d) * (q_pi_targ - self.alpha * logp_a2)

            if self.curiosity:
                backup += prediction_loss * self.curiosity_weighting # curiosity weighting needs to be relatively large as prediction error is pretty small.
            # only need to do this because of PER
            #self.replay_buffer.update_priorities(data['PER_tree_idxs'], backup.cpu().numpy())

        # MSE loss against Bellman backup
        loss_q1 = ((q1 - backup) ** 2).mean()
        loss_q2 = ((q2 - backup) ** 2).mean()
        loss_q = loss_q1 + loss_q2

        # Conservative Q Loss
        # We want to min the q value of random actions, while maximising the q values of data distrib
        a_random = (torch.FloatTensor(o.shape[0], a.shape[-1]).uniform_(-1, 1)*self.act_limit).cuda()
        q_random = torch.log(torch.exp( (self.ac.q1(o, a_random)  / a_random.norm(dim=-1)).mean() ) + torch.exp( (self.ac.q2(o, a_random)  / a_random.norm(dim=-1)).mean()))

        q_data = self.ac.q1(o, a).mean() + self.ac.q2(o, a).mean()
        CQL = q_random  - q_data # We are minimisng this

        if self.CQL:
            loss_q += CQL

        # Useful info for logging
        q_info = dict(Q1Vals=q1.detach().cpu().numpy(),
                      Q2Vals=q2.detach().cpu().numpy())

        return loss_q, q_info

    # Set up function for computing SAC pi loss
    def compute_loss_pi(self,data):
        o = data['obs']
        pi, logp_pi, distrib = self.ac.pi(o)

        if self.AWAC:
            a = data['act']
            logprob = distrib.log_prob(a).sum(axis=-1)
            with torch.no_grad():
                adv_weighting = torch.min( self.ac.q1(o, a), self.ac.q2(o,a) )
                Lambda = 0.3

            loss = -logprob
            loss = loss * torch.exp((1/Lambda) * adv_weighting)
            loss = loss.mean()
            return loss, None, loss
        else:
            q1_pi = self.ac.q1(o, pi)
            q2_pi = self.ac.q2(o, pi)
            q_pi = torch.min(q1_pi, q2_pi)

            # Entropy-regularized policy loss
            loss_pi = (self.alpha * logp_pi - q_pi).mean()

            # Useful info for logging
            pi_info = dict(LogPi=logp_pi.detach().cpu().numpy())

            return loss_pi, pi_info, q_pi.mean()

    def polyak_target_update(self):
        with torch.no_grad():
            for p, p_targ in zip(self.ac.parameters(), self.ac_targ.parameters()):
                # NB: We use an in-place operations "mul_", "add_" to update target
                # params, as opposed to "mul" and "add", which would make new tensors.
                p_targ.data.mul_(self.polyak)
                p_targ.data.add_((1 - self.polyak) * p.data)

    def q_update(self,data):
        self.q_optimizer.zero_grad()
        loss_q, q_info = self.compute_loss_q(data)
        loss_q.backward()
        self.q_optimizer.step()
        return loss_q

    def update(self,data):
        # First run one gradient descent step for Q1 and Q2
        loss_q = self.q_update(data)
        # # Record things
        # logger.store(LossQ=loss_q.item(), **q_info)
        # Freeze Q-networks so you don't waste computational effort
        # computing gradients for them during the policy learning step.
        for p in self.q_params:
            p.requires_grad = False

        # Next run one gradient descent step for pi.
        self.pi_optimizer.zero_grad()
        loss_pi, pi_info, q_pi = self.compute_loss_pi(data)
        loss_pi.backward()
        self.pi_optimizer.step()

        # Unfreeze Q-networks so you can optimize it at next step.
        for p in self.q_params:
            p.requires_grad = True

        # Finally, update target networks by polyak averaging.
        self.polyak_target_update()
        return loss_q, loss_pi, q_pi

    def supervised_update(self, data, supervised_loss,  use_RL=True):
        # First run one gradient descent step for Q1 and Q2

        if use_RL:
            self.q_update(data)
            for p in self.q_params:
                p.requires_grad = False
            loss_pi, pi_info = self.compute_loss_pi(data)
            loss_pi += supervised_loss
        else:
            loss_pi = supervised_loss
        loss_pi.backward()
        self.pi_optimizer.step()
        # Unfreeze Q-networks so you can optimize it at next step.
        for p in self.q_params:
            p.requires_grad = True

        # Finally, update target networks by polyak averaging.
        self.polyak_target_update()

    def load_weights(self, path='saved_models/'):
        self.ac.load_state_dict(torch.load(path+'_SAC_'+self.exp_name+'.pth'))
        self.ac.eval()
        print('Successfully loaded')

    def save_weights(self, path='saved_models/'):
        print('Saved at ', path+'_SAC_'+self.exp_name+'.pth')
        torch.save(self.ac.state_dict(), path+'_SAC_'+self.exp_name+'.pth')

def SAC(env_fn, ac_kwargs=dict(), seed=0,
        steps_per_epoch=2000, epochs=100, replay_size=int(1e6), gamma=0.99,
        polyak=0.995, lr=1e-3, alpha=0.2, batch_size=100, start_steps=200,
        max_ep_len=1000, save_freq=1, load=False, exp_name="Experiment_1", render=False):

    torch.manual_seed(seed) # TODO COMPAT
    np.random.seed(seed)
    env, test_env = env_fn(), env_fn()

    test_env.reset()

    # test_env.render(mode='human')
    # env = wrappers.FlattenDictWrapper(env, dict_keys=['observation', 'desired_goal'])
    # test_env = wrappers.FlattenDictWrapper(test_env, dict_keys=['observation', 'desired_goal'])
    # Get Env dimensions
    obs_dim = env.observation_space.shape[0]
    act_dim = env.action_space.shape[0]
    act_limit = env.action_space.high
    # Experience buffer
    replay_buffer = ReplayBuffer(obs_dim=obs_dim, act_dim=act_dim, size=replay_size)

    SAC = SAC_model(act_limit, obs_dim, act_dim, ac_kwargs['hidden_sizes'], lr, gamma, alpha, polyak, load,
                    exp_name)

    # Logging
    start_time = time.time()
    train_log_dir = 'logs/' + exp_name + str(int(start_time))
    summary_writer = SummaryWriter(train_log_dir)

    def update_models(model, replay_buffer, steps, batch_size):
        for j in range(steps):
            batch = replay_buffer.sample_batch(batch_size) #
            model.update(batch)

    # now collect epsiodes
    total_steps = steps_per_epoch * epochs
    steps_collected = 0

    if not load:
        # collect some initial random steps to initialise
        steps_collected += rollout_trajectories(n_steps=start_steps, env=env, max_ep_len=max_ep_len, actor='random',
                                                replay_buffer=replay_buffer, summary_writer=summary_writer,
                                                exp_name=exp_name)
        update_models(SAC, replay_buffer, steps=steps_collected, batch_size=batch_size)

    # now act with our actor, and alternately collect data, then train.
    while steps_collected < total_steps:
        # collect an episode
        steps_collected += rollout_trajectories(n_steps=max_ep_len, env=env, max_ep_len=max_ep_len,
                                                actor=SAC.actor.get_stochastic_action, replay_buffer=replay_buffer,
                                                summary_writer=summary_writer, current_total_steps=steps_collected,
                                                exp_name=exp_name)
        # take than many training steps
        update_models(SAC, replay_buffer, steps=max_ep_len, batch_size=batch_size)

        # if an epoch has elapsed, save and test.
        if steps_collected > 0 and steps_collected % steps_per_epoch == 0:
            SAC.save_weights()
            # Test the performance of the deterministic version of the agent.
            rollout_trajectories(n_steps=max_ep_len * 10, env=test_env, max_ep_len=max_ep_len,
                                 actor=SAC.actor.get_deterministic_action, summary_writer=summary_writer,
                                 current_total_steps=steps_collected, train=False, render=True, exp_name=exp_name)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--env', type=str, default='Pendulum-v0')
    parser.add_argument('--hid', type=int, default=128)
    parser.add_argument('--l', type=int, default=2)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--seed', '-s', type=int, default=0)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--max_ep_len', type=int, default=200)
    parser.add_argument('--exp_name', type=str, default='experiment_1')
    parser.add_argument('--load', type=str2bool, default=False)
    parser.add_argument('--render', type=str2bool, default=False)
    args = parser.parse_args()

    experiment_name = '_' + args.env + '_Hidden_' + str(args.hid) + 'l_' + str(args.l)

    SAC(lambda: gym.make(args.env),
        ac_kwargs=dict(hidden_sizes=[args.hid] * args.l),
        gamma=args.gamma, seed=args.seed, epochs=args.epochs, load=args.load, exp_name=experiment_name,
        max_ep_len=args.max_ep_len, render=args.render)


import os,sys
os.environ['PYOPENCL_CTX']='0:1'
os.environ['PYOPENCL_COMPILER_OUTPUT'] = '0'
#os.environ['CUDA_PROFILE'] = '1'
import chroma.api as api
api.use_opencl()
#api.use_cuda()
from unittest_find import unittest
import numpy as np
from chroma.sim import Simulation
from chroma.event import Photons
from chroma.geometry import Surface
import chroma.event as event
from chroma.uboone.uboonedet import ubooneDet
from chroma.gpu.photon import GPUPhotons

try:
    import ROOT as rt
    from rootpy.tree import Tree, TreeModel, FloatCol, IntCol
    from rootpy.io import root_open
    has_root = True
except:
    has_root = False
    raise ValueError("No ROOT")

if has_root:
    class PhotonData( TreeModel ):
        end_x = FloatCol()
        end_y = FloatCol()
        end_z = FloatCol()
        reflect_diffuse  = IntCol()
        reflect_specular = IntCol()
        bulk_scatter     = IntCol()
        bulk_absorb      = IntCol()
        surface_detect   = IntCol()
        surface_absorb   = IntCol()
        surface_reemit   = IntCol()

        def reset(self):
            self.reflect_diffuse  = 0
            self.reflect_specular = 0
            self.bulk_scatter     = 0
            self.bulk_absorb      = 0
            self.surface_detect   = 0
            self.surface_absorb   = 0
            self.surface_reemit   = 0

uboone_wireplane = Surface( 'uboone_wireplane' )
uboone_wireplane.nplanes = 3.0
uboone_wireplane.wire_pitch = 0.3
uboone_wireplane.wire_diameter = 0.015
uboone_wireplane.transmissive = 1
uboone_wireplane.model = Surface.SURFACE_WIREPLANE

def add_wireplane_surface( solid ):
    # function detector class will use to add a wireplane surface to the geometry
    # set surface for triangles on x=-1281.0 plane
    for n,triangle in enumerate(solid.mesh.triangles):
        nxplane = 0
        for ivert in triangle:
            if solid.mesh.vertices[ivert,0]==-1281.0:
                nxplane += 1
        if nxplane==3:
            print [ solid.mesh.vertices[x] for x in triangle ]
            solid.surface[ n ] = uboone_wireplane
            solid.unique_surfaces = np.unique( solid.surface )

class TestUbooneDetector(unittest.TestCase):
    def setUp(self):
        # UBOONE
        #daefile = "../gdml/microboone_nowires_chroma_simplified.dae"
        daefile = "dae/microboone_32pmts_nowires_cryostat.dae"
        #daefile = "dae/microboone_32pmts_nowires_cryostat_weldwireplanes.dae"
        self.geo = ubooneDet( daefile, detector_volumes=["vol_PMT_AcrylicPlate","volPaddle_PMT"],
                              wireplane_volumes=[('volTPCPlane_PV0x7f868ac5ef50',add_wireplane_surface)],
                              acrylic_detect=True, acrylic_wls=False,
                              read_bvh_cache=True, cache_dir="./uboone_cache",
                              dump_node_info=False)
        self.sim = Simulation(self.geo, geant4_processes=0, nthreads_per_block=1, max_blocks=1024)
        self.origin = self.geo.bvh.world_coords.world_origin

    @unittest.skip('skipping testDet')
    def testDet(self):

        # Run only one photon at a time
        nphotons = 1000000
        pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
        dir = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
        pol = np.zeros_like(pos)
        phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
        pol = np.cross( (np.cos(phi), np.sin(phi),0), dir )

        pol[:,0] = np.cos(phi)
        pol[:,1] = np.sin(phi)
        pol = np.cross( pol, dir )
        for n,p in enumerate(pol):
            norm = np.sqrt( p[0]*p[0] + p[1]*p[1] + p[2]*p[2] )
            p /= norm
        #print pol
        t = np.zeros(nphotons, dtype=np.float32) + 100.0 # Avoid negative photon times
        wavelengths = np.empty(nphotons, np.float32)
        wavelengths.fill(128.0)

        photons = Photons(pos=pos, dir=dir, pol=pol, t=t, wavelengths=wavelengths)
        hit_charges = []
        for ev in self.sim.simulate( (photons for i in xrange(1)), keep_photons_end=True, keep_photons_beg=False, ):
            ev.photons_end.dump_history()
            lht = ev.photons_end[0].last_hit_triangles

    #@unittest.skip('skipping testDet')
    def testPhotonBomb(self):

        # Run only one photon at a time
        #nphotons = 7200000
        nphotons = 256*100
        #nphotons = 256*10

        dphi = np.random.uniform(0,2.0*np.pi, nphotons)
        dcos = np.random.uniform(-1.0, 1.0, nphotons)
        dir = np.array( zip( np.sqrt(1-dcos[:]*dcos[:])*np.cos(dphi[:]), np.sqrt(1-dcos[:]*dcos[:])*np.sin(dphi[:]), dcos[:] ), dtype=np.float32 )

        pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
        pol = np.zeros_like(pos)
        phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
        pol[:,0] = np.cos(phi)
        pol[:,1] = np.sin(phi)
        pol = np.cross( pol, dir )
        for n,p in enumerate(pol):
            norm = np.sqrt( p[0]*p[0] + p[1]*p[1] + p[2]*p[2] )
            p /= norm

        t = np.zeros(nphotons, dtype=np.float32) + 100.0 # Avoid negative photon times
        wavelengths = np.empty(nphotons, np.float32)
        wavelengths.fill(128.0)

        photons = Photons(pos=pos, dir=dir, pol=pol, t=t, wavelengths=wavelengths)
        hit_charges = []

        if has_root:
            root_file = root_open("output_test_uboone.root", "recreate")
            root_tree = Tree("PhotonData", model=PhotonData )
            root_tree.reset()

        for ev in self.sim.simulate( (photons for i in xrange(1)), keep_photons_end=True, keep_photons_beg=False, max_steps=100):
            ev.photons_end.dump_history()
            lht = ev.photons_end[0].last_hit_triangles
            nhits = ev.channels.hit[ np.arange(0,32)[:] ]

            #print "nchannels: ",len(ev.channels.hit)
            #print nhits
            #print ev.channels.q
            #print ev.channels.t

            if False:
                # Fill Tree
                #print "save info for ",len(ev.photons_end)
                for photon in ev.photons_end:
                    root_tree.end_x = photon.pos[0]
                    root_tree.end_y = photon.pos[1]
                    root_tree.end_z = photon.pos[2]

                    root_tree.reflect_diffuse  = int( event.REFLECT_DIFFUSE & photon.flags )
                    root_tree.reflect_specular = int( event.REFLECT_SPECULAR & photon.flags )
                    root_tree.bulk_scatter     = int( event.RAYLEIGH_SCATTER & photon.flags )
                    root_tree.bulk_absorb      = int( event.BULK_ABSORB & photon.flags )
                    root_tree.surface_detect   = int( event.SURFACE_DETECT & photon.flags )
                    root_tree.surface_absorb   = int( event.SURFACE_ABSORB & photon.flags )
                    root_tree.surface_reemit   = int( event.SURFACE_REEMIT & photon.flags )
                    root_tree.fill()
        if has_root:
            root_tree.Write()

    @unittest.skip('skipping testDet')
    def testWorkQueue(self):

        # Run only one photon at a time
        nphotons = 32*12

        dphi = np.random.uniform(0,2.0*np.pi, nphotons)
        dcos = np.random.uniform(-1.0, 1.0, nphotons)
        dir = np.array( zip( np.sqrt(1-dcos[:]*dcos[:])*np.cos(dphi[:]),
                             np.sqrt(1-dcos[:]*dcos[:])*np.sin(dphi[:]),
                             dcos[:] ), dtype=np.float32 )
        pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
        pol = np.zeros_like(pos)
        phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
        pol[:,0] = np.cos(phi)
        pol[:,1] = np.sin(phi)
        pol = np.cross( pol, dir )
        for n,p in enumerate(pol):
            norm = np.sqrt( p[0]*p[0] + p[1]*p[1] + p[2]*p[2] )
            p /= norm

        t = np.zeros(nphotons, dtype=np.float32) + 100.0 # Avoid negative photon times
        wavelengths = np.empty(nphotons, np.float32)
        wavelengths.fill(128.0)

        photons = Photons(pos=pos, dir=dir, pol=pol, t=t, wavelengths=wavelengths)

        #rq = RayQueue( self.sim.context )
        #rq.checknodes.print_dev_info()
        #rq.simulate( photons, self.sim )

if __name__ == "__main__":
    unittest.main()
    pycuda.driver.stop_profiler()


#coding=utf-8
import numpy as np
import cv2
import json
from caffe2.python import workspace

from core.config import merge_cfg_from_file
import core.test_engine as infer_engine
import datasets.dummy_datasets as dummy_datasets
import utils.c2 as c2_utils
import core.test_engine as infer_engine

c2_utils.import_detectron_ops()
# OpenCL may be enabled by default in OpenCV3; disable it because it's not
# thread safe and causes unwanted GPU memory allocations.
cv2.ocl.setUseOpenCL(False)

class Model():

    def __init__(self,cfg_path,weights_path):

        #nms_same_class 0.3  ----  *.yaml/TEST:NMS:0.3 中设置 defalut 0.5

        self.gpu_id = 0               #gpu_id default 0

        self.score_thresh = 0.4       #score > score_thresh  default 0.3

        self.per_class_thresh = False         #score > class_score_thresh
        self.autotruck_score_thresh = 0.6
        self.forklift_score_thresh = 0.65
        self.digger_score_thresh = 0.65
        self.car_score_thresh = 0.45
        self.bus_score_thresh = 0.0
        self.tanker_score_thresh = 0.55
        self.person_score_thresh = 0.35
        self.minitruck_score_thresh = 0.0
        self.minibus_score_thresh = 0.59

        self.class_nms_thresh = 0.85   #nms_between_classes  IOU > class_nms_thresh    default 0.9
        merge_cfg_from_file(cfg_path)
        self.model = infer_engine.initialize_model_from_cfg(weights_path,self.gpu_id)
        self.dummy_coco_dataset = dummy_datasets.get_steal_oil_class14_dataset()
        print ("model is ok")

    def predict(self,im):

        #class_str_list = []
        data_list = []

        with c2_utils.NamedCudaScope(self.gpu_id):
            cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(self.model, im, None, None
            )

        #get box classes
        if isinstance(cls_boxes, list):
            boxes, segms, keypoints, classes = self.convert_from_cls_format(cls_boxes, cls_segms, cls_keyps)
        if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < self.score_thresh:
            return data_list
        #get score
        areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
        sorted_inds = np.argsort(-areas)

        #no nms between classes
        '''im1 = cv2.cvtColor(im,cv2.COLOR_RGB2BGR)
        result1= im1.copy()
        for i in sorted_inds:

            bbox = boxes[i, :4]
            score = boxes[i, -1]
            if score < self.score_thresh:
                continue
            #get class-str
            class_str = self.get_class_string(classes[i], score, self.dummy_coco_dataset)
            cv2.rectangle(result1,(int(bbox[0]),int(bbox[1])),(int(bbox[2]),int(bbox[3])),(255,255,0),1)
            font = cv2.FONT_HERSHEY_SIMPLEX
            ((txt_w, txt_h), _) = cv2.getTextSize(class_str, font, 0.35, 1)
            txt_tl = int(bbox[0]), int(bbox[1]) - int(0.3 * txt_h)
            cv2.putText(result1, class_str, txt_tl, font, 0.35, (218, 227, 218), lineType=cv2.LINE_AA)
            txt_tl = int(bbox[0])+txt_w, int(bbox[1]) - int(0.3 * txt_h)
            cv2.putText(result1, ('%.2f' % score), txt_tl, font, 0.35, (218, 227, 218), lineType=cv2.LINE_AA)
        cv2.imwrite("test1.jpg", result1)'''

        #nms between classes
        #im2 = cv2.cvtColor(im,cv2.COLOR_RGB2BGR)
        #result2= im2.copy()
        if (len(sorted_inds) > 0):
            nmsIndex = self.nms_between_classes(boxes, self.class_nms_thresh)  #阈值为0.9，阈值越大，过滤的越少
            for i in xrange(len(nmsIndex)):
                bbox = boxes[nmsIndex[i], :4]
                score = boxes[nmsIndex[i], -1]
                if score < self.score_thresh:
                    continue
                #get class-str
                class_str = self.get_class_string(classes[nmsIndex[i]], score, self.dummy_coco_dataset)

                #score thresd per class
                if self.per_class_thresh:
                    if 'autotruck' == class_str and score < self.autotruck_score_thresh:
                        continue
                    if 'forklift'  == class_str and score < self.forklift_score_thresh:
                        continue
                    if 'digger'    == class_str and score < self.digger_score_thresh:
                        continue
                    if 'car'       == class_str and score < self.car_score_thresh:
                        continue
                    if 'bus'       == class_str and score < self.bus_score_thresh:
                        continue
                    if 'tanker'    == class_str and score < self.tanker_score_thresh:
                        continue
                    if 'person'    == class_str and score < self.person_score_thresh:
                        continue
                    if 'minitruck' == class_str and score < self.minitruck_score_thresh:
                        continue
                    if 'minibus'   == class_str and score < self.minibus_score_thresh:
                        continue

                single_data = {"cls":class_str,"score":float('%.2f' % score),"bbox":{"xmin":int(bbox[0]),"ymin":int(bbox[1]),"xmax":int(bbox[2]),"ymax":int(bbox[3])}}
                data_list.append(single_data)

                '''cv2.rectangle(result2,(int(bbox[0]),int(bbox[1])),(int(bbox[2]),int(bbox[3])),(255,255,0),1)
                font = cv2.FONT_HERSHEY_SIMPLEX
                ((txt_w, txt_h), _) = cv2.getTextSize(class_str, font, 0.55, 1)
                txt_tl = int(bbox[0]), int(bbox[1]) - int(0.3 * txt_h)
                cv2.putText(result2, class_str, txt_tl, font, 0.55, (218, 227, 218), lineType=cv2.LINE_AA)
                txt_tl = int(bbox[0])+txt_w, int(bbox[1]) - int(0.3 * txt_h)
                cv2.putText(result2, ('%.2f' % score), txt_tl, font, 0.35, (218, 227, 218), lineType=cv2.LINE_AA)'''
        #cv2.imwrite("test2.jpg", result2)

        #construcrion - data_list
        return data_list

    def convert_from_cls_format(self,cls_boxes, cls_segms, cls_keyps):
        """Convert from the class boxes/segms/keyps format generated by the testing
        code.
        """
        box_list = [b for b in cls_boxes if len(b) > 0]
        if len(box_list) > 0:
            boxes = np.concatenate(box_list)
        else:
            boxes = None
        if cls_segms is not None:
            segms = [s for slist in cls_segms for s in slist]
        else:
            segms = None
        if cls_keyps is not None:
            keyps = [k for klist in cls_keyps for k in klist]
        else:
            keyps = None
        classes = []
        for j in range(len(cls_boxes)):
            classes += [j] * len(cls_boxes[j])
        return boxes, segms, keyps, classes

    def get_class_string(self,class_index, score, dataset):
        class_text = dataset.classes[class_index] if dataset is not None else \
            'id{:d}'.format(class_index)
        #return class_text + ' {:0.2f}'.format(score).lstrip('0')
        return class_text
    def nms_between_classes(self,boxes, threshold):
        if boxes.size==0:
            return np.empty((0,3))
        x1 = boxes[:,0]
        y1 = boxes[:,1]
        x2 = boxes[:,2]
        y2 = boxes[:,3]
        s = boxes[:,4]
        area = (x2-x1+1) * (y2-y1+1)
        I = np.argsort(s)
        pick = np.zeros_like(s, dtype=np.int16)
        counter = 0
        while I.size>0:
            i = I[-1]
            pick[counter] = i
            counter += 1
            idx = I[0:-1]
            xx1 = np.maximum(x1[i], x1[idx])
            yy1 = np.maximum(y1[i], y1[idx])
            xx2 = np.minimum(x2[i], x2[idx])
            yy2 = np.minimum(y2[i], y2[idx])
            w = np.maximum(0.0, xx2-xx1+1)
            h = np.maximum(0.0, yy2-yy1+1)
            inter = w * h
            o = inter / (area[i] + area[idx] - inter)
            I = I[np.where(o<=threshold)]
        pick = pick[0:counter]  #返回nms后的索引
        return pick