#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
File BejeweledEnvironment.py created on 15:39 2017/10/19

@author: Yichi Xiao
@version: 1.0
"""

import numpy as np
import cv2
import time
import win32api
import win32gui
from PIL import ImageGrab
import win32con

from itertools import product, tee
from collections import deque

from SpriteConvnetModel import SpriteConvnetModel, tf_flags
from ROISelector import selectROI
from img_utils import img_crop_to_array
from Tagging import Tagging
from Environment import Environment

def interpolate(a, b, t):
    if t <= 0:
        return a
    if t >= 1:
        return b
    return a + (b-a) * t

class BejeweledState():
    TYPE_NUM = 11
    def __init__(self, predictions):
        self.prediction = predictions
        assert len(predictions) == 64
        self.state = self.prediction_to_state_one_hot(predictions)

    @staticmethod
    def prediction_to_state_one_hot(prediction):
        a = np.zeros((8,8,BejeweledState.TYPE_NUM))
        for idx, p in enumerate(prediction):
            i, j = int(idx/8), idx%8
            p = int(p)
            if p <= 8:
                a[i,j,p] = 1
            elif 9 <= p and p <= 15:
                a[i,j,p-8] = a[i,j,9] = 1
            else:
                a[i,j,p-15] = a[i,j,10] = 1
        return a # np.eye(BejeweledState.TYPE_NUM)[prediction].reshape(8,8,-1)

    @staticmethod
    def one_hot_state_to_prediction(s):
        r = [0,1,2,3,4,5,6,7,8,8,15]
        p2 = np.tensordot(r, s, axes=([0], [2]))
        return p2.reshape((64,))

class BejeweledAction():
    def __init__(self):
        self.action_space = list(product([0,1,2,3,4,5,6,7], [0,1,2,3,4,5,6], ['H','V']))
        self.action_space = self.action_space + [(0, 0, 'W')]

    def random_action(self):
        return self.action_space[np.random.randint(len(self.action_space))]

    @staticmethod
    def matrix6():
        mat = np.zeros((6, 6, 13, 2*8*7+1), dtype=np.float32)
        for i in range(6):
            for j in range(6):
                # H => 6: (i+{0,1,2},j+{0,1})
                # V => 6: (j+{0,1,2},i+{0,1})
                for idx, (delta_i, delta_j) in enumerate(((0,0), (0,1), (1,0), (1,1), (2,0), (2,1))):
                    x = i + delta_i
                    y = j + delta_j
                    mat[i, j, idx, (x*7+y)*2] = 1
                    x = j + delta_i
                    y = i + delta_j
                    mat[i, j, idx+6, (x*7+y)*2+1] = 1
                mat[i, j, 12, 2*7*8] = 1
        return mat

class BejeweledEnvironment(Environment):
    def __init__(self, ratio=1.25):
        super(BejeweledEnvironment, self).__init__()
        self.ratio = ratio
        self.SPRITE_RATIO = (0.3305, 0.1175, 0.9572, 0.9088)
        self.hwnd = None
        self.game_rect = None
        self.get_hwnd("Bejeweled 3", "MainWindow")
        self.screen_size = self.get_screen_resolution()
        self.force_front_flag = True
        self.recognize_digit = True
        self.action_space = BejeweledAction().action_space
        self._state_queue = deque()
        for _ in range(4):
            self._state_queue.append(BejeweledState(np.zeros(64)).state)

        # state store
        self.last_image = None
        self.last_state = None
        self.last_score = 0

        self.state_iterator = self.gen_state()
        self.state_iterator.send(None)

        # others
        self.render_timestamp = time.time()

    def get_initial_state(self):
        return next(self.state_iterator)

    def reset(self):
        self.last_image = None
        self.last_state = None
        self.last_score = 0

        self.mouse_click_on_screen(self.game_ratio_to_screen_point((0.16, 0.88)))
        time.sleep(0.1)
        ## robust start
        self.mouse_click_on_screen(self.game_ratio_to_screen_point((0.57, 0.85)))
        self.mouse_click_on_screen(self.game_ratio_to_screen_point((0.57, 0.90)))
        time.sleep(0.7)
        ## robust end
        self.mouse_click_on_screen(self.game_ratio_to_screen_point((0.54, 0.76)))
        time.sleep(1.8)
        self.mouse_click_on_screen(self.game_ratio_to_screen_point((0.20, 0.28)))
        time.sleep(0.5)
        self.mouse_click_on_screen(self.game_ratio_to_screen_point((0.54, 0.44)))
        time.sleep(3)

        # do reset step
        ret = self.get_initial_state()

        return ret

    def step(self, action, wait=0):
        if type(action) != BejeweledAction:
            action = self.action_space[action]
        a, b, c = action
        row1, row2, col1, col2 = 0, 0, 0, 0
        w = False
        if c == 'H':
            row1, row2 = a, a
            col1, col2 = b, b+1
        elif c == 'V':
            col1, col2 = a, a
            row1, row2 = b, b+1
        elif c == 'W':
            w = True
        else:
            print("Invalid Action!")
            return None
        if not w:
            self.mouse_click_on_sprite(row1, col1)
            time.sleep(0.05)
            self.mouse_click_on_sprite(row2, col2)
        else:
            time.sleep(0.5)
        time.sleep(wait)
        # capture next state after wait
        cached_digits = self.last_score
        predictions, digits = next(self.state_iterator)
        self.last_score = digits
        reward = 0
        if cached_digits:
            reward = digits - cached_digits
        if reward < 0:
            reward = 0
        if reward >= 10:
            reward = 10
        #if reward > 0:
        #    print("[GOT REWARD]", int(reward*100))
        if reward == 0:
            reward = -0.1
        # print("Step Action: {}, Reward: {} ({} -> {})".format(action, reward, cached_digits, digits))
        return predictions, reward, False

    def render(self, show=True):
        duration = int(1000*(time.time() - self.render_timestamp))
        zero_img = np.zeros(self.last_image.shape, np.uint8)
        result = Tagging.attach(zero_img, BejeweledState.one_hot_state_to_prediction(self._state_queue[0]))

        cv2.putText(result, '%s ms' % duration, (15, 35), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (100, 100, 100), 3)
        if show:
            cv2.imshow('Sprites', result)
            cv2.moveWindow('Sprites', 0, 0)
            cv2.waitKey(1)
        self.render_timestamp = time.time()
        return result

    ''' Get current state representation:
        1. Grab Screen
        2. Use SpriteConvnetModel to predict
        3. restructure the prediction
    '''
    def gen_state(self):
        #self.prediction_iterator.send(None)
        #self.digit_iterator.send(None)
        model = SpriteConvnetModel(tf_flags(), False, False)
        model.predict_prepare()
        _last_score = 0
        _last_pred = [0] * 64
        while True:
            img = self.get_screen_image()
            sprite_feature = self.get_sprite_roi_data(img)

            digits = self.get_digit(img, min_score=_last_score)
            predictions = model.predict(sprite_feature)

            # if score updates, ignore this frame
            if _last_score < digits:
                #print("[SCORE]{} -> {}".format(int(_last_score*100), int(digits*100)))
                _last_score = digits
                _last_pred = predictions
                continue
            else:
                pass

            # if prediction updates, ignore this frame
            if np.count_nonzero(np.array(predictions-_last_pred)) > 0:
                _last_pred = predictions
                continue
            else:
                pass

            predictions = BejeweledState(predictions).state # Package the predictions

            self._state_queue.appendleft(predictions)
            if len(self._state_queue) > 4:
                self._state_queue.pop()
            self.last_state = predictions
            yield predictions, digits

    def get_screen_image(self):
        # ts = time.time()
        img = self.grab_screen(delay=0.01, force_front=self.force_front_flag)
        # d = int((time.time() - ts) * 1000)
        # print("Grab time:", d, 'ms.')
        return img

    def get_sprite_roi_data(self, image):
        # ts = time.time()
        img = selectROI(image, ratio=self.SPRITE_RATIO)
        self.last_image = img
        data = img_crop_to_array(img)
        # d = int((time.time() - ts) * 1000)
        # print("Predict time:", d, 'ms.')
        return data

    def get_digit(self, image, min_score):
        if self.recognize_digit:
            import pyocr
            import pyocr.tesseract as tess
            import pyocr.builders
            from PIL import Image

            digit_ratio = (0.1016, 0.1836, 0.2228, 0.2268)
            digits = selectROI(image, ratio=digit_ratio, round8=False)
            bw_img = digits
            cv2.imwrite('digit_sample.jpg', bw_img)
            txt = tess.image_to_string(
                Image.fromarray(bw_img),
                lang='eng',
                builder=pyocr.builders.TextBuilder()
            )
            txt = txt.replace(',','').replace('.','')
            if txt.isdigit() and int(txt) >= min_score: # reward should not decrease
                score = int(txt) / 100.0
            else:
                score = min_score
            # print("TXT:", txt, 'score:', score , 'last_score:', last_score)
            # scale score

            return score
        else:
            return 0

    def get_screen_resolution(self):
        width = win32api.GetSystemMetrics(0)
        height = win32api.GetSystemMetrics(1)
        return width, height

    def get_hwnd(self, caption, clazz):
        self.hwnd = win32gui.FindWindow(clazz, caption)
        if not self.hwnd:
            print('window not found!')
        return self.hwnd

    def grab_screen(self, delay=0.25, force_front=False):
        try:
            if force_front:
                win32gui.ShowWindow(self.hwnd, win32con.SW_RESTORE)  # 强行显示界面后才好截图
                win32gui.SetForegroundWindow(self.hwnd)  # 将窗口提到最前
                self.force_front_flag = False
        except:
            return None
        time.sleep(delay)
        #  裁剪得到全图
        game_rect = win32gui.GetWindowRect(self.hwnd)
        game_rect = tuple(int(self.ratio * x) for x in game_rect)
        self.game_rect = game_rect
        pil_image = ImageGrab.grab(game_rect).convert('RGB')
        image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_BGR2RGB)
        return image

    def get_cursor_ratio(self):
        cursor_info = win32gui.GetCursorInfo()[2]
        c = (cursor_info[0]*1.25, cursor_info[1]*1.25)
        gr = self.game_rect
        r1 = 1.0 * (c[0] - gr[0]) / (gr[2] - gr[0])
        r2 = 1.0 * (c[1] - gr[1]) / (gr[3] - gr[1])
        return r1, r2

    ''' input: (ratio_w, ratio_h)
        pos_w = game_rect_ws + game_rect_width * ratio_w
        pos_h = game_rect_hs + game_rect_height* ratio_h
    '''
    def game_ratio_to_screen_point(self, cur_ratio):
        gr = self.game_rect
        pos_w = gr[0] * (1-cur_ratio[0]) + gr[2] * cur_ratio[0]
        pos_h = gr[1] * (1-cur_ratio[1]) + gr[3] * cur_ratio[1]
        p1 = int(pos_w/self.ratio)
        p2 = int(pos_h/self.ratio)
        return p1, p2

    def mouse_click_on_sprite(self, row, col):
        r1 = interpolate(self.SPRITE_RATIO[0], self.SPRITE_RATIO[2], (col+0.5)/8)
        r2 = interpolate(self.SPRITE_RATIO[1], self.SPRITE_RATIO[3], (row+0.5)/8)
        pos = self.game_ratio_to_screen_point((r1, r2))
        self.mouse_click_on_screen(pos)

    def mouse_click_on_screen(self, pos):
        win32api.SetCursorPos(pos)
        win32api.mouse_event(win32con.MOUSEEVENTF_LEFTDOWN | win32con.MOUSEEVENTF_LEFTUP, 0, 0, 0, 0)

    def send_esc(self):
        win32api.keybd_event(27, 0, 0, 0)
        win32api.keybd_event(27, 0, win32con.KEYEVENTF_KEYUP, 0)


import numpy as np
import warnings

def wavelet_lowpass(img_in, cpd_cuttoff, stim_cpd, wavelet='Gaussian'):
    '''
    Lowpass filter an image at a given cpd cuttoff, using wavelet representation for a given cpd of the img stimulus

    Args:
        img_in (2d numpy array):   stimluius img
        cpd_cuttoff (float):  maximum CPD value present in output img
        stim_cpd (float):    CPD of stimulus (should be larger than cpd_cuttoff)

    Returns:
        stim (2d numpy array):   stimlulus image fourier filtered and no frequencies higher than cpd_cuttoff
    '''

    warnings.warn('Wavelets Not Yet Implemented! Returning Unfiltered Image.')
    return(img_in)


# Code in this file is copied and adapted from
# https://github.com/openai/evolution-strategies-starter and from
# https://github.com/modestyachts/ARS

from collections import namedtuple
import logging
import numpy as np
import random
import time
from typing import Optional

import ray
from ray.rllib.agents import Trainer, TrainerConfig
from ray.rllib.agents.ars.ars_tf_policy import ARSTFPolicy
from ray.rllib.agents.es import optimizers, utils
from ray.rllib.agents.es.es_tf_policy import rollout
from ray.rllib.env.env_context import EnvContext
from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
from ray.rllib.utils import FilterManager
from ray.rllib.utils.annotations import override
from ray.rllib.utils.deprecation import Deprecated
from ray.rllib.utils.metrics import (
    NUM_AGENT_STEPS_SAMPLED,
    NUM_AGENT_STEPS_TRAINED,
    NUM_ENV_STEPS_SAMPLED,
    NUM_ENV_STEPS_TRAINED,
)
from ray.rllib.utils.torch_utils import set_torch_seed
from ray.rllib.utils.typing import TrainerConfigDict

logger = logging.getLogger(__name__)

Result = namedtuple(
    "Result",
    [
        "noise_indices",
        "noisy_returns",
        "sign_noisy_returns",
        "noisy_lengths",
        "eval_returns",
        "eval_lengths",
    ],
)

class ARSConfig(TrainerConfig):
    """Defines an ARSTrainer configuration class from which an ARSTrainer can be built.

    Example:
        >>> from ray.rllib.agents.ars import ARSConfig
        >>> config = ARSConfig().training(sgd_stepsize=0.02, report_length=20)\
        ...     .resources(num_gpus=0)\
        ...     .rollouts(num_rollout_workers=4)
        >>> print(config.to_dict())
        >>> # Build a Trainer object from the config and run 1 training iteration.
        >>> trainer = config.build(env="CartPole-v1")
        >>> trainer.train()

    Example:
        >>> from ray.rllib.agents.ars import ARSConfig
        >>> from ray import tune
        >>> config = ARSConfig()
        >>> # Print out some default values.
        >>> print(config.action_noise_std)
        >>> # Update the config object.
        >>> config.training(rollouts_used=tune.grid_search([32, 64]), eval_prob=0.5)
        >>> # Set the config object's env.
        >>> config.environment(env="CartPole-v1")
        >>> # Use to_dict() to get the old-style python config dict
        >>> # when running with tune.
        >>> tune.run(
        ...     "ARS",
        ...     stop={"episode_reward_mean": 200},
        ...     config=config.to_dict(),
        ... )
    """

    def __init__(self):
        """Initializes a ARSConfig instance."""
        super().__init__(trainer_class=ARSTrainer)

        # fmt: off
        # __sphinx_doc_begin__

        # ARS specific settings:
        self.action_noise_std = 0.0
        self.noise_stdev = 0.02
        self.num_rollouts = 32
        self.rollouts_used = 32
        self.sgd_stepsize = 0.01
        self.noise_size = 250000000
        self.eval_prob = 0.03
        self.report_length = 10
        self.offset = 0

        # Override some of TrainerConfig's default values with ARS-specific values.
        self.num_workers = 2
        self.observation_filter = "MeanStdFilter"
        # ARS will use Trainer's evaluation WorkerSet (if evaluation_interval > 0).
        # Therefore, we must be careful not to use more than 1 env per eval worker
        # (would break ARSPolicy's compute_single_action method) and to not do
        # obs-filtering.
        self.evaluation_config["num_envs_per_worker"] = 1
        self.evaluation_config["observation_filter"] = "NoFilter"

        # __sphinx_doc_end__
        # fmt: on

    @override(TrainerConfig)
    def training(
        self,
        *,
        action_noise_std: Optional[float] = None,
        noise_stdev: Optional[float] = None,
        num_rollouts: Optional[int] = None,
        rollouts_used: Optional[int] = None,
        sgd_stepsize: Optional[float] = None,
        noise_size: Optional[int] = None,
        eval_prob: Optional[float] = None,
        report_length: Optional[int] = None,
        offset: Optional[int] = None,
        **kwargs,
    ) -> "ARSConfig":
        """Sets the training related configuration.

        Args:
            action_noise_std: Std. deviation to be used when adding (standard normal)
                noise to computed actions. Action noise is only added, if
                `compute_actions` is called with the `add_noise` arg set to True.
            noise_stdev: Std. deviation of parameter noise.
            num_rollouts: Number of perturbs to try.
            rollouts_used: Number of perturbs to keep in gradient estimate.
            sgd_stepsize: SGD step-size used for the Adam optimizer.
            noise_size: Number of rows in the noise table (shared across workers).
                Each row contains a gaussian noise value for each model parameter.
            eval_prob: Probability of evaluating the parameter rewards.
            report_length: How many of the last rewards we average over.
            offset: Value to subtract from the reward (e.g. survival bonus
                from humanoid) during rollouts.

        Returns:
            This updated TrainerConfig object.
        """
        # Pass kwargs onto super's `training()` method.
        super().training(**kwargs)

        if action_noise_std is not None:
            self.action_noise_std = action_noise_std
        if noise_stdev is not None:
            self.noise_stdev = noise_stdev
        if num_rollouts is not None:
            self.num_rollouts = num_rollouts
        if rollouts_used is not None:
            self.rollouts_used = rollouts_used
        if sgd_stepsize is not None:
            self.sgd_stepsize = sgd_stepsize
        if noise_size is not None:
            self.noise_size = noise_size
        if eval_prob is not None:
            self.eval_prob = eval_prob
        if report_length is not None:
            self.report_length = report_length
        if offset is not None:
            self.offset = offset

        return self

@ray.remote
def create_shared_noise(count):
    """Create a large array of noise to be shared by all workers."""
    seed = 123
    noise = np.random.RandomState(seed).randn(count).astype(np.float32)
    return noise

class SharedNoiseTable:
    def __init__(self, noise):
        self.noise = noise
        assert self.noise.dtype == np.float32

    def get(self, i, dim):
        return self.noise[i : i + dim]

    def sample_index(self, dim):
        return np.random.randint(0, len(self.noise) - dim + 1)

    def get_delta(self, dim):
        idx = self.sample_index(dim)
        return idx, self.get(idx, dim)

@ray.remote
class Worker:
    def __init__(self, config, env_creator, noise, worker_index, min_task_runtime=0.2):

        # Set Python random, numpy, env, and torch/tf seeds.
        seed = config.get("seed")
        if seed is not None:
            # Python random module.
            random.seed(seed)
            # Numpy.
            np.random.seed(seed)
            # Torch.
            if config.get("framework") == "torch":
                set_torch_seed(seed)

        self.min_task_runtime = min_task_runtime
        self.config = config
        self.config["single_threaded"] = True
        self.noise = SharedNoiseTable(noise)

        env_context = EnvContext(config["env_config"] or {}, worker_index)
        self.env = env_creator(env_context)
        # Seed the env, if gym.Env.
        if not hasattr(self.env, "seed"):
            logger.info("Env doesn't support env.seed(): {}".format(self.env))
        # Gym.env.
        else:
            self.env.seed(seed)

        from ray.rllib import models

        self.preprocessor = models.ModelCatalog.get_preprocessor(self.env)

        policy_cls = get_policy_class(config)
        self.policy = policy_cls(
            self.env.observation_space, self.env.action_space, config
        )

    @property
    def filters(self):
        return {DEFAULT_POLICY_ID: self.policy.observation_filter}

    def sync_filters(self, new_filters):
        for k in self.filters:
            self.filters[k].sync(new_filters[k])

    def get_filters(self, flush_after=False):
        return_filters = {}
        for k, f in self.filters.items():
            return_filters[k] = f.as_serializable()
            if flush_after:
                f.reset_buffer()
        return return_filters

    def rollout(self, timestep_limit, add_noise=False):
        rollout_rewards, rollout_fragment_length = rollout(
            self.policy,
            self.env,
            timestep_limit=timestep_limit,
            add_noise=add_noise,
            offset=self.config["offset"],
        )
        return rollout_rewards, rollout_fragment_length

    def do_rollouts(self, params, timestep_limit=None):
        # Set the network weights.
        self.policy.set_flat_weights(params)

        noise_indices, returns, sign_returns, lengths = [], [], [], []
        eval_returns, eval_lengths = [], []

        # Perform some rollouts with noise.
        while len(noise_indices) == 0:
            if np.random.uniform() < self.config["eval_prob"]:
                # Do an evaluation run with no perturbation.
                self.policy.set_flat_weights(params)
                rewards, length = self.rollout(timestep_limit, add_noise=False)
                eval_returns.append(rewards.sum())
                eval_lengths.append(length)
            else:
                # Do a regular run with parameter perturbations.
                noise_index = self.noise.sample_index(self.policy.num_params)

                perturbation = self.config["noise_stdev"] * self.noise.get(
                    noise_index, self.policy.num_params
                )

                # These two sampling steps could be done in parallel on
                # different actors letting us update twice as frequently.
                self.policy.set_flat_weights(params + perturbation)
                rewards_pos, lengths_pos = self.rollout(timestep_limit)

                self.policy.set_flat_weights(params - perturbation)
                rewards_neg, lengths_neg = self.rollout(timestep_limit)

                noise_indices.append(noise_index)
                returns.append([rewards_pos.sum(), rewards_neg.sum()])
                sign_returns.append(
                    [np.sign(rewards_pos).sum(), np.sign(rewards_neg).sum()]
                )
                lengths.append([lengths_pos, lengths_neg])

        return Result(
            noise_indices=noise_indices,
            noisy_returns=returns,
            sign_noisy_returns=sign_returns,
            noisy_lengths=lengths,
            eval_returns=eval_returns,
            eval_lengths=eval_lengths,
        )

def get_policy_class(config):
    if config["framework"] == "torch":
        from ray.rllib.agents.ars.ars_torch_policy import ARSTorchPolicy

        policy_cls = ARSTorchPolicy
    else:
        policy_cls = ARSTFPolicy
    return policy_cls

class ARSTrainer(Trainer):
    """Large-scale implementation of Augmented Random Search in Ray."""

    @classmethod
    @override(Trainer)
    def get_default_config(cls) -> TrainerConfigDict:
        return ARSConfig().to_dict()

    @override(Trainer)
    def validate_config(self, config: TrainerConfigDict) -> None:
        # Call super's validation method.
        super().validate_config(config)

        if config["num_gpus"] > 1:
            raise ValueError("`num_gpus` > 1 not yet supported for ARS!")
        if config["num_workers"] <= 0:
            raise ValueError("`num_workers` must be > 0 for ARS!")
        if config["evaluation_config"]["num_envs_per_worker"] != 1:
            raise ValueError(
                "`evaluation_config.num_envs_per_worker` must always be 1 for "
                "ARS! To parallelize evaluation, increase "
                "`evaluation_num_workers` to > 1."
            )
        if config["evaluation_config"]["observation_filter"] != "NoFilter":
            raise ValueError(
                "`evaluation_config.observation_filter` must always be "
                "`NoFilter` for ARS!"
            )

    @override(Trainer)
    def setup(self, config):
        # Setup our config: Merge the user-supplied config (which could
        # be a partial config dict with the class' default).
        if isinstance(config, dict):
            self.config = self.merge_trainer_configs(
                self.get_default_config(), config, self._allow_unknown_configs
            )
        else:
            self.config = config.to_dict()

        # Validate our config dict.
        self.validate_config(self.config)

        # Generate `self.env_creator` callable to create an env instance.
        self.env_creator = self._get_env_creator_from_env_id(self._env_id)
        # Generate the local env.
        env_context = EnvContext(self.config["env_config"] or {}, worker_index=0)
        env = self.env_creator(env_context)

        self.callbacks = self.config["callbacks"]()

        self._policy_class = get_policy_class(self.config)
        self.policy = self._policy_class(
            env.observation_space, env.action_space, self.config
        )
        self.optimizer = optimizers.SGD(self.policy, self.config["sgd_stepsize"])

        self.rollouts_used = self.config["rollouts_used"]
        self.num_rollouts = self.config["num_rollouts"]
        self.report_length = self.config["report_length"]

        # Create the shared noise table.
        logger.info("Creating shared noise table.")
        noise_id = create_shared_noise.remote(self.config["noise_size"])
        self.noise = SharedNoiseTable(ray.get(noise_id))

        # Create the actors.
        logger.info("Creating actors.")
        self.workers = [
            Worker.remote(self.config, self.env_creator, noise_id, idx + 1)
            for idx in range(self.config["num_workers"])
        ]

        self.episodes_so_far = 0
        self.reward_list = []
        self.tstart = time.time()

    @override(Trainer)
    def get_policy(self, policy=DEFAULT_POLICY_ID):
        if policy != DEFAULT_POLICY_ID:
            raise ValueError(
                "ARS has no policy '{}'! Use {} "
                "instead.".format(policy, DEFAULT_POLICY_ID)
            )
        return self.policy

    @override(Trainer)
    def step_attempt(self):
        config = self.config

        theta = self.policy.get_flat_weights()
        assert theta.dtype == np.float32
        assert len(theta.shape) == 1

        # Put the current policy weights in the object store.
        theta_id = ray.put(theta)
        # Use the actors to do rollouts, note that we pass in the ID of the
        # policy weights.
        results, num_episodes, num_timesteps = self._collect_results(
            theta_id, config["num_rollouts"]
        )
        # Update our sample steps counters.
        self._counters[NUM_AGENT_STEPS_SAMPLED] += num_timesteps
        self._counters[NUM_ENV_STEPS_SAMPLED] += num_timesteps

        all_noise_indices = []
        all_training_returns = []
        all_training_lengths = []
        all_eval_returns = []
        all_eval_lengths = []

        # Loop over the results.
        for result in results:
            all_eval_returns += result.eval_returns
            all_eval_lengths += result.eval_lengths

            all_noise_indices += result.noise_indices
            all_training_returns += result.noisy_returns
            all_training_lengths += result.noisy_lengths

        assert len(all_eval_returns) == len(all_eval_lengths)
        assert (
            len(all_noise_indices)
            == len(all_training_returns)
            == len(all_training_lengths)
        )

        self.episodes_so_far += num_episodes

        # Assemble the results.
        eval_returns = np.array(all_eval_returns)
        eval_lengths = np.array(all_eval_lengths)
        noise_indices = np.array(all_noise_indices)
        noisy_returns = np.array(all_training_returns)
        noisy_lengths = np.array(all_training_lengths)

        # keep only the best returns
        # select top performing directions if rollouts_used < num_rollouts
        max_rewards = np.max(noisy_returns, axis=1)
        if self.rollouts_used > self.num_rollouts:
            self.rollouts_used = self.num_rollouts

        percentile = 100 * (1 - (self.rollouts_used / self.num_rollouts))
        idx = np.arange(max_rewards.size)[
            max_rewards >= np.percentile(max_rewards, percentile)
        ]
        noise_idx = noise_indices[idx]
        noisy_returns = noisy_returns[idx, :]

        # Compute and take a step.
        g, count = utils.batched_weighted_sum(
            noisy_returns[:, 0] - noisy_returns[:, 1],
            (self.noise.get(index, self.policy.num_params) for index in noise_idx),
            batch_size=min(500, noisy_returns[:, 0].size),
        )
        g /= noise_idx.size
        # scale the returns by their standard deviation
        if not np.isclose(np.std(noisy_returns), 0.0):
            g /= np.std(noisy_returns)
        assert g.shape == (self.policy.num_params,) and g.dtype == np.float32
        # Compute the new weights theta.
        theta, update_ratio = self.optimizer.update(-g)

        # Update our train steps counters.
        self._counters[NUM_AGENT_STEPS_TRAINED] += num_timesteps
        self._counters[NUM_ENV_STEPS_TRAINED] += num_timesteps

        # Set the new weights in the local copy of the policy.
        self.policy.set_flat_weights(theta)
        # update the reward list
        if len(all_eval_returns) > 0:
            self.reward_list.append(eval_returns.mean())

        # Now sync the filters
        FilterManager.synchronize(
            {DEFAULT_POLICY_ID: self.policy.observation_filter}, self.workers
        )

        info = {
            "weights_norm": np.square(theta).sum(),
            "weights_std": np.std(theta),
            "grad_norm": np.square(g).sum(),
            "update_ratio": update_ratio,
            "episodes_this_iter": noisy_lengths.size,
            "episodes_so_far": self.episodes_so_far,
        }
        result = dict(
            episode_reward_mean=np.mean(self.reward_list[-self.report_length :]),
            episode_len_mean=eval_lengths.mean(),
            timesteps_this_iter=noisy_lengths.sum(),
            info=info,
        )

        return result

    @override(Trainer)
    def cleanup(self):
        # workaround for https://github.com/ray-project/ray/issues/1516
        for w in self.workers:
            w.__ray_terminate__.remote()

    @override(Trainer)
    def compute_single_action(self, observation, *args, **kwargs):
        action, _, _ = self.policy.compute_actions([observation], update=True)
        if kwargs.get("full_fetch"):
            return action[0], [], {}
        return action[0]

    @override(Trainer)
    def _sync_weights_to_workers(self, *, worker_set=None, workers=None):
        # Broadcast the new policy weights to all evaluation workers.
        assert worker_set is not None
        logger.info("Synchronizing weights to evaluation workers.")
        weights = ray.put(self.policy.get_flat_weights())
        worker_set.foreach_policy(lambda p, pid: p.set_flat_weights(ray.get(weights)))

    def _collect_results(self, theta_id, min_episodes):
        num_episodes, num_timesteps = 0, 0
        results = []
        while num_episodes < min_episodes:
            logger.debug(
                "Collected {} episodes {} timesteps so far this iter".format(
                    num_episodes, num_timesteps
                )
            )
            rollout_ids = [
                worker.do_rollouts.remote(theta_id) for worker in self.workers
            ]
            # Get the results of the rollouts.
            for result in ray.get(rollout_ids):
                results.append(result)
                # Update the number of episodes and the number of timesteps
                # keeping in mind that result.noisy_lengths is a list of lists,
                # where the inner lists have length 2.
                num_episodes += sum(len(pair) for pair in result.noisy_lengths)
                num_timesteps += sum(sum(pair) for pair in result.noisy_lengths)

        return results, num_episodes, num_timesteps

    def __getstate__(self):
        return {
            "weights": self.policy.get_flat_weights(),
            "filter": self.policy.observation_filter,
            "episodes_so_far": self.episodes_so_far,
        }

    def __setstate__(self, state):
        self.episodes_so_far = state["episodes_so_far"]
        self.policy.set_flat_weights(state["weights"])
        self.policy.observation_filter = state["filter"]
        FilterManager.synchronize(
            {DEFAULT_POLICY_ID: self.policy.observation_filter}, self.workers
        )

# Deprecated: Use ray.rllib.agents.ars.ARSConfig instead!
class _deprecated_default_config(dict):
    def __init__(self):
        super().__init__(ARSConfig().to_dict())

    @Deprecated(
        old="ray.rllib.agents.ars.ars.DEFAULT_CONFIG",
        new="ray.rllib.agents.ars.ars.ARSConfig(...)",
        error=False,
    )
    def __getitem__(self, item):
        return super().__getitem__(item)

DEFAULT_CONFIG = _deprecated_default_config()


from PIL import Image
import os
import cv2
from torchvision import transforms
import torch
import torch.nn as nn
from torchvision.models.resnet import resnet18, resnet34, resnet50, resnet101, resnet152
import numpy as np
from AI_Model_Loading import net, device, GestureClass,model_name
import torch.nn.functional as F
import time

from SSD300_Model_Loading import nvidia_ssd_processing_utils

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

# SSD 모델
class ResNet(nn.Module):
    def __init__(self, backbone='resnet50', backbone_path=None):
        super().__init__()
        if backbone == 'resnet18':
            backbone = resnet18(pretrained=not backbone_path)
            self.out_channels = [256, 512, 512, 256, 256, 128]
        elif backbone == 'resnet34':
            backbone = resnet34(pretrained=not backbone_path)
            self.out_channels = [256, 512, 512, 256, 256, 256]
        elif backbone == 'resnet50':
            backbone = resnet50(pretrained=not backbone_path)
            self.out_channels = [1024, 512, 512, 256, 256, 256]
        elif backbone == 'resnet101':
            backbone = resnet101(pretrained=not backbone_path)
            self.out_channels = [1024, 512, 512, 256, 256, 256]
        else:  # backbone == 'resnet152':
            backbone = resnet152(pretrained=not backbone_path)
            self.out_channels = [1024, 512, 512, 256, 256, 256]
        if backbone_path:
            backbone.load_state_dict(torch.load(backbone_path))

        self.feature_extractor = nn.Sequential(*list(backbone.children())[:7])

        conv4_block1 = self.feature_extractor[-1][0]

        conv4_block1.conv1.stride = (1, 1)
        conv4_block1.conv2.stride = (1, 1)
        conv4_block1.downsample[0].stride = (1, 1)

    def forward(self, x):
        x = self.feature_extractor(x)
        return x
class SSD300(nn.Module):
    def __init__(self, backbone=ResNet('resnet50')):
        super().__init__()

        self.feature_extractor = backbone

        self.label_num = 81  # number of COCO classes
        self._build_additional_features(self.feature_extractor.out_channels)
        self.num_defaults = [4, 6, 6, 6, 4, 4]
        self.loc = []
        self.conf = []

        for nd, oc in zip(self.num_defaults, self.feature_extractor.out_channels):
            self.loc.append(nn.Conv2d(oc, nd * 4, kernel_size=3, padding=1))
            self.conf.append(nn.Conv2d(oc, nd * self.label_num, kernel_size=3, padding=1))

        self.loc = nn.ModuleList(self.loc)
        self.conf = nn.ModuleList(self.conf)
        self._init_weights()

    def _build_additional_features(self, input_size):
        self.additional_blocks = []
        for i, (input_size, output_size, channels) in enumerate(
                zip(input_size[:-1], input_size[1:], [256, 256, 128, 128, 128])):
            if i < 3:
                layer = nn.Sequential(
                    nn.Conv2d(input_size, channels, kernel_size=1, bias=False),
                    nn.BatchNorm2d(channels),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(channels, output_size, kernel_size=3, padding=1, stride=2, bias=False),
                    nn.BatchNorm2d(output_size),
                    nn.ReLU(inplace=True),
                )
            else:
                layer = nn.Sequential(
                    nn.Conv2d(input_size, channels, kernel_size=1, bias=False),
                    nn.BatchNorm2d(channels),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(channels, output_size, kernel_size=3, bias=False),
                    nn.BatchNorm2d(output_size),
                    nn.ReLU(inplace=True),
                )

            self.additional_blocks.append(layer)

        self.additional_blocks = nn.ModuleList(self.additional_blocks)

    def _init_weights(self):
        layers = [*self.additional_blocks, *self.loc, *self.conf]
        for layer in layers:
            for param in layer.parameters():
                if param.dim() > 1: nn.init.xavier_uniform_(param)

    # Shape the classifier to the view of bboxes
    def bbox_view(self, src, loc, conf):
        ret = []
        for s, l, c in zip(src, loc, conf):
            ret.append((l(s).view(s.size(0), 4, -1), c(s).view(s.size(0), self.label_num, -1)))

        locs, confs = list(zip(*ret))
        locs, confs = torch.cat(locs, 2).contiguous(), torch.cat(confs, 2).contiguous()
        return locs, confs

    def forward(self, x):
        x = self.feature_extractor(x)

        detection_feed = [x]
        for l in self.additional_blocks:
            x = l(x)
            detection_feed.append(x)

        # Feature Map 38x38x4, 19x19x6, 10x10x6, 5x5x6, 3x3x4, 1x1x4
        locs, confs = self.bbox_view(detection_feed, self.loc, self.conf)

        # For SSD 300, shall return nbatch x 8732 x {nlabels, nlocs} results
        return locs, confs
class Loss(nn.Module):
    """
        Implements the loss as the sum of the followings:
        1. Confidence Loss: All labels, with hard negative mining
        2. Localization Loss: Only on positive labels
        Suppose input dboxes has the shape 8732x4
    """

    def __init__(self, dboxes):
        super(Loss, self).__init__()
        self.scale_xy = 1.0 / dboxes.scale_xy
        self.scale_wh = 1.0 / dboxes.scale_wh

        self.sl1_loss = nn.SmoothL1Loss(reduce=False)
        self.dboxes = nn.Parameter(dboxes(order="xywh").transpose(0, 1).unsqueeze(dim=0),
                                   requires_grad=False)
        # Two factor are from following links
        # http://jany.st/post/2017-11-05-single-shot-detector-ssd-from-scratch-in-tensorflow.html
        self.con_loss = nn.CrossEntropyLoss(reduce=False)

    def _loc_vec(self, loc):
        """
            Generate Location Vectors
        """
        gxy = self.scale_xy * (loc[:, :2, :] - self.dboxes[:, :2, :]) / self.dboxes[:, 2:, ]
        gwh = self.scale_wh * (loc[:, 2:, :] / self.dboxes[:, 2:, :]).log()
        return torch.cat((gxy, gwh), dim=1).contiguous()

    def forward(self, ploc, plabel, gloc, glabel):
        """
            ploc, plabel: Nx4x8732, Nxlabel_numx8732
                predicted location and labels
            gloc, glabel: Nx4x8732, Nx8732
                ground truth location and labels
        """
        mask = glabel > 0
        pos_num = mask.sum(dim=1)

        vec_gd = self._loc_vec(gloc)

        # sum on four coordinates, and mask
        sl1 = self.sl1_loss(ploc, vec_gd).sum(dim=1)
        sl1 = (mask.float() * sl1).sum(dim=1)

        # hard negative mining
        con = self.con_loss(plabel, glabel)

        # postive mask will never selected
        con_neg = con.clone()
        con_neg[mask] = 0
        _, con_idx = con_neg.sort(dim=1, descending=True)
        _, con_rank = con_idx.sort(dim=1)

        # number of negative three times positive
        neg_num = torch.clamp(3 * pos_num, max=mask.size(1)).unsqueeze(-1)
        neg_mask = con_rank < neg_num

        # print(con.shape, mask.shape, neg_mask.shape)
        closs = (con * (mask.float() + neg_mask.float())).sum(dim=1)

        # avoid no object detected
        total_loss = sl1 + closs
        num_mask = (pos_num > 0).float()
        pos_num = pos_num.float().clamp(min=1e-6)
        ret = (total_loss * num_mask / pos_num).mean(dim=0)
        return ret

# 텐서화
class ImageTransform():
    def __init__(self):
        self.data_transform = transforms.Compose([
            transforms.Resize((size, size)),  # 리사이즈
            transforms.ToTensor(),
        ])

    def __call__(self, img):
        return self.data_transform(img)

# SSD모델 파라미터 불러오기.
ssd_model = SSD300()
ssd_model.load_state_dict(torch.load('ssd300.pth'))
ssd_model.to('cuda')
ssd_model.eval()

if torch.cuda.is_available():
    print("Cuda mode Online")
else:
    print("****** Warning!!  Cuda mode Offline")

utils = nvidia_ssd_processing_utils()

class_names = open("category_names.txt").readlines()
class_names = [c.strip() for c in class_names]
classes_to_labels = class_names

GestureTypeClass = ['None', 'Right', 'Left']

#Video input
cap = cv2.VideoCapture(0)
cap.set(3, 1920)  # width=1920
cap.set(4, 1080)  # height=1080
#1920 1080이 최대인데, 1080 1080으로 하면 다운사이징이 아니라 centerCrop임.

print(cap.get(cv2.CAP_PROP_FPS))
print(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
print(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

#Video output
#fourcc = cv2.VideoWriter_fourcc(*'DIVX')
#out = cv2.VideoWriter('{} {} AI demo {}frame.avi'.format(Video_name.split('.')[0],model_name,int(cap.get(cv2.CAP_PROP_FPS))), fourcc, int(cap.get(cv2.CAP_PROP_FPS)), (int(cap.get(3)), int(cap.get(4))))
#PoseFrameCount = int(cap.get(cv2.CAP_PROP_FPS))

#Parameters
size = 150

SSD_confidence = 40         #SSD가 몇퍼센트를 기준으로 그 이상의 결과가 나오면 물체로 볼것인지. ex)어떤 물체를 봣는데 Confidence의 %기준으로 인식할지 말지 결정.
No_drawing = False #결과값 그리기.

PoseFrameCount = 8
Frame_list = [None for i in range(PoseFrameCount)]      #8개의 프레임을 기준으로 특정포즈 결과가 일부% 이상을 차지하면 결과 출력.

Image_w ,Image_h = int(cap.get(3)),int(cap.get(4))
frameCount = 0
timer = 1

while cap.isOpened():  # 영상 시작
    success, frame = cap.read()

    if success:
        start = time.time()
        Input_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

        # Input이미지를 1:1비율로 만들기
        longer = int(max(Input_image.size))
        filledImage = Image.new("RGB", (longer, longer))
        filledImage.paste(Input_image)  # 0,0 기준으로 붙임.

        # Todo 디버깅용
        #filledImage.save("Sqaured Image.jpg")

        image = filledImage
        OriginalImage = Image.fromarray(frame)

        # Image -> ndarray
        ndarrayImage = np.array(image)

        print(type(ndarrayImage))
        print(ndarrayImage.shape)

        inputs = [utils.prepare_input(ndarrayImage)]  # 여기서 300 300 3으로 바뀜. 또한 List로 들어가야함. #파일의 이름이 들어감.
        tensor = utils.prepare_tensor(inputs)

        with torch.no_grad():
            detections_batch = ssd_model(tensor)

        try:
            results_per_input = utils.decode_results(detections_batch)
        except:
            continue

        best_results_per_input = [utils.pick_best(results, SSD_confidence / 100) for results in
                                  results_per_input]  # 이부분에서 Confidence 조정.

        for image_idx in range(len(best_results_per_input)):
            filledImage = image

            image_min = min(image.size)

            image = inputs[image_idx] / 2 + 0.5

            img = np.array(OriginalImage)

            bboxes, classes, confidences = best_results_per_input[image_idx]  # 0번배열에 물체의 박스위치, 1번에 물체의 카테고리, 2번에 컨피던스.

            for idx in range(len(bboxes)):
                left, bot, right, top = bboxes[idx]
                local_positions = [left, bot, right, top]

                x, y, w, h = [int(val * max(OriginalImage.size)) for val in [left, bot, right - left, top - bot]]
                # x는 왼쪽위 x점. y는 왼쪽위 y점. h는 높이. w는 넓이.

                for i, v in enumerate(local_positions):
                    if v > 1:
                        local_positions[i] = 1

                local_positions = [v * image_min for v in local_positions]
                # 왼쪽위 구석을 0,0으로 기준으로 두고, 거기서 왼쪽부분은 Left고 오른쪽부분은 right고 위는 bot이고 아래는 top임. 그래서 bot보다 top이 큼.

                if classes[idx] == 1:
                    local_positions = [local_positions[0] - 10, local_positions[1] - 10, local_positions[2] + 10,
                                   local_positions[3] + 10]
                else:
                    local_positions = [local_positions[0], local_positions[1], local_positions[2],
                                       local_positions[3]]

                # 왼쪽, 위 ,오른쪽 아래
                # print(local_positions)

                for i, v in enumerate(local_positions):
                    if v < 0:
                        local_positions[i] = 0

                if (classes[idx] >=2 and classes[idx] <=8) and classes[idx] != 5:
                    if not No_drawing :
                        cv2.rectangle(img, (x, y), (x + w, y + h), (0,0,255), 2)

                        cv2.putText(img, "{}".format(classes_to_labels[classes[idx] - 1]), (x, y - 10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 3)
                        cv2.putText(img, "{:.2f}%".format(confidences[idx] * 100),
                                    (x, y + 20),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 3)

                        cv2.putText(img, "{}".format(classes_to_labels[classes[idx] - 1]), (x, y - 10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (1, 1, 1), 2)
                        cv2.putText(img, "{:.2f}%".format(confidences[idx] * 100),
                                    (x, y + 20),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (1, 1, 1), 2)

                elif classes[idx] == 1:     #사람이라면??
                    if local_positions[3] - local_positions[1] > 180 and confidences[idx] > 0.6:           #영상 크기에 맞게 반응형으로 바꿔야함

                        timer = 0

                        #여기서부터는 인식도 잘되고 어느정도 수신호라고 볼수있음. 수신호 인식 시작.
                        PersonRect = filledImage.crop(
                            (filledImage.size[0] / 2 - image_min / 2, 0, filledImage.size[0] / 2 + image_min / 2,
                             image_min))
                        PersonRectImage = PersonRect.crop((local_positions))

                        # Todo 디버깅용      이미지 저장
                        #PersonRectImage.save('debug/human picture {}.jpg'.format(frameCount))
                        frameCount += 1

                        transform = ImageTransform()

                        GestureInputImage = transform(PersonRectImage)
                        GestureInputImage = GestureInputImage.unsqueeze(0)
                        GestureInputImage = GestureInputImage.to(device)

                        outputs = net(GestureInputImage)

                        preds_percents = F.softmax(outputs, dim=1)

                        highest_chance = torch.max(preds_percents, 1)  # 가장 확률이 높은값
                        location = highest_chance.indices[0].tolist()  # 높은값의 위치.

                        if highest_chance.values[0] >= 0.38:        #38%이상의 정확도가 나와야 유의미한 결과로 결정.

                            # Frame_List에 값 넣어서 평균 확인하기.
                            if None in Frame_list:  # 한자리라도 비어있으면,
                                for i in range(len(Frame_list)):
                                    if Frame_list[i] == None:
                                        Frame_list[i] = GestureClass[location]
                                        break
                            else:  # 꽉 차있으면
                                del Frame_list[0]
                                Frame_list.append(GestureClass[location])

                        if not No_drawing:

                            BGR_white = (255, 255, 255)
                            BGR_Green = (0, 255, 0)

                            BGR_Array = [BGR_white, BGR_white, BGR_white]

                            BGR_Array[location] = BGR_Green

                            if GestureClass[location] == 'N':       #N이면 파란색, L과 R은 초록색
                                cv2.rectangle(img, (x - 10, y - 10), (x + w + 10, y + h + 10), (255, 0, 0),
                                              2)  # (top,left),(bot,right),

                            elif GestureClass[location] == 'L':
                                cv2.rectangle(img, (x - 10, y - 10), (x + w + 10, y + h + 10), (0, 255, 0),
                                              2)  # (top,left),(bot,right),
                            else:
                                cv2.rectangle(img, (x - 10, y - 10), (x + w + 10, y + h + 10), (0, 255, 0),
                                              2)  # (top,left),(bot,right),

                            cv2.putText(img, 'Person {:.2f}%'.format(confidences[idx] * 100), (x, y-15),
                                        cv2.FONT_HERSHEY_SIMPLEX,
                                        0.5, (0,0,0), 3)
                            cv2.putText(img, 'Person {:.2f}%'.format(confidences[idx] * 100), (x, y-15),
                                        cv2.FONT_HERSHEY_SIMPLEX,
                                        0.5, BGR_white, 2)

                            cv2.putText(img, 'Left {:.2f}%'.format(preds_percents.tolist()[0][2] * 100), (x, y + 20),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 3)
                            cv2.putText(img, 'Right {:.2f}%'.format(preds_percents.tolist()[0][1] * 100), (x, y + 40),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 3)
                            cv2.putText(img, 'None {:.2f}%'.format(preds_percents.tolist()[0][0] * 100), (x, y + 60),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 3)

                            cv2.putText(img, 'Left {:.2f}%'.format(preds_percents.tolist()[0][2] * 100), (x, y + 20),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, BGR_Array[2], 2)
                            cv2.putText(img, 'Right {:.2f}%'.format(preds_percents.tolist()[0][1] * 100), (x, y + 40),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, BGR_Array[1], 2)
                            cv2.putText(img, 'None {:.2f}%'.format(preds_percents.tolist()[0][0] * 100), (x, y + 60),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, BGR_Array[0], 2)

                        PersonRectImage = np.array(PersonRectImage)
                        PersonRectImage = cv2.cvtColor(PersonRectImage, cv2.COLOR_BGR2RGB)

                        if local_positions[3] - local_positions[1] > 200:           #영상 크기에 맞게 반응형으로 바꿔야함
                            #여기는 100퍼센트 수신호임. 그래서 결과를 내줘야함..

                            #만약 이 크기가 갑자기 어디서 나타났다고 가정하면, 화살표를 내주지못한다..
                            if Frame_list.count(None) < len(Frame_list)*0.4 and Frame_list.count('N') <len(Frame_list)*0.4:

                                if max(Frame_list, key=Frame_list.count) == 'R':
                                    cv2.arrowedLine(img, (int(Image_w / 2) - 75, 50), (int(Image_w / 2) + 75, 50),
                                                    (255, 255, 255), 9)
                                    cv2.arrowedLine(img, (int(Image_w/2) - 75, 50), (int(Image_w/2)+75, 50), (255, 0, 0), 8)

                                elif max(Frame_list, key=Frame_list.count) == 'L':
                                    cv2.arrowedLine(img, (int(Image_w / 2) + 75, 50), (int(Image_w / 2) - 75, 50),
                                                    (255, 255, 255), 9)
                                    cv2.arrowedLine(img, (int(Image_w/2)+75, 50), (int(Image_w/2) - 75, 50), (255, 0, 0), 8)

                                if not No_drawing:
                                    rows, cols, channels = PersonRectImage.shape
                                    PersonRectImage = cv2.addWeighted(img[0:0 + rows, 0:0 + cols], 0, PersonRectImage, 1, 0)

                                    img[0:0 + rows, 0:0 + cols] = PersonRectImage
                            else:
                                print("Frame_list에 None이 너무 많음. 의사결정 불가")

                    else:       #사람의 크기가 일정크기보다 작다면 그냥 행인.
                        if not No_drawing:
                            cv2.rectangle(img, (x, y), (x + w+10, y + h+10), (0, 255, 255), 2)

                            cv2.putText(img, "{}".format(classes_to_labels[classes[idx] - 1]), (x, y - 10),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 3)
                            cv2.putText(img, "{:.2f}%".format(confidences[idx] * 100),
                                        (x, y + 20),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 3)

                            cv2.putText(img, "{}".format(classes_to_labels[classes[idx] - 1]), (x, y - 10),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (1, 1, 1), 2)
                            cv2.putText(img, "{:.2f}%".format(confidences[idx] * 100),
                                        (x, y + 20),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (1, 1, 1), 2)

                else:       #사람과 차가 아니라면
                    if not No_drawing:
                        cv2.rectangle(img, (x, y), (x + w, y + h), (255, 255, 255), 2)

                        cv2.putText(img, "{}".format(classes_to_labels[classes[idx] - 1]), (x, y - 10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 3)
                        cv2.putText(img, "{:.2f}%".format(confidences[idx] * 100),
                                    (x, y + 20),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 3)

                        cv2.putText(img, "{}".format(classes_to_labels[classes[idx]-1]), (x, y -10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (1, 1, 1), 2)
                        cv2.putText(img, "{:.2f}%".format(confidences[idx] * 100),
                                    (x, y + 20),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (1, 1, 1), 2)
        end = time.time()
        cv2.putText(img, "fps : {}".format(int(1/(end-start))),(int(cap.get(3)-100), 30),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 2)

        #out.write(img)
        cv2.imshow("test",img)
        cv2.waitKey(1)

    else:
        cap.release()
        #out.release()
        cv2.destroyAllWindows()
        print("끝")
        break


#!/usr/bin/python

# This scripts scaffolds prototxt files defining models and training configs
# for various experiments

import sys
import random
import subprocess
import itertools
import argparse
import numpy as np
import os
import stat
sys.path.append('./external/caffe/python/caffe/proto'); import caffe_pb2

from caffe_pb2 import NetParameter, LayerParameter, DataParameter, SolverParameter, ParamSpec, TRAIN, TEST

def add_weight_filler(param):
  param.type = "gaussian"
  param.std = 0.01

def add_bias_filler(param, value=0):
  param.type = 'constant'
  param.value = value

class CreateNet(object):
  ''' Class to scaffold out caption nets in Caffe prototxt '''

  def get_data_layers(self, param, mode='train'):
    ''' Mode can be 'train','test' or 'scst' for self-critical sequence training. '''
    net = NetParameter()

    if mode == 'test':
      data = net.layer.add()
      data.type = "Python"
      data.name = "test_data"
      data.top.append("image_id")
      data.top.append("num_boxes") # (b, 1) - excludes the pooled box in count
      data.top.append("boxes") # (b, max_boxes, 4)
      data.top.append("features") # (b, max_boxes, 2048)
      data.python_param.module = "rcnn_layers"
      data.python_param.layer = "RCNNTestDataLayer"
      data.python_param.param_str = str({
        'batch_size': param['test_batch_size'],
        'feature_sources' : param['test_feature_sources'],
        'pool': True,
        'max_boxes' : param['max_att_features']
      })

    elif mode == 'train':
      data = net.layer.add()
      data.type = "Python"
      data.name = "sequence_data"
      data.top.append("image_id")
      data.top.append("input_sentence")
      data.top.append("target_sentence")
      data.top.append("num_boxes") # (b, 1) - excludes the pooled box in count
      data.top.append("boxes") # (b, max_boxes, 4)
      data.top.append("features") # (b, max_boxes, 2048)
      data.python_param.module = "rcnn_layers"
      data.python_param.layer = "RCNNCaptionTrainDataLayer"
      data.python_param.param_str = str({
        'end_of_sequence': param['end_of_sequence'],
        'ignore_label': param['ignore_label'],
        'sequence_length': param['max_length'],
        'batch_size': param['train_batch_size'],
        'caption_sources' : param['train_caption_sources'],
        'feature_sources' : param['train_feature_sources'],
        'pool': True,
        'max_boxes' : param['max_att_features']
      })

    elif mode == 'scst':
      data = net.layer.add()
      data.type = "Python"
      data.name = "train_data"
      data.top.append("image_id")
      data.top.append("num_boxes") # (b, 1) - excludes the pooled box in count
      data.top.append("boxes") # (b, max_boxes, 4)
      data.top.append("features") # (b, max_boxes, 2048)
      data.python_param.module = "rcnn_layers"
      data.python_param.layer = "RCNNTestDataLayer"
      data.python_param.param_str = str({
        'batch_size': param['train_batch_size'],
        'feature_sources' : param['train_feature_sources'],
        'pool': True,
        'max_boxes' : param['max_att_features']
      })

    else:
      raise ValueError('unknown data layer mode')

    slice_layer = net.layer.add()
    slice_layer.name = "feature_slice_layer"
    slice_layer.type = "Slice"
    slice_layer.slice_param.slice_dim = 1
    slice_layer.slice_param.slice_point.append(1)
    slice_layer.bottom.append("features")
    slice_layer.top.append("avg_pool") # batch_size x 1 x cnn_filters
    slice_layer.top.append("spatial_features") # batch_size x att_features x cnn_filters

    flatten_layer = net.layer.add()
    flatten_layer.name = "context"
    flatten_layer.type = "Flatten"
    flatten_layer.bottom.append("avg_pool") # batch_size 2048
    flatten_layer.top.append(flatten_layer.name) # batch_size x 2048
    flatten_layer.flatten_param.axis = 1

    fc_layer = net.layer.add()
    fc_layer.name = "fc"
    fc_layer.type = "InnerProduct"
    fc_layer.bottom.append("spatial_features") # batch_size x att_features x cnn_filters
    fc_layer.top.append(fc_layer.name) # batch_size x att_features x att_hidden_units
    fc_layer.inner_product_param.num_output = param['att_hidden_units']
    fc_layer.inner_product_param.axis = 2
    fc_layer.inner_product_param.bias_term = False
    add_weight_filler(fc_layer.inner_product_param.weight_filler)

    if mode == 'test' or mode == 'scst':
      # Tile inputs for the number of beams and reshape into batchsize
      tile_layer = net.layer.add()
      tile_layer.name = "tile_context"
      tile_layer.top.append(tile_layer.name)
      tile_layer.bottom.append("context")
      tile_layer.type = "Tile"
      tile_layer.tile_param.tiles = param['test_beam_size']
      tile_layer.tile_param.axis = 1

      reshape_layer = net.layer.add()
      reshape_layer.name = "beam_context"
      reshape_layer.type = "Reshape"
      reshape_layer.bottom.append("tile_context")
      reshape_layer.top.append(reshape_layer.name)
      reshape_layer.reshape_param.shape.dim.append(-1)
      reshape_layer.reshape_param.shape.dim.append(param['cnn_filters'])

      tile_layer = net.layer.add()
      tile_layer.name = "tile_spatial_features"
      tile_layer.top.append(tile_layer.name)
      tile_layer.bottom.append("spatial_features")
      tile_layer.type = "Tile"
      tile_layer.tile_param.tiles = param['test_beam_size']
      tile_layer.tile_param.axis = 1

      reshape_layer = net.layer.add()
      reshape_layer.name = "beam_spatial_features"
      reshape_layer.type = "Reshape"
      reshape_layer.bottom.append(tile_layer.name)
      reshape_layer.top.append(reshape_layer.name)
      reshape_layer.reshape_param.shape.dim.append(-1)
      reshape_layer.reshape_param.shape.dim.append(param['max_att_features'])
      reshape_layer.reshape_param.shape.dim.append(param['cnn_filters'])

      tile_layer = net.layer.add()
      tile_layer.name = "tile_fc"
      tile_layer.top.append(tile_layer.name)
      tile_layer.bottom.append("fc")
      tile_layer.type = "Tile"
      tile_layer.tile_param.tiles = param['test_beam_size']
      tile_layer.tile_param.axis = 1

      reshape_layer = net.layer.add()
      reshape_layer.name = "beam_fc"
      reshape_layer.type = "Reshape"
      reshape_layer.bottom.append(tile_layer.name)
      reshape_layer.top.append(reshape_layer.name)
      reshape_layer.reshape_param.shape.dim.append(-1)
      reshape_layer.reshape_param.shape.dim.append(param['max_att_features'])
      reshape_layer.reshape_param.shape.dim.append(param['att_hidden_units'])

      tile_layer = net.layer.add()
      tile_layer.name = "tile_num_boxes"
      tile_layer.top.append(tile_layer.name)
      tile_layer.bottom.append("num_boxes")
      tile_layer.type = "Tile"
      tile_layer.tile_param.tiles = param['test_beam_size']
      tile_layer.tile_param.axis = 1

      reshape_layer = net.layer.add()
      reshape_layer.name = "beam_num_boxes"
      reshape_layer.type = "Reshape"
      reshape_layer.bottom.append(tile_layer.name)
      reshape_layer.top.append(reshape_layer.name)
      reshape_layer.reshape_param.shape.dim.append(-1)
      reshape_layer.reshape_param.shape.dim.append(1)
    return net

  def get_scst_net(self, param):
    net = NetParameter()

    beam_layer = net.layer.add()
    beam_layer.name = "beam"
    beam_layer.type = "BeamSearch"
    beam_layer.bottom.append("num_boxes")
    beam_layer.bottom.append("spatial_features")
    beam_layer.bottom.append("fc")
    beam_layer.bottom.append("context")
    beam_layer.top.append("caption")
    beam_layer.top.append("log_prob")
    beam_layer.top.append("log_prob_sequence")
    bs = beam_layer.beam_search_param
    bs.beam_size = param['test_beam_size']
    bs.sequence_length = param['max_length']
    bs.end_of_sequence = param['end_of_sequence']
    for word in param['allowed_multiple']:
      bs.allowed_multiple.append(word)
    for i in range(param['num_lstm_stacks']):
      # Previous hidden state
      rc = bs.recurrent_connection.add()
      rc.src = 'lstm%d_hidden0' % i
      rc.dest = 'lstm%d_hidden_prev' % i
      # Previous mem cell
      rc = bs.recurrent_connection.add()
      rc.src = 'lstm%d_mem_cell0' % i
      rc.dest = 'lstm%d_mem_cell_prev' % i
    bs.beam_search_connection.src = 'logp_0'
    bs.beam_search_connection.dest = 'input'
    for pname in ["embed_param", "lstm0_param_0", "lstm0_param_1", "hidden_att_param_0",
          "predict_att_param_0", "lstm1_param_0", "lstm1_param_1", "predict_param_0",
          "predict_param_1"]:
      p = beam_layer.param.add()
      p.name = pname # Share weights

    inner_net = bs.net_param
    input_layer = inner_net.layer.add()
    input_layer.name = "input"
    input_layer.type = "Input"
    input_layer.top.append("num_boxes")
    input_layer.top.append("spatial_features")
    input_layer.top.append("fc")
    input_layer.top.append("context")
    input_layer.top.append(input_layer.name)
    blob_shape = input_layer.input_param.shape.add()
    blob_shape.dim.append(param['train_batch_size'])
    blob_shape.dim.append(1)
    blob_shape = input_layer.input_param.shape.add()
    blob_shape.dim.append(param['train_batch_size'])
    blob_shape.dim.append(param['max_att_features'])
    blob_shape.dim.append(param['cnn_filters'])
    blob_shape = input_layer.input_param.shape.add()
    blob_shape.dim.append(param['train_batch_size'])
    blob_shape.dim.append(param['max_att_features'])
    blob_shape.dim.append(param['att_hidden_units'])
    blob_shape = input_layer.input_param.shape.add()
    blob_shape.dim.append(param['train_batch_size'])
    blob_shape.dim.append(param['cnn_filters'])
    blob_shape = input_layer.input_param.shape.add()
    blob_shape.dim.append(param['train_batch_size'])
    blob_shape.dim.append(1)

    max_length = param['max_length']
    param['max_length'] = 1
    inner_net = self.get_net(param, param['test_batch_size'],
        net=inner_net, mode="scst_decode")

    silence_layer = net.layer.add()
    silence_layer.name = "silence_bs"
    silence_layer.type = "Silence"
    silence_layer.bottom.append("log_prob")
    silence_layer.bottom.append("log_prob_sequence")

    scst_layer = net.layer.add()
    scst_layer.type = "Python"
    scst_layer.name = "scst"
    scst_layer.bottom.append("image_id")
    scst_layer.bottom.append("caption")
    scst_layer.propagate_down.append(False)
    scst_layer.propagate_down.append(False)
    scst_layer.top.append("score_weights")
    scst_layer.top.append("input_sentence")
    scst_layer.top.append("target_sentence")
    scst_layer.top.append("mean_score")
    scst_layer.python_param.module = "scst_layers"
    scst_layer.python_param.layer = "SCSTLayer"
    scst_layer.python_param.param_str = str({
      'vocab_path' : param['data_dir'] + param['vocab_file'],
      'gt_caption_paths' : param['gt_caption_paths'],
      'end_of_sequence' : param['end_of_sequence'],
      'ignore_label' : param['ignore_label']
    })

    # Add rest of training network
    param['max_length'] = max_length
    net = self.get_net(param, param['test_batch_size']*param['test_beam_size'],
        net=net, mode="scst_train")

    return net

  def get_net(self, param, batch_size, loss_weight=None, net=None, mode="train"):
    ''' Mode can be 'train','test','scst_decode' or 'scst_train'. '''

    if not net:
      net = NetParameter()

    if mode=="test":
      dummy_layer = net.layer.add()
      dummy_layer.name = "input"
      dummy_layer.top.append(dummy_layer.name)
      dummy_layer.type = "DummyData"
      filler = dummy_layer.dummy_data_param.data_filler.add()
      # Note that DataLayer uses the end_of_sequence to start sequences,
      # so we will do the same
      filler.value = param['end_of_sequence']
      blob_shape = dummy_layer.dummy_data_param.shape.add()
      blob_shape.dim.append(batch_size)
      blob_shape.dim.append(1)
    if mode == "train" or mode == "scst_train":
      input_slice_layer = net.layer.add()
      input_slice_layer.name = "input_slice"
      input_slice_layer.type = "Slice"
      input_slice_layer.slice_param.slice_dim = 1
      input_slice_layer.bottom.append('input_sentence')
      for i in range(param['max_length']):
        input_slice_layer.top.append("input" if i == 0 else "input_%d" % i)
        if i != 0:
          input_slice_layer.slice_param.slice_point.append(i)

    for i in range(param['max_length']):
      if i == 0:

        for j in range(param['num_lstm_stacks']):
          dummy_layer = net.layer.add()
          dummy_layer.name = 'lstm%d_hidden_prev' % j
          dummy_layer.top.append(dummy_layer.name)
          dummy_layer.type = "DummyData"
          blob_shape = dummy_layer.dummy_data_param.shape.add()
          blob_shape.dim.append(batch_size)
          blob_shape.dim.append(param['lstm_num_cells'])

          dummy_mem_cell = net.layer.add()
          dummy_mem_cell.name = 'lstm%d_mem_cell_prev' %j
          dummy_mem_cell.top.append(dummy_mem_cell.name)
          dummy_mem_cell.type = "DummyData"
          blob_shape = dummy_mem_cell.dummy_data_param.shape.add()
          blob_shape.dim.append(batch_size)
          blob_shape.dim.append(param['lstm_num_cells'])

      for j in range(param['num_lstm_stacks']):

        if j == 0:
          embed_layer = net.layer.add()
          embed_layer.name = "embedding" if i == 0 else "embedding_%d" % i
          embed_layer.type = "Embed"
          embed_layer.bottom.append("input" if i == 0 else "input_%d" % i)
          embed_layer.propagate_down.append(False)
          embed_layer.top.append(embed_layer.name)
          embed_layer.embed_param.bias_term = False
          embed_layer.embed_param.input_dim = param['vocab_size']
          embed_layer.embed_param.num_output = param['lstm_num_cells']
          add_weight_filler(embed_layer.embed_param.weight_filler)
          p = embed_layer.param.add()
          p.name = 'embed_param' # Share weights

        if j != 0:
          # Set up attention mechanism

          inner_product_layer = net.layer.add()
          inner_product_layer.name = "hidden_att_%d" % i
          inner_product_layer.bottom.append(lstm_output_blob) # batch_size x lstm_num_cells
          inner_product_layer.top.append(inner_product_layer.name) # batch_size x att_hidden_units
          inner_product_layer.type = "InnerProduct"
          inner_product_layer.inner_product_param.num_output = param['att_hidden_units']
          inner_product_layer.inner_product_param.bias_term = False
          add_weight_filler(inner_product_layer.inner_product_param.weight_filler)
          p = inner_product_layer.param.add()
          p.name = 'hidden_att_param_0' # Share weights

          tile_layer = net.layer.add()
          tile_layer.name = "tile_hidden_att_%d" % i
          tile_layer.top.append(tile_layer.name) # batch_size x (att_features x att_hidden_units)
          tile_layer.bottom.append(inner_product_layer.name)
          tile_layer.type = "Tile"
          tile_layer.tile_param.axis = 1
          tile_layer.tile_param.tiles = param['max_att_features']

          reshape_layer = net.layer.add()
          reshape_layer.name = "tile_hidden_reshape_%d" % i
          reshape_layer.type = "Reshape"
          reshape_layer.bottom.append(tile_layer.name)
          reshape_layer.top.append(reshape_layer.name)
          reshape_layer.reshape_param.shape.dim.append(0) # Batch size
          reshape_layer.reshape_param.shape.dim.append(-1)
          reshape_layer.reshape_param.shape.dim.append(param['att_hidden_units'])

          sum_layer = net.layer.add()
          sum_layer.name = "sum_hidden_att_%d" % i
          sum_layer.top.append(sum_layer.name) # batch_size x att_features x att_hidden_units
          if mode == 'test' or mode == 'scst_train':
            sum_layer.bottom.append("beam_fc") # batch_size x att_features x att_hidden_units
          else:
            sum_layer.bottom.append("fc") # batch_size x att_features x att_hidden_units
          sum_layer.bottom.append(reshape_layer.name) # batch_size x att_features x att_hidden_units
          sum_layer.type = "Eltwise"
          sum_layer.eltwise_param.operation = sum_layer.eltwise_param.SUM

          tanh_layer = net.layer.add()
          tanh_layer.name = "hidden_tanh_%d" % i
          tanh_layer.top.append(sum_layer.name)
          tanh_layer.bottom.append(sum_layer.name)
          tanh_layer.type = "TanH"

          proj_layer = net.layer.add()
          proj_layer.name = "predict_att_%d" % i
          proj_layer.type = "InnerProduct"
          proj_layer.bottom.append(sum_layer.name) # batch_size x att_features x att_hidden_units
          proj_layer.top.append(proj_layer.name) # batch_size x att_features x 1
          proj_layer.inner_product_param.num_output = 1
          proj_layer.inner_product_param.axis = 2
          proj_layer.inner_product_param.bias_term = False
          add_weight_filler(proj_layer.inner_product_param.weight_filler)
          p = proj_layer.param.add()
          p.name = 'predict_att_param_0' # Share weights

          reshape_layer = net.layer.add()
          reshape_layer.name = "reshape_predict_att_%d" % i
          reshape_layer.type = "Reshape"
          reshape_layer.bottom.append(proj_layer.name)
          reshape_layer.top.append(reshape_layer.name) # batch_size x att_features
          reshape_layer.reshape_param.shape.dim.append(0) # Batch size
          reshape_layer.reshape_param.shape.dim.append(-1)

          softmax_layer = net.layer.add()
          softmax_layer.name = "att_weight_%d" % i
          softmax_layer.type = "Softmax"
          softmax_layer.bottom.append(reshape_layer.name) # batch_size x att_features
          if mode == 'test' or mode == 'scst_train':
            softmax_layer.bottom.append("beam_num_boxes") # batch_size x 1
          else:
            softmax_layer.bottom.append("num_boxes") # batch_size x 1
          softmax_layer.top.append(softmax_layer.name) # batch_size x att_features
          softmax_layer.softmax_param.axis = 1
          softmax_layer.softmax_param.engine = softmax_layer.softmax_param.CAFFE # Not implemented in CUDNN

          scale_layer = net.layer.add()
          scale_layer.name = "att_product_%d" % i
          scale_layer.top.append(scale_layer.name) # batch_size x att_features x input_feature_size
          if mode == 'test' or mode == 'scst_train':
            scale_layer.bottom.append("beam_spatial_features") # batch_size x att_features x input_feature_size
          else:
            scale_layer.bottom.append("spatial_features") # batch_size x att_features x input_feature_size
          scale_layer.bottom.append(softmax_layer.name) # batch_size x att_features
          scale_layer.type = "Scale"
          scale_layer.scale_param.axis = 0

          permute_layer = net.layer.add()
          permute_layer.name = "permute_att_%d" % i
          permute_layer.bottom.append(scale_layer.name) # batch_size x att_features x input_feature_size
          permute_layer.top.append(permute_layer.name) # batch_size x input_feature_size x att_features
          permute_layer.type = "Permute"
          permute_layer.permute_param.order.append(0)
          permute_layer.permute_param.order.append(2)
          permute_layer.permute_param.order.append(1)

          reduction_layer = net.layer.add()
          reduction_layer.name = "fc8_%d" % i
          reduction_layer.type = "Reduction"
          reduction_layer.bottom.append(permute_layer.name) # batch_size x input_feature_size x att_features
          reduction_layer.top.append(reduction_layer.name) # batch_size x input_feature_size
          reduction_layer.reduction_param.axis = 2

        concat_layer = net.layer.add()
        concat_layer.name = 'concat%d_t%d' % (j, i)
        concat_layer.top.append(concat_layer.name)
        concat_layer.type = "Concat"
        # Data input is either the last word or the layer below output
        if j == 0:
          concat_layer.bottom.append(embed_layer.name)
          if mode == 'test' or mode == 'scst_train':
            concat_layer.bottom.append("beam_context") # Add CNN context
          else:
            concat_layer.bottom.append("context") # Add CNN context
          # Add copy down from lstm above
          if i == 0:
            concat_layer.bottom.append('lstm%d_hidden_prev' % (j+1))
          else:
            if mode == 'test' or mode == 'scst_decode':
              concat_layer.bottom.append('lstm%d_hidden_prev%d' % (j+1, i))
            else:
              concat_layer.bottom.append('lstm%d_hidden%d' % (j+1, i-1))
        else:
          concat_layer.bottom.append('lstm%d_hidden%d' % (j - 1, i))
          concat_layer.bottom.append(reduction_layer.name) # Include attended cnn feature as next LSTM input
        # Plus either dummy or the diagonal connection from before
        if i == 0:
          concat_layer.bottom.append('lstm%d_hidden_prev' % (j))
        else:
          if mode == 'test' or mode == 'scst_decode':
            concat_layer.bottom.append('lstm%d_hidden_prev%d' % (j, i))
          else:
            concat_layer.bottom.append('lstm%d_hidden%d' % (j, i - 1))

        lstm_layer = net.layer.add()
        lstm_layer.name = 'lstm%d' % (j+1) if i==0 else 'lstm%d_t%d' % (j+1, i)
        lstm_layer.type = "LSTMNode"
        lstm_layer.lstm_param.num_cells = param['lstm_num_cells']

        add_weight_filler(lstm_layer.lstm_param.input_weight_filler)
        add_weight_filler(lstm_layer.lstm_param.input_gate_weight_filler)
        add_weight_filler(lstm_layer.lstm_param.forget_gate_weight_filler)
        add_weight_filler(lstm_layer.lstm_param.output_gate_weight_filler)

        add_bias_filler(lstm_layer.lstm_param.input_bias_filler)
        add_bias_filler(lstm_layer.lstm_param.input_gate_bias_filler)
        add_bias_filler(lstm_layer.lstm_param.forget_gate_bias_filler, 1)
        add_bias_filler(lstm_layer.lstm_param.output_gate_bias_filler)

        for k in range(2):
          param_spec = lstm_layer.param.add()
          param_spec.name = 'lstm%d_param_%d' % (j, k)
        lstm_output_blob = 'lstm%d_hidden%d' % (j, i)
        lstm_layer.top.append(lstm_output_blob)
        lstm_layer.top.append('lstm%d_mem_cell%d' % (j, i))
        lstm_layer.bottom.append('concat%d_t%d' % (j, i))
        lstm_layer.propagate_down.append(True)
        if i == 0:
          lstm_layer.bottom.append('lstm%d_mem_cell_prev' % j)
          lstm_layer.propagate_down.append(False)
        else:
          if mode == 'test' or mode == 'scst_decode':
            lstm_layer.bottom.append('lstm%d_mem_cell_prev%d' % (j, i))
          else:
            lstm_layer.bottom.append('lstm%d_mem_cell%d' % (j, i - 1))
          lstm_layer.propagate_down.append(True)

        if j == param['num_lstm_stacks']-1:

          inner_product_layer = net.layer.add()
          inner_product_layer.name = "predict" if i == 0 else "predict_%d" % i
          inner_product_layer.top.append(inner_product_layer.name)
          inner_product_layer.bottom.append(lstm_output_blob)
          inner_product_layer.type = "InnerProduct"
          inner_product_layer.inner_product_param.num_output = param['vocab_size']
          inner_product_layer.inner_product_param.axis = 1
          p = inner_product_layer.param.add()
          p.lr_mult = 1
          p.decay_mult = 1
          p.name = 'predict_param_0'
          p = inner_product_layer.param.add()
          p.lr_mult = 2
          p.decay_mult = 0
          p.name = 'predict_param_1'
          add_weight_filler(inner_product_layer.inner_product_param.weight_filler)
          add_bias_filler(inner_product_layer.inner_product_param.bias_filler)

          if mode == 'test' or mode == 'scst_decode':
            softmax_layer = net.layer.add()
            softmax_layer.name = "probs_%d" % i
            softmax_layer.type = "Softmax"
            softmax_layer.bottom.append(inner_product_layer.name)
            softmax_layer.top.append(softmax_layer.name)
            softmax_layer.softmax_param.axis = 1

            log_layer = net.layer.add()
            log_layer.name = "logp_%d" % i
            log_layer.type = "Log"
            log_layer.bottom.append(softmax_layer.name)
            log_layer.top.append(log_layer.name)

          if mode == 'test':
            # Beam search test
            bs_layer = net.layer.add()
            bs_layer.name = "beam_search_%d" % i
            bs_layer.type = "BeamSearchNode"
            if i > 0:
              bs_layer.bottom.append("bs_scores_%d" % (i-1))
              bs_layer.bottom.append("bs_sentence_%d" % (i-1))
            bs_layer.bottom.append(log_layer.name)
            for k in range(param['num_lstm_stacks']):
              bs_layer.bottom.append("lstm%d_hidden%d" % (k,i))
              bs_layer.bottom.append("lstm%d_mem_cell%d" % (k,i))

            bs_layer.top.append("log_prob" if i+1==param['max_length'] else "bs_scores_%d" % i)
            bs_layer.top.append("caption" if i+1==param['max_length'] else "bs_sentence_%d" % i)
            bs_layer.top.append("input_%d" % (i+1))
            for k in range(param['num_lstm_stacks']):
              bs_layer.top.append("lstm%d_hidden_prev%d" % (k,i+1))
              bs_layer.top.append("lstm%d_mem_cell_prev%d" % (k,i+1))

            bs = bs_layer.beam_search_param
            bs.beam_size = param['test_beam_size']
            bs.end_of_sequence = param['end_of_sequence']
            bs.ignore_label = param['ignore_label']
            for word in param['allowed_multiple']:
              bs.allowed_multiple.append(word)

    if mode == 'train' or mode == "scst_train":
      hidden_concat_layer = net.layer.add()
      hidden_concat_layer.type = "Concat"
      hidden_concat_layer.name = 'predict_concat'
      hidden_concat_layer.top.append(hidden_concat_layer.name)
      hidden_concat_layer.concat_param.concat_dim = 1
      for i in range(param['max_length']):
        hidden_concat_layer.bottom.append("predict" if i == 0 else "predict_%d" % i)

      reshape_layer = net.layer.add()
      reshape_layer.name = 'predict_reshape'
      reshape_layer.type = "Reshape"
      reshape_layer.bottom.append('predict_concat')
      reshape_layer.top.append(reshape_layer.name)
      reshape_layer.reshape_param.shape.dim.append(0) # Batch size
      reshape_layer.reshape_param.shape.dim.append(param['max_length'])
      reshape_layer.reshape_param.shape.dim.append(param['vocab_size'])

      word_loss_layer = net.layer.add()
      word_loss_layer.name = "cross_entropy_loss"
      word_loss_layer.type = "SoftmaxWithLoss"
      word_loss_layer.bottom.append("predict_reshape")
      word_loss_layer.bottom.append("target_sentence")
      word_loss_layer.propagate_down.append(True)
      word_loss_layer.propagate_down.append(False)
      if mode == "scst_train":
        word_loss_layer.bottom.append("score_weights")
        word_loss_layer.propagate_down.append(False)
      word_loss_layer.top.append(word_loss_layer.name)
      if loss_weight is None:
        word_loss_layer.loss_weight.append(param['max_length'])
      else:
        word_loss_layer.loss_weight.append(loss_weight)
      word_loss_layer.loss_param.ignore_label = param['ignore_label']
      word_loss_layer.softmax_param.axis = 2

      accuracy_layer = net.layer.add()
      accuracy_layer.name = "accuracy"
      accuracy_layer.type = "Accuracy"
      accuracy_layer.bottom.append("predict_reshape")
      accuracy_layer.bottom.append("target_sentence")
      accuracy_layer.top.append(accuracy_layer.name)
      accuracy_layer.accuracy_param.ignore_label = param['ignore_label']
      accuracy_layer.accuracy_param.axis = 2

    if mode != 'scst_decode':
      silence_layer = net.layer.add()
      silence_layer.name = "silence"
      silence_layer.type = "Silence"
      for j in range(param['num_lstm_stacks']):
        silence_layer.bottom.append("lstm%d_mem_cell%d" % (j, param['max_length'] - 1))
      silence_layer.bottom.append("image_id")
      silence_layer.bottom.append("boxes")
      if mode == 'test':
        silence_layer.bottom.append("input_%d" % param['max_length'])
        for j in range(param['num_lstm_stacks']):
          silence_layer.bottom.append("lstm%d_hidden_prev%d" % (j,param['max_length']))
          silence_layer.bottom.append("lstm%d_mem_cell_prev%d" % (j,param['max_length']))

    return net

  def write_solver(self, param, file_name):
    with open(file_name, 'w') as f:
      f.write(str(param))

  def write_net(self, param, file_name):
    with open(file_name, 'w') as f:
      f.write('name: "%s"\n' % param['net_name'])
      f.write(str(self.get_data_layers(param, mode='train')))
      f.write(str(self.get_net(param, mode='train',
        batch_size = param['train_batch_size'])))

  def write_decoder(self, param, file_name):
    with open(file_name, 'w') as f:
      f.write('name: "%s"\n' % param['net_name'])
      f.write(str(self.get_data_layers(param, mode='test')))
      batch_size = param['test_beam_size']*param['test_batch_size']
      f.write(str(self.get_net(param, batch_size, mode='test')))

  def write_scst_net(self, param, file_name):
    with open(file_name, 'w') as f:
      f.write('name: "%s"\n' % param['net_name'])
      param['train_batch_size'] = param['test_batch_size']
      f.write(str(self.get_data_layers(param, mode='scst')))
      f.write(str(self.get_scst_net(param)))

  def write_train_script(self, param):
    file_name = '%s%s/train.sh' % (param['base_dir'],param['net_name'])
    with open(file_name, 'w') as f:
      f.write('''#!/bin/bash

GPU_ID=%s
BASE_DIR=%s
DATA_DIR=%s
LOG_DIR=%s
SNAPSHOT_DIR=%s
NET_NAME=%s
OUT_DIR=%s
VOCAB_FILE=%s
MAX_IT=%s
SCST_MAX_IT=%s

mkdir -p ${LOG_DIR}${NET_NAME}
mkdir -p ${SNAPSHOT_DIR}${NET_NAME}
mkdir -p ${OUT_DIR}${NET_NAME}

python -u ./scripts/train.py \
    --solver ${BASE_DIR}${NET_NAME}/solver.prototxt \
    --gpus ${GPU_ID//,/ } \
    > ${LOG_DIR}${NET_NAME}/solver.log 2<&1

# Decode the cross entropy trained model
python ./scripts/beam_decode.py   --gpu ${GPU_ID:0:1}   \
    --model ${BASE_DIR}${NET_NAME}/decoder.prototxt \
    --weights=${SNAPSHOT_DIR}${NET_NAME}/lstm_iter_${MAX_IT}.caffemodel.h5 \
    --vocab ${DATA_DIR}${VOCAB_FILE} \
    --outfile ${OUT_DIR}/${NET_NAME}/iter_${MAX_IT}.json

# Self-critical sequence training
python -u ./scripts/train.py \
    --solver ${BASE_DIR}${NET_NAME}/scst_solver.prototxt \
    --gpus ${GPU_ID//,/ } \
    --weights=${SNAPSHOT_DIR}${NET_NAME}/lstm_iter_${MAX_IT}.caffemodel.h5 \
    > ${LOG_DIR}${NET_NAME}/scst.log 2<&1

# Decode the finished model
python ./scripts/beam_decode.py   --gpu ${GPU_ID:0:1}   \
    --model ${BASE_DIR}${NET_NAME}/decoder.prototxt \
    --weights=${SNAPSHOT_DIR}${NET_NAME}/lstm_scst_iter_${SCST_MAX_IT}.caffemodel.h5 \
    --vocab ${DATA_DIR}${VOCAB_FILE} \
    --outfile ${OUT_DIR}/${NET_NAME}/scst_iter_${SCST_MAX_IT}.json

''' % (param['gpu_ids'], param['base_dir'], param['data_dir'], param['log_dir'],
        param['snapshot_dir'], param['net_name'], param['out_dir'],
        param['vocab_file'], param['max_iter'], param['scst_max_iter']))
    st = os.stat(file_name) # Make executable
    os.chmod(file_name, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)

  def get_solver(self, param):
    solver = SolverParameter()
    solver.net = param['base_dir'] + param['net_name'] + "/net.prototxt"

    solver.base_lr = 0.01
    solver.weight_decay = 0.0005
    solver.lr_policy = "poly"
    solver.power = 1
    solver.momentum = 0.9
    solver.type = "SGD"
    solver.clip_gradients = 10

    solver.display = 100
    solver.max_iter = param['max_iter']
    solver.average_loss = 100
    solver.snapshot = param['solver_snapshot_interval']
    solver.snapshot_prefix = param['snapshot_dir']+param['net_name']+"/lstm"
    solver.snapshot_format = solver.HDF5
    solver.random_seed = param['random_seed']
    solver.iter_size = 1
    solver.layer_wise_reduce = False
    return solver

  def get_scst_solver(self, param):
    solver = self.get_solver(param)
    solver.base_lr = 0.001
    solver.lr_policy = "exp"
    solver.iter_size = 2
    solver.gamma = 0.99975
    solver.snapshot = 500
    solver.max_iter = param['scst_max_iter']
    solver.net = param['base_dir'] + param['net_name'] + "/scst_net.prototxt"
    solver.snapshot_prefix = param['snapshot_dir']+param['net_name']+"/lstm_scst"
    return solver

  def get_base_param(self, test_submission=False):
    param = {}
    param['net_name'] = 'caption_lstm'
    param['random_seed'] = 1701
    param['train_caption_sources'] = ['data/coco_splits/train_captions.txt']
    param['train_feature_sources'] = ['data/tsv/trainval/karpathy_train_resnet101_faster_rcnn_genome.tsv.%d' % i for i in range(2)]
    if test_submission:
      # Test on coco test2014
      param['test_feature_sources'] = ['data/tsv/test2014/test2014_resnet101_faster_rcnn_genome.tsv']
    else:
      # Test on karpathy test
      param['test_feature_sources'] = ['data/tsv/trainval/karpathy_test_resnet101_faster_rcnn_genome.tsv']
    param['max_iter'] = 60000
    param['scst_max_iter'] = 1000
    param['max_length'] = 20
    param['gt_caption_paths'] = ['data/coco/captions_train2014.json',
                                 'data/coco/captions_val2014.json'] # For SCST training
    param['vocab_file']= 'trainval_vocab.txt' if test_submission else 'train_vocab.txt'
    param['vocab_size'] = 10387 if test_submission else 10010
    # Only typical stopwords are allowed to be repeated in captions:
    # a the of , on in with and an " to some at are it that ' by
    param['allowed_multiple'] = [2,5,4,15,3,6,8,7,9,13,277,11,30,16,19,27,25,119,48]
    param['end_of_sequence'] = 0 # This is the period token in the vocab file
    param['ignore_label'] = 1 # This is the unknown word token in vocab file
    param['train_batch_size'] = 50 # Remember, this is per gpu
    param['test_batch_size'] = 12
    param['test_beam_size'] = 5

    param['cnn_filters'] = 2048
    param['num_lstm_stacks'] = 2
    param['lstm_num_cells'] = 1000
    param['max_att_features'] = 100
    param['att_hidden_units'] = 512

    param['gpu_ids'] = '0'
    param['base_dir'] = 'experiments/'
    param['data_dir'] = 'data/coco_splits/'
    param['snapshot_dir'] = 'snapshots/'
    param['log_dir'] = 'logs/'
    param['out_dir'] = 'outputs/'
    param['solver_snapshot_interval'] = 5000

    return param

def make_dir(path):
  if not os.path.exists(path):
    os.makedirs(path)

def main(builder):

  param = builder.get_base_param()
  make_dir(param['base_dir']+param['net_name'])
  print 'Creating net: %s' % param['net_name']
  builder.write_train_script(param)

  solver_param = builder.get_solver(param)
  file_name = '%s%s/solver.prototxt' % (param['base_dir'], param['net_name'])
  builder.write_solver(solver_param, file_name)

  solver_param = builder.get_scst_solver(param)
  file_name = '%s%s/scst_solver.prototxt' % (param['base_dir'], param['net_name'])
  builder.write_solver(solver_param, file_name)

  file_name = '%s%s/net.prototxt' % (param['base_dir'],param['net_name'])
  builder.write_net(param, file_name)

  file_name = '%s%s/decoder.prototxt' % (param['base_dir'],param['net_name'])
  builder.write_decoder(param, file_name)

  file_name = '%s%s/scst_net.prototxt' % (param['base_dir'],param['net_name'])
  builder.write_scst_net(param, file_name)

if __name__ == '__main__':
  builder = CreateNet()
  main(builder)