#!/usr/bin/env python
# -*-coding:utf-8 -*-
'''
@author: daishilong
@contact: daishilong1236@gmail.com
'''

import numpy as np
import matplotlib.pyplot as plt
import gym
from gym import spaces
from crl_kino.env.differential_env import DifferentialDriveEnv, normalize_angle
import itertools
import transforms3d.euler as euler
import pickle
from crl_kino.utils.draw import *
from tianshou.data import Batch

class DifferentialDriveGym(gym.Env):
    """
    The gym wrapper for DifferentailDriveEnv
    """

    def __init__(self,
                 robot_env=DifferentialDriveEnv(1.0, -0.1, np.pi, 1.0, np.pi),
                 reward_param=np.array(
                     [50.0, -1.5, -0.5, -300.0, 1.0, 1.0, -.5, -2.5]),
                 obs_list_list=[]):
        """
        :param robot_env: simulation environment
        """
        super(DifferentialDriveGym, self).__init__()
        #
        self.robot_env = robot_env
        self.state_bounds = self.robot_env.get_bounds()['state_bounds'].copy()

        self.state = np.zeros(len(self.state_bounds))
        self.goal = np.zeros(len(self.state))

        # actual range of v, w
        self.v_space = spaces.Box(
            low=self.state_bounds[-2:, 0], high=self.state_bounds[-2:, 1])
        # output of network, need to be scaled to v_space
        self.action_space = self.v_space

        self.sample_positions = self._init_sample_positions()
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(
            len(self.sample_positions)+4,))

        self.n_step = 0
        self.max_step = 200
        self.current_time = 0.0
        self.max_time = 100.0

        self.obs_list_list = obs_list_list
        self.n_case = 0

        self.reward_param = reward_param.copy()
        self.curriculum = {'obs_num': 7, 'ori': False}
        # self.reset()

    def _init_sample_positions(self):
        """
        the robot sense the local environment by sampling the local map.
        sample_positions contain the positions need to be sampled in the body coordinate system.
        """
        percept_region = np.array([-4, 4, -2, 6])  # left, right, behind, ahead
        sample_reso = 0.3
        lr = np.arange(percept_region[0], percept_region[1], sample_reso)
        ba = np.arange(percept_region[2], percept_region[3], sample_reso)
        sample_positions = list(itertools.product(ba, lr))
        self.percept_region = percept_region
        self.percept_shape = (1, len(ba), len(lr))  # (channel, width, height)
        self.sample_reso = sample_reso
        return np.array(sample_positions)

    def sample_local_map(self, wTb):
        """
        sampling points of local map, with value representing occupancy (0 is non-free, 255 is free)
        wTb: world transform robot
        """
        # world sample position
        tmp_wPos = self.T_transform2d(wTb, self.sample_positions)
        local_map = np.zeros(len(self.sample_positions))
        local_map[:] = self.robot_env.valid_point_check(tmp_wPos)
        # for ind, pos in enumerate(tmp_wPos):
        #     if self.robot_env.valid_point_check(pos):
        #         local_map[ind] = 255
        #     else:
        #         local_map[ind] = 0
        return np.array(local_map)

    def plot_observation(self):
        wRb = euler.euler2mat(0, 0, self.state[2])[:2, :2]
        wTb = np.block([[wRb, self.state[:2].reshape((-1, 1))],
                        [np.zeros((1, 2)), 1]])
        tmp_wPos = self.T_transform2d(wTb, self.sample_positions)
        local_map = self.sample_local_map(wTb)
        ind_non_free = local_map == 0
        ind_free = local_map > 0

        # plot boundary
        left_bot = [self.percept_region[2], self.percept_region[0]]
        left_top = [self.percept_region[2], self.percept_region[1]]
        rigit_top = [self.percept_region[3], self.percept_region[1]]
        rigit_bot = [self.percept_region[3], self.percept_region[0]]
        boundry = np.array([left_bot, left_top, rigit_top,
                            rigit_bot, left_bot], dtype=np.float32)
        boundry = self.T_transform2d(wTb, boundry)
        plt.plot(boundry[:, 0], boundry[:, 1], c='cyan')
        plt.plot(tmp_wPos[ind_free, 0], tmp_wPos[ind_free, 1],
                 '.', c='g', markersize=1)
        plt.plot(tmp_wPos[ind_non_free, 0],
                 tmp_wPos[ind_non_free, 1], '.', c='black', markersize=1)

    def a2v(self, action):
        '''
        action space to v space
        '''
        v = np.zeros(self.action_space.shape)
        for i in range(len(action)):
            v[i] = (action[i]-self.action_space.low[i])/(self.action_space.high[i] -
                                                         self.action_space.low[i]) * (self.v_space.high[i]-self.v_space.low[i]) + self.v_space.low[i]
        return v

    def v2a(self, v):
        '''
        v space to action space
        '''
        action = np.zeros(self.action_space.shape)
        for i in range(len(v)):
            action[i] = (v[i]-self.v_space.low[i])/(self.v_space.high[i]-self.v_space.low[i]) * (
                self.action_space.high[i]-self.action_space.low[i]) + self.action_space.low[i]
        return action

    def step(self, action):
        v = action
        dt = 1.0/5.0  # 5 Hz
        self.state = self.robot_env.motion_velocity(self.state, v, dt)
        self.current_time += dt
        obs = self._obs()
        info = self._info()
        reward = self._reward(info)
        done = info['goal'] or info['collision'] or self.current_time > self.max_time
        return obs, reward, done, info

    def _info(self):
        info = {'goal': False,
                'goal_dis': 0.0,
                'heading': 0.0,
                'collision': False,
                'clearance': 0.0,
                'v': self.state[3],
                'w': abs(self.state[4]),
                'step': 1.0
                }
        info['goal_dis'] = np.linalg.norm(self.state[:2]-self.goal[:2])
        if info['goal_dis'] <= 1.0:
            info['goal'] = True
        info['heading'] = np.abs(normalize_angle(np.arctan2(
            self.goal[1]-self.state[1], self.goal[0]-self.state[0]) - self.state[2]))
        info['clearance'] = min(self.robot_env.get_clearance(self.state), 1.0)
        info['collision'] = not self.robot_env.valid_state_check(self.state)
        return info

    def _reward(self, info):
        # clearance_reward = 0.5*info['clearance'] if info['clearance']>0.3 else -1.0
        # reward = 50*info['goal']\
        #     -0.5*(info['goal_dis'])\
        #         -0.3*info['heading']*(3**(info['goal_dis']<=2.0))\
        #         -100*info['collision']\
        #         +0.5*(info['clearance']>=0.5)\
        #         +0.1*info['v']\
        #         -0.1*np.tanh(info['w'])\
        #         -0.6
        info_arr = np.array([info[a] for a in info])
        reward = info_arr @ self.reward_param
        return reward

    def _obs(self):
        wRb = euler.euler2mat(0, 0, self.state[2])[:2, :2]
        wTb = np.block([[wRb, self.state[:2].reshape((-1, 1))],
                        [np.zeros((1, 2)), 1]])

        local_map = self.sample_local_map(wTb)
        bTw = np.linalg.inv(wTb)
        b_goal_pos = self.T_transform2d(bTw, self.goal[:2])
        # goal position in the robot coordinate frame, robot velocity, local map
        obs = np.block([b_goal_pos, self.state[-2:], local_map])
        return obs

    @staticmethod
    def T_transform2d(aTb, bP):
        """
        aTb: transform with rotation, translation
        bP: non-holonomic coordinates in b
        return: non-holonomic coordinates in a
        """
        if len(bP.shape) == 1:  # vector
            bP_ = np.concatenate((bP, np.ones(1)))
            aP_ = aTb @ bP_
            aP = aP_[:2]
        elif bP.shape[1] == 2:
            bP = bP.T
            bP_ = np.vstack((bP, np.ones((1, bP.shape[1]))))
            aP_ = aTb @ bP_
            aP = aP_[:2, :].T
        return aP

    def set_curriculum(self, **kwargs):
        for k, v in kwargs.items():
            if self.curriculum[k] is not None:
                self.curriculum[k] = v

    def reset(self, low=5, high=8):
        assert len(self.obs_list_list)>0, 'No defined training environments'
        ind_obs = np.random.randint(0, len(self.obs_list_list))
        self.robot_env.set_obs(self.obs_list_list[ind_obs])

        # sample a random start goal configuration
        start = np.zeros(len(self.state_bounds))
        goal = np.zeros(len(self.state_bounds))
        while True:
            # random sample start and goal configuration

            # sample a valid state
            start[:] = np.random.uniform(
                self.state_bounds[:, 0], self.state_bounds[:, 1])
            if self.robot_env.get_clearance(start) <= 0.5:
                continue

            # sample a valid goal
            for _ in range(5):
                r = np.random.uniform(low, high)
                theta = np.random.uniform(-np.pi, np.pi)
                goal[0] = np.clip(start[0] + r*np.cos(theta), *self.state_bounds[0,:])
                goal[1] = np.clip(start[1] + r*np.sin(theta), *self.state_bounds[1,:])
                if self.robot_env.get_clearance(goal) > 0.5:
                    break

            # start point to goal
            if self.curriculum['ori']:
                start[2] = normalize_angle(np.arctan2(
                    goal[1]-start[1], goal[0]-start[0]))
            if self.robot_env.get_clearance(start) > 0.5 and self.robot_env.get_clearance(goal) > 0.5 and low < np.linalg.norm(start[:2]-goal[:2]) < high:
                break

        self.state = start
        self.goal = goal

        self.current_time = 0
        obs = self._obs()
        return obs

    def render(self, mode='human', plot_localwindow=True, pause=True):
        ax = plt.gca()
        ax.cla()  # clear things

        plot_obs_list(ax, self.robot_env.obs_list)
        if plot_localwindow:
            self.plot_observation()
        plot_problem_definition(ax, self.robot_env.obs_list, self.robot_env.rigid_robot, self.state, self.goal)
        plot_robot(ax, self.robot_env.rigid_robot, self.state[:3])

        plt.axis('equal')
        # plt.ylim(-20.0, 20.0)
        plt.tight_layout()
        if pause:
            plt.pause(0.0001)

        return np.array([[[1, 1, 1]]
                         ], dtype=np.uint8)

################### motion primitives ######################
class DifferentialDriveGymPrimitive(DifferentialDriveGym):
    def __init__(self,
                 robot_env=DifferentialDriveEnv(1.0, -0.1, np.pi, 1.0, np.pi),
                 reward_param=np.zeros(5),
                 obc_list=np.zeros((10, 0, 2))):
        super().__init__(robot_env, reward_param, obc_list)
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(
            3,))
        self.max_time = 10.0

    def step(self, action):
        v = action
        dt = 1.0/5.0  # 5 Hz
        self.state = self.robot_env.motion_velocity(self.state, v, dt)
        self.current_time += dt
        obs = self._obs()
        info = self._info()
        reward = self._reward()
        done = self.current_time > self.max_time
        return obs, reward, done, info

    def _reward(self):
        vx = self.state[3] * np.cos(self.state[2])
        vy = self.state[3] * np.sin(self.state[2])
        w = self.state[4]
        reward = vx - 0.2*abs(vy) - 1.0*abs(np.tanh(w))
        return reward

    def _obs(self):
        obs = self.state[2:]
        return obs

    def reset(self):
        # sample a random start goal configuration
        start = np.random.uniform(
                self.state_bounds[:, 0], self.state_bounds[:, 1])
        self.state = start
        self.current_time = 0
        obs = self._obs()
        return obs

class DifferentialDriveGymForward(DifferentialDriveGymPrimitive):
    def _reward(self):
        vx = self.state[3] * np.cos(self.state[2])
        vy = self.state[3] * np.sin(self.state[2])
        w = self.state[4]
        reward = vx - 0.5*abs(vy) - 1.0*abs(np.tanh(w))
        return reward

class DifferentialDriveGymBackward(DifferentialDriveGymPrimitive):
    def _reward(self):
        vx = self.state[3] * np.cos(self.state[2])
        vy = self.state[3] * np.sin(self.state[2])
        w = self.state[4]
        reward = -vx - 0.5*abs(vy) - 1.0*abs(np.tanh(w))
        return reward

class DifferentialDriveGymUpward(DifferentialDriveGymPrimitive):
    def _reward(self):
        vx = self.state[3] * np.cos(self.state[2])
        vy = self.state[3] * np.sin(self.state[2])
        w = self.state[4]
        reward = vy - 0.5*abs(vx) - 1.0*abs(np.tanh(w))
        return reward

class DifferentialDriveGymDownward(DifferentialDriveGymPrimitive):
    def _reward(self):
        vx = self.state[3] * np.cos(self.state[2])
        vy = self.state[3] * np.sin(self.state[2])
        w = self.state[4]

        reward = -vy - 0.5*abs(vx) - 1.0*abs(np.tanh(w))
        return reward

####################### compositional policy #####################
class DifferentialDriveGymCompose(DifferentialDriveGym):
    def __init__(self, base_policies, obc_list):
        super().__init__(obc_list=obc_list)
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(
            5+len(self.sample_positions)+8,))
        self.action_space = spaces.Box(low=0.0, high=1.0, shape=(len(base_policies),))
        self.base_policies = base_policies

    def primitives(self):
        obs_proprioceptive = Batch(obs=self.state[2:].reshape((1, -1)), info=None)

        v_list = []
        for k,v_ in enumerate(self.base_policies):
            act_batch = v_(obs_proprioceptive, eps=0.05)
            act = act_batch.act.detach().cpu().numpy()[0]
            v_primitive = act
            assert len(v_primitive)==2
            v_list.append(v_primitive)
        return v_list

    def step(self, action):
        action /= np.sum(action)
        assert abs(np.sum(action) - 1.0) < 1e-5,\
            'the summation of weights'+str(np.sum(action))+'does not equal to 1'
        weights = action
        v_list = self.primitives()

        v = np.array(v_list).T @ weights
        dt = 1.0/5.0  # 5 Hz

        self.state = self.robot_env.motion_velocity(self.state, v, dt)
        self.current_time += dt
        obs = self._obs()
        info = self._info()
        reward = self._reward(info)
        done = info['goal'] or info['collision'] or self.current_time > self.max_time
        return obs, reward, done, info

    def _obs(self):
        wRb = euler.euler2mat(0, 0, self.state[2])[:2, :2]
        wTb = np.block([[wRb, self.state[:2].reshape((-1, 1))],
                        [np.zeros((1, 2)), 1]])

        local_map = self.sample_local_map(wTb)
        obs = np.block([self.state[2:], self.goal[:2]-self.state[:2], local_map]+self.primitives())
        return obs


import argparse
import torch
import torch.nn.functional as F
import pickle
from pathlib import Path

from torch.utils import data
import numpy as np
import tqdm
import logging

import random
import torch
from cswm import utils
from cswm.models.modules import RewardPredictor, CausalTransitionModel
from cswm.utils import OneHot

import sys
import datetime
import os
import pickle
from pathlib import Path

import numpy as np

from itertools import chain

torch.backends.cudnn.deterministic = True

parser = argparse.ArgumentParser()
parser.add_argument('--save-folder', type=Path,
                    default='checkpoints',
                    help='Path to checkpoints.')

parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='Disable CUDA training.')
parser.add_argument('--finetune', action='store_true', default=False,
                    help='Whether to use finetuned model')
parser.add_argument('--random', action='store_true', default=False)
parser.add_argument('--epochs', type=int, default=50)
parser.add_argument('--recurrent', action = 'store_true')
args_eval = parser.parse_args()

def set_seed(seed):
    """Set seed"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if args.cuda:
        torch.cuda.manual_seed_all(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.backends.cudnn.deterministic = True

meta_file = args_eval.save_folder / 'metadata.pkl'
if args_eval.finetune:
    model_file = args_eval.save_folder / 'finetuned_model.pt'
    reward_model_file = args_eval.save_folder / 'finetune_reward_model.pt'
    log_file = args_eval.save_folder / 'finetune_reward_log.txt'
elif args_eval.random:
    model_file = args_eval.save_folder / 'random_model.pt'
    reward_model_file = args_eval.save_folder / 'random_reward_model.pt'
    log_file = args_eval.save_folder / 'random_reward_log.txt'
else:
    model_file = args_eval.save_folder / 'model.pt'
    reward_model_file = args_eval.save_folder / 'reward_model.pt'
    log_file = args_eval.save_folder / 'reward_log.txt'

with open(meta_file, 'rb') as f:
    args = pickle.load(f)['args']

handlers = [logging.FileHandler(log_file, 'a')]
if args.silent:
    handlers.append(logging.StreamHandler(sys.stdout))
logging.basicConfig(level=logging.INFO, format='%(message)s', handlers=handlers)
logger = logging.getLogger()
print = logger.info

set_seed(args.seed)

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

print("Loading data...")
dataset = utils.StateTransitionsDataset(
    hdf5_file=args.dataset, action_transform=OneHot(args.num_objects * args.action_dim))
valid_dataset = utils.StateTransitionsDataset(
    hdf5_file=args.valid_dataset, action_transform=OneHot(args.num_objects * args.action_dim))

train_loader = data.DataLoader(
    dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)
valid_loader = data.DataLoader(
    valid_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)

print("Loading model...")

# Get data sample
obs = next(iter(train_loader))[0]
input_shape = obs[0].size()

print(f"VAE: {args.vae}")
print(f"Modular: {args.modular}")
print(f"Encoder: {args.encoder}")
print(f"Num Objects: {args.num_objects}")
print(f"Dataset: {args.dataset}")

if not args_eval.recurrent:
    model = CausalTransitionModel(
        embedding_dim_per_object=args.embedding_dim_per_object,
        hidden_dim=args.hidden_dim,
        action_dim=args.action_dim,
        input_dims=input_shape,
        input_shape=input_shape,
        modular=args.modular,
        predict_diff=args.predict_diff,
        vae=args.vae,
        num_objects=args.num_objects,
        encoder=args.encoder,
        gnn=args.gnn,
        multiplier=args.multiplier,
        ignore_action=args.ignore_action,
        copy_action=args.copy_action).to(device)

else:
    model = CausalTransitionModelLSTM(
    embedding_dim_per_object=args.embedding_dim_per_object,
    hidden_dim=args.hidden_dim,
    action_dim=args.action_dim,
    input_dims=(3, 50, 50),
    input_shape=(3, 50, 50),
    modular=args.modular,
    predict_diff=args.predict_diff,
    vae=args.vae,
    num_objects=args.num_objects,
    encoder=args.encoder,
    rim = args.rim,
    scoff = args.scoff).to(device)

num_enc = sum(p.numel() for p in model.encoder_parameters())
num_dec = sum(p.numel() for p in model.decoder_parameters())
num_tr = sum(p.numel() for p in model.transition_parameters())

print(f'Number of parameters in Encoder: {num_enc}')
print(f'Number of parameters in Decoder: {num_dec}')
print(f'Number of parameters in Transition: {num_tr}')
print(f'Number of parameters: {num_enc+num_dec+num_tr}')

if not args_eval.random:
    model.load_state_dict(torch.load(model_file))
else:
    torch.save(model.state_dict(), model_file)

model.eval()

Reward_Model = RewardPredictor(args.embedding_dim_per_object * args.num_objects).to(device)

def evaluate(valid_loader):
    valid_loss = 0.0
    Reward_Model.eval()

    for batch_idx, data_batch in enumerate(valid_loader):
        data_batch = [tensor.to(device).float() for tensor in data_batch]
        _, _, obs, reward, target = data_batch

        state, _ = model.encode(obs)
        state = state.view(state.shape[0], args.num_objects * args.embedding_dim_per_object)
        reward_state, _ = model.encode(target)
        reward_state = reward_state.view(state.shape[0], args.num_objects * args.embedding_dim_per_object)

        state_emb = torch.cat([state, reward_state], dim=1)
        reward_pred = Reward_Model(state_emb).view(reward.shape)

        loss = F.l1_loss(reward_pred, reward)

        valid_loss += loss.item()

    avg_loss = valid_loss / len(valid_loader)
    print('====> Average valid loss: {:.6f}'.format(avg_loss))
    return avg_loss

def train(max_epochs, lr):

    optimizer = torch.optim.Adam(Reward_Model.parameters(), lr = lr)

    print('Starting model training...')
    best_loss = 1e9
    for epoch in range(1, max_epochs + 1):
        train_loss = 0

        iterator = tqdm.tqdm(train_loader, desc=f'Epoch {epoch}',
                             disable=args.silent)

        for batch_idx, data_batch in enumerate(iterator):
            Reward_Model.train()
            data_batch = [tensor.to(device).float() for tensor in data_batch]
            _, _, obs, reward, target = data_batch

            optimizer.zero_grad()
            #print(obs.size())
            state, _ = model.encode(obs)
            state = state.view(state.shape[0], args.num_objects * args.embedding_dim_per_object)
            #print(target.size())
            reward_state, _ = model.encode(target)
            reward_state = reward_state.view(state.shape[0], args.num_objects * args.embedding_dim_per_object)

            state_emb = torch.cat([state, reward_state], dim=1)
            reward_pred = Reward_Model(state_emb).view(reward.shape)

            loss = F.l1_loss(reward_pred, reward)

            loss.backward()
            train_loss += loss.item()

            optimizer.step()

            if batch_idx % args.log_interval == 0 and batch_idx > 0:
                iterator.set_postfix(loss=f'{train_loss / (1 + batch_idx):.6f}')
                print(
                    'Epoch: {} [ {}/{} ] \t Loss: {:.6f}'.format(
                        epoch, (batch_idx+1),
                        len(train_loader.dataset),
                        loss.item()))

        avg_loss = train_loss / len(train_loader)
        print('====> Epoch: {} Average train loss: {:.6f}'.format(
            epoch, avg_loss))

        avg_loss = evaluate(valid_loader)

        if avg_loss < best_loss:
            best_loss = avg_loss
            torch.save(Reward_Model.state_dict(), reward_model_file)

train(args_eval.epochs, lr = args.lr)


from keras import Input, Model, Sequential
from keras.layers import Flatten, Conv2D, Concatenate, Dense, Activation
from keras.optimizers import Adam
from rl.agents import DDPGAgent, CEMAgent
from rl.memory import SequentialMemory, EpisodeParameterMemory
from rl.policy import EpsGreedyQPolicy

from hexagon_agent import *
from random import shuffle
from multi_agent import *
import sys
import hexagon_ui_api
import os
from square_grid import *
import numpy as np
import warnings
import argparse

# ______________________________________________________________________________________________________________________________
class EnvDef:
  centaur_name = 'centaur'
  game_name = '1'
  HASH_POOL = 10000
  NODE_FEATURE_COUNT = 5
  DECISION_ACTION_SPACE = 2
  SHORT_MEMORY_SIZE = 1
  MAX_ROUND = 2000
  CELL_FEATURE = 1
  MAX_GRID_LENGTH = 9
  SPATIAL_INPUT = (MAX_GRID_LENGTH, MAX_GRID_LENGTH)
  SPATIAL_OUTPUT = (MAX_GRID_LENGTH * MAX_GRID_LENGTH, )
  EPISODE_REWARD = 1000
  MOVE_REWARD_MULTIPLIER = 10
  DONT_OWN_MOVE_REWARD = -5
  CANT_ATTACK_MOVE_REWARD = -3

# __________________________________________________________________________________________________________________________

class NoneZeroEpsGreedyQPolicy(EpsGreedyQPolicy):
  """Implement the epsilon greedy policy

  Eps Greedy policy either:

  - takes a random action with probability epsilon from Non-Zero Q-values
  - takes current best action with prob (1 - epsilon)
  """

  def __init__(self, eps=.1):
    super(EpsGreedyQPolicy, self).__init__()
    self.eps = eps

  def select_action(self, q_values):
    """Return the selected action

    # Arguments
        q_values (np.ndarray): List of the estimations of Q for each action

    # Returns
        Selection action
    """
    assert q_values.ndim == 1
    nb_actions = q_values.shape[0]
    copy_q_values = np.copy(q_values)
    for i in range(0, nb_actions):
      if copy_q_values[i] == 0:
        copy_q_values[i] = -1e9
    if np.random.uniform() < self.eps:
      idx = np.argmax(copy_q_values)
      copy_q_values[idx] = 0
      for i in range(0, nb_actions):
        copy_q_values[i] *= np.random.uniform()
      action = np.argmax(copy_q_values)
      if action == 0:  # there was only one possible value
        action = idx  # make it same old
    else:
      action = np.argmax(copy_q_values)
    if action == 0:
      x = 1
    return action

  def get_config(self):
    """Return configurations of EpsGreedyPolicy

    # Returns
        Dict of config
    """
    config = super(EpsGreedyQPolicy, self).get_config()
    config['eps'] = self.eps
    return config

# __________________________________________________________________________________________________________________________

class MaskableDDPGAgent(DDPGAgent):

  def __init__(self, nb_actions, actor, critic, critic_action_input, memory,
               gamma=.99, batch_size=32, nb_steps_warmup_critic=1000, nb_steps_warmup_actor=1000,
               train_interval=1, memory_interval=1, delta_range=None, delta_clip=np.inf,
               random_process=None, custom_model_objects={}, target_model_update=.001,
               mask_processor=None, **kwargs):

    DDPGAgent.__init__(self, nb_actions=nb_actions, actor=actor, critic=critic,
                       critic_action_input=critic_action_input, memory=memory,
                       gamma=gamma, batch_size=batch_size, nb_steps_warmup_critic=nb_steps_warmup_critic,
                       nb_steps_warmup_actor=nb_steps_warmup_actor,train_interval=train_interval,
                       memory_interval=memory_interval, delta_range=delta_range, delta_clip=delta_clip,
                       random_process=random_process, custom_model_objects=custom_model_objects,
                       target_model_update=target_model_update, **kwargs)

    self.mask_processor = mask_processor

  def forward(self, observation):
    # Select an action.
    state = self.memory.get_recent_state(observation)
    action = self.select_action(state)  # TODO: move this into policy

    if self.mask_processor is not None:
      action = self.mask_processor.mask(action)

    # Book-keeping.
    self.recent_observation = observation
    self.recent_action = action

    return action

# __________________________________________________________________________________________________________________________

class AgentType:
  BoostDecision = 'BoostDecision'
  Attack = 'Attack'
  Boost = 'Boost'

# __________________________________________________________________________________________________________________________
class SuperCentaurPlayer(Aliostad):
  def __init__(self, name, boost_likelihood=0.23, boosting_off=False, attack_off=False):
    Aliostad.__init__(self, name)
    self.reset_state()
    self.boost_likelihood = boost_likelihood
    self.boosting_off = boosting_off
    self.attack_off = attack_off

  def reset_state(self):
    self.actions = {}
    self.current_move = None
    self.was_called = {}
    self.illegal_move_reward_by_agents = {}

  def timeForBoost(self, world):
    """

    :type world: World
    :return:
    """
    if self.boosting_off:
      return Aliostad.timeForBoost(self, world)

    #self.was_called[AgentType.BoostDecision] = True
    #return self.actions[AgentType.BoostDecision] == 1
    isTimeForBoost = np.random.uniform() < self.boost_likelihood
    self.boost_stats.append(isTimeForBoost)
    return isTimeForBoost

  def move(self, playerView):
    return self.current_move

  def getAttackFromCellId(self, world):
    if self.attack_off:
      return Aliostad.getAttackFromCellId(self, world)

    self.was_called[AgentType.Attack] = True
    cellId = self.actions[AgentType.Attack]
    if cellId not in world.uberCells:
      self.illegal_move_reward_by_agents[AgentType.Attack] = EnvDef.DONT_OWN_MOVE_REWARD
      print('{} - illegal move (dont own): {}'.format(self.round_no, cellId))
      return None
    if not world.uberCells[cellId].canAttackOrExpand:
      self.illegal_move_reward_by_agents[AgentType.Attack] = EnvDef.CANT_ATTACK_MOVE_REWARD
      print('{} - illegal move (cant attack): {}'.format(self.round_no, cellId))
      return None
    #print ('{} - legal!!'.format(world.round_no))
    return cellId

# __________________________________________________________________________________________________________________________
class HierarchicalCentaurEnv(Env):
  def __init__(self, centaur_boost_likelihood, opponent_randomness=None, boosting_off=False, attack_off=False):
    self.boosting_off = boosting_off
    self.attack_off = attack_off
    self.opponent_randomness = opponent_randomness
    self.players = []
    self.game = None
    self._seed = 0
    self.centaur = None
    self.cell_count = 1
    self.resources = 100
    self.world = None
    self.leaderBoard = {}
    self.cellLeaderBoard = {}
    self.shortMemory = []
    self.centaur_boost_likelihood = centaur_boost_likelihood

  def configure(self, *args, **kwargs):
    pass

  def seed(self, seed=None):
    self._seed = seed
    return [self._seed]

  def render(self, mode='human', close=False):
    pass

  def step(self, actions):
    for name in actions:
      self.centaur.actions[name] = actions[name]  # we can also deep copy
    self.centaur.current_move = self.centaur.movex(self.world)
    stats, isFinished , extraInfo = self.game.run_sync()
    info = self.game.board.get_cell_infos_for_player(EnvDef.centaur_name)
    reward = (len(info) - self.cell_count) * EnvDef.MOVE_REWARD_MULTIPLIER
    if len(info) == 0 or self.game.round_no > EnvDef.MAX_ROUND:
      isFinished = True  # it has no cells anymore
    else:
      self.cell_count = len(info)
    if self.game.round_no % 100 == 0:
      print(self.cell_count)
    if isFinished:
      winner = stats[0]
      reward = EnvDef.EPISODE_REWARD if winner.playerName == EnvDef.centaur_name else -EnvDef.EPISODE_REWARD
      if winner.playerName in self.leaderBoard:
        self.leaderBoard[winner.playerName] += 1
      else:
        self.leaderBoard[winner.playerName] = 1
      for stat in stats:
        if stat.playerName not in self.cellLeaderBoard:
          self.cellLeaderBoard[stat.playerName] = 0
        self.cellLeaderBoard[stat.playerName] += stat.cellsOwned
        print('{} {} ({}) - {}'.format(stat.playerName, stat.cellsOwned, stat.totalResources, str(extraInfo[stat.playerName])))

      for name in self.leaderBoard:
        print(' - {}: {} ({})'.format(name, self.leaderBoard[name], self.cellLeaderBoard[name]))

    playerView = PlayerView(self.game.round_no, info)
    wrld = self.centaur.build_world(playerView.ownedCells)
    self.push_world(wrld)
    rewards = {name: (reward + self.centaur.illegal_move_reward_by_agents[name]) if name in self.centaur.illegal_move_reward_by_agents else
      reward for name in self.centaur.was_called}
    self.centaur.reset_state()
    return wrld, rewards, isFinished, {}

  def push_world(self, world):
    """

    :type world: World
    :return:
    """
    self.world = world
    self.shortMemory.append(world)
    if len(self.shortMemory) > EnvDef.SHORT_MEMORY_SIZE:
      del self.shortMemory[0]

  def close(self):
    print('closing CentaurEnv')

  def reset(self):
    self.cell_count = 1
    self.resources = 100
    if self.game is not None:
      self.game.finish()

    self.shortMemory = []
    for i in range(0, EnvDef.SHORT_MEMORY_SIZE):
      self.shortMemory.append(World([], 0))

    self.centaur = SuperCentaurPlayer(EnvDef.centaur_name,
                      boost_likelihood=self.centaur_boost_likelihood,
                      attack_off=self.attack_off, boosting_off=self.boosting_off)
    self.players = [self.centaur, Aliostad('aliostad', randomBoostFactor=self.opponent_randomness)]
    shuffle(self.players)
    self.game = Game(EnvDef.game_name, self.players, radius=5)
    hexagon_ui_api.games[EnvDef.game_name] = self.game
    self.game.start()
    playerView = PlayerView(self.game.round_no, self.game.board.get_cell_infos_for_player(EnvDef.centaur_name))
    wrld = self.centaur.build_world(playerView.ownedCells)
    self.push_world(wrld)
    return wrld

# ____________________________________________________________________________________________________________________________
class CentaurDecisionProcessor(Processor):

  @staticmethod
  def calculate_hash_index(cellName):
    # type: (str) -> int
    return int(abs(hash(cellName))) % EnvDef.HASH_POOL

  def buildInput(self, world):
    """
    Each input has 5 values:
      - first 4 items is 1-hot-vector:
        - No cell (we do not see)
        - My cell
        - Enemy cell
        - Not-captured cell
      - Resources

    :type world:
    :return:
    """
    inpt = [np.array([1, 0, 0, 0, 0]) for i in range(0, EnvDef.HASH_POOL)]

    for cell in world.cells.values():
      id = self.calculate_hash_index(str(cell.id))
      inpt[id] = np.array([0, 1, 0, 0, cell.resources])
      for n in cell.neighbours:
        id = int(abs(hash(n.id))) % EnvDef.HASH_POOL
        if n.isOwned is None:  # neutral
          inpt[id] = np.array(
            [0, 0, 0, 1, cell.resources])  # resources would be 0 but better just to use resources property
        elif n.isOwned is False:  # enemy
          inpt[id] = np.array([0, 0, 1, 0, cell.resources])
    inpt = np.array(inpt).flatten()
    return inpt

  def process_action(self, action):
    return action  # this is due to a BUG in keras-rl when it tries to calc mean

  def process_observation(self, observation):
    """

    :type observation: World
    :return:
    """
    return self.buildInput(observation)

# ______________________________________________________________________________________________________________________________
class CentaurAttackProcessor(Processor):
  def __init__(self, masking=True, random_action=False):
    self.masking = masking
    self.last_world = None
    self.policy = NoneZeroEpsGreedyQPolicy()
    self.random_action = random_action

  def buildInput(self, world):
    """
    returns a MxN map of the world with hexagon grid transposed to square grid

    :type world: World
    :return:
    """
    self.last_world = world
    hector = np.zeros(EnvDef.SPATIAL_INPUT)
    for cid in world.worldmap:
      hid = GridCellId.fromHexCellId(cid)
      thid = hid.transpose(EnvDef.SPATIAL_INPUT[0] / 2, EnvDef.SPATIAL_INPUT[1] / 2)
      hector[thid.x][thid.y] = world.worldmap[cid]
    return hector

  def buildOutput(self, world):
    """
    returns a MxN map of the world with hexagon grid transposed to square grid

    :type world: World
    :return:
    """
    hector = np.zeros(EnvDef.SPATIAL_INPUT)
    for cid in world.uberCells:
      hid = GridCellId.fromHexCellId(cid)
      thid = hid.transpose(EnvDef.SPATIAL_INPUT[0] / 2, EnvDef.SPATIAL_INPUT[1] / 2)
      hector[thid.x][thid.y] = 0 if not world.uberCells[cid].canAttackOrExpand else world.uberCells[cid].attackPotential
    return hector

  def process_action_old(self, action):
    """

    :type action: ndarray
    :return:
    """
    flat = action.flatten()
    idx = np.argmax(flat, 0)
    y = idx % action.shape[1]
    x = idx / action.shape[1]
    if idx == 0:
      print('zero!')
    thid = GridCellId(x, y).transpose(-(EnvDef.SPATIAL_INPUT[0] / 2), -(EnvDef.SPATIAL_INPUT[1] / 2))
    return thid.to_cell_id()

  def process_action(self, action):
    """

    :type action: int
    :return:
    """
    idx = self.policy.select_action(action)
    y = idx % EnvDef.SPATIAL_INPUT[1]
    x = idx / EnvDef.SPATIAL_INPUT[1]
    thid = GridCellId(x, y).transpose(-(EnvDef.SPATIAL_INPUT[0] / 2), -(EnvDef.SPATIAL_INPUT[1] / 2))
    cid = thid.to_cell_id()
    return cid

  @staticmethod
  def add_noise(arr, noise_range=0.3):
    """

    :type arr: ndarray
    :type noise_range: float
    :return:
    """
    shp = arr.shape
    flat = arr.flatten()
    noisy = np.array([value * np.random.uniform(1. - noise_range, 1. + noise_range) for value in flat])
    return noisy.reshape(shp)

  def mask(self, Y):
    """

    :type Y: ndarray
    :return:
    """
    assert len(Y.shape) == 1  # it is flat
    if self.last_world is None:
      warnings.warn("Last world is None. Could not mask.")
      return Y
    mask = self.buildOutput(self.last_world).flatten()
    assert mask.shape == Y.shape

    for i in range(0, len(Y)):
      if mask[i] == 0:
        Y[i] = 0
      elif self.random_action:
        Y[i] = np.random.uniform()
    return Y

  def process_and_mask(self, Y):
    """

    :type Y: ndarray
    :return:
    """
    shibo = AttackModel.process_y(Y)
    real_shibo = shibo
    if self.masking and self.last_world is not None:
      if min(shibo.flat) == sum(shibo.flat):
        print('all zero output')
        return self.add_noise(self.buildOutput(self.last_world))
      mask = np.zeros(shibo.shape)
      for cid in self.last_world.uberCells:
        if self.last_world.uberCells[cid].canAttackOrExpand:
          hid = GridCellId.fromHexCellId(cid)
          thid = hid.transpose(EnvDef.SPATIAL_INPUT[0] / 2, EnvDef.SPATIAL_INPUT[1] / 2)
          mask[thid.x][thid.y] = 1
      minimum = min(shibo.flatten())
      if minimum < 0:  # rescale to zero
        shibo += -minimum
      real_shibo = shibo * mask
      if sum(real_shibo.flat) == 0:
        print('all zero output after masking')
        return self.add_noise(self.buildOutput(self.last_world))
    return real_shibo

  def process_observation(self, observation):
    """

    :type observation: World
    :return:
    """
    return np.reshape(self.buildInput(observation), EnvDef.SPATIAL_INPUT + (1, ))

  def process_state_batch(self, batch):
    """

    :type batch: ndarray
    :return:
    """
    return np.reshape(batch, (batch.shape[0], ) + EnvDef.SPATIAL_INPUT + (1, ))

class CentaurBoostProcessor(CentaurAttackProcessor):

  def __init__(self, masking=True):
    CentaurAttackProcessor.__init__(self, masking)

  def process_and_mask(self, Y):
    """

    :type Y: ndarray
    :return:
    """
    shibo = AttackModel.process_y(Y)
    real_shibo = shibo
    if self.masking and self.last_world is not None:
      mask = np.zeros(Y.shape)
      for cid in self.last_world.uberCells:
        hid = GridCellId.fromHexCellId(cid)
        thid = hid.transpose(EnvDef.SPATIAL_INPUT[0] / 2, EnvDef.SPATIAL_INPUT[1] / 2)
        mask[thid.x][thid.y] = 1
      real_shibo = shibo * mask
      if max(real_shibo.flatten()) == 0:
        print('what??')
    return real_shibo
# ______________________________________________________________________________________________________________________________

class DecisionModel:
  def __init__(self, modelName=None):
    """

    :type theMethod: str
    """
    self.modelName = modelName if modelName is not None else 'Decision_model_params.h5f' + str(r.uniform(0, 10000))
    model = Sequential()
    model.add(Flatten(input_shape=(1,) + (EnvDef.HASH_POOL * EnvDef.NODE_FEATURE_COUNT * EnvDef.SHORT_MEMORY_SIZE,)))
    model.add(Dense(32, activation="relu"))
    model.add(Dense(16, activation="relu"))
    model.add(Dense(EnvDef.DECISION_ACTION_SPACE))
    model.add(Activation('softmax'))
    print(model.summary())
    model.compile(loss="categorical_crossentropy",
                  optimizer='adadelta', metrics=['accuracy'])
    self.model = model

# ______________________________________________________________________________________________________________________________

class AttackModel:
  def __init__(self, modelName=None):
    """

    :type theMethod: str
    """
    self.modelName = modelName if modelName is not None else 'Attack_model_params.h5f' + str(r.uniform(0, 10000000))

    model = Sequential()
    model.add(Conv2D(128, (5, 5), padding='same', activation='relu',
              input_shape=EnvDef.SPATIAL_INPUT + (1, ), name='INPUT_ATTACK'))
    model.add(Conv2D(64, (3, 3), padding='same', activation='relu'))
    model.add(Conv2D(16, (3, 3), padding='same', activation='relu'))
    model.add(Conv2D(4, (3, 3), padding='same', activation='relu'))
    model.add(Conv2D(1, (3, 3), padding='same', activation='tanh'))
    model.add(Flatten())
    model.add(Dense(EnvDef.SPATIAL_OUTPUT[0], activation='tanh'))

    self.model = model

    action_input = Input(shape=EnvDef.SPATIAL_OUTPUT, name='action_input')
    observation_input = Input(shape=EnvDef.SPATIAL_INPUT + (1, ), name='observation_input')
    self.critic_action_input = action_input
    conv = Conv2D(256, (5, 5), padding='same', activation='relu')(observation_input)
    conv = Conv2D(128, (3, 3), padding='same', activation='relu')(conv)
    conv = Conv2D(32, (3, 3), padding='same', activation='relu')(conv)
    conv = Conv2D(8, (3, 3), padding='same', activation='relu')(conv)
    conv = Conv2D(1, (1, 1), padding='same', activation='tanh')(conv)
    conv = Flatten()(conv)
    conv = Dense(EnvDef.SPATIAL_OUTPUT[0], activation='tanh')(conv)
    x = Concatenate()([action_input, conv])
    x = Dense(32, activation='relu')(x)
    x = Dense(16, activation='relu')(x)
    x = Dense(8, activation='relu')(x)
    x = Dense(1, activation='tanh')(x)
    critic = Model(inputs=[action_input, observation_input], outputs=x)
    self.critic = critic

  def get_critic(self):
    return self.critic

  def get_critic_action_input(self):
    return self.critic_action_input

  @staticmethod
  def prepare_x(X, batch=False):
    """

    :type X: ndarray
    :type batch: bool
    :return:
    """
    shape = EnvDef.SPATIAL_INPUT + (1, )
    if batch:
      shape = (X.shape[0], ) + shape
    return np.reshape(X, shape)

  @staticmethod
  def prepare_y(Y, batch=False):
    """

    :type Y: ndarray
    :type batch: bool
    :return:
    """
    shape = EnvDef.SPATIAL_OUTPUT + (1, )
    if batch:
      shape = (Y.shape[0], ) + shape
    return np.reshape(Y, shape)

  @staticmethod
  def process_y(Y):
    """

    :type Y: ndarray
    :return:
    """
    return np.reshape(Y, EnvDef.SPATIAL_OUTPUT)

# ______________________________________________________________________________________________________________________________

if __name__ == '__main__':

  parser = argparse.ArgumentParser()
  parser.add_argument('what', help="what to do", type=str)
  parser.add_argument('--model_name', '-m', help="model name", type=str)
  parser.add_argument('--randomness', '-r', help="randomness of aliostad", type=float)
  parser.add_argument('--randomaction', '-x', help="action completely random but valid", type=bool, nargs='?', const=True)
  parser.add_argument('--boostingoff', '-y', help="don't use boosting method of centaur",
                      type=bool, nargs='?', const=True, default=True)
  parser.add_argument('--attackoff', '-z', help="dont use attack method of centaur",
                      type=bool, nargs='?', const=True, default=False)
  parser.add_argument('--testrounds', '-t', help="number of epochs when testing", type=int, default=100)
  parser.add_argument('--centaur_boost_likelihood', '-b', help="likelihood of random boost for centaur", type=float, default=0.23)

  randomness = None
  attack_model_name = 'Attack_model_params.h5f'

  args = parser.parse_args(sys.argv[1:])
  if args.model_name is not None:
    attack_model_name = args.model_name

  if args.randomness is not None:
    randomness = args.randomness

  env = HierarchicalCentaurEnv(opponent_randomness=randomness,
                               centaur_boost_likelihood=args.centaur_boost_likelihood,
                               boosting_off=args.boostingoff, attack_off=args.attackoff)
  np.random.seed(42)
  env.seed(42)

  prc = CentaurDecisionProcessor()
  dec_model = DecisionModel()
  attack_model = AttackModel(attack_model_name)

  prc = MultiProcessor({AgentType.BoostDecision: prc, AgentType.Attack: CentaurAttackProcessor(random_action=args.randomaction)})
  memory = EpisodeParameterMemory(limit=1000, window_length=1)
  decision_agent = CEMAgent(model=dec_model.model, nb_actions=EnvDef.DECISION_ACTION_SPACE, memory=memory,
                            batch_size=50, nb_steps_warmup=200, train_interval=50, elite_frac=0.05)

  decision_agent.compile()
  memory2 = SequentialMemory(limit=100000, window_length=1)
  policy = NoneZeroEpsGreedyQPolicy()
  attack_agent = MaskableDDPGAgent(nb_actions=EnvDef.SPATIAL_OUTPUT[0],
                                   actor=attack_model.model,
                                   processor=prc.inner_processors[AgentType.Attack],
                                   critic=attack_model.get_critic(),
                                   memory=memory2, nb_steps_warmup_actor=200,
                                   nb_steps_warmup_critic=200,
                                   critic_action_input=attack_model.get_critic_action_input(),
                                   mask_processor=prc.inner_processors[AgentType.Attack])

  agent = MultiAgent({AgentType.BoostDecision: decision_agent, AgentType.Attack: attack_agent}, processor=prc, save_frequency=0.05)
  agent.inner_agents[AgentType.Attack].compile(Adam(lr=0.001), metrics=['mean_squared_logarithmic_error'])
  if args.model_name is not None:
    agent.inner_agents[AgentType.Attack].load_weights(attack_model_name)

  hexagon_ui_api.run_in_background()
  if len(sys.argv) == 1:
    print('Usage: python centaur_ai_gym.py (train|test)')
  elif args.what == 'train':
    agent.fit(env, nb_steps=300 * 1000, visualize=False, verbose=2, interim_filenames={AgentType.Attack: attack_model.modelName})
    agent.save_weights({AgentType.BoostDecision: dec_model.modelName, AgentType.Attack: attack_model.modelName}, overwrite=True)
  elif args.what == 'test':
    agent.test(env, nb_episodes=args.testrounds)
  else:
    print('argument not recognised: ' + sys.argv[1])


#!/usr/bin/env python

import rospy
import json
from wr8_ai.msg import BoundingBox
import wr8_ai.utils as ut
import wr8_ai.ncs as ncs
from wr8_ai.yolo import yolo
from wr8_ai.yolo import bbox
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image, CompressedImage

class DetectorNCS:
    def __init__(self):
        pass

    def init(self, stick_idx, graph_path, config_path):
        model = ncs.InferNCS(graph_path, device_idx=stick_idx, fp16=False)

        rospy.loginfo('Opened params: {} / {}'.format(graph_path, config_path))

        if not model.is_opened():
            rospy.logerr('Failed to init device')
            return False

        rospy.loginfo('Model created')

        with open(config_path) as config_buffer:
            config = json.load(config_buffer)

        rospy.loginfo('Config opened')

        self.labels = config['model']['labels']
        anchors = config['model']['anchors']

        net_h, net_w = config['model']['infer_shape']
        self.obj_thresh, nms_thresh = 0.5, 0.45

        self.inferer = yolo.YOLO(model, net_h, net_w, anchors, self.obj_thresh, nms_thresh)

        return True

    def infer(self, cv_img):
        # Return (ROS signs boxes, render_img)

        boxes = self.inferer.make_infer([cv_img])[0]

        ros_bboxes = []
        for box in boxes:
            # Fixed - current NMS removes class porbabilities,
            if box.get_score() > 0:
                ros_bboxes += [ut.yolo_bbox_2_ros_bbox(box, self.labels)]

        yolo_bboxes = boxes

        return ros_bboxes, yolo_bboxes

    def draw_boxes(self, img, ros_boxes):
        bbox.draw_ros_boxes(img, ros_boxes, self.labels, self.obj_thresh)

import wr8_ai.tsr.image as tsr_im
import numpy as np
import keras

class RecognizerCPU:
    def __init__(self):
        pass

    def init(self, weights_path, config_path):
        from keras.models import load_model

        rospy.loginfo('Opened params: {} / {}'.format(weights_path, config_path))

        self.model = load_model(weights_path)

        rospy.loginfo('Model created')

        with open(config_path) as config_buffer:
            config = json.load(config_buffer)

        rospy.loginfo('Config opened')

        self.labels = config['model']['labels']
        self.net_h, self.net_w = config['model']['infer_shape']

        return True

    def infer(self, cv_img):
        # Return (ROS signs boxes, render_img)
        image = tsr_im.image_preprocess(cv_img, self.net_h, self.net_w)

        inf_img = np.expand_dims(image, axis=0)
        result = self.model.predict(inf_img)[0]
        # print(result)
        max_idx = np.argmax(result)

        return self.labels[max_idx]

class RecognizerNCS:
    def __init__(self):
        pass

    def init(self, stick_idx, graph_path, config_path):
        self.model = ncs.InferNCS(graph_path, device_idx=stick_idx, fp16=False)

        rospy.loginfo('Opened params: {} / {}'.format(graph_path, config_path))

        if not self.model.is_opened():
            rospy.logerr('Failed to init device')
            return False

        rospy.loginfo('Model created')

        with open(config_path) as config_buffer:
            config = json.load(config_buffer)

        rospy.loginfo('Config opened')

        self.labels = config['model']['labels']
        self.net_h, self.net_w = config['model']['infer_shape']

        return True

    def infer(self, cv_img):
        # Return (ROS signs boxes, render_img)
        image = tsr_im.image_preprocess(cv_img, self.net_h, self.net_w)

        result = self.model.infer(image)
        print(result)
        max_idx = np.argmax(result)

        return self.labels[max_idx]

class DoubleDetector:
    def __init__(self):
        pass

    def init(self, config_rec, weights_rec, graph_rec, config_det, graph_det):
        self.nn_det = DetectorNCS()
        if not self.nn_det.init(0, graph_det, config_det):
            print('Failed to init NCS detector')
            return False

        if 1:
            self.nn_model = RecognizerCPU()
            if not self.nn_model.init(weights_rec, config_rec):
                print('Failed to init CPU recognizer')
                return False

        else:
            self.nn_model = RecognizerNCS()
            if not self.nn_model.init(1, graph_rec, config_rec):
                print('Failed to init CPU recognizer')
                return False

        return True

    def infer(self, cv_img, yolo_img=None, corr_img=None):
        # render_img = img.copy()
        ros_bboxes, yolo_bboxes = self.nn_det.infer(cv_img)

        yolo_result_img = cv_img.copy()

        if yolo_img is not None:
            self.nn_det.draw_boxes(yolo_img, ros_bboxes)

        corrected_ros_bboxes = []

        for ros_box in ros_bboxes:
            # print('Expected class: {}'.format(ros_box.Class))
            small_img = cv_img[ros_box.ymin:ros_box.ymax, ros_box.xmin:ros_box.xmax]
            result_cl = self.nn_model.infer(small_img)

            if result_cl == 'negative':
                continue

            ros_box.Class = result_cl

            corrected_ros_bboxes += [ros_box]
            # print('Predicted class: {}'.format(result_cl))

        if corr_img is not None:
            self.nn_det.draw_boxes(corr_img, corrected_ros_bboxes)

        return corrected_ros_bboxes

from wr8_ai.srv import ObjectDetection, ObjectDetectionResponse

class SignsDetector:

    def __init__(self):
        pass

    def init(self):
        rec_weights_path = rospy.get_param('~rec_weights_path')
        rec_config_path = rospy.get_param('~rec_config_path')
        rec_graph_path = rospy.get_param('~rec_graph_path')

        det_graph_path = rospy.get_param('~det_graph_path')
        det_config_path = rospy.get_param('~det_config_path')

        self.detector = DoubleDetector()
        if not self.detector.init(rec_config_path, rec_weights_path, rec_graph_path, det_config_path, det_graph_path):
            rospy.logerr('Failed to init detectors')
            self.initialized = False
            return False

        self.initialized = True

        self.bridge = CvBridge()
        self.image_sub = rospy.Subscriber("image", Image, self.callback_img, queue_size=1)
        # self.image_sub = rospy.Subscriber("nn_input_compr", CompressedImage, self.callback_comp_img, queue_size=1)

        self.image_pub_yolo = rospy.Publisher("nn/yolo/image/compressed", CompressedImage, queue_size=10)
        self.image_pub_yolo_corr = rospy.Publisher("nn/image/compressed", CompressedImage, queue_size=10)

        rospy.Service('nn/detect_signs', ObjectDetection, self.get_signs_srv_hndlr)

        self.cv_image = None
        self.cv_image_comp = None

        self.last_boxes_ros = None

        return True

    def callback_comp_img(self, data):
        np_arr = np.fromstring(data.data, np.uint8)
        self.cv_image_comp = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

        self.last_boxes_ros = self.get_signs_compr()

    def callback_img(self, data):
        try:
            self.cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as e:
            rospy.logwarn(e)

    def get_signs_srv_hndlr(self, req):
        resp = ObjectDetectionResponse()
        if self.last_boxes_ros is None:
            resp.bboxes = []
        else:
            resp.bboxes = self.last_boxes_ros
        return resp

    def update(self):
        if self.cv_image is None:
            self.last_boxes_ros = []
        else:
            self.last_boxes_ros = self.get_signs()

    def get_signs(self):
        if self.cv_image is None:
            return []

        # yolo_img = self.cv_image.copy()
        corr_yolo_img = self.cv_image.copy()

        ros_boxes = self.detector.infer(self.cv_image, None, corr_yolo_img)

        msg = CompressedImage()
        msg.header.stamp = rospy.Time.now()
        msg.format = "png"
        # msg.data = np.array(cv2.imencode('.png', yolo_img)[1]).tostring()
        # self.image_pub_yolo.publish(msg)

        msg.data = np.array(cv2.imencode('.png', corr_yolo_img)[1]).tostring()
        self.image_pub_yolo_corr.publish(msg)

        return ros_boxes

    def get_signs_compr(self):
        if self.cv_image_comp is None:
            return []

        yolo_img = self.cv_image_comp.copy()
        corr_yolo_img = self.cv_image_comp.copy()

        ros_boxes = self.detector.infer(self.cv_image_comp, yolo_img, corr_yolo_img)

        msg = CompressedImage()
        msg.header.stamp = rospy.Time.now()
        msg.format = "png"
        msg.data = np.array(cv2.imencode('.png', yolo_img)[1]).tostring()
        self.image_pub_yolo.publish(msg)

        msg.data = np.array(cv2.imencode('.png', corr_yolo_img)[1]).tostring()
        self.image_pub_yolo_corr.publish(msg)

        return ros_boxes

    def get_image(self):
        return self.cv_image

import cv2
import argparse
from wr8_ai.yolo import fps
import wr8_ai.nn_utils as nn_ut
import time

def full_test(config_rec, weights_rec, config_det, graph_det, input_path):
    print('Full test!')

    detector = DoubleDetector()
    if not detector.init(config_rec, weights_rec, config_det, graph_det):
        exit(1)

    data_generator = nn_ut.data_generator(input_path)
    fps_meter = fps.FPSMeter()

    for _, img in data_generator:
        start = time.time()

        yolo_img = img.copy()
        corr_yolo_img = img.copy()

        ros_boxes = detector.infer(img, yolo_img, corr_yolo_img)

        fps_meter.update(time.time() - start)
        if fps_meter.milliseconds > 5000:
            fps_meter.print_statistics()
            fps_meter.reset()

        cv2.imshow('result', yolo_img)
        cv2.imshow('corr_result', corr_yolo_img)

        if 27 == cv2.waitKey(0):
            break

        fps_meter.print_statistics()

if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='Test NCS NNs')
    argparser.add_argument('-m', '--mode', help='yolo/rec/reccpu')
    argparser.add_argument('-cr', '--configrec', help='path to configuration file')
    argparser.add_argument('-cd', '--configdet', help='path to configuration file')
    argparser.add_argument('-i', '--input', help='path to an image, a directory of images, a video, or webcam')
    argparser.add_argument('-g', '--graph', help='graph path')
    argparser.add_argument('-w', '--weights', help='weights path')

    args = argparser.parse_args()

    input_path = args.input
    graph_path = args.graph
    weights_path = args.weights
    config_rec_path = args.configrec
    config_det_path = args.configdet
    nn_mode = args.mode

    if nn_mode == 'yolo':
        nn_model = DetectorNCS()
        if not nn_model.init(0, graph_path, config_det_path):
            print('Failed to init NCS detector')
            exit(1)
    elif nn_mode == 'reccpu':
        nn_model = RecognizerCPU()
        if not nn_model.init(weights_path, config_rec_path):
            print('Failed to init CPU recognizer')
            exit(1)
    elif nn_mode == 'rec':
        nn_model = RecognizerNCS()
        if not nn_model.init(0, graph_path, config_rec_path):
            print('Failed to init NCS recognizer')
            exit(1)
    elif nn_mode == 'full':
        full_test(config_rec_path, weights_path, config_det_path, graph_path, input_path)
        exit(0)
    else:
        print('Invalid mode, check help!')

    data_generator = nn_ut.data_generator(input_path)
    fps_meter = fps.FPSMeter()

    for _, img in data_generator:
        start = time.time()

        if nn_mode == 'yolo':
            render_img = img.copy()
            result_bxs = nn_model.infer(img, render_img=render_img)
        else:
            result_cl = nn_model.infer(img)

        fps_meter.update(time.time() - start)
        if fps_meter.milliseconds > 5000:
            fps_meter.print_statistics()
            fps_meter.reset()

        if nn_mode == 'yolo':
            cv2.imshow('result', render_img)
        else:
            print(result_cl)
            cv2.imshow('orig', img)

        if 27 == cv2.waitKey(0):
            break

        fps_meter.print_statistics()