import numpy as np
import torch
import argparse
import os
import math
import gym
import sys
import random
import time
import json
import dmc2gym
import copy

import utils
from logger import Logger
from video import VideoRecorder

from sac_ae import SacAeAgent

def parse_args():
    parser = argparse.ArgumentParser()
    # environment
    parser.add_argument('--domain_name', default='cheetah')
    parser.add_argument('--task_name', default='run')
    parser.add_argument('--image_size', default=84, type=int)
    parser.add_argument('--action_repeat', default=1, type=int)
    parser.add_argument('--frame_stack', default=3, type=int)
    parser.add_argument('--env_spec', default='full')
    # replay buffer
    parser.add_argument('--replay_buffer_capacity', default=200000, type=int)
    # train
    parser.add_argument('--agent', default='sac_ae', type=str)
    parser.add_argument('--init_steps', default=1000, type=int)
    parser.add_argument('--num_train_steps', default=200000, type=int)
    parser.add_argument('--batch_size', default=128, type=int)
    parser.add_argument('--hidden_dim', default=1024, type=int)
    # eval
    parser.add_argument('--eval_freq', default=10000, type=int)
    parser.add_argument('--num_eval_episodes', default=10, type=int)
    # critic
    parser.add_argument('--critic_lr', default=1e-3, type=float)
    parser.add_argument('--critic_beta', default=0.9, type=float)
    parser.add_argument('--critic_tau', default=0.01, type=float)
    parser.add_argument('--critic_target_update_freq', default=2, type=int)
    # actor
    parser.add_argument('--actor_lr', default=1e-3, type=float)
    parser.add_argument('--actor_beta', default=0.9, type=float)
    parser.add_argument('--actor_log_std_min', default=-10, type=float)
    parser.add_argument('--actor_log_std_max', default=2, type=float)
    parser.add_argument('--actor_update_freq', default=2, type=int)
    # encoder/decoder
    parser.add_argument('--encoder_type', default='pixel', type=str)
    parser.add_argument('--encoder_feature_dim', default=50, type=int)
    parser.add_argument('--encoder_lr', default=1e-3, type=float)
    parser.add_argument('--encoder_tau', default=0.05, type=float)
    parser.add_argument('--decoder_type', default='identity', type=str)
    parser.add_argument('--decoder_lr', default=1e-3, type=float)
    parser.add_argument('--decoder_update_freq', default=1, type=int)
    parser.add_argument('--decoder_latent_lambda', default=1e-6, type=float)
    parser.add_argument('--decoder_weight_lambda', default=1e-7, type=float)
    parser.add_argument('--num_layers', default=4, type=int)
    parser.add_argument('--num_filters', default=32, type=int)
    # sac
    parser.add_argument('--discount', default=0.99, type=float)
    parser.add_argument('--init_temperature', default=0.1, type=float)
    parser.add_argument('--alpha_lr', default=1e-4, type=float)
    parser.add_argument('--alpha_beta', default=0.5, type=float)
    # misc
    parser.add_argument('--seed', default=1, type=int)
    parser.add_argument('--work_dir', default='.', type=str)
    parser.add_argument('--save_tb', default=False, action='store_true')
    parser.add_argument('--save_model', default=True, action='store_true')
    parser.add_argument('--save_buffer', default=True, action='store_true')
    parser.add_argument('--save_video', default=False, action='store_true')

    args = parser.parse_args()
    return args

def evaluate(env, agent, video, num_episodes, L, step):
    for i in range(num_episodes):
        obs, state = env.reset()
        video.init(enabled=(i == 0))
        done = False
        episode_reward = 0
        while not done:
            with utils.eval_mode(agent):
                if agent.encoder_type == 'identity':
                    action = agent.select_action(state)
                else:
                    action = agent.select_action(obs)
            obs, state, reward, done, _ = env.step(action)
            video.record(env)
            episode_reward += reward

        video.save('%d.mp4' % step)
        L.log('eval/episode_reward', episode_reward, step)
    L.dump(step)

def make_agent(obs_shape, action_shape, args, device):
    if args.agent == 'sac_ae':
        return SacAeAgent(
            obs_shape=obs_shape,
            action_shape=action_shape,
            device=device,
            hidden_dim=args.hidden_dim,
            discount=args.discount,
            init_temperature=args.init_temperature,
            alpha_lr=args.alpha_lr,
            alpha_beta=args.alpha_beta,
            actor_lr=args.actor_lr,
            actor_beta=args.actor_beta,
            actor_log_std_min=args.actor_log_std_min,
            actor_log_std_max=args.actor_log_std_max,
            actor_update_freq=args.actor_update_freq,
            critic_lr=args.critic_lr,
            critic_beta=args.critic_beta,
            critic_tau=args.critic_tau,
            critic_target_update_freq=args.critic_target_update_freq,
            encoder_type=args.encoder_type,
            encoder_feature_dim=args.encoder_feature_dim,
            encoder_lr=args.encoder_lr,
            encoder_tau=args.encoder_tau,
            decoder_type=args.decoder_type,
            decoder_lr=args.decoder_lr,
            decoder_update_freq=args.decoder_update_freq,
            decoder_latent_lambda=args.decoder_latent_lambda,
            decoder_weight_lambda=args.decoder_weight_lambda,
            num_layers=args.num_layers,
            num_filters=args.num_filters
        )
    else:
        assert 'agent is not supported: %s' % args.agent

def main():
    args = parse_args()
    utils.set_seed_everywhere(args.seed)

    env = dmc2gym.make(
        domain_name=args.domain_name,
        task_name=args.task_name,
        seed=args.seed,
        visualize_reward=False,
        spec=args.env_spec,
        height=args.image_size,
        width=args.image_size,
        frame_skip=args.action_repeat
    )
    env.seed(args.seed)

    # stack several consecutive frames together
    #if args.encoder_type == 'pixel':
    env = utils.FrameStack(env, k=args.frame_stack)

    utils.make_dir(args.work_dir)
    video_dir = utils.make_dir(os.path.join(args.work_dir, 'video'))
    model_dir = utils.make_dir(os.path.join(args.work_dir, 'model'))
    buffer_dir = utils.make_dir(os.path.join(args.work_dir, 'buffer'))

    video = VideoRecorder(video_dir if args.save_video else None)

    with open(os.path.join(args.work_dir, 'args.json'), 'w') as f:
        json.dump(vars(args), f, sort_keys=True, indent=4)

    from soco_device import DeviceCheck

    dc = DeviceCheck()
    # will return a device name ('cpu'/'cuda') and a list of gpu ids, if any
    device_name, device_ids = dc.get_device(n_gpu=1)

    if len(device_ids) == 1:
        device_name = '{}:{}'.format(device_name, device_ids[0])
        device = torch.device(device_name)

    print(device)
    # the dmc2gym wrapper standardizes actions
    assert env.action_space.low.min() >= -1
    assert env.action_space.high.max() <= 1

    replay_buffer = utils.ReplayBuffer(
        obs_shape=env.observation_space.shape,
        state_shape=env.state_space.shape,
        action_shape=env.action_space.shape,
        capacity=args.replay_buffer_capacity,
        batch_size=args.batch_size,
        device=device
    )

    if args.encoder_type == 'pixel':

        agent = make_agent(
            obs_shape=env.observation_space.shape,
            action_shape=env.action_space.shape,
            args=args,
            device=device
        )

    else:
        agent = make_agent(
            obs_shape=env.state_space.shape,
            action_shape=env.action_space.shape,
            args=args,
            device=device
        )

    agent.load(os.path.join(args.work_dir, 'model'),990000)
    #if args.expert_encoder_type == 'pixel':
    #    agent.warm_start_from(expert_agent)
    print("expert loaded.")

    #replay_buffer.load(os.path.join(args.work_dir, 'buffer'))
    #print("buffer loaded.")
    L = Logger(args.work_dir, use_tb=args.save_tb)

    episode, episode_reward, done = 0, 0, True
    start_time = time.time()
    for step in range(args.num_train_steps):
        if done:
            if step > 0:
                L.log('train/duration', time.time() - start_time, step)
                start_time = time.time()
                L.dump(step)

            # evaluate agent periodically
            if step % args.eval_freq == 0:
                L.log('eval/episode', episode, step)
                evaluate(env, agent, video, args.num_eval_episodes, L, step)
                #if args.save_model:
                #    agent.save(model_dir, step)
                #if args.save_buffer:
                #    replay_buffer.save(buffer_dir)

            L.log('train/episode_reward', episode_reward, step)

            obs, state = env.reset()
            done = False
            episode_reward = 0
            episode_step = 0
            episode += 1

            L.log('train/episode', episode, step)

        # sample action for data collection

        with utils.eval_mode(agent):
            if agent.encoder_type == 'identity':
                action = agent.sample_action(state)
            else:
                action = agent.sample_action(obs)

        # run training update
        if step >= args.init_steps:
            num_updates = args.init_steps if step == args.init_steps else 1
            for _ in range(num_updates):
                agent.update(replay_buffer, L, step)

        next_obs, next_state, reward, done, _ = env.step(action)

        # allow infinit bootstrap
        done_bool = 0 if episode_step + 1 == env._max_episode_steps else float(
            done
        )
        episode_reward += reward

        replay_buffer.add(obs, state, action, reward, next_obs, next_state, done_bool)

        obs = next_obs
        state = next_state
        episode_step += 1

if __name__ == '__main__':
    main()


"""
Mask R-CNN
The main Mask R-CNN model implemenetation.

Version of Detection_Target_layer without MASK processing
"""

import os
import sys
# import glob
# import random
# import math
# import datetime
# import itertools
# import json
# import re
# import logging
import numpy as np
import tensorflow as tf
# from collections import OrderedDict
# import keras.backend as K
# import keras.layers as KL
# import keras.initializers as KI
import keras.engine as KE
import keras.backend as KB
sys.path.append('..')
import mrcnn.utils as utils

############################################################
#  Detection Target Layer
############################################################

def overlaps_graph_mod(boxes1, boxes2):
    '''
    Computes IoU overlaps between two sets of boxes.in normalized coordinates

    boxes1 - proposals :  [batch_size,  proposal_counts, 4 (y1, x1, y2, x2)] <-- Region proposals
    boxes2 - gt_boxes  :  [batch_size, max_gt_instances, 4 (y1, x1, y2, x2)] <-- input_normlzd_gt_boxes

    proposal_counts : 1000 or 2000 based on training or inference
    max_gt_instances: 100

    returns :
    ---------
    overlaps :          [ proposal_counts, max_gt_instances]
                        IoU of all proposal box / gt_box pairs
                        The dimensionality :
                            row:  number of non_zero proposals
                            cols: number of non_zero gt_bboxes
    '''
    # 1. Tile boxes2 and repeat boxes1. This allows us to compare every boxes1 against every boxes2 without loops.
    #    TF doesn't have an equivalent to np.repeat() so simulate it using tf.tile() and tf.reshape.

    # print('\t>>> detection_targets_graph - calculate Overlaps_graph')
    # print('\t     overlaps_graph: shape of boxes1 before reshape: ',tf.shape(boxes1).eval())  # (?,?)
    # print('\t     overlaps_graph: shape of boxes2 before reshape: ',tf.shape(boxes1).eval())  # (?,?)

    # tf.expand_dims(boxes1, 1) : makes b1:[1, batch_size, proposal_count_sz, 4]

    b1 = tf.reshape(tf.tile(tf.expand_dims(boxes1, 1), [1, 1, tf.shape(boxes2)[0]]), [-1, 4])
    b2 = tf.tile(boxes2, [tf.shape(boxes1)[0], 1])
    # print('\t     overlaps_graph: shape of boxes1 after reshape: ',tf.shape(b1).eval())  # (?,4)
    # print('\t     overlaps_graph: shape of boxes2 after reshape: ',tf.shape(b2).eval())  # (?,4)

    # 2. Compute intersections
    b1_y1, b1_x1, b1_y2, b1_x2 = tf.split(b1, 4, axis=1)
    b2_y1, b2_x1, b2_y2, b2_x2 = tf.split(b2, 4, axis=1)

    # print('     overlaps_graph: shape of b1_y1 after split: ',b2_y1.shape)  # (?,4)

    y1 = tf.maximum(b1_y1, b2_y1)
    x1 = tf.maximum(b1_x1, b2_x1)
    y2 = tf.minimum(b1_y2, b2_y2)
    x2 = tf.minimum(b1_x2, b2_x2)
    intersection = tf.maximum(x2 - x1, 0) * tf.maximum(y2 - y1, 0)

    # 3. Compute unions
    b1_area = (b1_y2 - b1_y1) * (b1_x2 - b1_x1)
    b2_area = (b2_y2 - b2_y1) * (b2_x2 - b2_x1)
    union = b1_area + b2_area - intersection

    # 4. Compute IoU and reshape to [boxes1, boxes2]
    iou = intersection / union
    overlaps = tf.reshape(iou, [tf.shape(boxes1)[0], tf.shape(boxes2)[0]])
    # print('\t     Overlaps_graph(): Shape of output overlaps', tf.shape(overlaps), overlaps.get_shape())
    return overlaps

def detection_targets_graph_mod(proposals, gt_class_ids, gt_boxes,  config):
    '''
    Generates detection targets for one image. Subsamples proposals and
    generates target class IDs, bounding box deltas, and masks for each.

    Inputs:
    -------
    proposals:              [N, 2000, (y1, x1, y2, x2)] in normalized coordinates.
                            Proposals passed from the proposal layer.
                            Might be zero padded if there are not enough proposals.

    gt_class_ids:           [MAX_GT_INSTANCES]
                            ground truth object class ids

    gt_boxes:               [MAX_GT_INSTANCES, (y1, x1, y2, x2)]
                            ground truth bbox normalized coordinates

    gt_masks:               [height, width, MAX_GT_INSTANCES] of boolean type.
                            ground truth mask information

    Returns:            Target ROIs and corresponding class IDs, bounding box shifts, and masks.
    --------
    output_rois:            [TRAIN_ROIS_PER_IMAGE, (y1, x1, y2, x2)] in normalized coordinates
                            Contains BOTH POSITIVE and NEGATIVE examples

    target_class_ids:       [TRAIN_ROIS_PER_IMAGE]. Integer class IDs. Zero padded.
                            only contains POSTITIVE exmaples

    target_bbox_deltas:     [TRAIN_ROIS_PER_IMAGE, NUM_CLASSES, (dy, dx, log(dh), log(dw))]
                            Class-specific bbox refinments between the  POSITIVE examples
                            (in target_gt_bboxes) and their ground truth counter parts (in output_rois)

    target_gt_bboxes        [TRAIN_ROIS_PER_IMAGE, (y1, x1, y2, x2)]
    (roi_gt_bboxes)         Ground Truth bbox coordinates - Only for POSITIVE examples

    target_gt_masks:        [TRAIN_ROIS_PER_IMAGE, height, width). Masks cropped to bbox
                            boundaries and resized to neural network output size.
                            only for POSITIVE examples.

    Note: Returned arrays might be zero padded if not enough target ROIs.

    '''
    # Assertions
    asserts = [
        tf.Assert(tf.greater(tf.shape(proposals)[0], 0), [proposals], name="roi_assertion"),
    ]

    with tf.control_dependencies(asserts):
        proposals = tf.identity(proposals)
    # print('>>> detection_targets_graph ')
    # print('     propsals.shape        :',  proposals.shape, proposals.get_shape(), KB.int_shape(proposals) )
    # print('     gt_boxes.shape        :',  gt_boxes.shape ,    KB.int_shape(gt_boxes)   )
    # print('     gt_class_ids.shape    :',  gt_class_ids.shape, KB.int_shape(gt_class_ids))
    # print('     gt_masks.shape        :',  gt_masks.shape ,    KB.int_shape(gt_masks)   )

    # Remove zero padding
    # non_zeros returns indicies to valid bboxes, which we use to index gt_class_ids, and gt_masks
    proposals, _        = utils.trim_zeros_graph(proposals, name="trim_proposals")
    gt_boxes, non_zeros = utils.trim_zeros_graph(gt_boxes , name="trim_gt_boxes")
    gt_class_ids        = tf.boolean_mask(gt_class_ids, non_zeros, name="trim_gt_class_ids")
    # gt_masks            = tf.gather(gt_masks, tf.where(non_zeros)[:, 0], axis=2,name="trim_gt_masks")

    #------------------------------------------------------------------------------------------
    # Handle COCO crowds
    # A crowd box in COCO is a bounding box around several instances. Exclude
    # them from training. A crowd box is given a negative class ID.
    #------------------------------------------------------------------------------------------
    # tf.where : returns the coordinates of true elements of  the specified conditon.
    #            The coordinates are returned in a 2-D tensor where the first dimension (rows)
    #            represents the number of true elements, and the second dimension (columns)
    #            represents the coordinates of the true elements.
    #            Keep in mind, the shape of the output tensor can vary depending on how many
    #            true values there are in input. Indices are output in row-major order.
    #
    # tf.gather: Gather slices from params axis (default = 0) according to indices.
    #            indices must be an integer tensor of any dimension (usually 0-D or 1-D).
    #            Produces an output tensor with shape:
    #                   params.shape[:axis] + indices.shape + params.shape[axis + 1:]
    #
    # tf.squeeze: Removes dimensions of size 1 from the shape of a tensor.
    #            Given a tensor input, this operation returns a tensor of the same type with
    #            all dimensions of size 1 removed. If you don't want to remove all size 1
    #            dimensions, you can remove specific size 1 dimensions by specifying axis.
    #------------------------------------------------------------------------------------------

    #------------------------------------------------------------------------------------------
    ##  Separate GT boxes and masks by 'crowd' and 'non-crowd' classifications
    #------------------------------------------------------------------------------------------
    crowd_ix        = tf.where(gt_class_ids < 0)[:, 0]
    non_crowd_ix    = tf.where(gt_class_ids > 0)[:, 0]
    crowd_boxes     = tf.gather(gt_boxes, crowd_ix)
    # crowd_masks     = tf.gather(gt_masks, crowd_ix, axis=2)
    gt_class_ids    = tf.gather(gt_class_ids, non_crowd_ix)
    gt_boxes        = tf.gather(gt_boxes, non_crowd_ix)
    # gt_masks        = tf.gather(gt_masks, non_crowd_ix, axis=2)

    #------------------------------------------------------------------------------------------
    # Compute overlaps with crowd boxes [anchors, crowds]
    #------------------------------------------------------------------------------------------
    crowd_overlaps  = overlaps_graph_mod(proposals, crowd_boxes)
    crowd_iou_max   = tf.reduce_max(crowd_overlaps, axis=1)
    no_crowd_bool   = (crowd_iou_max < 0.001)

    #------------------------------------------------------------------------------------------
    # Compute overlaps matrix [proposals, gt_boxes] - The IoU between
    # proposals and gt_boxes (non-crowd gt boxes, designated by classId < 0 in Coco)
    # overlaps is
    # compute max of elements across axis 1 of overlaps tensor.
    #------------------------------------------------------------------------------------------
    overlaps        = overlaps_graph_mod(proposals, gt_boxes)
    roi_iou_max     = tf.reduce_max(overlaps, axis=1)
    # print('     overlaps.shape        :',  overlaps.shape, KB.int_shape(overlaps)   )

    #------------------------------------------------------------------------------------------
    ## 1. Determine indices of postive ROI propsal boxes
    #------------------------------------------------------------------------------------------
    #    Identify ROI proposal boxes that have an IoU >= 05 overlap with some gt_box, and store
    #    indices into positive_indices
    positive_roi_bool     = (roi_iou_max >= 0.5)
    positive_indices      = tf.where(positive_roi_bool)[:, 0]

    #------------------------------------------------------------------------------------------
    ## 2. Determine indices of negative ROI proposal boxes
    #------------------------------------------------------------------------------------------
    #    those with < 0.5 with every GT box and are not crowds bboxes
    # the where creates a array with shape [# of answers, 1] so we use [:, 0] after
    ## current method
    negative_indices      = tf.where(tf.logical_and(roi_iou_max < 0.5, no_crowd_bool))[:, 0]

    ## new method
    # this modification will determine negative ROI proposal boxes but in addition,
    # will suppress the zero RoIs from the indicies
    # note that   ( negative_bool         = ~positive_roi_bool)
    # negative_nonzero_bool = tf.logical_and(~positive_roi_bool, (roi_iou_max > 0))
    # negative_nonzero_bool = tf.logical_and(negative_nonzero_bool, no_crowd_bool)
    # negative_indices2     = tf.where(negative_nonzero_bool) [:, 0]

    #------------------------------------------------------------------------------------------
    ## 3. Subsample positive ROIs based on ROI_POSITIVE_RATIO
    #------------------------------------------------------------------------------------------
    #    Aim for 33% positive (config.ROI_POSITIVE_RATIO = 0.33)
    #    Positive ROIs   33% of config.TRAIN_ROIS_PER_IMAGE ~  11
    positive_count        = int(config.TRAIN_ROIS_PER_IMAGE * config.ROI_POSITIVE_RATIO)
    positive_indices      = tf.random_shuffle(positive_indices)[:positive_count]
    positive_count        = tf.shape(positive_indices)[0]

    #------------------------------------------------------------------------------------------
    ## 4. Add Negative ROIs. Add enough to maintain positive:negative ratio
    #------------------------------------------------------------------------------------------
    #     negative_count = int((positive_count / config.ROI_POSITIVE_RATIO) - positive_count)
    r = 1.0 / config.ROI_POSITIVE_RATIO
    negative_count        = tf.cast(r * tf.cast(positive_count, tf.float32), tf.int32) - positive_count
    negative_indices      = tf.random_shuffle(negative_indices)[:negative_count]

    #------------------------------------------------------------------------------------------
    ## 5.   Gather selected positive and negative ROIs
    #------------------------------------------------------------------------------------------
    positive_rois         = tf.gather(proposals, positive_indices)
    negative_rois         = tf.gather(proposals, negative_indices)

    #------------------------------------------------------------------------------------------
    ## 6.   Assign positive ROIs to GT boxes.
    #------------------------------------------------------------------------------------------
    #      roi_gt_box_assignment shows for each positive overlap, which class has the maximum overlap
    positive_overlaps     = tf.gather(overlaps, positive_indices)
    roi_gt_box_assignment = tf.argmax(positive_overlaps, axis=1)
    roi_gt_boxes          = tf.gather(gt_boxes    , roi_gt_box_assignment)
    roi_gt_class_ids      = tf.gather(gt_class_ids, roi_gt_box_assignment)
    # print('     shape of positive overlaps is :', positive_overlaps.get_shape())

    #------------------------------------------------------------------------------------------
    ## 7.   Compute bbox delta (between the selected positive RPN proposals and corrspoing gt bboxes)
    #------------------------------------------------------------------------------------------
    # calculate refinement (difference b/w positive rois and gt_boxes) for positive ROIs
    roi_gt_deltas  = utils.box_refinement_graph(positive_rois, roi_gt_boxes)
    roi_gt_deltas /= config.BBOX_STD_DEV

    #------------------------------------------------------------------------------------------
    ## 8.  prepare gt_masks
    #------------------------------------------------------------------------------------------
    #      transpose gt_masks from [h, w, N] to [N, height, width] and add 4th dim at end [N, height, width, 1]
    #      Pick the right mask for each ROI
    # transposed_masks = tf.expand_dims(tf.transpose(gt_masks, [2, 0, 1]), -1)
    # roi_masks = tf.gather(transposed_masks, roi_gt_box_assignment)

    # Compute mask targets
    # boxes = positive_rois

    # if config.USE_MINI_MASK:
        # Transform ROI corrdinates from normalized image space
        # to normalized mini-mask space.
        # y1, x1, y2, x2 = tf.split(positive_rois, 4, axis=1)
        # gt_y1, gt_x1, gt_y2, gt_x2 = tf.split(roi_gt_boxes, 4, axis=1)
        # gt_h = gt_y2 - gt_y1
        # gt_w = gt_x2 - gt_x1
        # y1 = (y1 - gt_y1) / gt_h
        # x1 = (x1 - gt_x1) / gt_w
        # y2 = (y2 - gt_y1) / gt_h
        # x2 = (x2 - gt_x1) / gt_w
        # boxes = tf.concat([y1, x1, y2, x2], 1)

    # box_ids = tf.range(0, tf.shape(roi_masks)[0])
    # masks   = tf.image.crop_and_resize(tf.cast(roi_masks, tf.float32),
                                       # boxes,
                                       # box_ids,
                                       # config.MASK_SHAPE)

    # Remove the extra dimension from masks.
    # masks = tf.squeeze(masks, axis=3)

    # Threshold mask pixels at 0.5 to have GT masks be 0 or 1 to use with
    # binary cross entropy loss.
    # masks = tf.round(masks)

    #------------------------------------------------------------------------------------------
    ## 9. Prepare final outputs
    # Append negative ROIs and pad bbox roi_gt_deltas and masks that
    # are not used for negative ROIs with zeros.
    #------------------------------------------------------------------------------------------
    rois             = tf.concat([positive_rois, negative_rois], axis=0)
    N                = tf.shape(negative_rois)[0]
    P                = tf.maximum(config.TRAIN_ROIS_PER_IMAGE - tf.shape(rois)[0], 0)
    rois             = tf.pad(rois            , [(0, P), (0, 0)])

    roi_gt_class_ids = tf.pad(roi_gt_class_ids, [(0, N + P)])
    roi_gt_deltas    = tf.pad(roi_gt_deltas   , [(0, N + P), (0, 0)])
    roi_gt_boxes     = tf.pad(roi_gt_boxes    , [(0, N + P), (0, 0)])
    # masks            = tf.pad(masks           , [[0, N + P], (0, 0), (0, 0)])

    # print(' roi_gt_boxes :  ' , tf.shape(roi_gt_boxes) )
    # print(' P:  ' , P,  ' N :    ', N)
    # print('     roi.shape             :',  rois.shape            , tf.shape(rois))
    # print('     roi_gt_class_ids.shape:',  roi_gt_class_ids.shape, tf.shape(roi_gt_class_ids))
    # print('     roi_gt_deltas.shape   :',  roi_gt_deltas.shape   , tf.shape(roi_gt_deltas))
    # print('     masks.shape           :',  masks.shape           , tf.shape(masks))
    # print('     roi_gt_boxes.shape    :',  roi_gt_boxes.shape    , tf.shape(roi_gt_boxes))

    return rois, roi_gt_class_ids,  roi_gt_deltas, roi_gt_boxes

class DetectionTargetLayer_mod(KE.Layer):
    '''
    Subsamples proposals and generates target box refinment, class_ids,
    and masks for each.

    Inputs:
    -------
    proposals:              [N, 2000, (y1, x1, y2, x2)] in normalized coordinates.
                            Proposals passed from the proposal layer.
                            Might be zero padded if there are not enough proposals.

    gt_class_ids:           [MAX_GT_INSTANCES]
                            ground truth object class ids

    gt_boxes:               [MAX_GT_INSTANCES, (y1, x1, y2, x2)]
                            ground truth bbox normalized coordinates

    gt_masks:               [height, width, MAX_GT_INSTANCES] of boolean type.
                            ground truth mask information

    Returns:            Target ROIs and corresponding class IDs, bounding box shifts, and masks.
    --------
    output_rois:            [TRAIN_ROIS_PER_IMAGE, (y1, x1, y2, x2)] in normalized coordinates
                            Contains BOTH POSITIVE and NEGATIVE examples

    target_class_ids:       [TRAIN_ROIS_PER_IMAGE]. Integer class IDs. Zero padded.
                            only contains POSTITIVE exmaples

    target_bbox_deltas:     [TRAIN_ROIS_PER_IMAGE, NUM_CLASSES, (dy, dx, log(dh), log(dw))]
                            Class-specific bbox refinments between the  POSITIVE examples
                            (in target_gt_bboxes) and their ground truth counter parts (in output_rois)

    target_gt_bboxes        [TRAIN_ROIS_PER_IMAGE, (y1, x1, y2, x2)]
    (roi_gt_bboxes)         Ground Truth bbox coordinates - Only for POSITIVE examples

    target_gt_masks:        [TRAIN_ROIS_PER_IMAGE, height, width). Masks cropped to bbox
                            boundaries and resized to neural network output size.
                            only for POSITIVE examples.

    Note: Returned arrays might be zero padded if not enough target ROIs.
    '''

    def __init__(self, config, **kwargs):
        # super(DetectionTargetLayer, self).__init__(**kwargs)
        super().__init__(**kwargs)
        print('\n>>> Detection Target Layer (Training Mode)')
        self.config = config

    def call(self, inputs):
        print('    Detection Target Layer : call() ', type(inputs), len(inputs))
        print('     proposals.shape    :',  inputs[0].shape, inputs[0].get_shape(), KB.int_shape(inputs[0]) )
        print('     gt_class_ids.shape :',  inputs[1].shape, inputs[1].get_shape(), KB.int_shape(inputs[1]) )
        print('     gt_bboxes.shape    :',  inputs[2].shape, inputs[2].get_shape(), KB.int_shape(inputs[2]) )
        # print('     gt_masks.shape     :',  inputs[3].shape, inputs[3].get_shape(), KB.int_shape(inputs[3]) )
        proposals    = inputs[0]    # target_rois           --  proposals generated by the RPN (or artificially generated proposals)
        gt_class_ids = inputs[1]    # input_gt_class_ids
        gt_boxes     = inputs[2]    # input_normlzd_gt_boxes
        # gt_masks     = inputs[3]    # input_gt_masks

        # Slice the batch and run a graph for each slice
        # TODO: Rename target_bbox to target_deltas for clarity
        # detection_target_graph() returns:
        #         rois,    roi_gt_class_ids,  deltas,         masks,       roi_gt_boxes

        names = ["output_rois", "target_class_ids", "target_bbox_deltas", "roi_gt_boxes"]

        outputs = utils.batch_slice([proposals, gt_class_ids, gt_boxes],             # inputs
                                    lambda w, x, y: detection_targets_graph_mod(w, x, y, self.config),        # batch function
                                    self.config.IMAGES_PER_GPU,                                                 # batch_size, name
                                    names=names)

        print('\n    Detection Target Layer : return ', type(outputs) , len(outputs))
        for i,out in enumerate(outputs):
            print('     output {}  shape {}  type {} '.format(i, out.shape, type(out)))

        return outputs

    def compute_output_shape(self, input_shape):
        return [
            (None, self.config.TRAIN_ROIS_PER_IMAGE, 4),                                                        # rois boxes
            (None, self.config.TRAIN_ROIS_PER_IMAGE),                                                                                # gt class_ids
            (None, self.config.TRAIN_ROIS_PER_IMAGE, 4),                                                        # roi/gt deltas
            # (None, self.config.TRAIN_ROIS_PER_IMAGE, self.config.MASK_SHAPE[0], self.config.MASK_SHAPE[1]),     # masks
            (None, self.config.TRAIN_ROIS_PER_IMAGE, 4)                                                         # roi_gt_bboxes

        ]

    def compute_mask(self, inputs, mask=None):
        return [None, None, None,  None]


#!/usr/bin/env python3
import sys
import time
from threading import Thread, Lock
import tensorflow as tf
from gi.repository import Playerctl, GLib
import numpy as np
import cv2
from ml import encoder, decoder, detect_face, model_fn, get_parser_fn

class WebcamVideoStream:

    def __init__(self, src=0):
        self.stream = cv2.VideoCapture(src)
        (self.grabbed, self.frame) = self.stream.read()
        self.started = False
        self.read_lock = Lock()

    def start(self):
        if self.started:
            print("already started!!")
            return None
        self.started = True
        self.thread = Thread(target=self.update, args=())
        self.thread.start()
        return self

    def update(self):
        while self.started:
            (grabbed, frame) = self.stream.read()
            self.read_lock.acquire()
            self.grabbed, self.frame = grabbed, frame
            self.read_lock.release()

    def read(self):
        self.read_lock.acquire()
        frame = self.frame.copy()
        self.read_lock.release()
        return frame

    def stop(self):
        self.started = False
        self.thread.join()

    def __exit__(self, exc_type, exc_value, traceback):
        self.stream.release()

def facer(player):
    face_cascade_classifier = cv2.CascadeClassifier(
        "resources/haarcascade_frontalface_default.xml")
    thres = 0.56
    #cap = cv2.VideoCapture(0)
    cap = WebcamVideoStream().start()

    tick = 20

    estimator = tf.estimator.Estimator(model_fn, "./model_dir/")
    input_ = tf.placeholder(tf.int8, shape=(None, None, None, 3))
    parser = get_parser_fn(tf.estimator.ModeKeys.PREDICT)
    parsed_input = parser(input_)
    reconstructions = decoder(encoder(parsed_input), parsed_input.shape[-1])
    loss = tf.losses.mean_squared_error(parsed_input, reconstructions)
    saver = tf.train.Saver()

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    with tf.Session(config=config) as sess:
        saver.restore(sess, estimator.latest_checkpoint())

        paused = False
        squared_error_vals = []
        not_detected = 0
        while True:
            #_, frame = cap.read()
            frame = cap.read()
            face_rect = detect_face(face_cascade_classifier, frame)
            if face_rect is not None:
                not_detected = 0
                x, y, w, h = face_rect
                x -= 30
                y -= 30
                w += 30
                h += 30
                if x < 0 or y < 0 or w > frame.shape[0] or h > frame.shape[1]:
                    continue
                face = frame[y:y + h, x:x + w]
                squared_error_val = sess.run(
                    loss, feed_dict={input_: np.expand_dims(face, axis=0)})
                squared_error_vals.append(squared_error_val)
                if len(squared_error_vals) >= tick or paused:
                    mse = np.mean(np.array(squared_error_vals))
                    print(mse)
                    squared_error_vals.clear()
                    if mse > thres:
                        player.pause()
                    else:
                        player.play()
                        paused = False
            else:
                # there's no one in front of the camera for more than 2*tick
                not_detected += 1
                if not_detected > 2 * tick:
                    not_detected = 0
                    player.pause()
                    paused = True
            time.sleep(0.15)

def main():
    player = Playerctl.Player(player_name='spotify')
    vision = Thread(target=facer, args=(player,))
    main_loop = GLib.MainLoop()
    vision.start()
    main_loop.run()
    vision.join()

if __name__ == "__main__":
    sys.exit(main())


# original author: Zylo117
# modified by: Marek Sicha

"""
Simple Inference Script of EfficientDet-Pytorch for detecting objects on webcam
"""
import time
import torch
import cv2
import numpy as np
import tensorflow as tf

from torch.backends import cudnn

from backbone import EfficientDetBackbone
from efficientdet.utils import BBoxTransform, ClipBoxes
from utils.utils import preprocess_video, invert_affine, postprocess, STANDARD_COLORS, standard_to_bgr, plot_one_box

compound_coef = 2
num_classes_detector = 1
force_input_size = None  # set None to use default size

# Video's path
video_src = 0  # set int to use webcam, set str to read from a video file

# Classifier´s model path
model_classifier = tf.keras.models.load_model("saved_model/")

# Detector´s model path
model_detector = 'efficientdet-d2_72_36000.pth'

anchor_ratios = [(1.0, 1.0), (1.0, 1.0), (1.0, 0.9)]
anchor_scales = [2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)]

# Detector
threshold = 0.5
iou_threshold = 0.5

# Classifier
classifier_threshold = 0.8

use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True

anchor_ratios = [(1.0, 1.0), (1.0, 1.0), (1.0, 0.9)]
anchor_scales = [2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)]

color_list = standard_to_bgr(STANDARD_COLORS)

input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1536]
input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size

model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=num_classes_detector,
                             ratios=anchor_ratios, scales=anchor_scales)
model.load_state_dict(torch.load(model_detector))
model.requires_grad_(False)
model.eval()

if use_cuda:
    model = model.cuda()
if use_float16:
    model = model.half()

def classifier_img(img):
    img = cv2.equalizeHist(cv2.cvtColor(cv2.resize(img, (32, 32)), cv2.COLOR_BGR2GRAY))
    img = img.reshape(1, 32, 32, 1)
    prediction = np.amax(model_classifier.predict(img))
    class_ID = np.argmax(model_classifier.predict(img), axis=-1)

    return(img, prediction, class_ID)

def get_class_name(class_ID):
    classes = open('classes_cz.txt', 'r')
    class_line = classes.read()
    class_name = class_line.splitlines()

    return(class_name[class_ID[0]])

def display(preds, imgs):
    for i in range(len(imgs)):
        if len(preds[i]['rois']) == 0:
            return imgs[i]

        for j in range(len(preds[i]['rois'])):
            x1, y1, x2, y2 = preds[i]['rois'][j].astype(np.int32)
            _, prediction, class_ID = classifier_img((imgs[i][y1:y2, x1:x2]))
            class_name = get_class_name(class_ID)
            if prediction <= classifier_threshold:
                continue
            plot_one_box(imgs[i], [x1, y1, x2, y2], label=class_name,
                         score=prediction, color=color_list[(int(class_ID))])

        return imgs[i]

regressBoxes = BBoxTransform()
clipBoxes = ClipBoxes()

# Video capture
cap = cv2.VideoCapture(video_src)
prev_frame_time = 0
while True:
    ret, frame = cap.read()
    if not ret:
        break
    new_frame_time = time.time()
    ori_imgs, framed_imgs, framed_metas = preprocess_video(frame, max_size=input_size)

    if use_cuda:
        x = torch.stack([torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)
    else:
        x = torch.stack([torch.from_numpy(fi) for fi in framed_imgs], 0)

    x = x.to(torch.float32 if not use_float16 else torch.float16).permute(0, 3, 1, 2)

    with torch.no_grad():
        features, regression, classification, anchors = model(x)

        out = postprocess(x,
                          anchors, regression, classification,
                          regressBoxes, clipBoxes,
                          threshold, iou_threshold)

    out = invert_affine(framed_metas, out)
    img_show = display(out, ori_imgs)
    cv2.putText(img_show, str(int(1/(new_frame_time-prev_frame_time))),
                (4, 30), cv2.FONT_HERSHEY_SIMPLEX,
                1, (100, 255, 0), 1, cv2.LINE_AA)
    prev_frame_time = new_frame_time
    cv2.imshow('frame', img_show)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()


#!/usr/bin/env python

from __future__ import print_function

import argparse
import glob
import os.path as osp

import itertools, pkg_resources, sys
from distutils.version import LooseVersion
if LooseVersion(pkg_resources.get_distribution("chainer").version) >= LooseVersion('7.0.0') and \
   sys.version_info.major == 2:
   print('''Please install chainer <= 7.0.0:

    sudo pip install chainer==6.7.0

c.f https://github.com/jsk-ros-pkg/jsk_recognition/pull/2485
''', file=sys.stderr)
   sys.exit(1)
if [p for p in list(itertools.chain(*[pkg_resources.find_distributions(_) for _ in sys.path])) if "cupy-" in p.project_name ] == []:
   print('''Please install CuPy

    sudo pip install cupy-cuda[your cuda version]
i.e.
    sudo pip install cupy-cuda91

''', file=sys.stderr)
   sys.exit(1)
import chainer
from chainer import cuda
import numpy as np
import rospkg
import skimage.io
from sklearn.metrics import classification_report
from sklearn.neighbors import KNeighborsClassifier
import tqdm
import yaml

# TODO(wkentaro): Support Resnet50/101
from jsk_recognition_utils.chainermodels import ResNet152Feature

def get_templates(template_dir):
    target_names = yaml.load(open(osp.join(template_dir, 'target_names.yaml')))
    for cls_id, cls_name in enumerate(target_names):
        obj_dir = osp.join(template_dir, cls_name)
        for img_file in glob.glob(osp.join(obj_dir, '*.jpg')):
            dirname, basename = osp.split(img_file)
            mask_file = osp.join(dirname, 'masks', basename)
            img = skimage.io.imread(img_file)
            mask = skimage.io.imread(mask_file) >= 127
            yield cls_id, img, mask

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('template_dir', help='Template dir')
    parser.add_argument('db_file', help='DB file which will be created')
    parser.add_argument('-g', '--gpu', type=int, default=0)
    args = parser.parse_args()

    template_dir = args.template_dir
    db_file = args.db_file
    gpu = args.gpu

    pkg_path = rospkg.RosPack().get_path('jsk_perception')
    mean_file = osp.join(pkg_path, 'trained_data/resnet_lsvrc2012_mean.npy')
    pretrained_model = osp.join(
        pkg_path, 'trained_data/resnet152_from_caffe.npz')

    target_names = yaml.load(open(osp.join(template_dir, 'target_names.yaml')))

    mean = np.load(mean_file)

    model = ResNet152Feature()
    chainer.serializers.load_npz(pretrained_model, model)
    if gpu >= 0:
        chainer.cuda.get_device_from_id(gpu).use()
        model.to_gpu()

    chainer.global_config.train = False
    chainer.global_config.enable_backprop = False

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

    X = []
    y = []
    for cls_id, img, mask in tqdm.tqdm(get_templates(template_dir)):
        img = img[:, :, ::-1]  # RGB- > BGR
        img = img.astype(np.float64)
        img[mask] -= mean[mask]
        img[~mask] = 0

        img = img.transpose(2, 0, 1)
        img = img.astype(np.float32)
        x_data = np.asarray([img])
        x_data = cuda.to_gpu(x_data)
        x = chainer.Variable(x_data)
        feat = model(x)

        feat = cuda.to_cpu(feat.data)
        feat = feat.squeeze(axis=(2, 3))
        for f in feat:
            X.append(f)
            y.append(cls_id)
    X = np.asarray(X)
    y = np.asarray(y)
    np.savez_compressed(db_file, X=X, y=y, target_names=target_names)

    knn = KNeighborsClassifier(n_neighbors=1)
    knn.fit(X, y)
    y_pred = knn.predict(X)
    # validation: must be all 1.0
    print(classification_report(y, y_pred, labels=range(len(target_names)),
                                target_names=target_names))

if __name__ == '__main__':
    main()


import copy
from typing import Callable, List, Literal, Optional, Tuple, Union, overload

import numpy
import openmm
import openmm.app
from openff.toolkit.topology import Topology
from openmm import unit

SystemGenerator = Callable[
    [Topology, unit.Quantity, Literal["solvent-a", "solvent-b"]], openmm.System
]

OpenMMPlatform = Literal["Reference", "OpenCL", "CUDA", "CPU"]

@overload
def array_to_vectors(box_vectors: numpy.ndarray) -> List[openmm.Vec3]:
    ...

@overload
def array_to_vectors(box_vectors: unit.Quantity) -> unit.Quantity:
    ...

def array_to_vectors(array):
    """A utility to convert numpy arrays into lists of OpenMM vectors.

    Args:
        array: The array to convert with shape=(n_items, 3)

    Returns:
        A list of three ``Vec3`` objects
    """

    original_units = None if not isinstance(array, unit.Quantity) else array.unit

    if original_units is not None:
        array = array.value_in_unit(original_units)

    vectors = [openmm.Vec3(row[0], row[1], row[2]) for row in array]

    return vectors if original_units is None else vectors * original_units

def build_context(
    system: openmm.System,
    coordinates: unit.Quantity,
    box_vectors: Optional[unit.Quantity],
    temperature: unit.Quantity,
    pressure: Optional[unit.Quantity],
    platform: OpenMMPlatform = "Reference",
    timestep: unit.Quantity = 2.0 * unit.femtoseconds,
    thermostat_friction: unit.Quantity = 1.0 / unit.picosecond,
) -> openmm.Context:
    """Builds a new OpenMM context.

    Args:
        system: The system encoding the potential energy function.
        coordinates: The atom coordinates stored in a unit wrapped array with
            shape=(n_atoms, 3).
        box_vectors: The (optional) box vectors stored in a unit wrapped array with
            shape=(3, 3).
        temperature: The temperature to simulate at.
        pressure: The (optional) pressure to simulate at.
        platform: The platform to simulate on.
        timestep: The default timestep to use.
        thermostat_friction: The default thermostat friction to set.
    """

    system = copy.deepcopy(system)

    assert pressure is None or (
        pressure is not None and box_vectors is not None
    ), "box vectors must be provided when the pressure is specified."

    assert (
        len(
            [
                force
                for force in system.getForces()
                if isinstance(force, openmm.MonteCarloBarostat)
            ]
        )
        == 0
    ), "the system should not already contain a barostat"

    if pressure is not None:
        system.addForce(openmm.MonteCarloBarostat(pressure, temperature, 25))

    integrator = openmm.LangevinIntegrator(temperature, thermostat_friction, timestep)

    context = openmm.Context(
        system,
        integrator,
        openmm.Platform.getPlatformByName(platform),
    )

    set_coordinates(context, coordinates, box_vectors)

    context.setVelocitiesToTemperature(temperature * unit.kelvin)

    return context

def set_alchemical_lambdas(
    context: openmm.Context,
    lambda_sterics: Optional[float] = None,
    lambda_electrostatics: Optional[float] = None,
):
    """Set the values of the alchemical lambdas on an OpenMM context.

    Args:
        context: The context to update.
        lambda_sterics: The (optional) value of the steric lambda.
        lambda_electrostatics: The (optional) value of the electrostatics lambda.
    """

    if lambda_sterics is not None:
        assert 0.0 <= lambda_sterics <= 1.0, "`lambda_sterics` must be between 0 and 1"
        context.setParameter("lambda_sterics", lambda_sterics)

    if lambda_electrostatics is not None:

        assert (
            0.0 <= lambda_electrostatics <= 1.0
        ), "`lambda_electrostatics` must be between 0 and 1"

        context.setParameter("lambda_electrostatics", lambda_electrostatics)

def set_coordinates(
    context: openmm.Context,
    coordinates: unit.Quantity,
    box_vectors: Optional[unit.Quantity],
):
    """Updates an OpenMM context with a new set of coordinates and box vectors.

    Args:
        context: The context to update.
        coordinates: The new coordinates stored in a unit wrapped array with
            shape=(n_atoms, 3).
        box_vectors: The (optional) box vectors stored in a unit wrapped array with
            shape=(3, 3) or list of three ``Vec3`` objects.
    """

    if box_vectors is not None:

        if isinstance(box_vectors.value_in_unit(unit.nanometers), numpy.ndarray):
            box_vectors = array_to_vectors(box_vectors)

        context.setPeriodicBoxVectors(*box_vectors)

    system: openmm.System = context.getSystem()

    if len(coordinates) != system.getNumParticles():

        coordinates = coordinates.value_in_unit(unit.nanometers)

        n_v_sites = sum(
            1 for i in range(system.getNumParticles()) if system.isVirtualSite(i)
        )
        n_atoms = system.getNumParticles() - n_v_sites

        assert len(coordinates) == n_atoms, (
            "coordinates must either have shape=(n_atoms, 3) "
            "or (n_atoms + n_v_sites, 3)"
        )

        full_coordinates = numpy.zeros((system.getNumParticles(), 3))
        counter = 0

        for i in range(system.getNumParticles()):

            if system.isVirtualSite(i):
                continue

            full_coordinates[i] = coordinates[counter]
            counter += 1

        coordinates = full_coordinates * unit.nanometers

    context.setPositions(coordinates)
    context.computeVirtualSites()

def extract_coordinates(
    state: Union[openmm.State, openmm.Context]
) -> Tuple[unit.Quantity, unit.Quantity]:
    """Extracts the current coordinates and box vectors from an OpenMM context.

    Args:
        state: The state (or context) to extract from.

    Returns:
        The current coordinates and box vectors stored in unit wrapped arrays with
        shape=(n_atoms, 3) and shape=(3, 3) respectively.
    """

    if isinstance(state, openmm.Context):
        state = state.getState(getPositions=True)

    box_vectors = state.getPeriodicBoxVectors().value_in_unit(unit.nanometers)

    return (
        state.getPositions(asNumpy=True),
        numpy.array(
            [
                [box_vectors[0].x, box_vectors[0].y, box_vectors[0].z],
                [box_vectors[1].x, box_vectors[1].y, box_vectors[1].z],
                [box_vectors[2].x, box_vectors[2].y, box_vectors[2].z],
            ]
        )
        * unit.nanometers,
    )

def minimize(
    context: openmm.Context,
    tolerance: unit.Quantity = 10 * unit.kilojoules_per_mole / unit.nanometer,
    max_iterations: int = 0,
):
    """Energy minimize an OpenMM context.

    Args:
        context: The context to minimize.
        tolerance: How precisely the energy minimum must be located
        max_iterations: The maximum number of iterations to perform. If this is 0,
            minimization is continued until the results converge.
    """
    openmm.LocalEnergyMinimizer.minimize(context, tolerance, max_iterations)

def evaluate_energy(
    system: openmm.System,
    coordinates: unit.Quantity,
    box_vectors: Optional[unit.Quantity] = None,
    lambda_sterics: Optional[float] = None,
    lambda_electrostatics: Optional[float] = None,
    platform: OpenMMPlatform = "Reference",
) -> unit.Quantity:
    """Evaluates the energy of a given system at a particular set of coordinates.

    Args:
        system: The openmm system that should be used to evaluate the energies.
        coordinates: The coordinates that should be used when evaluating the energies.
        box_vectors: The (optional) periodic box vectors that should be used when
            evaluating the energies.
        lambda_sterics: The value of `lambda_sterics` to evaluate the energies at.
        lambda_electrostatics: The value of `lambda_electrostatics` to evaluate the
            energies at.
        platform: The OpenMM platform to simulate using.

    Returns:
        The energy.
    """

    context = build_context(
        system, coordinates, box_vectors, 1.0 * unit.kelvin, None, platform
    )
    set_alchemical_lambdas(context, lambda_sterics, lambda_electrostatics)

    return context.getState(getEnergy=True).getPotentialEnergy()

def create_system_generator(
    force_field: openmm.app.ForceField,
    solvent_a_nonbonded_method: int,
    solvent_b_nonbonded_method: int,
    nonbonded_cutoff: unit.Quantity = 1.0 * unit.nanometer,
    constraints: Optional[int] = None,
    rigid_water: Optional[bool] = None,
    remove_cmm_motion: bool = True,
    hydrogen_mass: Optional[unit.Quantity] = None,
    switch_distance: Optional[unit.Quantity] = None,
) -> SystemGenerator:
    """Creates a 'system generator' that can be used when setting up an alchemical
    free energy calculation from an OpenMM force field.

    Args:
        force_field: The OpenMM force field to parameterize the topology using.
        solvent_a_nonbonded_method: The non-bonded method to use in solvent a.
        solvent_b_nonbonded_method: The non-bonded method to use in solvent b.
        nonbonded_cutoff: The non-bonded cutoff to use.
        constraints: The type of constraints to apply to the system.
        rigid_water: Whether to force rigid water.
        remove_cmm_motion: Whether to remove any CMM motion.
        hydrogen_mass: The mass to use for hydrogens.
        switch_distance: The switch distance to use.

    Returns:
        A callable that will create an OpenMM system from an OpenFF topology and the
        name of the solvent (i.e. ``"solvent-a"`` or ``"solvent-b"``) the system will
        be used for.
    """

    def system_generator(
        topology: Topology,
        coordinates: unit.Quantity,
        solvent_index: Literal["solvent-a", "solvent-b"],
    ) -> openmm.System:

        openmm_topology = topology.to_openmm()

        if topology.box_vectors is not None:
            openmm_topology.setPeriodicBoxVectors(topology.box_vectors)

        # We need to fix the special case of water in order for OMM to correctly apply
        # a constraint between H atoms.
        for chain in openmm_topology.chains():

            for residue in chain.residues():

                if len(residue) != 3:
                    continue

                symbols = sorted(atom.element.symbol for atom in residue.atoms())

                if symbols != ["H", "H", "O"]:
                    continue

                residue.name = "HOH"

        from openmm.app import Modeller

        modeller = Modeller(
            openmm_topology,
            [
                openmm.Vec3(coordinate[0], coordinate[1], coordinate[2])
                for coordinate in coordinates.value_in_unit(unit.nanometers)
            ]
            * unit.nanometers,
        )
        modeller.addExtraParticles(force_field)

        system = force_field.createSystem(
            modeller.getTopology(),
            nonbondedMethod=(
                solvent_a_nonbonded_method
                if solvent_index == "solvent-a"
                else solvent_b_nonbonded_method
            ),
            nonbondedCutoff=nonbonded_cutoff,
            constraints=constraints,
            rigidWater=rigid_water,
            removeCMMotion=remove_cmm_motion,
            hydrogenMass=hydrogen_mass,
            switchDistance=switch_distance,
        )

        return system

    return system_generator

def disable_long_range_corrections(system: openmm.System):
    """Disables all long range corrections found in a system in-place.

    This is performed by calling

        * ``NonbondedForce.setUseDispersionCorrection(False)``
        * ``CustomNonbondedForce.setUseLongRangeCorrection(False)``

    on all relevant forces.
    """

    for force in (
        force
        for force in system.getForces()
        if isinstance(force, openmm.NonbondedForce)
    ):
        force.setUseDispersionCorrection(False)
    for force in (
        force
        for force in system.getForces()
        if isinstance(force, openmm.CustomNonbondedForce)
    ):
        force.setUseLongRangeCorrection(False)


# Copyright (c) OpenMMLab. All rights reserved.
import bisect
import collections
import copy
import math
from collections import defaultdict

import numpy as np
from mmcv.utils import build_from_cfg, print_log
from torch.utils.data.dataset import ConcatDataset as _ConcatDataset

from .builder import DATASETS, PIPELINES
from .coco import CocoDataset

@DATASETS.register_module()
class ConcatDataset(_ConcatDataset):
    """A wrapper of concatenated dataset.

    Same as :obj:`torch.utils.data.dataset.ConcatDataset`, but
    concat the group flag for image aspect ratio.

    Args:
        datasets (list[:obj:`Dataset`]): A list of datasets.
        separate_eval (bool): Whether to evaluate the results
            separately if it is used as validation dataset.
            Defaults to True.
    """

    def __init__(self, datasets, separate_eval=True):
        super(ConcatDataset, self).__init__(datasets)
        self.CLASSES = datasets[0].CLASSES
        self.separate_eval = separate_eval
        if not separate_eval:
            if any([isinstance(ds, CocoDataset) for ds in datasets]):
                raise NotImplementedError(
                    'Evaluating concatenated CocoDataset as a whole is not'
                    ' supported! Please set "separate_eval=True"')
            elif len(set([type(ds) for ds in datasets])) != 1:
                raise NotImplementedError(
                    'All the datasets should have same types')

        if hasattr(datasets[0], 'flag'):
            flags = []
            for i in range(0, len(datasets)):
                flags.append(datasets[i].flag)
            self.flag = np.concatenate(flags)

    def get_cat_ids(self, idx):
        """Get category ids of concatenated dataset by index.

        Args:
            idx (int): Index of data.

        Returns:
            list[int]: All categories in the image of specified index.
        """

        if idx < 0:
            if -idx > len(self):
                raise ValueError(
                    'absolute value of index should not exceed dataset length')
            idx = len(self) + idx
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
        return self.datasets[dataset_idx].get_cat_ids(sample_idx)

    def get_ann_info(self, idx):
        """Get annotation of concatenated dataset by index.

        Args:
            idx (int): Index of data.

        Returns:
            dict: Annotation info of specified index.
        """
        if idx < 0:
            if -idx > len(self):
                raise ValueError(
                    'absolute value of index should not exceed dataset length')
            idx = len(self) + idx
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
        return self.datasets[dataset_idx].get_ann_info(sample_idx)

    def evaluate(self, results, logger=None, **kwargs):
        """Evaluate the results.

        Args:
            results (list[list | tuple]): Testing results of the dataset.
            logger (logging.Logger | str | None): Logger used for printing
                related information during evaluation. Default: None.

        Returns:
            dict[str: float]: AP results of the total dataset or each separate
            dataset if `self.separate_eval=True`.
        """
        assert len(results) == self.cumulative_sizes[-1], \
            ('Dataset and results have different sizes: '
             f'{self.cumulative_sizes[-1]} v.s. {len(results)}')

        # Check whether all the datasets support evaluation
        for dataset in self.datasets:
            assert hasattr(dataset, 'evaluate'), \
                f'{type(dataset)} does not implement evaluate function'

        if self.separate_eval:
            dataset_idx = -1
            total_eval_results = dict()
            for size, dataset in zip(self.cumulative_sizes, self.datasets):
                start_idx = 0 if dataset_idx == -1 else \
                    self.cumulative_sizes[dataset_idx]
                end_idx = self.cumulative_sizes[dataset_idx + 1]

                results_per_dataset = results[start_idx:end_idx]
                print_log(
                    f'\nEvaluateing {dataset.ann_file} with '
                    f'{len(results_per_dataset)} images now',
                    logger=logger)

                eval_results_per_dataset = dataset.evaluate(
                    results_per_dataset, logger=logger, **kwargs)
                dataset_idx += 1
                for k, v in eval_results_per_dataset.items():
                    total_eval_results.update({f'{dataset_idx}_{k}': v})

            return total_eval_results
        elif any([isinstance(ds, CocoDataset) for ds in self.datasets]):
            raise NotImplementedError(
                'Evaluating concatenated CocoDataset as a whole is not'
                ' supported! Please set "separate_eval=True"')
        elif len(set([type(ds) for ds in self.datasets])) != 1:
            raise NotImplementedError(
                'All the datasets should have same types')
        else:
            original_data_infos = self.datasets[0].data_infos
            self.datasets[0].data_infos = sum(
                [dataset.data_infos for dataset in self.datasets], [])
            eval_results = self.datasets[0].evaluate(
                results, logger=logger, **kwargs)
            self.datasets[0].data_infos = original_data_infos
            return eval_results

@DATASETS.register_module()
class RepeatDataset:
    """A wrapper of repeated dataset.

    The length of repeated dataset will be `times` larger than the original
    dataset. This is useful when the data loading time is long but the dataset
    is small. Using RepeatDataset can reduce the data loading time between
    epochs.

    Args:
        dataset (:obj:`Dataset`): The dataset to be repeated.
        times (int): Repeat times.
    """

    def __init__(self, dataset, times):
        self.dataset = dataset
        self.times = times
        self.CLASSES = dataset.CLASSES
        if hasattr(self.dataset, 'flag'):
            self.flag = np.tile(self.dataset.flag, times)

        self._ori_len = len(self.dataset)

    def __getitem__(self, idx):
        return self.dataset[idx % self._ori_len]

    def get_cat_ids(self, idx):
        """Get category ids of repeat dataset by index.

        Args:
            idx (int): Index of data.

        Returns:
            list[int]: All categories in the image of specified index.
        """

        return self.dataset.get_cat_ids(idx % self._ori_len)

    def get_ann_info(self, idx):
        """Get annotation of repeat dataset by index.

        Args:
            idx (int): Index of data.

        Returns:
            dict: Annotation info of specified index.
        """

        return self.dataset.get_ann_info(idx % self._ori_len)

    def __len__(self):
        """Length after repetition."""
        return self.times * self._ori_len

# Modified from https://github.com/facebookresearch/detectron2/blob/41d475b75a230221e21d9cac5d69655e3415e3a4/detectron2/data/samplers/distributed_sampler.py#L57 # noqa
@DATASETS.register_module()
class ClassBalancedDataset:
    """A wrapper of repeated dataset with repeat factor.

    Suitable for training on class imbalanced datasets like LVIS. Following
    the sampling strategy in the `paper <https://arxiv.org/abs/1908.03195>`_,
    in each epoch, an image may appear multiple times based on its
    "repeat factor".
    The repeat factor for an image is a function of the frequency the rarest
    category labeled in that image. The "frequency of category c" in [0, 1]
    is defined by the fraction of images in the training set (without repeats)
    in which category c appears.
    The dataset needs to instantiate :func:`self.get_cat_ids` to support
    ClassBalancedDataset.

    The repeat factor is computed as followed.

    1. For each category c, compute the fraction # of images
       that contain it: :math:`f(c)`
    2. For each category c, compute the category-level repeat factor:
       :math:`r(c) = max(1, sqrt(t/f(c)))`
    3. For each image I, compute the image-level repeat factor:
       :math:`r(I) = max_{c in I} r(c)`

    Args:
        dataset (:obj:`CustomDataset`): The dataset to be repeated.
        oversample_thr (float): frequency threshold below which data is
            repeated. For categories with ``f_c >= oversample_thr``, there is
            no oversampling. For categories with ``f_c < oversample_thr``, the
            degree of oversampling following the square-root inverse frequency
            heuristic above.
        filter_empty_gt (bool, optional): If set true, images without bounding
            boxes will not be oversampled. Otherwise, they will be categorized
            as the pure background class and involved into the oversampling.
            Default: True.
    """

    def __init__(self, dataset, oversample_thr, filter_empty_gt=True):
        self.dataset = dataset
        self.oversample_thr = oversample_thr
        self.filter_empty_gt = filter_empty_gt
        self.CLASSES = dataset.CLASSES

        repeat_factors = self._get_repeat_factors(dataset, oversample_thr)
        repeat_indices = []
        for dataset_idx, repeat_factor in enumerate(repeat_factors):
            repeat_indices.extend([dataset_idx] * math.ceil(repeat_factor))
        self.repeat_indices = repeat_indices

        flags = []
        if hasattr(self.dataset, 'flag'):
            for flag, repeat_factor in zip(self.dataset.flag, repeat_factors):
                flags.extend([flag] * int(math.ceil(repeat_factor)))
            assert len(flags) == len(repeat_indices)
        self.flag = np.asarray(flags, dtype=np.uint8)

    def _get_repeat_factors(self, dataset, repeat_thr):
        """Get repeat factor for each images in the dataset.

        Args:
            dataset (:obj:`CustomDataset`): The dataset
            repeat_thr (float): The threshold of frequency. If an image
                contains the categories whose frequency below the threshold,
                it would be repeated.

        Returns:
            list[float]: The repeat factors for each images in the dataset.
        """

        # 1. For each category c, compute the fraction # of images
        #   that contain it: f(c)
        category_freq = defaultdict(int)
        num_images = len(dataset)
        for idx in range(num_images):
            cat_ids = set(self.dataset.get_cat_ids(idx))
            if len(cat_ids) == 0 and not self.filter_empty_gt:
                cat_ids = set([len(self.CLASSES)])
            for cat_id in cat_ids:
                category_freq[cat_id] += 1
        for k, v in category_freq.items():
            category_freq[k] = v / num_images

        # 2. For each category c, compute the category-level repeat factor:
        #    r(c) = max(1, sqrt(t/f(c)))
        category_repeat = {
            cat_id: max(1.0, math.sqrt(repeat_thr / cat_freq))
            for cat_id, cat_freq in category_freq.items()
        }

        # 3. For each image I, compute the image-level repeat factor:
        #    r(I) = max_{c in I} r(c)
        repeat_factors = []
        for idx in range(num_images):
            cat_ids = set(self.dataset.get_cat_ids(idx))
            if len(cat_ids) == 0 and not self.filter_empty_gt:
                cat_ids = set([len(self.CLASSES)])
            repeat_factor = 1
            if len(cat_ids) > 0:
                repeat_factor = max(
                    {category_repeat[cat_id]
                     for cat_id in cat_ids})
            repeat_factors.append(repeat_factor)

        return repeat_factors

    def __getitem__(self, idx):
        ori_index = self.repeat_indices[idx]
        return self.dataset[ori_index]

    def get_ann_info(self, idx):
        """Get annotation of dataset by index.

        Args:
            idx (int): Index of data.

        Returns:
            dict: Annotation info of specified index.
        """
        ori_index = self.repeat_indices[idx]
        return self.dataset.get_ann_info(ori_index)

    def __len__(self):
        """Length after repetition."""
        return len(self.repeat_indices)

@DATASETS.register_module()
class MultiImageMixDataset:
    """A wrapper of multiple images mixed dataset.

    Suitable for training on multiple images mixed data augmentation like
    mosaic and mixup. For the augmentation pipeline of mixed image data,
    the `get_indexes` method needs to be provided to obtain the image
    indexes, and you can set `skip_flags` to change the pipeline running
    process. At the same time, we provide the `dynamic_scale` parameter
    to dynamically change the output image size.

    Args:
        dataset (:obj:`CustomDataset`): The dataset to be mixed.
        pipeline (Sequence[dict]): Sequence of transform object or
            config dict to be composed.
        dynamic_scale (tuple[int], optional): The image scale can be changed
            dynamically. Default to None. It is deprecated.
        skip_type_keys (list[str], optional): Sequence of type string to
            be skip pipeline. Default to None.
    """

    def __init__(self,
                 dataset,
                 pipeline,
                 dynamic_scale=None,
                 skip_type_keys=None):
        if dynamic_scale is not None:
            raise RuntimeError(
                'dynamic_scale is deprecated. Please use Resize pipeline '
                'to achieve similar functions')
        assert isinstance(pipeline, collections.abc.Sequence)
        if skip_type_keys is not None:
            assert all([
                isinstance(skip_type_key, str)
                for skip_type_key in skip_type_keys
            ])
        self._skip_type_keys = skip_type_keys

        self.pipeline = []
        self.pipeline_types = []
        for transform in pipeline:
            if isinstance(transform, dict):
                self.pipeline_types.append(transform['type'])
                transform = build_from_cfg(transform, PIPELINES)
                self.pipeline.append(transform)
            else:
                raise TypeError('pipeline must be a dict')

        self.dataset = dataset
        self.CLASSES = dataset.CLASSES
        if hasattr(self.dataset, 'flag'):
            self.flag = dataset.flag
        self.num_samples = len(dataset)

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        results = copy.deepcopy(self.dataset[idx])
        for (transform, transform_type) in zip(self.pipeline,
                                               self.pipeline_types):
            if self._skip_type_keys is not None and \
                    transform_type in self._skip_type_keys:
                continue

            if hasattr(transform, 'get_indexes'):
                indexes = transform.get_indexes(self.dataset)
                if not isinstance(indexes, collections.abc.Sequence):
                    indexes = [indexes]
                mix_results = [
                    copy.deepcopy(self.dataset[index]) for index in indexes
                ]
                results['mix_results'] = mix_results

            results = transform(results)

            if 'mix_results' in results:
                results.pop('mix_results')

        return results

    def update_skip_type_keys(self, skip_type_keys):
        """Update skip_type_keys. It is called by an external hook.

        Args:
            skip_type_keys (list[str], optional): Sequence of type
                string to be skip pipeline.
        """
        assert all([
            isinstance(skip_type_key, str) for skip_type_key in skip_type_keys
        ])
        self._skip_type_keys = skip_type_keys


import torch
import numpy as np
import numpy.random as rd
from RL.Agent.Algo.Need.Base.AgentBase import AgentBase
from RL.Agent.Algo.Need.Base.net import CriticTwin, ActorSAC

class AgentSAC(AgentBase):  # [ElegantRL.2021.11.11]
    """
    Bases: ``AgentBase``

    Soft Actor-Critic algorithm. “Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor”. Tuomas Haarnoja et al.. 2018.

    :param net_dim[int]: the dimension of networks (the width of neural networks)
    :param state_dim[int]: the dimension of state (the number of state vector)
    :param action_dim[int]: the dimension of action (the number of discrete action)
    :param learning_rate[float]: learning rate of optimizer
    :param if_per_or_gae[bool]: PER (off-policy) or GAE (on-policy) for sparse reward
    :param env_num[int]: the env number of VectorEnv. env_num == 1 means don't use VectorEnv
    :param agent_id[int]: if the visible_gpu is '1,9,3,4', agent_id=1 means (1,9,4,3)[agent_id] == 9
    """

    def __init__(self):
        AgentBase.__init__(self)
        self.ClassCri = CriticTwin
        self.ClassAct = ActorSAC
        self.if_use_cri_target = True
        self.if_use_act_target = False

        self.alpha_log = None
        self.alpha_optim = None
        self.target_entropy = None
        self.obj_critic = (-np.log(0.5)) ** 0.5  # for reliable_lambda

    def init(self, net_dim=256, state_dim=8, action_dim=2, reward_scale=1.0, gamma=0.99,
             learning_rate=1e-4, if_per_or_gae=False, env_num=1, gpu_id=0):
        """
        Explict call ``self.init()`` to overwrite the ``self.object`` in ``__init__()`` for multiprocessing.
        """
        AgentBase.init(self, net_dim=net_dim, state_dim=state_dim, action_dim=action_dim,
                       reward_scale=reward_scale, gamma=gamma,
                       learning_rate=learning_rate, if_per_or_gae=if_per_or_gae,
                       env_num=env_num, gpu_id=gpu_id, )

        self.alpha_log = torch.tensor((-np.log(action_dim) * np.e,), dtype=torch.float32,
                                      requires_grad=True, device=self.device)  # trainable parameter
        self.alpha_optim = torch.optim.Adam((self.alpha_log,), lr=learning_rate)
        self.target_entropy = np.log(action_dim)

        if if_per_or_gae:  # if_use_per
            self.criterion = torch.nn.SmoothL1Loss(reduction='none')
            self.get_obj_critic = self.get_obj_critic_per
        else:
            self.criterion = torch.nn.SmoothL1Loss(reduction='mean')
            self.get_obj_critic = self.get_obj_critic_raw

    def select_actions(self, state: torch.Tensor) -> torch.Tensor:
        """
        Select actions given an array of states.

        :param state: an array of states in a shape (batch_size, state_dim, ).
        :return: an array of actions in a shape (batch_size, action_dim, ) where each action is clipped into range(-1, 1).
        """
        state = state.to(self.device)
        if rd.rand() < self.explore_rate:  # epsilon-greedy
            actions = self.act.get_action(state)
        else:
            actions = self.act(state)
        return actions.detach().cpu()

    def update_net(self, buffer, batch_size, repeat_times, soft_update_tau):
        """
        Update the neural networks by sampling batch data from ``ReplayBuffer``.

        :param buffer: the ReplayBuffer instance that stores the trajectories.
        :param batch_size: the size of batch data for Stochastic Gradient Descent (SGD).
        :param repeat_times: the re-using times of each trajectory.
        :param soft_update_tau: the soft update parameter.
        :return: a tuple of the log information.
        """
        buffer.update_now_len()

        obj_critic = None
        obj_actor = None
        alpha = None
        for _ in range(int(buffer.now_len * repeat_times / batch_size)):
            alpha = self.alpha_log.exp()

            '''objective of critic (loss function of critic)'''
            obj_critic, state = self.get_obj_critic(buffer, batch_size, alpha)
            self.optim_update(self.cri_optim, obj_critic)
            if self.if_use_cri_target:
                self.soft_update(self.cri_target, self.cri, soft_update_tau)

            '''objective of alpha (temperature parameter automatic adjustment)'''
            action_pg, logprob = self.act.get_action_logprob(state)  # policy gradient
            obj_alpha = (self.alpha_log * (logprob - self.target_entropy).detach()).mean()
            self.optim_update(self.alpha_optim, obj_alpha)

            '''objective of actor'''
            obj_actor = -(self.cri(state, action_pg) + logprob * alpha).mean()
            self.optim_update(self.act_optim, obj_actor)
            if self.if_use_act_target:
                self.soft_update(self.act_target, self.act, soft_update_tau)
        return obj_critic, obj_actor.item(), alpha.item()

    def get_obj_critic_raw(self, buffer, batch_size, alpha):
        """
        Calculate the loss of networks with **uniform sampling**.

        :param buffer: the ReplayBuffer instance that stores the trajectories.
        :param batch_size: the size of batch data for Stochastic Gradient Descent (SGD).
        :param alpha: the trade-off coefficient of entropy regularization.
        :return: the loss of the network and states.
        """
        with torch.no_grad():
            reward, mask, action, state, next_s = buffer.sample_batch(batch_size)

            next_a, next_log_prob = self.act_target.get_action_logprob(next_s)  # stochastic policy
            next_q = torch.min(*self.cri_target.get_q1_q2(next_s, next_a))  # twin critics

            q_label = reward + mask * (next_q + next_log_prob * alpha)
        q1, q2 = self.cri.get_q1_q2(state, action)
        obj_critic = (self.criterion(q1, q_label) + self.criterion(q2, q_label)) / 2.
        return obj_critic, state

    def get_obj_critic_per(self, buffer, batch_size, alpha):
        """
        Calculate the loss of the network with **Prioritized Experience Replay (PER)**.

        :param buffer: the ReplayBuffer instance that stores the trajectories.
        :param batch_size: the size of batch data for Stochastic Gradient Descent (SGD).
        :param alpha: the trade-off coefficient of entropy regularization.
        :return: the loss of the network and states.
        """
        with torch.no_grad():
            reward, mask, action, state, next_s, is_weights = buffer.sample_batch(batch_size)

            next_a, next_log_prob = self.act_target.get_action_logprob(next_s)  # stochastic policy
            next_q = torch.min(*self.cri_target.get_q1_q2(next_s, next_a))  # twin critics

            q_label = reward + mask * (next_q + next_log_prob * alpha)
        q1, q2 = self.cri.get_q1_q2(state, action)
        td_error = (self.criterion(q1, q_label) + self.criterion(q2, q_label)) / 2.
        obj_critic = (td_error * is_weights).mean()

        buffer.td_error_update(td_error.detach())
        return obj_critic, state