from scipy.misc import imresize
from skimage.color import rgb2gray
from multiprocessing import *
from collections import deque
import queue
import gym
import numpy as np
import argparse

# -----
parser = argparse.ArgumentParser(description='Training model')
parser.add_argument('--game', default='Breakout-v0', help='OpenAI gym environment name', dest='game', type=str)
parser.add_argument('--processes', default=4, help='Number of processes that generate experience for agent',
                    dest='processes', type=int)
parser.add_argument('--lr', default=0.0001, help='Learning rate', dest='learning_rate', type=float)
parser.add_argument('--batch_size', default=20, help='Batch size to use during training', dest='batch_size', type=int)
parser.add_argument('--swap_freq', default=10000, help='Number of frames before swapping network weights',
                    dest='swap_freq', type=int)
parser.add_argument('--checkpoint', default=0, help='Iteration to resume training', dest='checkpoint', type=int)
parser.add_argument('--save_freq', default=250000, help='Number of frame before saving weights', dest='save_freq',
                    type=int)
parser.add_argument('--eps_decay', default=4000000,
                    help='Number of frames needed to decay epsilon to the lowest value', dest='eps_decay', type=int)
parser.add_argument('--lr_decay', default=80000000,
                    help='Number of frames needed to decay lr to the lowest value', dest='lr_decay', type=int)
parser.add_argument('--queue_size', default=256, help='Size of queue holding agent experience', dest='queue_size',
                    type=int)
parser.add_argument('--th_comp_fix', default=True,
                    help='Sets different Theano compiledir for each process', dest='th_fix', type=bool)
# -----
args = parser.parse_args()

# -----

def build_network(input_shape, output_shape):
    from keras.models import Model
    from keras.layers import Input, Conv2D, Flatten, Dense

    x = Input(shape=input_shape)
    h = Conv2D(16, kernel_size=(8, 8), strides=(4, 4), activation='relu', data_format='channels_first')(x)
    h = Conv2D(32, kernel_size=(4, 4), strides=(2, 2), activation='relu', data_format='channels_first')(h)
    h = Flatten()(h)
    h = Dense(256, activation='relu')(h)
    v = Dense(output_shape, activation='linear')(h)
    return Model(inputs=x, outputs=v)

class LearningAgent(object):
    def __init__(self, action_space, batch_size=32, screen=(84, 84), swap_freq=200):
        from keras.optimizers import RMSprop
        # -----
        self.screen = screen
        self.input_depth = 1
        self.past_range = 3
        self.observation_shape = (self.input_depth * self.past_range,) + self.screen
        self.batch_size = batch_size

        self.action_value = build_network(self.observation_shape, action_space.n)
        self.action_value_freeze = build_network(self.observation_shape, action_space.n)

        self.action_value.compile(optimizer='rmsprop', loss='mse')
        self.action_value_freeze.compile(optimizer='rmsprop', loss='mse')

        self.losses = deque(maxlen=25)
        self.q_values = deque(maxlen=25)
        self.swap_freq = swap_freq
        self.swap_counter = self.swap_freq
        self.unroll = np.arange(self.batch_size)
        self.frames = 0

    def learn(self, last_observations, actions, rewards, observations, not_terminals, discount=0.99,
              learning_rate=0.001):
        self.action_value.optimizer.lr.set_value(learning_rate)
        frames = len(last_observations)
        self.frames += frames
        # -----
        targets = self.action_value.predict_on_batch(last_observations)
        q_values = self.action_value_freeze.predict_on_batch(observations)
        # -----
        # equation = rewards + not_terminals * discount * np.argmax(q_values)
        rewards = np.clip(rewards, -1., 1.)
        equation = not_terminals
        equation *= np.max(q_values, axis=1)
        equation *= discount
        targets[self.unroll, actions] = rewards + equation
        # -----
        loss = self.action_value.train_on_batch(last_observations, targets)
        self.losses.append(loss)
        self.q_values.append(np.mean(targets))
        print(
            '\rFrames: %8d; Lr: %8.7f; Loss: %7.4f; Min: %7.4f; Max: %7.4f; Avg: %7.4f --- Q-value; Min: %7.4f; Max: %7.4f; Avg: %7.4f' % (
                self.frames, learning_rate, loss, min(self.losses), max(self.losses), np.mean(self.losses),
                np.min(self.q_values), np.max(self.q_values), np.mean(self.q_values)), end='')
        self.swap_counter -= frames
        if self.swap_counter < 0:
            self.swap_counter += self.swap_freq
            self.action_value_freeze.set_weights(self.action_value.get_weights())
            return True
        return False

def learn_proc(global_frame, mem_queue, weight_dict):
    import os
    pid = os.getpid()
    if args.th_fix:
        os.environ['THEANO_FLAGS'] = 'floatX=float32,device=gpu,nvcc.fastmath=False,lib.cnmem=0,' + \
                                     'compiledir=th_comp_learn'
    # -----
    save_freq = args.save_freq
    learning_rate = args.learning_rate
    batch_size = args.batch_size
    checkpoint = args.checkpoint
    lr_decay = args.lr_decay
    # -----
    env = gym.make(args.game)
    agent = LearningAgent(env.action_space, batch_size=args.batch_size, swap_freq=args.swap_freq)
    # -----
    if checkpoint > 0:
        agent.action_value.load_weights('model-%d.h5' % (checkpoint,))
        agent.action_value_freeze.set_weights(agent.action_value.get_weights())
    print(' %5d> Setting weights in dict' % (pid,))
    # -----
    weight_dict['update'] = 0
    weight_dict['weights'] = agent.action_value.get_weights()
    # -----
    last_obs = np.zeros((batch_size,) + agent.observation_shape)
    actions = np.zeros(batch_size, dtype=np.int32)
    rewards = np.zeros(batch_size)
    obs = np.zeros((batch_size,) + agent.observation_shape)
    not_term = np.zeros(batch_size)
    # -----
    index = 0
    agent.frames = checkpoint
    save_counter = checkpoint % save_freq + save_freq
    while True:
        last_obs[index, ...], actions[index], rewards[index], obs[index, ...], not_term[index] = mem_queue.get()
        # -----
        index = (index + 1) % batch_size
        if index == 0:
            lr = max(0.00000001, learning_rate * (1. - agent.frames * batch_size / lr_decay))
            updated = agent.learn(last_obs, actions, rewards, obs, not_term, learning_rate=lr)
            global_frame.value = agent.frames
            if updated:
                # print(' %5d> Updating weights in dict' % (pid,))
                weight_dict['weights'] = agent.action_value_freeze.get_weights()
                weight_dict['update'] += 1
        # -----
        save_counter -= 1
        if save_counter < 0:
            save_counter += save_freq
            agent.action_value_freeze.save_weights('model-%d.h5' % (agent.frames,), overwrite=True)

class ActingAgent(object):
    def __init__(self, action_space, screen=(84, 84)):
        from keras.optimizers import RMSprop
        # -----
        self.screen = screen
        self.input_depth = 1
        self.past_range = 3
        self.observation_shape = (self.input_depth * self.past_range,) + self.screen

        self.action_value = build_network(self.observation_shape, action_space.n)
        self.action_value.compile(optimizer='rmsprop', loss='mse')

        self.action_space = action_space
        self.observations = np.zeros(self.observation_shape)
        self.last_observations = np.zeros_like(self.observations)

    def init_episode(self, observation):
        for _ in range(self.past_range):
            self.save_observation(observation)

    def sars_data(self, action, reward, observation, not_terminal):
        self.save_observation(observation)
        return self.last_observations, action, reward, self.observations, not_terminal

    def choose_action(self, epsilon=0.0):
        if np.random.random() < epsilon:
            return self.action_space.sample()
        else:
            return np.argmax(self.action_value.predict(self.observations[None, ...]))

    def save_observation(self, observation):
        self.last_observations = self.observations[...]
        self.observations = np.roll(self.observations, -self.input_depth, axis=0)
        self.observations[-self.input_depth:, ...] = self.transform_screen(observation)

    def transform_screen(self, data):
        return rgb2gray(imresize(data, self.screen))[None, ...]

def generate_experience_proc(global_frame, mem_queue, weight_dict, no, epsilon):
    import os
    pid = os.getpid()
    if args.th_fix:
        os.environ['THEANO_FLAGS'] = 'floatX=float32,device=gpu,nvcc.fastmath=True,lib.cnmem=0,' + \
                                     'compiledir=th_comp_act_' + str(no)
    # -----
    print(' %5d> Process started with %6.3f' % (pid, epsilon))
    # -----
    env = gym.make(args.game)
    agent = ActingAgent(env.action_space)

    if args.checkpoint > 0:
        print(' %5d> Loaded weights from file' % (pid,))
        agent.action_value.load_weights('model-%d.h5' % (args.checkpoint,))
    else:
        import time
        while 'weights' not in weight_dict:
            time.sleep(0.1)
        agent.action_value.set_weights(weight_dict['weights'])
        print(' %5d> Loaded weights from dict' % (pid,))

    best_score, last_update, frames = 0, 0, 0
    avg_score = deque(maxlen=20)
    stop_decay = global_frame.value > args.eps_decay

    while True:
        done = False
        episode_reward, noops, last_op = 0, 0, 0
        observation = env.reset()
        agent.init_episode(observation)

        # -----
        while not done:
            frames += 1
            if not stop_decay:
                frame_tmp = global_frame.value
                decayed_epsilon = max(epsilon, epsilon + (1. - epsilon) * (
                                        args.eps_decay - frame_tmp) / args.eps_decay)
                stop_decay = frame_tmp > args.eps_decay
            # -----
            action = agent.choose_action(decayed_epsilon)
            observation, reward, done, _ = env.step(action)
            episode_reward += reward
            best_score = max(best_score, episode_reward)
            # -----
            if action == last_op:
                noops += 1
            else:
                last_op, noops = action, 0
            # -----
            if noops > 100:
                break
            # -----
            mem_queue.put(agent.sars_data(action, reward, observation, not done))
            # -----
            if frames % 2000 == 0:
                print(' %5d> Epsilon: %9.6f; Best: %4d; Avg: %6.2f' % (
                    pid, decayed_epsilon, best_score, np.mean(avg_score)))
            if frames % args.batch_size == 0:
                update = weight_dict.get('update', 0)
                if update > last_update:
                    last_update = update
                    # print(' %5d> Getting weights from dict' % (pid,))
                    agent.action_value.set_weights(weight_dict['weights'])
        # -----
        avg_score.append(episode_reward)

def init_worker():
    import signal
    signal.signal(signal.SIGINT, signal.SIG_IGN)

def main():
    manager = Manager()
    weight_dict = manager.dict()
    global_frame = manager.Value('i', args.checkpoint)
    mem_queue = manager.Queue(args.queue_size)

    eps = [0.1, 0.01, 0.5]
    pool = Pool(args.processes + 1, init_worker)

    try:
        for i in range(args.processes):
            pool.apply_async(generate_experience_proc,
                             args=(global_frame, mem_queue, weight_dict, i, eps[i % len(eps)]))

        pool.apply_async(learn_proc, args=(global_frame, mem_queue, weight_dict))

        pool.close()
        pool.join()

    except KeyboardInterrupt:
        pool.terminate()
        pool.join()

if __name__ == "__main__":
    main()


import os
import cv2
import numpy as np

start_idx = 1
PREDEFINE_LEN = 6
anno_num = 0
total_anno = 0

class generating:
    def __init__(self, path, label_path, dst_dir):
        self.Path = path  #####视频路径
        self.base_name = os.path.basename(self.Path)
        self.base_name = os.path.splitext(self.base_name)[0]
        self.label_path = os.path.join(label_path, self.base_name + '.txt')#bbox 文件
        self.dst_dir = dst_dir
        self.target_actions = []
        os.makedirs(self.dst_dir,exist_ok=True)
        global anno_num

        if os.path.exists(self.label_path):
            self._all_boxes = np.loadtxt(self.label_path,dtype='float32')
            self._all_boxes = self._all_boxes.reshape([-1,5])
            self._all_boxes = self._all_boxes[np.argsort(self._all_boxes[:, 0])]
            self.target_actions = self.extract_whole_action(self._all_boxes)
            with open(os.path.join(self.dst_dir,self.base_name + '.txt'),'w') as outfile:
                for action in self.target_actions:
                    anno_num += 1
                    outfile.write('{} {}'.format(int(action[0]),int(action[1])) + " " + " ".join([str(a) for a in action[2:]]) + '\n')
        else:
            print('this video has none box annotation')

    def bbox_iou(self,box1, box2, x1y1x2y2=True, GIoU=False):
        # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4

        # Get the coordinates of bounding boxes
        if x1y1x2y2:
            # x1, y1, x2, y2 = box1
            b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
            b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
        else:
            # x, y, w, h = box1
            b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
            b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
            b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
            b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2

        # Intersection area
        inter_area = (min(b1_x2, b2_x2) - max(b1_x1, b2_x1)) * (min(b1_y2, b2_y2) - max(b1_y1, b2_y1))

        # Union Area
        union_area = ((b1_x2 - b1_x1) * (b1_y2 - b1_y1) + 1e-16) + \
                     (b2_x2 - b2_x1) * (b2_y2 - b2_y1) - inter_area

        iou = inter_area / union_area  # iou
        return iou

    def extract_whole_action(self,anno_boxes):
        sorted_anno = anno_boxes[np.argsort(anno_boxes[:, 0])]
        target_actions = np.ones([0,6],'float32')
        global total_anno
        for anno in sorted_anno:
            total_anno += 1
            loc = np.where(target_actions[:,0] == anno[0])#起始帧数值一样
            loc2 = np.where(target_actions[:,1] == anno[0])#结束帧数值相同
            if loc[0].shape[0] == 0:
                if loc2[0].shape[0] == 0:
                    target_action = np.ones([1,6],'float32')
                    target_action[0,0] = anno[0]
                    target_action[0,1] = anno[0]+PREDEFINE_LEN
                    target_action[0,2:6] = anno[1:]
                    target_actions = np.concatenate((target_actions,target_action),axis=0)
                else:
                    cur_actions = target_actions[loc2]
                    is_exist = False#是否为同一个action
                    for cur_action in cur_actions:
                        iou = self.bbox_iou(anno[1:], cur_action[2:])
                        if (iou > 0.1):
                            cur_action[1] = anno[0] + PREDEFINE_LEN
                            cur_action[2], cur_action[3] = min(anno[1], cur_action[2]), min(anno[2], cur_action[3])
                            cur_action[4], cur_action[5] = max(anno[3], cur_action[4]), max(anno[4], cur_action[5])
                            is_exist = True
                    if not is_exist:
                        target_action = np.ones([1,6], 'float32')
                        target_action[0, 0] = anno[0]
                        target_action[0, 1] = anno[0] + PREDEFINE_LEN
                        target_action[0, 2:6] = anno[1:]
                        target_actions = np.concatenate((target_actions, target_action), axis=0)
                    else:
                        target_actions[loc2] = cur_actions
            else:
                target_action = np.ones([1,6], 'float32')
                target_action[0, 0] = anno[0]
                target_action[0, 1] = anno[0] + PREDEFINE_LEN
                target_action[0, 2:6] = anno[1:]
                target_actions = np.concatenate((target_actions, target_action), axis=0)
        return target_actions
    def generate_single_sample(self,img_list,labels):
        standard_with = img_list[0].shape[1]
        standard_height = img_list[0].shape[0]

        global start_idx
        if(len(labels) != 0):
            for label in labels:
                fourcc = cv2.VideoWriter_fourcc(*'XVID')
                pos_path = os.path.join(dst_dir, 'pos','%05d.mp4' % start_idx)
                os.makedirs(os.path.join(dst_dir, 'pos'),exist_ok=True)
                print(pos_path)
                vw = cv2.VideoWriter(pos_path, fourcc, 6,(standard_with, standard_height))
                for image in img_list:
                    vw.write(image)
                vw.release()
                start_idx += 1

    def start(self):
        if os.path.exists(self.label_path) and len(self.target_actions) != 0:
            video_cap = cv2.VideoCapture(self.Path)
            frame_list = []
            frame_idx = 0
            global start_idx

            action_id = 0
            self.target_actions = self.target_actions[np.argsort(self.target_actions[:, 0])]
            cur_action = self.target_actions[action_id]
            while True:
                is_read,img = video_cap.read()
                if is_read:
                    if frame_idx >= cur_action[0] and frame_idx < cur_action[1]:
                        frame_list.append(img)
                    if len(frame_list) == (cur_action[1]-cur_action[0]) and frame_idx == cur_action[1]:
                        self.generate_single_sample(frame_list,cur_action)
                        frame_list = []
                        action_id += 1
                        cur_action = self.target_actions[action_id]
                    frame_idx += 1
                else:
                    break

if __name__ == '__main__':
    total_tatgets = [r'/media/hzh/NO.6/Q3_RIGHT_newest/Q3/fall/fall20200617/fall20200617',
                     r'/media/hzh/NO.6/Q3_RIGHT_newest/Q3/fall/fall20200618/fall20200618',
                     r'/media/hzh/NO.6/Q3_RIGHT_newest/Q3/fall/fall20200619/fall20200619',
                     r'/media/hzh/NO.6/Q3_RIGHT_newest/Q3/fall/fall20200620/fall20200620',
                     r'/media/hzh/NO.6/Q3_RIGHT_newest/Q3/fall/fall20200621/fall20200621',
                     r'/media/hzh/NO.6/Q3_RIGHT_newest/Q3/fall/fall20200622/fall20200622']
    for tar_path in total_tatgets:
        # tar_path = r'/media/hzh/ssd_disk/摔倒标注数据/fall/fall20200622/fall20200622'
        label_path = tar_path+'_label'
        # label_path = r'/media/hzh/ssd_disk/摔倒标注数据/fall/fall20200622/fall20200622_label'
        # new_label_path = tar_path+'_newlabel'
        dst_dir = tar_path+'_newlabel'

        sub_video_list = os.listdir(tar_path)
        suffix = '.mp4'
        for i, sub_video in enumerate(sub_video_list):
            if sub_video.endswith(suffix) or sub_video.endswith('.flv'):
                print("process " + sub_video + "..")
                full_sub_video = os.path.join(tar_path, sub_video)
                instance_ = generating(full_sub_video,label_path,dst_dir)
                # instance_.start()
    print('total fall action:',anno_num)
    print('total src fall action:',total_anno)


import numpy as np
import pandas as pd

from .baseSlicer import BaseSlicer
from rubin_sim.maf.plots.moPlotters import MetricVsH, MetricVsOrbit

from .orbits import Orbits

__all__ = ['MoObjSlicer']

class MoObjSlicer(BaseSlicer):
    """ Slice moving object _observations_, per object and optionally clone/per H value.

    Iteration over the MoObjSlicer will go as:
    * iterate over each orbit;
    * if Hrange is not None, for each orbit, iterate over Hrange.

    Parameters
    ----------
    Hrange : numpy.ndarray or None
        The H values to clone the orbital parameters over. If Hrange is None, will not clone orbits.
    """
    def __init__(self, Hrange=None, verbose=True, badval=0):
        super(MoObjSlicer, self).__init__(verbose=verbose, badval=badval)
        self.Hrange = Hrange
        self.slicer_init = {'Hrange': Hrange, 'badval': badval}
        # Set default plotFuncs.
        self.plotFuncs = [MetricVsH(),
                          MetricVsOrbit(xaxis='q', yaxis='e'),
                          MetricVsOrbit(xaxis='q', yaxis='inc')]

    def setupSlicer(self, orbitFile, delim=None, skiprows=None, obsFile=None):
        """Set up the slicer and read orbitFile and obsFile from disk.

        Sets self.orbits (with orbit parameters), self.allObs, and self.obs
        self.orbitFile and self.obsFile

        Parameters
        ----------
        orbitFile : str
            The file containing the orbit information.
            This is necessary, in order to be able to generate plots.
        obsFile : str, optional
            The file containing the observations of each object, optional.
            If not provided (default, None), then the slicer will not be able to 'slice', but can still plot.
        """
        self.readOrbits(orbitFile, delim=delim, skiprows=skiprows)
        if obsFile is not None:
            self.readObs(obsFile)
        else:
            self.obsFile = None
            self.allObs = None
            self.obs = None
        # Add these filenames to the slicer init values, to preserve in output files.
        self.slicer_init['orbitFile'] = self.orbitFile
        self.slicer_init['obsFile'] = self.obsFile

    def readOrbits(self, orbitFile, delim=None, skiprows=None):
        # Use sims_movingObjects to read orbit files.
        orb = Orbits()
        orb.readOrbits(orbitFile, delim=delim, skiprows=skiprows)
        self.orbitFile = orbitFile
        self.orbits = orb.orbits
        # Then go on as previously. Need to refactor this into 'setupSlicer' style.
        self.nSso = len(self.orbits)
        self.slicePoints = {}
        self.slicePoints['orbits'] = self.orbits
        # And set the slicer shape/size.
        if self.Hrange is not None:
            self.shape = [self.nSso, len(self.Hrange)]
            self.slicePoints['H'] = self.Hrange
        else:
            self.shape = [self.nSso, 1]
            self.slicePoints['H'] = self.orbits['H']
        # Set the rest of the slicePoint information once
        self.nslice = self.shape[0] * self.shape[1]

    def readObs(self, obsFile):
        """Read observations of the solar system objects (such as created by sims_movingObjects).

        Parameters
        ----------
        obsFile: str
            The file containing the observation information.
        """
        # For now, just read all the observations (should be able to chunk this though).
        self.allObs = pd.read_csv(obsFile, delim_whitespace=True, comment='#')
        self.obsFile = obsFile
        # We may have to rename the first column from '#objId' to 'objId'.
        if self.allObs.columns.values[0].startswith('#'):
            newcols = self.allObs.columns.values
            newcols[0] = newcols[0].replace('#', '')
            self.allObs.columns = newcols
        if 'velocity' not in self.allObs.columns.values:
            self.allObs['velocity'] = np.sqrt(self.allObs['dradt']**2 + self.allObs['ddecdt']**2)
        if 'visitExpTime' not in self.allObs.columns.values:
            self.allObs['visitExpTime'] = np.zeros(len(self.allObs['objId']), float) + 30.0
        # If we created intermediate data products by pandas, we may have an inadvertent 'index'
        #  column. Since this creates problems later, drop it here.
        if 'index' in self.allObs.columns.values:
            self.allObs.drop('index', axis=1, inplace=True)
        self.subsetObs()

    def subsetObs(self, pandasConstraint=None):
        """
        Choose a subset of all the observations, such as those in a particular time period.
        """
        if pandasConstraint is None:
            self.obs = self.allObs
        else:
            self.obs = self.allObs.query(pandasConstraint)

    def _sliceObs(self, idx):
        """Return the observations of a given ssoId.

        For now this works for any ssoId; in the future, this might only work as ssoId is
        progressively iterated through the series of ssoIds (so we can 'chunk' the reading).

        Parameters
        ----------
        idx : integer
            The integer index of the particular SSO in the orbits dataframe.
        """
        # Find the matching orbit.
        orb = self.orbits.iloc[idx]
        # Find the matching observations.
        if self.obs['objId'].dtype == 'object':
            obs = self.obs.query('objId == "%s"' %(orb['objId']))
        else:
            obs = self.obs.query('objId == %d' %(orb['objId']))
        # Return the values for H to consider for metric.
        if self.Hrange is not None:
            Hvals = self.Hrange
        else:
            Hvals = np.array([orb['H']], float)
        # Note that ssoObs / obs is a recarray not Dataframe!
        # But that the orbit IS a Dataframe.
        return {'obs': obs.to_records(),
                'orbit': orb,
                'Hvals': Hvals}

    def __iter__(self):
        """
        Iterate through each of the ssoIds.
        """
        self.idx = 0
        return self

    def __next__(self):
        """
        Returns result of self._getObs when iterating over moSlicer.
        """
        if self.idx >= self.nSso:
            raise StopIteration
        idx = self.idx
        self.idx += 1
        return self._sliceObs(idx)

    def __getitem__(self, idx):
        # This may not be guaranteed to work if/when we implement chunking of the obsfile.
        return self._sliceObs(idx)

    def __eq__(self, otherSlicer):
        """
        Evaluate if two slicers are equal.
        """
        result = False
        if isinstance(otherSlicer, MoObjSlicer):
            if otherSlicer.orbitFile == self.orbitFile:
                if otherSlicer.obsFile == self.obsFile:
                    if np.array_equal(otherSlicer.slicePoints['H'], self.slicePoints['H']):
                        result = True
        return result


import sys
from dataset import Dataset
import tqdm
import time
from cnn import Net32
import torch
import torch.nn as nn
import numpy as np
import logging
import argparse
import torchvision.models as models
from torchvision import datasets
from tensorboardX import SummaryWriter
import sendAMail as send

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def getArgs():
    parser = argparse.ArgumentParser(description='Setting model')
    parser.add_argument('-l','--learning',type=str, choices=['images','landmarks','images_landmarks'], default='images', help='Select the learning type')
    args = parser.parse_args()
    return args.learning

def save_checkpoint(state, is_best, checkpoint_path,best_model_path):
    torch.save(state, checkpoint_path)
    if is_best:
        torch.save(state, best_model_path)

def startLearning(mode:str,bs,me,f,p,l):
    print(device)
    #cudnn.Benchmark = True
    writer = SummaryWriter()

    batch_size = bs
    max_epochs = me
    factor = f
    patience = p
    lr = l
    print("load dataset")
    training_loader = torch.utils.data.DataLoader(dataset=Dataset('training_dataset_pack.h5',mode,"std_training.png","mean_training.png"), batch_size=batch_size, shuffle=True,num_workers=0)
    validation_loader = torch.utils.data.DataLoader(dataset=Dataset('validation_dataset_pack.h5',mode,"std_training.png","mean_training.png"), batch_size=batch_size, shuffle=True,num_workers=0)

    print("Done")

    model = Net32(training_loader.dataset.getInputSize()).cuda()

    criterion = nn.CrossEntropyLoss().cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9,0.999))
    #déclarer scheduler reduce plateau
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,mode='min',factor=factor,patience=patience)
    best_loss = None
    ######

    ######
    for epoch in range(1,max_epochs+1):  # loop over the dataset multiple times

        is_best= False
        print("Training...")
        pbart = tqdm.tqdm(total=int(len(training_loader.dataset)/batch_size),postfix={"loss":None,"accuracy":None},desc="Epoch: {}/{}".format(epoch,max_epochs))
        loss = 0.
        acc = 0.
        val_loss = 0.
        val_acc =0.

        for i, data in enumerate(training_loader, 0):
            # get the inputs
            images, expressions = data
            images = images.to(device)
            expressions = expressions.to(device).long()
            # zero the parameter gradients
            optimizer.zero_grad()
            outputs = model(images)
            current_loss = criterion(outputs, expressions)
            current_loss.backward()
            optimizer.step()

            # print statistics
            loss += current_loss.item()
            _,predicted = torch.max(outputs,1)
            acc += (predicted == expressions).sum().float()/batch_size

            pbart.update(1)
            pbart.set_postfix({"loss": loss/(i+1),"accuracy":acc.item()/(i+1)})

        pbart.close()
        loss /= (len(training_loader.dataset)/batch_size)
        acc /= (len(training_loader.dataset)/batch_size)

        with open("log/training_"+mode+".log","a") as f:
            f.write("epoch: {} / {} loss: {} accuracy: {}\n".format(epoch,max_epochs,loss,acc))
        f.close()

        print("Validation...")
        pbarv = tqdm.tqdm(total=int(len(validation_loader.dataset)/batch_size),postfix={"loss":None,"accuracy":None},desc="Epoch: {}/{}".format(epoch,max_epochs))

        with torch.no_grad():
            for i, data in enumerate(validation_loader, 0):
                # get the inputs
                images, expressions = data
                images = images.to(device)
                expressions = expressions.to(device).long()

                outputs = model(images)
                loss = criterion(outputs, expressions)
                val_loss += loss.item()

                _,predicted = torch.max(outputs,1)
                val_acc  +=(predicted == expressions).sum().float()/batch_size

                pbarv.update(1)
                pbarv.set_postfix({"loss": val_loss/(i+1),"accuracy":val_acc.item()/(i+1)})

        pbarv.close()
        val_loss /= (len(validation_loader.dataset)/batch_size)
        val_acc /= (len(validation_loader.dataset)/batch_size)

        if best_loss == None or val_loss < best_loss:
            best_loss = val_loss
            is_best = True
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'loss': val_loss,
            'accuracy': val_acc,
            'optimizer' : optimizer.state_dict()}, is_best,"save_model/checkpoint_"+mode+".pth","save_model/best_model_"+mode+".pth")
        is_best = False

        scheduler.step(val_loss)
        #loss accuracy training et validation dans les logs

        with open("log/validation_"+mode+".log","a") as f:
            f.write("epoch: {} / {} loss: {} accuracy: {}\n".format(epoch,max_epochs,val_loss,val_acc))
        f.close()

        writer.add_scalar('data/Loss training', loss, epoch)
        writer.add_scalar('data/Loss validation', val_loss, epoch)
        writer.add_scalar('data/Accuracy training', acc, epoch)
        writer.add_scalar('data/Accuracy validation', val_acc, epoch)

        if (epoch%30 == 1 or epoch == max_epochs):
            desc = dict()
            desc["bs"] = batch_size
            desc["lr"] = lr
            desc["f"] = factor
            desc["p"] = patience
            desc["d"] = 0
            desc["weights"] = [1,1]
            desc["epoch"] = epoch
            desc["nbepochs"] = max_epochs

            send.sendInfos("1 (3 models)",desc,loss,acc,"...","...","...","...",val_loss,val_acc,"...","...","...","...")

if __name__ == "__main__":
    print(sys.version)
    args =getArgs()
    mode = args
    with open("log/training_"+mode+".log","w") as f:
        f.close()
    with open("log/validation_"+mode+".log","w") as f:
        f.close()

    batch_size = 64
    max_epochs = 150
    factor = 0.1
    patience = 10
    lr = 0.0001
    startLearning(mode,batch_size,max_epochs,factor,patience,lr)


import pycurl
from StringIO import StringIO
import multiprocessing as mp
from workloads import cost_driver, spc_traces, gdwheel, developer_ubuntu
from workloads.virality_trees import vtree
import sys
import time
import numpy as np
import random

d_lambdas = [(lambda seed, nclients :
              cost_driver.ZipfFixedDriver(seed = seed,
                                          item_range_max = 10**5)),
             (lambda seed, nclients:
              spc_traces.SPCTraceDriver(trace_file = "workloads/spc1/WebSearch1.spc",
                                        name = "SPCWebSearch",
                                        column_map = spc_traces.SPCMap1,
                                        seed = (seed, nclients))),
             (lambda seed, nclients :
              cost_driver.ZipfFixedDriver(seed = seed, zipf_param = 0.75,
                                          item_range_max = 10**6)), # = 2
             (lambda seed, nclients:
              cost_driver.DynamicPromote(period = 100).wrap_driver(cost_driver.ZipfFixedDriver(seed = seed,
                                                                                item_range_max = 10**5))), # = 3
             (lambda seed, nclients:
              cost_driver.Introducing(period = 100,
                                      move_to = 0, name = "IntHigh.%d").wrap_driver(cost_driver.ZipfFixedDriver(seed = seed,
                                                                                item_range_max = 10**5))), # = 4
             (lambda seed, nclients:
              vtree.ViralityTreeDriver(seed = seed, name = "VTree(40k)", filename = "workloads/virality_trees/output-40k.json")),
             (lambda seed, nclients:
              gdwheel.GD2(seed = seed, max_cost = (2**32 - 2), item_range_max = 10**5)),
             (lambda seed, nclients:
              spc_traces.SPCTraceDriver(trace_file = "workloads/arc/S1.lis",
                                        name = "Arc.S1",
                                        column_map = spc_traces.ARCMap,)), #=7
             (lambda seed, nclients :
              cost_driver.ZipfFixedDriver(seed = seed,
                                          item_range_max = 10**6)), # = 8
             (lambda seed, nclients :
              cost_driver.ZipfFixedDriver(seed = seed, zipf_param = 0.9,
                                          item_range_max = 10**6)), # = 9
             (lambda seed, nclients :
              developer_ubuntu.URLDriver(seed = seed)
              ), # = 10
             (lambda seed, nclients :
              developer_ubuntu.WikiZURLDriver(seed = seed, item_range_max = 10**6,)
              ), # = 11
             ]

PROGRESS_CTR_T = mp.Array('i', [0 for i in range(1024)], lock=False)
MEASURING = mp.Value('i', 0, lock = True)

REQ_ARRS = None
COST_ARRS = None

COST_DRIVERS = [5, 6]

SEED = random.randint(0, 10000)

def update_progress_bar(progress, total):
    col_size = 65

    count = (col_size * progress) / total

    sys.stderr.write("\r")
    sys.stderr.write("[" + "".join(["=" for i in range(count)]))
    sys.stderr.write(">")
    sys.stderr.write("".join([" " for i in range(col_size - count - 1)]))
    sys.stderr.write("] %.1f%%" % (float(100 * progress) / total))
    sys.stderr.flush()

def generate_req_arr(atuple):
    driver_lambda, seed, nreqs, nclients = atuple
    driver = d_lambdas[driver_lambda](seed+SEED, nclients)
    cost_out = None
    gen_costs = driver_lambda in COST_DRIVERS
    if gen_costs:
        cost_out = np.zeros(nreqs, np.uint32)

    using_np = True

    for i in range(nreqs):
        item, cost = driver.sample_item_w_cost()
        if i == 0:
            using_np = not isinstance(item, str)
            if using_np:
                out = np.zeros(nreqs, np.uint32)
            else:
                out = []

        if using_np:
            out[i] = item
        else:
            out.append(item)

        if gen_costs:
            cost_out[i] = int(cost + 1)

    return (out, cost_out)

class WorkloadRunner(object):
    def __init__(self, driver_lambda, host ="sns49.cs.princeton.edu:3590", # "128.112.7.141:3590",
                 nclients = 8, nreqs = 10**6):
        self.driver_lambda = driver_lambda

        self.cost = self.driver_lambda in COST_DRIVERS

        self.nclients = nclients
        if self.cost:
            self.host_str = "http://%s/%%s/%%s" % host
        else:
            self.host_str = "http://%s/%%s" % host

        self.workload = [(nreqs, self.host_str, seed, self.cost)
                         for seed in range(nclients)]
        self.nreqs = nreqs

    def bench_run(self):
        global PROGRESS_CTR_T, MEASURING
        global REQ_ARRS
        global COST_ARRS

        if self.nclients == 1:
            results = map(generate_req_arr,
                          [(self.driver_lambda, i, self.nreqs, self.nclients)
                           for i in range(self.nclients)])
            REQ_ARRS = []
            COST_ARRS = []
            for (out_arr, cost_arr) in results:
                REQ_ARRS.append(out_arr)
                COST_ARRS.append(cost_arr)
            t_start = time.time()
            results = map(work_driver, self.workload)
            t_end = time.time()
            requests = sum([reqs for reqs, _ in results])
            misses = sum([misses for _, misses in results])

            throughput = float(requests) / (t_end - t_start)
            miss_rate = float(misses) / float(requests)

            return throughput, miss_rate, requests, (t_end - t_start)

        t_start = time.time()
        pool = mp.Pool(processes = 16)
        results = pool.map(generate_req_arr,
                           [(self.driver_lambda, i, self.nreqs, self.nclients)
                            for i in range(self.nclients)])

        REQ_ARRS = []
        COST_ARRS = []

        for (out_arr, cost_arr) in results:
            REQ_ARRS.append(out_arr)
            COST_ARRS.append(cost_arr)
        pool.close()
        t_end = time.time()
        sys.stderr.write("finished building request arrays in %f s" % ( t_end - t_start ))
        sys.stderr.write("\n")
        sys.stderr.flush()

        pool = mp.Pool(processes = self.nclients)
        result_async = pool.map_async(work_driver, self.workload)
        t_start = time.time()

        update_time = 5
        start_window = 30 / update_time
        end_window = 1 * 30 / update_time
        end_window += update_time
        updates = 0

        while not result_async.ready():
            result_async.wait(5)
            updates += 1
            update_progress_bar(sum(PROGRESS_CTR_T),
                                self.nclients * self.nreqs)

        t_end = time.time()

        sys.stderr.write("\n")
        sys.stderr.flush()
        results = result_async.get()

        requests = sum([reqs for reqs, _ in results])
        misses = sum([misses for _, misses in results])

        throughput = float(requests) / (t_end - t_start)
        miss_rate = float(misses) / float(requests)

        pool.close()
        return throughput, miss_rate, requests, (t_end - t_start)

def work_driver(arg_tup):
    global PROGRESS_CTR_T, MEASURING
    global REQ_ARRS
    global COST_ARRS
    nreqs, host_str, req_arr_ix, send_cost = arg_tup
    req_arr = REQ_ARRS[req_arr_ix]
    cost_arr = COST_ARRS[req_arr_ix]

    success = 0
    miss = 0

    NCURL = 128
    curl_hands = []
    for i in range(NCURL):
        c = pycurl.Curl()
        curl_hands.append(c)
    n = 0
    for ix, item in enumerate(req_arr):
        c = curl_hands[ix % NCURL]
        try:
            if send_cost:
                url = host_str % (item, cost_arr[ix])
            else:
                url = host_str % item
            c.setopt(c.URL, url)
            resp = StringIO()
            headers = StringIO()
            c.setopt(c.WRITEFUNCTION, resp.write)
            c.setopt(c.HEADERFUNCTION, headers.write)
            c.setopt(pycurl.CONNECTTIMEOUT, 20)
            c.setopt(pycurl.TIMEOUT, 20)
            c.perform()

            if c.getinfo(c.RESPONSE_CODE) == 200:
                success += 1
                is_hit = handle_response(resp, headers)
                if not is_hit:
                    miss += 1

        except Exception as e:
            pass

        PROGRESS_CTR_T[req_arr_ix] = ix

    return (success, miss)

HEADER_STRING = "X-CACHED-MIDDLEWARE: "
def handle_response(resp, headers):
    """
    return TRUE = HIT, FALSE = MISS
    """
    val = resp.getvalue()
    head_val = headers.getvalue()
    index = head_val.find("X-CACHED-MIDDLEWARE: ")
    if index == -1:
        return False
    response = head_val[index + len(HEADER_STRING)]
    return response == "H"

def unitc_zpop(d_lam, nr, nc, host = False):
    if host:
        WL = WorkloadRunner(d_lam, nclients = nc, nreqs = nr, host = host)
    else:
        WL = WorkloadRunner(d_lam, nclients = nc, nreqs = nr)
    return WL.bench_run()

if __name__ in ('main', '__main__'):
    if len(sys.argv) < 3:
        unitc_zpop(0, 10**3, 1)
    else:
        nclients = int(sys.argv[2])
        nreqs = int(float(sys.argv[3]))
        host = False
        if len(sys.argv) > 4:
            host = sys.argv[4]
        throughput, missrate, reqs, time = unitc_zpop(int(sys.argv[1]), nreqs, nclients, host)

        if len(sys.argv) > 3:
            print "%s, %s, %f, %f, %d, %s" % (sys.argv[5], sys.argv[6], throughput, missrate, reqs, time)