import cv2
import numpy as np
import os
import operator
from torch.utils import data
from collections import OrderedDict
from preprocess.data_preprocess import TrainAugmentation
from postprocess.visualization import cv_draw_bbox_3d
from datasets.dataset_utils import AspectRatioGroupedDataset
from datasets.dataset_utils import Label
from datasets.distributed_sampler import *
from utils.ParamList import ParamList
from Archs_3D.configs.mobilev2_retina3d_configs import CONFIGS

class KITTI(data.Dataset):
    def __init__(self, cfg, transform=None, is_training=True):
        super(KITTI).__init__()
        self.root = cfg.DATASET.PATH
        self.image_dir = os.path.join(self.root, "image_2")
        self.label_dir = os.path.join(self.root, "label_2")
        self.calib_dir = os.path.join(self.root, "calib")

        self.split = cfg.DATASET.TRAIN_SPLIT if is_training else cfg.DATASET.TEST_SPLIT
        self.is_training = is_training
        self.transform = transform
        self.K = np.array(cfg.DATASET.INTRINSIC, dtype=np.float32)
        self.K = self.K.reshape((3, 4))

        if self.split == "train":
            imageset_txt = os.path.join(self.root, "ImageSets", "train.txt")
        elif self.split == "val":
            imageset_txt = os.path.join(self.root, "ImageSets", "val.txt")
        elif self.split == "trainval":
            imageset_txt = os.path.join(self.root, "ImageSets", "trainval.txt")
        elif self.split == "test":
            imageset_txt = os.path.join(self.root, "ImageSets", "test.txt")
        else:
            raise ValueError("Invalid split!")

        image_files = []
        for line in open(imageset_txt, "r"):
            base_name = line.replace("\n", "")
            image_name = base_name + ".png"
            image_files.append(image_name)
        self.image_files = image_files
        self.label_files = [i.replace(".png", ".txt") for i in self.image_files]
        self.num_samples = len(self.image_files)
        self.classes = cfg.DATASET.DETECT_CLASSES

        self.label_map = OrderedDict({'{}'.format(c): i for i, c in enumerate(self.classes)})

    # should return (Input, Label)
    def __getitem__(self, index):

        img = cv2.imread(os.path.join(self.image_dir, self.image_files[index]))
        targets = ParamList(image_size=(img.shape[1], img.shape[0]), is_train=self.is_training)
        targets.add_field('K', np.copy(self.K))
        if not self.is_training:
            if self.transform is not None:
                img, targets = self.transform(img, targets=targets)
            return img, targets
        labels = self.parse_label(os.path.join(self.label_dir, self.label_files[index]))
        targets.add_field('class', labels[0])
        targets.add_field('bbox', labels[1])
        targets.add_field('dimension', labels[2])
        targets.add_field('alpha', labels[3])
        targets.add_field('Ry', labels[4])
        targets.add_field('mask', labels[5])

        if self.transform is not None:
            img, targets = self.transform(img, targets=targets)
        return img,  targets

    def __len__(self):
        return self.num_samples

    def parse_label(self, label_path):
        gt_classes = []
        gt_bboxes = []
        gt_dims = []
        gt_alphas = []
        gt_rays = []

        with open(label_path, 'r') as f:
            for line in f:
                line = line[:-1].split(' ')

                Class = line[0]
                if Class not in self.classes:
                    continue

                for i in range(1, len(line)):
                    line[i] = float(line[i])

                Box_2D = [float(line[4]), float(line[5]), float(line[6]), float(line[7])]
                gt_classes.append(self.label_map[Class])
                gt_bboxes.append(Box_2D)
                gt_dims.append([float(line[8]), float(line[9]), float(line[10])]) # H, W, L
                gt_alphas.append(float(line[3])) #alpha
                gt_rays.append(float(line[-1])) #ry
            gt_classes = np.array(gt_classes, dtype=np.float32)
            gt_bboxes = np.array(gt_bboxes, dtype=np.float32)
            gt_dims = np.array(gt_dims, dtype=np.float32)
            gt_alphas = np.array(gt_alphas, dtype=np.float32)
            reg_masks = np.zeros((gt_classes.shape[0],), dtype=np.float32) + 0.
            gt_rays = np.array(gt_rays, dtype=np.float32)

        return gt_classes, gt_bboxes, gt_dims, gt_alphas, gt_rays, reg_masks

    def calc_mean_of_dataset(self):
        rgb_mean = np.zeros((3,), dtype=np.float64)
        num = self.num_samples
        for i, file in enumerate(self.image_files):
            img = cv2.imread(os.path.join(self.image_dir, file))
            rgb_mean += np.mean(img, axis=(0, 1))
            print("progress : {:.2%}".format(i/num))
        print('rgb mean: ', rgb_mean/num)

from postprocess.postprocessing import calc_regressed_bbox_3d
from postprocess.visualization import KITTI_COLOR_MAP
if __name__ == '__main__':

    dataset = KITTI(CONFIGS,transform=TrainAugmentation(CONFIGS.INTENSOR_SIZE, mean=CONFIGS.DATASET.MEAN), is_training=True)
    # dataset.calc_mean_of_dataset()
    sampler = TrainingSampler(len(dataset))
    params = {'sampler': sampler,
              'batch_sampler': None,
              'collate_fn': operator.itemgetter(0),  # don't batch, but yield individual elements
              'num_workers': 1}

    generator = data.DataLoader(dataset, **params)
    generator = AspectRatioGroupedDataset(generator, 1)

    for d in generator:
        # print(d)
        image = d[0][0]
        # image = image.permute(1,2,0).numpy().astype(np.uint8)
        img = image.astype(np.uint8)
        h, w, _ = img.shape
        targets = d[0][1]
        classes = targets.get_field('class').cpu().numpy().astype(np.long)
        bboxes = targets.get_field('bbox').cpu().numpy()
        alphas = targets.get_field('alpha').cpu().numpy()
        K = targets.get_field('K').cpu().numpy()
        dims = targets.get_field('dimension').cpu().numpy()
        centers = np.concatenate([(bboxes[:, None, 0] + bboxes[:, None, 2]) * .5,
                                 (bboxes[:, None, 1] + bboxes[:, None, 3]) * .5], axis=1)
        rays = np.arctan((centers[:, 0] - K[0, 2])/K[0, 0])
        num = alphas.shape[0]
        K[:, 3] = 0.
        for i in range(num):
            loc, _ = calc_regressed_bbox_3d(alphas[i], theta_ray=rays[i], dimension=dims[i], bboxes=bboxes[i], proj_matrix=K)
            proj_matrix = np.copy(K)
            proj_matrix[0, :] *= w
            proj_matrix[1, :] *= h
            cv_draw_bbox_3d(img, proj_matrix,
                        alphas[i] + rays[i], dims[i], loc, classes[i], 1., KITTI_COLOR_MAP[classes[i]])
        # img = cv_draw_bboxes_2d(img, bboxes, classes, dataset.classes)
        cv2.imshow("preprocessed image", img)
        cv2.waitKey(1000)


import numpy as np
import torch.nn as nn
import torch

from .anchor_head_template import AnchorHeadTemplate

class AnchorHeadSingle(AnchorHeadTemplate):
    def __init__(self, model_cfg, input_channels, num_class, class_names, grid_size, point_cloud_range,
                 predict_boxes_when_training=True):
        super().__init__(
            model_cfg=model_cfg, num_class=num_class, class_names=class_names, grid_size=grid_size, point_cloud_range=point_cloud_range,
            predict_boxes_when_training=predict_boxes_when_training
        )

        self.num_anchors_per_location = sum(self.num_anchors_per_location)

        self.num_spatial_features_before_head = input_channels

        self.conv_cls = nn.Conv2d(
            input_channels, self.num_anchors_per_location * self.num_class,
            kernel_size=1
        )
        self.conv_box = nn.Conv2d(
            input_channels, self.num_anchors_per_location * self.box_coder.code_size,
            kernel_size=1
        )

        if self.model_cfg.get('USE_DIRECTION_CLASSIFIER', None) is not None:
            self.conv_dir_cls = nn.Conv2d(
                input_channels,
                self.num_anchors_per_location * self.model_cfg.NUM_DIR_BINS,
                kernel_size=1
            )
        else:
            self.conv_dir_cls = None
        self.init_weights()

    def init_weights(self):
        pi = 0.01
        nn.init.constant_(self.conv_cls.bias, -np.log((1 - pi) / pi))
        nn.init.normal_(self.conv_box.weight, mean=0, std=0.001)

    def forward(self, data_dict):
        spatial_features_2d = data_dict['spatial_features_2d']

        # 保存下来用于后面插值
        data_dict.update({'spatial_features_before_head': spatial_features_2d})

        # log something for visualization
        # 只可视化第0个batch
        vis_dict = {}
        vis_dict['vis_spatial_features'] = data_dict['spatial_features'][0].detach().clone()
        vis_dict['vis_res_info'] = spatial_features_2d[0].detach().clone() # 这个是可视化head之前所用的特征图

        cls_preds = self.conv_cls(spatial_features_2d)
        box_preds = self.conv_box(spatial_features_2d)

        cls_preds_split = torch.split(torch.sigmoid(cls_preds[0].detach().clone()), self.num_class, dim=0) # 这个是用于可视化对cls的预测
        cls_pred_decode_list = []
        for cls_pred in cls_preds_split:
            cls_pred_decode_list.append( torch.max(cls_pred, dim=0, keepdim=True)[0] )
        cls_pred_decode = torch.cat(cls_pred_decode_list, dim=0)
        vis_dict['vis_cls_pred'] = cls_pred_decode

        self.forward_ret_dict['vis_dict'] = vis_dict

        cls_preds = cls_preds.permute(0, 2, 3, 1).contiguous()  # [N, H, W, C]
        box_preds = box_preds.permute(0, 2, 3, 1).contiguous()  # [N, H, W, C]

        self.forward_ret_dict['cls_preds'] = cls_preds
        self.forward_ret_dict['box_preds'] = box_preds

        if self.conv_dir_cls is not None:
            dir_cls_preds = self.conv_dir_cls(spatial_features_2d)
            dir_cls_preds = dir_cls_preds.permute(0, 2, 3, 1).contiguous()
            self.forward_ret_dict['dir_cls_preds'] = dir_cls_preds
        else:
            dir_cls_preds = None

        if self.training:
            targets_dict = self.assign_targets(
                gt_boxes=data_dict['gt_boxes']
            )
            self.forward_ret_dict.update(targets_dict)

        if not self.training or self.predict_boxes_when_training:
            batch_cls_preds, batch_box_preds = self.generate_predicted_boxes(
                batch_size=data_dict['batch_size'],
                cls_preds=cls_preds, box_preds=box_preds, dir_cls_preds=dir_cls_preds
            )
            data_dict['batch_cls_preds'] = batch_cls_preds
            data_dict['batch_box_preds'] = batch_box_preds
            data_dict['cls_preds_normalized'] = False

        return data_dict


# %%
import os
import cv2
import torch
import numpy as np
import xml.etree.ElementTree as ET
import torch.utils.data as data
from ssd_utils import boxes_visual

VOC_CLASSES = (  # always index 0
    'aeroplane', 'bicycle', 'bird', 'boat',
    'bottle', 'bus', 'car', 'cat', 'chair',
    'cow', 'diningtable', 'dog', 'horse',
    'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor')

def get_class_by_id(ids):
    class_names = []
    for idx in range(len(ids)):
        class_names.append(VOC_CLASSES[ids[idx]])
    return class_names

class VOCDataset(data.Dataset):
    """ voc dataset defination
    Args:
        mode: 'train' or 'eval' mode
        transform: tranform used for raw image augmentation
        dataset_dir: dataset directory
        datasets: list with dataset name as 'VOC2007', 'VOC2012'
        include_difficult: whether or not include difficult detect target
        class_to_ind: dict of class to index map
    """

    def __init__(self, dataset_dir, transform=None, mode='train', datasets=['VOC2007', 'VOC2012'],
                 include_difficult=True, class_to_ind=None):
        self.dataset_dir = dataset_dir
        self.datasets = datasets
        self.transform = transform
        self.mode = mode
        self.image_info = []
        self.name = 'VOC'
        self.class_to_ind = class_to_ind or dict(
            zip(VOC_CLASSES, range(len(VOC_CLASSES))))
        for dataset in self.datasets:
            path = os.path.join(dataset_dir, 'VOCdevkit', dataset)
            trainval = open(os.path.join(
                path, 'ImageSets/Main', 'trainval.txt'))
            for line in trainval:
                id = line.strip()
                self.image_info.append(
                    {'id': id,
                     'path': os.path.join(path, 'JPEGImages/{}.jpg'.format(id)),
                     'annotation': os.path.join(path, 'Annotations/{}.xml'.format(id))})

    def __len__(self):
        return len(self.image_info)

    def __getitem__(self, index):
        _, image, gt_boxes, _, _ = self.get_data(index)
        return image, gt_boxes

    def get_data(self, index):
        """ Get data by index\n
        Args:
            index: index of dataset
        Returns:
            id: identifer of a image
            image: image data with formation NWH
            gt_boxes: [-1,xmin,ymin, xmax, ymax, label]
        """
        info = self.image_info[index]
        image_id = info['id']
        image = cv2.imread(info['path'])
        image = image[:, :, ::-1]
        tree = ET.ElementTree(file=info['annotation'])
        size = tree.find('size')
        width = int(size.find('width').text)
        height = int(size.find('height').text)
        depth = int(size.find('depth').text)
        bboxes = []
        labels = []
        for elem in tree.iter('object'):

            #     box=[]
            bndbox = elem.find('bndbox')
            pts = ['xmin', 'ymin', 'xmax', 'ymax']
            for idx, pt in enumerate(pts):
                bboxes.append((int(bndbox.find(pt).text)-1) /
                              (width if idx % 2 == 0 else height))
            labels.append(self.class_to_ind[elem.find('name').text])
        bboxes = np.array(bboxes, dtype=np.float32).reshape(-1, 4)
        labels = np.array(labels).reshape(-1, 1)
        if self.transform:
            image, bboxes, labels = self.transform(image, bboxes, labels)
        bboxes = np.concatenate([bboxes, labels], axis=1)
        return image_id, torch.as_tensor(image.copy()).permute(2, 0, 1), bboxes, width, height

    def get_image(self, index):
        info = self.image_info[index]
        image_id = info['id']
        image = cv2.imread(info['path'])
        image = image[:, :, ::-1]
        # tree = ET.ElementTree(file=info['annotation'])
        # size = tree.find('size')
        # width = int(size.find('width').text)
        # height = int(size.find('height').text)
        # depth = int(size.find('depth').text)
        return image

    def __str__(self):
        return 'dict of "id","path","annotation"'\
            ' elements\nlen = {}'.format(len(self.image_info))

if __name__ == "__main__":
    HOME = os.path.expanduser('~')
    DATA_DIR = HOME+'/dataset'
    dataset = VOCDataset(DATA_DIR)
    data = dataset[0]
    print(data)


## adapted from https://github.com/rail-berkeley/softlearning/blob/master/softlearning/algorithms/sac.py

import os
import math
import pickle
from collections import OrderedDict
from collections import defaultdict
from numbers import Number
from itertools import count
import gtimer as gt
import pdb
import random
import sys
import warnings

import numpy as np
import tensorflow as tf
from tensorflow.python.training import training_util

from softlearning.algorithms.rl_algorithm import RLAlgorithm
from softlearning.replay_pools.simple_replay_pool import SimpleReplayPool
from softlearning.replay_pools.mjc_state_replay_pool import MjcStateReplayPool
from softlearning.replay_pools.modelbuffer import ModelBuffer
from softlearning.replay_pools.cpobuffer import CPOBuffer
from softlearning.samplers.model_sampler import ModelSampler
from softlearning.policies.safe_utils.logx import EpochLogger
from softlearning.policies.safe_utils.utils import EPS

from mbpo.models.constructor import construct_model, format_samples_for_dyn, reset_model
from mbpo.models.fake_env import FakeEnv
from mbpo.models.perturbed_env import PerturbedEnv
from mbpo.utils.writer import Writer
from mbpo.utils.visualization import visualize_policy
from mbpo.utils.logging import Progress, update_dict
import mbpo.utils.filesystem as filesystem

def td_target(reward, discount, next_value):
    return reward + discount * next_value

class CMBPO(RLAlgorithm):
    """Model-Based Policy Optimization (MBPO)

    References
    ----------
        Michael Janner, Justin Fu, Marvin Zhang, Sergey Levine.
        When to Trust Your Model: Model-Based Policy Optimization.
        arXiv preprint arXiv:1906.08253. 2019.
    """

    def __init__(
            self,
            training_environment,
            evaluation_environment,
            mjc_model_environment,
            policy,
            Qs,
            pool,
            static_fns,
            plotter=None,
            tf_summaries=False,
            n_env_interacts = 1e7,
            lr=3e-4,
            reward_scale=1.0,
            target_entropy='auto',
            discount=0.99,
            tau=5e-3,
            target_update_interval=1,
            action_prior='uniform',
            reparameterize=False,
            store_extra_policy_info=False,
            eval_every_n_steps=5e3,

            deterministic=False,
            model_train_freq=250,
            num_networks=7,
            num_elites=5,
            model_retain_epochs=20,
            rollout_batch_size=100e3,
            real_ratio=0.1,
            dyn_model_train_schedule=[20, 100, 1, 5],
            cost_model_train_schedule=[20, 100, 1, 30],
            cares_about_cost = False,
            policy_alpha = 1,
            max_uncertainty_rew = None,
            max_uncertainty_c = None,
            rollout_mode = 'schedule',
            rollout_schedule=[20,100,1,1],
            maxroll = 80,
            max_tddyn_err = 1e-5,
            max_tddyn_err_decay = .995,
            min_real_samples_per_epoch = 1000,
            batch_size_policy = 5000,
            hidden_dims=(200, 200, 200, 200),
            max_model_t=None,

            use_mjc_state_model = False,
            model_std_inc = 0.02,
            **kwargs,
    ):
        """
        Args:
            env (`SoftlearningEnv`): Environment used for training.
            policy: A policy function approximator.
            initial_exploration_policy: ('Policy'): A policy that we use
                for initial exploration which is not trained by the algorithm.
            Qs: Q-function approximators. The min of these
                approximators will be used. Usage of at least two Q-functions
                improves performance by reducing overestimation bias.
            pool (`PoolBase`): Replay pool to add gathered samples to.
            plotter (`QFPolicyPlotter`): Plotter instance to be used for
                visualizing Q-function during training.
            lr (`float`): Learning rate used for the function approximators.
            discount (`float`): Discount factor for Q-function updates.
            tau (`float`): Soft value function target update weight.
            target_update_interval ('int'): Frequency at which target network
                updates occur in iterations.
            reparameterize ('bool'): If True, we use a gradient estimator for
                the policy derived using the reparameterization trick. We use
                a likelihood ratio based estimator otherwise.
        """

        super(CMBPO, self).__init__(**kwargs)

        self.obs_space = training_environment.observation_space
        self.act_space = training_environment.action_space
        self.obs_dim = np.prod(training_environment.observation_space.shape)

        self.act_dim = np.prod(training_environment.action_space.shape)
        self.n_env_interacts = n_env_interacts
        #### determine unstacked obs dim
        self.num_stacks = training_environment.stacks
        self.stacking_axis = training_environment.stacking_axis
        self.active_obs_dim = int(self.obs_dim/self.num_stacks)

        self.policy_alpha = policy_alpha
        self.cares_about_cost = cares_about_cost

        ## create fake environment for model
        self.fake_env = FakeEnv(training_environment,
                                    static_fns, num_networks=7,
                                    num_elites=3,
                                    hidden_dims=hidden_dims,
                                    cares_about_cost=cares_about_cost,
                                    session = self._session)

        self.use_mjc_state_model = use_mjc_state_model

        self._rollout_schedule = rollout_schedule
        self._max_model_t = max_model_t

        self._model_retain_epochs = model_retain_epochs
        self.eval_every_n_steps= eval_every_n_steps

        self._dyn_model_train_schedule = dyn_model_train_schedule
        self._cost_model_train_schedule = cost_model_train_schedule

        self._dyn_model_train_freq = 1
        self._cost_model_train_freq = 1
        self._rollout_batch_size = int(rollout_batch_size)
        self._max_uncertainty_rew = max_uncertainty_rew
        self._max_uncertainty_c = max_uncertainty_c
        self._deterministic = deterministic
        self._real_ratio = real_ratio

        self._log_dir = os.getcwd()
        self._writer = Writer(self._log_dir)

        self._training_environment = training_environment
        self._evaluation_environment = evaluation_environment
        self._mjc_model_environment = mjc_model_environment
        self.perturbed_env = PerturbedEnv(self._mjc_model_environment, std_inc=model_std_inc)
        self._policy = policy
        self._initial_exploration_policy = policy   #overwriting initial _exploration policy, not implemented for cpo yet

        self._pool = pool
        self._plotter = plotter
        self._tf_summaries = tf_summaries

        # set up pool
        pi_info_shapes = {k: v.shape.as_list()[1:] for k,v in self._policy.pi_info_phs.items()}
        self._pool.initialize(pi_info_shapes,
                                gamma = self._policy.gamma,
                                lam = self._policy.lam,
                                cost_gamma = self._policy.cost_gamma,
                                cost_lam = self._policy.cost_lam)

        self._policy_lr = lr

        self._reward_scale = reward_scale
        self._target_entropy = (
            -np.prod(self._training_environment.action_space.shape)
            if target_entropy == 'auto'
            else target_entropy)
        print('[ MBPO ] Target entropy: {}'.format(self._target_entropy))

        self._discount = discount
        self._tau = tau
        self._target_update_interval = target_update_interval
        self._action_prior = action_prior

        self._reparameterize = reparameterize
        self._store_extra_policy_info = store_extra_policy_info

        observation_shape = self._training_environment.active_observation_shape
        action_shape = self._training_environment.action_space.shape

        assert len(observation_shape) == 1, observation_shape
        self._observation_shape = observation_shape
        assert len(action_shape) == 1, action_shape
        self._action_shape = action_shape

        ### model sampler and buffer
        self.rollout_mode = rollout_mode
        self.max_tddyn_err = max_tddyn_err
        self.max_tddyn_err_decay = max_tddyn_err_decay
        self.min_real_samples_per_epoch = min_real_samples_per_epoch
        self.batch_size_policy = batch_size_policy

        self.model_pool = ModelBuffer(batch_size=self._rollout_batch_size,
                                        max_path_length=maxroll,
                                        env = self.fake_env,
                                        ensemble_size=num_networks,
                                        rollout_mode = self.rollout_mode,
                                        cares_about_cost = cares_about_cost,
                                        max_uncertainty_c = self._max_uncertainty_c,
                                        max_uncertainty_r = self._max_uncertainty_rew,
                                        )
        self.model_pool.initialize(pi_info_shapes,
                                    gamma = self._policy.gamma,
                                    lam = self._policy.lam,
                                    cost_gamma = self._policy.cost_gamma,
                                    cost_lam = self._policy.cost_lam,
                                    )
        #@anyboby debug
        self.model_sampler = ModelSampler(max_path_length=maxroll,
                                            batch_size=self._rollout_batch_size,
                                            store_last_n_paths=10,
                                            cares_about_cost = cares_about_cost,
                                            max_uncertainty_c = self._max_uncertainty_c,
                                            max_uncertainty_r = self._max_uncertainty_rew,
                                            logger=None,
                                            rollout_mode = self.rollout_mode,
                                            )

        # provide policy and sampler with the same logger
        self.logger = EpochLogger()
        self._policy.set_logger(self.logger)
        self.sampler.set_logger(self.logger)
        #self.model_sampler.set_logger(self.logger)

    def _train(self):

        """Return a generator that performs RL training.

        Args:
            env (`SoftlearningEnv`): Environment used for training.
            policy (`Policy`): Policy used for training
            initial_exploration_policy ('Policy'): Policy used for exploration
                If None, then all exploration is done using policy
            pool (`PoolBase`): Sample pool to add samples to
        """

        #### pool is e.g. simple_replay_pool
        training_environment = self._training_environment
        evaluation_environment = self._evaluation_environment
        policy = self._policy
        pool = self._pool

        if not self._training_started:
            #### perform some initial steps (gather samples) using initial policy
            ######  fills pool with _n_initial_exploration_steps samples
            self._initial_exploration_hook(
                training_environment, self._policy, pool)

        #### set up sampler with train env and actual policy (may be different from initial exploration policy)
        ######## note: sampler is set up with the pool that may be already filled from initial exploration hook
        self.sampler.initialize(training_environment, policy, pool)
        self.model_sampler.initialize(self.fake_env, policy, self.model_pool)
        rollout_dkl_lim = self.model_sampler.compute_dynamics_dkl(obs_batch=self._pool.rand_batch_from_archive(5000, fields=['observations'])['observations'], depth=self._rollout_schedule[2])
        self.model_sampler.set_rollout_dkl(rollout_dkl_lim)
        self.initial_model_dkl = self.model_sampler.dyn_dkl
        #### reset gtimer (for coverage of project development)
        gt.reset_root()
        gt.rename_root('RLAlgorithm')
        gt.set_def_unique(False)
        self.policy_epoch = 0       ### count policy updates
        self.new_real_samples = 0
        self.last_eval_step = 0
        self.diag_counter = 0
        running_diag = {}
        self.approx_model_batch = self.batch_size_policy-self.min_real_samples_per_epoch    ### some size to start off

        #### not implemented, could train policy before hook
        self._training_before_hook()

        #### iterate over epochs, gt.timed_for to create loop with gt timestamps
        for self._epoch in gt.timed_for(range(self._epoch, self._n_epochs)):

            #### do something at beginning of epoch (in this case reset self._train_steps_this_epoch=0)
            self._epoch_before_hook()
            gt.stamp('epoch_before_hook')

            #### util class Progress, e.g. for plotting a progress bar
            #######   note: sampler may already contain samples in its pool from initial_exploration_hook or previous epochs
            self._training_progress = Progress(self._epoch_length * self._n_train_repeat/self._train_every_n_steps)

            samples_added = 0
            #=====================================================================#
            #            Rollout model                                            #
            #=====================================================================#
            model_samples = None
            keep_rolling = True
            model_metrics = {}
            #### start model rollout
            if self._real_ratio<1.0: #if self._timestep % self._model_train_freq == 0 and self._real_ratio < 1.0:
                #=====================================================================#
                #                           Model Rollouts                            #
                #=====================================================================#
                if self.rollout_mode == 'schedule':
                    self._set_rollout_length()

                while keep_rolling:
                    ep_b = self._pool.epoch_batch(batch_size=self._rollout_batch_size, epochs=self._pool.epochs_list, fields=['observations','pi_infos'])
                    kls = np.clip(self._policy.compute_DKL(ep_b['observations'], ep_b['mu'], ep_b['log_std']), a_min=0, a_max=None)
                    btz_dist = self._pool.boltz_dist(kls, alpha=self.policy_alpha)
                    btz_b = self._pool.distributed_batch_from_archive(self._rollout_batch_size, btz_dist, fields=['observations','pi_infos'])
                    start_states, mus, logstds = btz_b['observations'], btz_b['mu'], btz_b['log_std']
                    btz_kl = np.clip(self._policy.compute_DKL(start_states, mus, logstds), a_min=0, a_max=None)

                    self.model_sampler.reset(start_states)
                    if self.rollout_mode=='uncertainty':
                        self.model_sampler.set_max_uncertainty(self.max_tddyn_err)

                    for i in count():
                        # print(f'Model Sampling step Nr. {i+1}')

                        _,_,_,info = self.model_sampler.sample(max_samples=int(self.approx_model_batch-samples_added))

                        if self.model_sampler._total_samples + samples_added >= .99*self.approx_model_batch:
                            keep_rolling = False
                            break

                        if info['alive_ratio']<= 0.1: break

                    ### diagnostics for rollout ###
                    rollout_diagnostics = self.model_sampler.finish_all_paths()
                    if self.rollout_mode == 'iv_gae':
                        keep_rolling = self.model_pool.size + samples_added <= .99*self.approx_model_batch

                    ######################################################################
                    ### get model_samples, get() invokes the inverse variance rollouts ###
                    model_samples_new, buffer_diagnostics_new = self.model_pool.get()
                    model_samples = [np.concatenate((o,n), axis=0) for o,n in zip(model_samples, model_samples_new)] if model_samples else model_samples_new

                    ######################################################################
                    ### diagnostics
                    new_n_samples = len(model_samples_new[0])+EPS
                    diag_weight_old = samples_added/(new_n_samples+samples_added)
                    diag_weight_new = new_n_samples/(new_n_samples+samples_added)
                    model_metrics = update_dict(model_metrics, rollout_diagnostics, weight_a= diag_weight_old,weight_b=diag_weight_new)
                    model_metrics = update_dict(model_metrics, buffer_diagnostics_new,  weight_a= diag_weight_old,weight_b=diag_weight_new)
                    ### run diagnostics on model data
                    if buffer_diagnostics_new['poolm_batch_size']>0:
                        model_data_diag = self._policy.run_diagnostics(model_samples_new)
                        model_data_diag = {k+'_m':v for k,v in model_data_diag.items()}
                        model_metrics = update_dict(model_metrics, model_data_diag, weight_a= diag_weight_old,weight_b=diag_weight_new)

                    samples_added += new_n_samples
                    model_metrics.update({'samples_added':samples_added})
                    ######################################################################

                ## for debugging
                model_metrics.update({'cached_var':np.mean(self.fake_env._model.scaler_out.cached_var)})
                model_metrics.update({'cached_mu':np.mean(self.fake_env._model.scaler_out.cached_mu)})

                print(f'Rollouts finished')
                gt.stamp('epoch_rollout_model')

            #=====================================================================#
            #  Sample                                                             #
            #=====================================================================#
            n_real_samples = self.model_sampler.dyn_dkl/self.initial_model_dkl * self.min_real_samples_per_epoch
            n_real_samples = max(n_real_samples, 1000)
            # n_real_samples = self.min_real_samples_per_epoch ### for ablation

            model_metrics.update({'n_real_samples':n_real_samples})
            start_samples = self.sampler._total_samples
            ### train for epoch_length ###
            for i in count():

                #### _timestep is within an epoch
                samples_now = self.sampler._total_samples
                self._timestep = samples_now - start_samples

                #### not implemented atm
                self._timestep_before_hook()
                gt.stamp('timestep_before_hook')

                ##### śampling from the real world ! #####
                _,_, _, _ = self._do_sampling(timestep=self.policy_epoch)
                gt.stamp('sample')

                self._timestep_after_hook()
                gt.stamp('timestep_after_hook')

                if self.ready_to_train or self._timestep>n_real_samples:
                    self.sampler.finish_all_paths(append_val=True, append_cval=True, reset_path=False)
                    self.new_real_samples += self._timestep
                    break

            #=====================================================================#
            #  Train model                                                        #
            #=====================================================================#
            if self.new_real_samples>2048 and self._real_ratio<1.0:
                model_diag = self.train_model(min_epochs=1, max_epochs=10)
                self.new_real_samples = 0
                model_metrics.update(model_diag)

            #=====================================================================#
            #  Get Buffer Data                                                    #
            #=====================================================================#
            real_samples, buf_diag = self._pool.get()

            ### run diagnostics on real data
            policy_diag = self._policy.run_diagnostics(real_samples)
            policy_diag = {k+'_r':v for k,v in policy_diag.items()}
            model_metrics.update(policy_diag)
            model_metrics.update(buf_diag)

            #=====================================================================#
            #  Update Policy                                                      #
            #=====================================================================#
            train_samples = [np.concatenate((r,m), axis=0) for r,m in zip(real_samples, model_samples)] if model_samples else real_samples
            self._policy.update_real_c(real_samples)
            self._policy.update_policy(train_samples)
            self._policy.update_critic(train_samples, train_vc=(train_samples[-3]>0).any())    ### only train vc if there are any costs

            if self._real_ratio<1.0:
                self.approx_model_batch = self.batch_size_policy-n_real_samples #self.model_sampler.dyn_dkl/self.initial_model_dkl * self.min_real_samples_per_epoch

            self.policy_epoch += 1
            self.max_tddyn_err *= self.max_tddyn_err_decay
            #### log policy diagnostics
            self._policy.log()

            gt.stamp('train')
            #=====================================================================#
            #  Log performance and stats                                          #
            #=====================================================================#

            self.sampler.log()
            # write results to file, ray prints for us, so no need to print from logger
            logger_diagnostics = self.logger.dump_tabular(output_dir=self._log_dir, print_out=False)
            #=====================================================================#

            if self._total_timestep // self.eval_every_n_steps > self.last_eval_step:
                evaluation_paths = self._evaluation_paths(
                    policy, evaluation_environment)
                gt.stamp('evaluation_paths')

                self.last_eval_step = self._total_timestep // self.eval_every_n_steps
            else:
                evaluation_paths = []

            if evaluation_paths:
                evaluation_metrics = self._evaluate_rollouts(
                    evaluation_paths, evaluation_environment)
                gt.stamp('evaluation_metrics')
                diag_obs_batch = np.concatenate(([evaluation_paths[i]['observations'] for i in range(len(evaluation_paths))]), axis=0)
            else:
                evaluation_metrics = {}
                diag_obs_batch = []

            gt.stamp('epoch_after_hook')

            new_diagnostics = {}

            time_diagnostics = gt.get_times().stamps.itrs

            # add diagnostics from logger
            new_diagnostics.update(logger_diagnostics)

            new_diagnostics.update(OrderedDict((
                *(
                    (f'evaluation/{key}', evaluation_metrics[key])
                    for key in sorted(evaluation_metrics.keys())
                ),
                *(
                    (f'times/{key}', time_diagnostics[key][-1])
                    for key in sorted(time_diagnostics.keys())
                ),
                *(
                    (f'model/{key}', model_metrics[key])
                    for key in sorted(model_metrics.keys())
                ),
            )))

            if self._eval_render_mode is not None and hasattr(
                    evaluation_environment, 'render_rollouts'):
                training_environment.render_rollouts(evaluation_paths)

            #### updateing and averaging
            old_ts_diag = running_diag.get('timestep', 0)
            new_ts_diag = self._total_timestep-self.diag_counter-old_ts_diag
            w_olddiag = old_ts_diag/(new_ts_diag+old_ts_diag)
            w_newdiag = new_ts_diag/(new_ts_diag+old_ts_diag)
            running_diag = update_dict(running_diag, new_diagnostics, weight_a=w_olddiag, weight_b=w_newdiag)
            running_diag.update({'timestep':new_ts_diag + old_ts_diag})
            ####

            if new_ts_diag + old_ts_diag > self.eval_every_n_steps:
                running_diag.update({
                    'epoch':self._epoch,
                    'timesteps_total':self._total_timestep,
                    'train-steps':self._num_train_steps,
                })
                self.diag_counter = self._total_timestep
                diag = running_diag.copy()
                running_diag = {}
                yield diag

            if self._total_timestep >= self.n_env_interacts:
                self.sampler.terminate()

                self._training_after_hook()

                self._training_progress.close()
                print("###### DONE ######")
                yield {'done': True, **running_diag}

                break

    def train(self, *args, **kwargs):
        return self._train(*args, **kwargs)

    def _initial_exploration_hook(self, env, initial_exploration_policy, pool):
        if self._n_initial_exploration_steps < 1: return

        if not initial_exploration_policy:
            raise ValueError(
                "Initial exploration policy must be provided when"
                " n_initial_exploration_steps > 0.")

        self.sampler.initialize(env, initial_exploration_policy, pool)
        while True:
            self.sampler.sample(timestep=0)
            if self.sampler._total_samples >= self._n_initial_exploration_steps:
                self.sampler.finish_all_paths(append_val=True, append_cval=True, reset_path=False)
                pool.get()  # moves policy samples to archive
                break

        ### train model
        if self._real_ratio<1.0:
            self.train_model(min_epochs=150, max_epochs=500)

    def train_model(self, min_epochs=5, max_epochs=100, batch_size=2048):
        self._dyn_model_train_freq = self._set_model_train_freq(
            self._dyn_model_train_freq,
            self._dyn_model_train_schedule
            ) ## set current train freq.
        self._cost_model_train_freq = self._set_model_train_freq(
            self._cost_model_train_freq,
            self._cost_model_train_schedule
            )

        print('[ MBPO ] log_dir: {} | ratio: {}'.format(self._log_dir, self._real_ratio))
        print('[ MBPO ] Training model at epoch {} | freq {} | timestep {} (total: {}) (total: {})'.format(
            self._epoch, self._dyn_model_train_freq, self._timestep, self._total_timestep, self._num_train_steps)
        )

        model_samples = self._pool.get_archive(['observations',
                                                'actions',
                                                'next_observations',
                                                'rewards',
                                                'costs',
                                                'terminals',
                                                'epochs',
                                                ])

        if self._epoch%self._dyn_model_train_freq==0:
            diag_dyn = self.fake_env.train_dyn_model(
                model_samples,
                batch_size=batch_size, #512
                max_epochs=max_epochs, # max_epochs
                min_epoch_before_break=min_epochs, # min_epochs,
                holdout_ratio=0.2,
                max_t=self._max_model_t
                )

        if self._epoch%self._cost_model_train_freq==0 and self.fake_env.cares_about_cost:
            diag_c = self.fake_env.train_cost_model(
                model_samples,
                batch_size= batch_size, #batch_size, #512,
                min_epoch_before_break= min_epochs,#min_epochs,
                max_epochs=max_epochs, # max_epochs,
                holdout_ratio=0.2,
                max_t=self._max_model_t
                )
            diag_dyn.update(diag_c)

        return diag_dyn

    @property
    def _total_timestep(self):
        total_timestep = self.sampler._total_samples
        return total_timestep

    def _set_rollout_length(self):
        min_epoch, max_epoch, min_length, max_length = self._rollout_schedule
        if self._epoch <= min_epoch:
            y = min_length
        else:
            dx = (self._epoch - min_epoch) / (max_epoch - min_epoch)
            dx = min(dx, 1)
            y = dx * (max_length - min_length) + min_length

        self._rollout_length = int(y)
        self.model_sampler.set_max_path_length(self._rollout_length)
        print('[ Model Length ] Epoch: {} (min: {}, max: {}) | Length: {} (min: {} , max: {})'.format(
            self._epoch, min_epoch, max_epoch, self._rollout_length, min_length, max_length
        ))

    def _do_sampling(self, timestep):
        return self.sampler.sample(timestep = timestep)

    def _set_model_train_freq(self, var, schedule):
        min_epoch, max_epoch, min_freq, max_freq = schedule
        if self._epoch <= min_epoch:
            y = min_freq
        else:
            dx = (self._epoch - min_epoch) / (max_epoch - min_epoch)
            dx = min(dx, 1)
            y = dx * (max_freq - min_freq) + min_freq

        var = int(y)
        print('[ Model Train Frequency ] Epoch: {} (min: {}, max: {}) | Frequency: {} (min: {} , max: {})'.format(
            self._epoch, min_epoch, max_epoch, var, min_freq, max_freq
        ))
        return var

    def _evaluate_rollouts(self, paths, env):
        """Compute evaluation metrics for the given rollouts."""

        total_returns = [path['rewards'].sum() for path in paths]
        episode_lengths = [len(p['rewards']) for p in paths]
        total_cost = [path['cost'].sum() for path in paths]
        diagnostics = OrderedDict((
            ('return-average', np.mean(total_returns)),
            ('return-min', np.min(total_returns)),
            ('return-max', np.max(total_returns)),
            ('return-std', np.std(total_returns)),
            ('episode-length-avg', np.mean(episode_lengths)),
            ('episode-length-min', np.min(episode_lengths)),
            ('episode-length-max', np.max(episode_lengths)),
            ('episode-length-std', np.std(episode_lengths)),
            ('creturn-average', np.mean(total_cost)),
            ('creturn-fullep-average', np.mean(total_cost)/np.mean(episode_lengths)*self.sampler.max_path_length),
            ('creturn-min', np.min(total_cost)),
            ('creturn-max', np.max(total_cost)),
            ('creturn-std', np.std(total_cost)),
        ))

        env_infos = env.get_path_infos(paths)
        for key, value in env_infos.items():
            diagnostics[f'env_infos/{key}'] = value

        return diagnostics

    def _visualize_model(self, env, timestep):
        ## save env state
        state = env.unwrapped.state_vector()
        qpos_dim = len(env.unwrapped.sim.data.qpos)
        qpos = state[:qpos_dim]
        qvel = state[qpos_dim:]

        print('[ Visualization ] Starting | Epoch {} | Log dir: {}\n'.format(self._epoch, self._log_dir))
        visualize_policy(env, self.fake_env, self._policy, self._writer, timestep)
        print('[ Visualization ] Done')
        ## set env state
        env.unwrapped.set_state(qpos, qvel)

    def get_diagnostics(self,
                        iteration,
                        obs_batch = None,
                        training_paths = None,
                        evaluation_paths = None):
        """Return diagnostic information as ordered dictionary.

        Records state value function, and TD-loss (mean squared Bellman error)
        for the sample batch.

        Also calls the `draw` method of the plotter, if plotter defined.
        """

        # @anyboby
        warnings.warn('diagnostics not implemented yet!')

        # diagnostics = OrderedDict({
        #     })

        # policy_diagnostics = self._policy.get_diagnostics(                  #use eval paths
        #     obs_batch[:,-self.active_obs_dim:])
        # diagnostics.update({
        #     f'policy/{key}': value
        #     for key, value in policy_diagnostics.items()
        # })

        # if self._plotter:
        #     self._plotter.draw()

        diagnostics = {}

        return diagnostics

    def save(self, savedir):
        self.fake_env._model.save(savedir, self._epoch)

    @property
    def tf_saveables(self):
        saveables = {
            self._policy.tf_saveables
        }

        # saveables = {
        #     '_policy_optimizer': self._policy_optimizer,
        #     **{
        #         f'Q_optimizer_{i}': optimizer
        #         for i, optimizer in enumerate(self._Q_optimizers)
        #     },
        #     '_log_alpha': self._log_alpha,
        # }

        # if hasattr(self, '_alpha_optimizer'):
        #     saveables['_alpha_optimizer'] = self._alpha_optimizer

        return saveables


# OpenRTiST
#   - Real-time Style Transfer
#
#   Author: Zhuo Chen <zhuoc@cs.cmu.edu>
#           Shilpa George <shilpag@andrew.cmu.edu>
#           Thomas Eiszler <teiszler@andrew.cmu.edu>
#           Padmanabhan Pillai <padmanabhan.s.pillai@intel.com>
#           Roger Iyengar <iyengar@cmu.edu>
#
#   Copyright (C) 2011-2019 Carnegie Mellon University
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
#
# Portions of this code borrow from sample code distributed as part of
# Intel OpenVino, which is also distributed under the Apache License.
#
# Portions of this code were modified from sampled code distributed as part of
# the fast_neural_style example that is part of the pytorch repository and is
# distributed under the BSD 3-Clause License.
# https://github.com/pytorch/examples/blob/master/LICENSE

from openrtist_adapter import OpenrtistAdapter
from torch.autograd import Variable
from transformer_net import TransformerNet
from torchvision import transforms
from distutils.version import LooseVersion
import numpy as np
import torch
import os
import logging

logger = logging.getLogger(__name__)

STARTUP_ONES_SIZE = (360, 240, 3)

class TorchAdapter(OpenrtistAdapter):
    def __init__(self, cpu_only, default_style):
        super().__init__(default_style)

        self.cpu_only = cpu_only

        # We do not need to compute gradients. This saves memory.
        torch.set_grad_enabled(False)

        self.style_model = TransformerNet()

        models_dir = "models"
        self.path = os.path.join(os.getcwd(), "..", models_dir)
        # self._update_model_style(default_style)

        self.content_transform = transforms.Compose([transforms.ToTensor()])

        for name in os.listdir(self.path):
            if name.endswith(".model"):
                self.add_supported_style(name[:-6])

        # Feed network an array of all ones. This makes it run faster on the
        # first real image.
        ones = np.ones(STARTUP_ONES_SIZE, np.uint8)
        preprocessed = self.preprocessing(ones)
        _ = self.inference(preprocessed)

    def set_style(self, new_style):
        if super().set_style(new_style):
            self._update_model_style(new_style)

    def preprocessing(self, img):
        content_image = self.content_transform(img)
        if not self.cpu_only:
            content_image = content_image.cuda()
        content_image = content_image.unsqueeze(0)
        return Variable(content_image)

    def inference(self, preprocessed):
        output = self.style_model(preprocessed)
        return output.data[0].clamp(0, 255).cpu().numpy()

    def postprocessing(self, post_inference):
        return post_inference.transpose(1, 2, 0)

    def _update_model_style(self, new_style):
        model = os.path.join(self.path, "{}.model".format(new_style))
        self.style_model.load_state_dict(torch.load(model))
        if not self.cpu_only:
            self.style_model.cuda()


from time import sleep
import pyautogui
import mss
import cv2
import mss.tools
import time
import numpy as np
from os import walk
import imutils
from analyzer import ColorLabeler
import keyboard
from random import randint
import matplotlib.pylab as plt
from PIL import Image
from lab.tradeanalyzer import Main
from collections import Counter
import os, shutil
from collections import defaultdict

# with mss.mss() as sct:
#     monitor = {"top": 480, "left": 1160, "width": 100, "height": 40}
#     img = sct.grab(monitor)
#     mss.tools.to_png(img.rgb, img.size, output="tradeconfirmation.png")
#     im = Image.open("tradeconfirmation.png")
#     pix = im.load()
#     for i in range(im.size[0]):
#         if i >= 39:
#             color = pix[i, 39]
#             print(color)
#             if color[0] == 113 and color[1] == 139 and color[2] == 80:
#                 print(True)
#         else:
#             color = pix[i, i]
#             print(color)
#             if color[0] == 113 and color[1] == 139 and color[2] == 80:
#                 print(True)
#
# mouseX = 1161
# mousex = 1312
# mouseY = 685
# offsetX = 13
# offsetY = 10
#
# mousex = 1312
# mouseY = 685
# offsetX = 50
# offsetY = 10
# mouseX = 1161
#
# sleep(2)
# pyautogui.keyDown('w')
# sleep(2)
# pyautogui.keyUp('w')
# quit = False
# while not quit:
#     with mss.mss() as sct:
#         pyautogui.keyDown('d')
#         sleep(0.01)
#         pyautogui.keyUp('d')
#         monitor = {"top": 662, "left": 1232, "width": 70, "height": 34}
#         img = sct.grab(monitor)
#         mss.tools.to_png(img.rgb, img.size, output="vault.png")
#         file = "vault.png"
#         img = Image.open(file)
#         colors = img.getcolors(256)
#         max_occurence, most_present = 0, 0
#         check = 0
#         try:
#             for c in colors:
#                 if c[0] > max_occurence:
#                     (max_occurence, most_present) = c
#             for value in most_present:
#                 if value == 255:
#                     check += 1
#             if check == 3:
#                 quit = True
#         except TypeError:
#             pass
#
# pyautogui.moveTo(mouseX, mouseY)
# pyautogui.moveTo(mousex - offsetX, mouseY - offsetY, 1)
# pyautogui.click(1325 - offsetX, 675)
# sleep(8)
# pyautogui.keyDown('w')
# sleep(0.2)
# pyautogui.keyUp('w')

with mss.mss() as sct:
    monitor = {"top": 600, "left": 1160, "width": 197, "height": 114}
    found_spots = defaultdict(list)

    sct_image = sct.grab(monitor)
    mss.tools.to_png(sct_image.rgb, sct_image.size, output="invault.png")

    imgpath = "invault.png"

    # load the image and resize it to a smaller factor so that
    # the shapes can be approximated better
    image = cv2.imread(imgpath)
    resized = imutils.resize(image, width=300)
    ratio = image.shape[0] / float(resized.shape[0])

    # blur the resized image slightly, then convert it to both
    # grayscale and the L*a*b* color spaces
    blurred = cv2.GaussianBlur(resized, (5, 5), 0)
    gray = cv2.cvtColor(blurred, cv2.COLOR_BGR2GRAY)
    lab = cv2.cvtColor(blurred, cv2.COLOR_BGR2LAB)
    thresh = cv2.threshold(gray, 60, 255, cv2.THRESH_BINARY)[1]

    # find contours in the thresholded image
    cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL,
                            cv2.CHAIN_APPROX_SIMPLE)
    cnts = cnts[0] if imutils.is_cv2() else cnts[1]

    # initialize the shape detector and color labeler
    cl = ColorLabeler()
    i = 0
    for c in cnts:
        # compute the center of the contour
        if i >= 1 and i <= 8:
            M = cv2.moments(c)
            cX = int((M["m10"] / M["m00"]) * ratio)
            cY = int((M["m01"] / M["m00"]) * ratio)

            # detect the shape of the contour and label the color
            color = cl.label(lab, c, True)

            # multiply the contour (x, y)-coordinates by the resize ratio,
            # then draw the contours and the name of the shape and labeled
            # color on the image
            c = c.astype("float")
            c *= ratio
            c = c.astype("int")
            text = "{}".format(color)
            cv2.drawContours(image, [c], -1, (0, 255, 0), 2)
            cv2.putText(image, text, (cX - 15, cY - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)

            # show the output image
            cv2.imshow("Image", image)
            if color != None:
                if color == "empty":
                    found_spots[color].append([cX, cY])
            print(found_spots["empty"][0][0])
        i += 1


import numpy as np
from Box2D.b2 import edgeShape, circleShape, fixtureDef, polygonShape, revoluteJointDef, contactListener

from teachDRL.gym_flowers.envs.bodies.walkers.WalkerAbstractBody import WalkerAbstractBody
from teachDRL.gym_flowers.envs.utils.custom_user_data import CustomBodyUserData, CustomMotorUserData

HULL_POLYGONS = [
        [(-46, +13), (+6, +13), (+50, +5),
        (+50, -12), (-46, -12)]
    ]
HULL_BOTTOM_WIDTH = 96
SPEED_HIP     = 4
SPEED_KNEE    = 6

class BigQuadruBody(WalkerAbstractBody):
    def __init__(self, scale, motors_torque = 300, nb_steps_under_water = 600):
        super(BigQuadruBody, self).__init__(scale, motors_torque, nb_steps_under_water)
        self.LEG_DOWN = 3 / self.SCALE  # 0 = center of hull
        self.LEG_W, self.LEG_H = 10 / self.SCALE, 51 / self.SCALE
        self.TORQUE_PENALTY = 0.00035 / 2 # 2 paris of legs

        self.AGENT_WIDTH = HULL_BOTTOM_WIDTH / self.SCALE
        self.AGENT_HEIGHT = 25 / self.SCALE + \
                            self.LEG_H * 2 - self.LEG_DOWN
        self.AGENT_CENTER_HEIGHT = self.LEG_H * 2 + self.LEG_DOWN

    def draw(self, world, init_x, init_y, force_to_center):
        HULL_FIXTURES = [
            fixtureDef(
                shape=polygonShape(vertices=[(x / self.SCALE, y / self.SCALE) for x, y in polygon]),
                density=5.0,
                friction=0.1,
                categoryBits=0x20,
                maskBits=0x000F)
            for polygon in HULL_POLYGONS
        ]

        LEG_FD = fixtureDef(
            shape=polygonShape(box=(self.LEG_W / 2, self.LEG_H / 2)),
            density=1.0,
            restitution=0.0,
            categoryBits=0x20,
            maskBits=0x000F)

        LOWER_FD = fixtureDef(
            shape=polygonShape(box=(0.8 * self.LEG_W / 2, self.LEG_H / 2)),
            density=1.0,
            restitution=0.0,
            categoryBits=0x20,
            maskBits=0x000F)

        hull = world.CreateDynamicBody(
            position=(init_x, init_y),
            fixtures=HULL_FIXTURES
        )

        hull.color1 = (0.5, 0.4, 0.9)
        hull.color2 = (0.3, 0.3, 0.5)
        hull.ApplyForceToCenter((force_to_center, 0), True)

        hull.userData = CustomBodyUserData(True, is_contact_critical=True, name="hull")
        self.body_parts.append(hull)
        self.reference_head_object = hull

        for x_anchor in [-0.7, 0.5]:
            absolute_x = init_x + np.interp(x_anchor, [-1,1], [-HULL_BOTTOM_WIDTH / 2 / self.SCALE, HULL_BOTTOM_WIDTH / 2 / self.SCALE])
            for i in [-1, +1]:
                leg = world.CreateDynamicBody(
                    position=(absolute_x, init_y - self.LEG_H / 2 - self.LEG_DOWN),
                    #angle=(i * 0.05),
                    fixtures=LEG_FD
                )
                leg.color1 = (0.6 - i / 10., 0.3 - i / 10., 0.5 - i / 10.)
                leg.color2 = (0.4 - i / 10., 0.2 - i / 10., 0.3 - i / 10.)
                rjd = revoluteJointDef(
                    bodyA=hull,
                    bodyB=leg,
                    anchor=(absolute_x, init_y - self.LEG_DOWN),
                    enableMotor=True,
                    enableLimit=True,
                    maxMotorTorque=self.MOTORS_TORQUE,
                    motorSpeed=i,
                    lowerAngle=-0.8,
                    upperAngle=1.1,
                )

                leg.userData = CustomBodyUserData(False, name="leg")
                self.body_parts.append(leg)

                joint_motor = world.CreateJoint(rjd)
                joint_motor.userData = CustomMotorUserData(SPEED_HIP, False)
                self.motors.append(joint_motor)

                lower = world.CreateDynamicBody(
                    position=(absolute_x, init_y - self.LEG_H * 3 / 2 - self.LEG_DOWN),
                    #angle=(i * 0.05),
                    fixtures=LOWER_FD
                )
                lower.color1 = (0.6 - i / 10., 0.3 - i / 10., 0.5 - i / 10.)
                lower.color2 = (0.4 - i / 10., 0.2 - i / 10., 0.3 - i / 10.)
                rjd = revoluteJointDef(
                    bodyA=leg,
                    bodyB=lower,
                    anchor=(absolute_x, init_y - self.LEG_DOWN - self.LEG_H),
                    enableMotor=True,
                    enableLimit=True,
                    maxMotorTorque=self.MOTORS_TORQUE,
                    motorSpeed=1,
                    lowerAngle=-1.6,
                    upperAngle=-0.1,
                )

                lower.userData = CustomBodyUserData(True, name="lower")
                self.body_parts.append(lower)

                joint_motor = world.CreateJoint(rjd)
                joint_motor.userData = CustomMotorUserData(
                    SPEED_KNEE,
                    True,
                    contact_body=lower,
                    angle_correction=1.0)
                self.motors.append(joint_motor)


"""
    Copyright 2020 Huawei Technologies Co., Ltd

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
    Typical usage example:
"""

import os
import numpy as np
import argparse

def read_class_names(class_file_name):
    """
    loads class name from a file
    """
    names = {}
    with open(class_file_name, 'r') as data:
        for id, name in enumerate(data):
            names[id] = name.strip('\n')
    return names

def parse_line(line):
    """
    parse line information
    """
    temp = line.split(" ")
    index = temp[1].split("/")[-1].split(".")[0]
    return index, (int(temp[3]), int(temp[2]))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Postprocessing of Yolo-caffe model')
    parser.add_argument("--bin_data_path", default="./pic_1024/", help='Benchamrk output path')
    parser.add_argument("--test_annotation", default="./VOC2017.info", help='Test set annotation file')
    parser.add_argument("--det_results_path", default="./input/detection-results/", help='Output TXT directory')
    parser.add_argument("--coco_class_names", default="./input/coco.names", help='coco dataset classes')
    parser.add_argument("--voc_class_names", default="./input/voc.names", help='vocal dataset classes')
    parser.add_argument("--net_input_width", default=416, help='width of the input image')
    parser.add_argument("--net_input_height", default=416, help='height of the input image')
    flags = parser.parse_args()

    # load the test image info
    img_size_dict = dict()
    with open(flags.test_annotation)as f:
        for line in f.readlines():
            temp = parse_line(line)
            img_size_dict[temp[0]] = temp[1]

    # load label class name
    coco_class_map = read_class_names(flags.coco_class_names)
    coco_set = set(coco_class_map.values())
    voc_class_map = read_class_names(flags.voc_class_names)
    voc_set = set(voc_class_map.values())

    # load the output files
    # "_1.bin" for detinfo
    # "_2.bin" for numbox
    bin_path = flags.bin_data_path
    net_input_width = flags.net_input_width
    net_input_height = flags.net_input_height
    det_results_path = flags.det_results_path
    if not os.path.exists(det_results_path):
        os.makedirs(det_results_path)
    for img_name in img_size_dict:
        boxbuf = []
        boxnum = []
        det_results_str = ""
        path_det_output1 = os.path.join(bin_path, img_name + "_1.bin")
        if os.path.exists(path_det_output1):
            boxbuf = np.fromfile(path_det_output1, dtype="float32")
        else:
            print("[ERROR] file not exist", path_det_output1)
            break
        path_det_output2 = os.path.join(bin_path, img_name + "_2.bin")
        if os.path.exists(path_det_output2):
            boxnum = np.fromfile(path_det_output2, dtype="float32")
        else:
            print("[ERROR] file not exist", path_det_output2)
            break
        boxcnt = boxnum[0].astype(int)
        if(boxcnt == 0):
            det_results_str = ""
        else:
            j = 0
            expect_data_list = []
            pre_box_list = []
            for i in boxbuf[0:boxcnt * 6]:
                expect_data_list.append(round(float(i), 3))
            for i in range(6):
                pre_box_list.append(list(expect_data_list[j:j + boxcnt]))
                j += boxcnt
            for i in range(int(boxcnt)):
                class_ind = int(pre_box_list[5][i])
                if class_ind < 0:
                    print("[ERROR] predicted object class error:", class_ind)
                    continue
                else:
                    class_name = coco_class_map[class_ind]
                    if class_name in voc_set:
                        det_results_str += "{} {} {} {} {} {}\n".format(
                            class_name, str(pre_box_list[4][i]), str(pre_box_list[0][i]),
                            str(pre_box_list[1][i]), str(pre_box_list[2][i]),
                            str(pre_box_list[3][i]))
        det_results_file = os.path.join(det_results_path, img_name + ".txt")
        with open(det_results_file, "w") as detf:
            detf.write(det_results_str)


import numpy as np
table = np.full(3268760, 255, dtype=np.uint8, order='C')
table[0]=0

table.tofile("v9_fevertime_table")


import cv2.cv2
import numpy as np #For image to matrix conversion
from PIL import ImageGrab #To get access to the screenshot
import win32gui #To access the Windows GUI
import win32api #To get api of the elements
import win32con #To perform keyboard press action
import time
#To access the windows running
toplist = []
windowList = []
def enum_win(windowHandle, result):
    windowText = win32gui.GetWindowText(windowHandle)
    windowList.append((windowHandle,windowText))

#Get the list of open windows
win32gui.EnumWindows(enum_win, toplist)

#To get the handle of the game
gameHandle = 0
for (handle,text) in windowList:
    if 'chrome://dino/' in text:
        gameHandle = handle
#To make sure the window is available or loaded
if gameHandle != 0:
    pass
else:
    print('Window not available')
    exit(-1)

#Kernel for the mask
kernelOpen = np.ones((3,3))
kernelClose = np.ones((12,12))
isFirst = True
#Set the game window to forground
win32gui.SetForegroundWindow(gameHandle)
#Resize the window
win32gui.MoveWindow(gameHandle, 0, 0,788,498, True)

#Reload the browser
win32api.SendMessage(gameHandle, win32con.WM_KEYDOWN, win32con.VK_F5)
win32api.SendMessage(gameHandle, win32con.WM_KEYUP, win32con.VK_F5)
time.sleep(1)

#Start the game
win32api.SendMessage(gameHandle, win32con.WM_KEYDOWN, win32con.VK_SPACE)
win32api.SendMessage(gameHandle, win32con.WM_KEYUP, win32con.VK_SPACE)
time.sleep(2)

while True:
    #Get the position of the window
    position = win32gui.GetWindowRect(gameHandle)
    x,y,w,h = position
    #Take screenshot
    screenshot = ImageGrab.grab(bbox=(x,y,w-200,h)) #give the location to take screenshot except the score
    #Convert the screenshot to numpy array
    #Convert the screenshot to an array
    screenshot_main = np.array(screenshot)
    #Convert from RGB to BGR for CV2
    screenshot_array = cv2.cvtColor(screenshot_main,cv2.COLOR_RGB2HSV)
    #Convert from RGB to BGR for display only
    screenshot_array2 = cv2.cvtColor(screenshot_main,cv2.COLOR_RGB2BGR)
    #To get the color of the images for masking
    low = np.array([0,0,166])
    high = np.array([0,0,172])
    mask = cv2.inRange(screenshot_array,low,high)
    #To hide noise
    maskOpen = cv2.morphologyEx(mask,cv2.MORPH_OPEN,kernelOpen)
    maskClose = cv2.morphologyEx(maskOpen,cv2.MORPH_CLOSE,kernelClose)
    contours, h = cv2.findContours(maskClose.copy(),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)

    #To store the location of each cactus and crow
    xs = dict()
    for i in range(len(contours)):
        x,y,w,h = cv2.boundingRect(contours[i])
        #Store the width and y coordinate of dino
        if isFirst:
            dinoWidth = x+w
            dinoY = y
            isFirst = False
        if x+w == dinoWidth:
            #Bound the dino by rectangle and write its name
            cv2.rectangle(screenshot_array2,(x,y),(x+w,y+h),(0,255,0),2)
            cv2.putText(screenshot_array2,"Dino",(x,y+h+20),cv2.FONT_HERSHEY_COMPLEX,1,(233,244,255))
        elif(len(contours)>1):
                #Bound the others in red rectangle and store their location
                xs[i] = [x,y,w]
                cv2.rectangle(screenshot_array2,(x,y),(x+w,y+h),(0,0,255),4)
    for i in xs.values():
        if((i[0] <= dinoWidth+110) and((dinoY < i[1]) or (((dinoY - i[1])<=8)and((dinoY - i[1])>=0)))):
                #Jump if it's closer than 110 and is at low position than dino
                win32api.SendMessage(gameHandle, win32con.WM_KEYDOWN, win32con.VK_SPACE)
                time.sleep(0.1)
                win32api.SendMessage(gameHandle, win32con.WM_KEYUP, win32con.VK_SPACE)
                cv2.putText(screenshot_array2,str("Jump"),(20,(dinoY - 100)),cv2.FONT_HERSHEY_TRIPLEX,1,(0,250,120))
        elif((i[0] <= dinoWidth+110) and(((i[1]-dinoY)<0))):
                #Bow if it's closer than 110 and is at higher position than dino
                cv2.putText(screenshot_array2,str("Bow"),(20,(dinoY - 100)),cv2.FONT_HERSHEY_TRIPLEX,1,(0,250,120))
                win32api.SendMessage(gameHandle, win32con.WM_KEYDOWN, win32con.VK_DOWN)
                time.sleep(0.4)
                win32api.SendMessage(gameHandle, win32con.WM_KEYUP, win32con.VK_DOWN)
    #cv2.WINDOW_NORMAL makes the output window resizealbe
    cv2.namedWindow('Screen', cv2.WINDOW_NORMAL)

    #resize the window according to the screen resolution
    cv2.resizeWindow('Screen', 640, 400)
    #Show the image
    cv2.imshow("Screen",screenshot_array2)
    #Wait for 25 milisecond to take another
    key = cv2.waitKey(1)

    #Break when 'A' is pressed on the keybo ard
    if key == 65:
        break
#Close all the windows
cv2.destroyAllWindows()


import torch
import numpy as np
from fairseq import checkpoint_utils, utils, options, tasks
from fairseq.logging import progress_bar
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
import ogb
import sys
import os
from pathlib import Path
from sklearn.metrics import roc_auc_score, mean_squared_error, mean_absolute_percentage_error, mean_absolute_error

from os import path
sys.path.append( path.dirname( path.dirname( path.abspath(__file__) ) ) )
from pretrain import load_pretrained_model
import logging
from data_class import geo_Omsk, single_geo_Omsk, GraphormerPYGDataset_predict
import os.path as osp
from torch_geometric.data import Dataset
from functools import lru_cache
import torch_geometric.datasets
from ogb.graphproppred import PygGraphPropPredDataset
from ogb.lsc.pcqm4m_pyg import PygPCQM4MDataset
import pyximport
from torch_geometric.data import InMemoryDataset, download_url
import pandas as pd
from sklearn import preprocessing
pyximport.install(setup_args={'include_dirs': np.get_include()})
from torch_geometric.data import Data
import time
from torch_geometric.utils import add_self_loops, negative_sampling
import copy
from fairseq.data import (
    NestedDictionaryDataset,
    NumSamplesDataset,
)
sys.path.insert(2, '/home/jovyan/graphormer_v2/graphormer')
from data.pyg_datasets.pyg_dataset import GraphormerPYGDataset
from data.dataset import (
    BatchedDataDataset,
    TargetDataset,
    GraphormerDataset)

def eval(args, use_pretrained, checkpoint_path=None, logger=None, data_name = None):
    cfg = convert_namespace_to_omegaconf(args)
    np.random.seed(cfg.common.seed)
    utils.set_torch_seed(cfg.common.seed)
    seed = 71

    ### data loading
    # name = 'omsk'
    if data_name == 'omsk':
        root = osp.join('dataset', data_name)
        raw_dir = osp.join(root, 'processed', 'data_omsk_1')
        data = single_geo_Omsk(root = raw_dir)
        GPYG = GraphormerPYGDataset_predict(data,seed,None,data,'omsk')
    batched_data = BatchedDataDataset(GPYG)
    data_sizes = np.array([128] * len(batched_data))
    dataset_total = NestedDictionaryDataset(
            {
                "nsamples": NumSamplesDataset(),
                "net_input": {"batched_data": batched_data},
                "target": batched_data,
            },
        sizes=data_sizes,
        )
    ###

    ### initialize task
    task = tasks.setup_task(cfg.task)
    model = task.build_model(cfg.model)
    batch_iterator = task.get_batch_iterator(
        dataset=dataset_total
    )
    itr = batch_iterator.next_epoch_itr(shuffle=False, set_dataset_epoch=False)
    progress = progress_bar.progress_bar(itr)
    ###

    ### load checkpoint
    model_state = torch.load(checkpoint_path)["model"]
    model.load_state_dict(model_state, strict=True, model_cfg=cfg.model)
    model.to(torch.cuda.current_device())
    del model_state
    ###

    ### prediction
    y_pred = []
    with torch.no_grad():
        model.eval()
        for i, sample in enumerate(progress):
            sample = utils.move_to_cuda(sample)
            y = model(**sample["net_input"])[:, 0, :].reshape(-1)
            y_pred.extend(y.detach().cpu())
            torch.cuda.empty_cache()
    ###

    # save predictions
    y_pred = torch.Tensor(y_pred)
    print(y_pred)

    return y_pred

def main():
    parser = options.get_training_parser()
    parser.add_argument(
        "--split",
        type=str,
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
    )
    parser.add_argument(
        "--metric",
        type=str,
    )
    args = options.parse_args_and_arch(parser, modify_parser=None)
    logger = logging.getLogger(__name__)
    if args.pretrained_model_name != "none":
        eval(args, True, logger=logger)
    elif hasattr(args, "save_dir"):
        checkpoint_fname = 'checkpoint_best.pt'
        checkpoint_path = Path(args.save_dir) / checkpoint_fname
        logger.info(f"evaluating checkpoint file {checkpoint_path}")
        eval(args, False, checkpoint_path, logger, args.dataset_name)

if __name__ == '__main__':
    main()


import os
import h5py
import numpy as np
import tensorflow as tf
from data_iter import next_batch, train_test_data
from network import NetworkModel

class Agent:

    def __init__(self, sess, params):
        '''Agent - powered by neural network, can trian and infer
        '''
        self.sess = sess
        self.params = params
        self.test_file_path = '/.firejobdata/data1/banlu/ViZDoom-SL/ViZDoomSL_baseline_lstm/data/dataset-1-test.hdf5'
        self.test_index = 0
        self.model = NetworkModel(params=params)
        self.build_model(self.params.log_dir)

    def choose_action(self, screen):
        predict_action, output_gf = self.sess.run([self.output_sc, self.output_gf], feed_dict={self.img: screen, self.keep_prob: 1})
        print predict_action
        action_id = int(np.argmax(predict_action, 1)[0])
        return action_id, output_gf

    def build_model(self, log_dir):
        self.img = tf.placeholder(tf.float32, shape=[None, self.params.seq_len, 60, 108, 3], name="img")
        # the latest action and game feature label
        self.label_sc = tf.placeholder(tf.int32, shape=[None, 1], name="label_sc")
        self.label_gf = tf.placeholder(tf.float32, shape=[None, 2], name="label_gf")
        self.keep_prob = tf.placeholder(tf.float32)

        train_op, output_sc, loss_sc, acc_sc, output_gf, loss_gf, acc_gf= self.model.build_train_graph(
                                         self.img, self.label_sc, self.label_gf, self.keep_prob)

        self.train_op = train_op
        self.output_sc, self.output_gf = output_sc, output_gf
        self.loss_sc, self.loss_gf = loss_sc, loss_gf
        self.acc_sc, self.acc_gf = acc_sc, acc_gf

        self.train_writer = tf.summary.FileWriter(log_dir + '/train', self.sess.graph)
        self.test_writer = tf.summary.FileWriter(log_dir + '/test')
        self.train_writer.add_graph(self.sess.graph)

        tf.summary.scalar('Loss_sc', self.loss_sc)
        tf.summary.scalar('Acc_sc', self.acc_sc)
        tf.summary.scalar('Loss_gf', self.loss_gf)
        tf.summary.scalar('Acc_gf', self.acc_gf)

        self.merged = tf.summary.merge_all()

        saver = tf.train.Saver()
        self.saver = saver
        self.sess.run(tf.global_variables_initializer())
        #self.saver.restore(self.sess, './checkpoint/model.ckpt')

    def get_test_data(self):
        test_file_list = os.listdir(self.test_file_path)
        test_file = test_file_list[self.test_index]
        self.test_index += 1
        if self.test_index == len(test_file_list):
            self.test_index = 0
        # read hdf5 test file
        ff = h5py.File(os.path.join(self.test_file_path, test_file), 'r')
        test_img = np.array([ff['img'+str(i)] for i in range(4)]).reshape(-1, 4, 60, 108, 3)
        test_label_gf = np.array(ff['label_gf'][:].reshape(-1, 2))
        test_label_sc = np.array(ff['action_id'][:]).reshape(-1, 1)
        return test_img, test_label_sc, test_label_gf

    def train(self, checkpoint_dir, data_dir):
        assert os.path.isfile(data_dir)

        train_img, train_label, train_label_gf = train_test_data(data_dir, self.params.seq_len)
        test_img, test_label, test_label_gf = train_test_data(self.test_file_path, self.params.seq_len)
        test_img = test_img.astype(np.float32) / 255.

        step = 0
        for e in range(self.params.epoch):
            train_loss_sum = 0
            train_acc_sum = 0
            train_gf_loss_sum = 0
            train_gf_acc_sum = 0
            count = 0

            for batch_img, batch_label_sc, batch_label_gf in next_batch(train_img, train_label, train_label_gf, self.params.batch_size):
                if batch_img.shape[0] != self.params.batch_size:
                    continue
                batch_img = batch_img.astype(np.float32) / 255.

                _, loss, acc, loss_gf, acc_gf = self.sess.run([self.train_op, self.loss_sc, self.acc_sc, self.loss_gf, self.acc_gf],
                                              feed_dict={self.img: batch_img,
                                                         self.label_sc: batch_label_sc[:, -1, :].reshape((-1, 1)),
                                                         self.label_gf: batch_label_gf[:, -1, :].reshape((-1, 2)),
                                                         self.keep_prob: 0.5})
                train_loss_sum += loss
                train_acc_sum += acc
                train_gf_loss_sum += loss_gf
                train_gf_acc_sum += acc_gf
                count += 1
                step += 1

                if step % 50 == 0:
                    # Get Train Summary for one batch and add summary to TensorBoard
                    summary = self.sess.run(self.merged, feed_dict={self.img: batch_img,
                                                                    self.label_sc: batch_label_sc[:, -1, :].reshape((-1, 1)),
                                                                    self.label_gf: batch_label_gf[:, -1, :].reshape((-1, 2)),
                                                                    self.keep_prob: 1})
                    self.train_writer.add_summary(summary, step)
                    self.train_writer.flush()

                    # Get Test Summary on random 32 test images and add summary to TensorBoard
                    test_idx = np.arange(test_img.shape[0])
                    np.random.shuffle(test_idx)
                    test_idx = test_idx[:32]
                    summary = self.sess.run(self.merged, feed_dict={self.img: test_img[test_idx,...],
                                                                    self.label_sc: test_label[test_idx, ...][:, -1, :].reshape((-1, 1)),
                                                                    self.label_gf: test_label_gf[test_idx, ...][:, -1, :].reshape((-1, 2)),
                                                                    self.keep_prob: 1})
                    self.test_writer.add_summary(summary, step)
                    self.test_writer.flush()

            test_acc_sc_sum = 0
            test_acc_gf_sum = 0
            test_count = 0
            #Feed forward all test images into graph and log accuracy
            for batch_test_img, batch_test_label_sc, batch_test_label_gf in next_batch(test_img, test_label, test_label_gf, 1):
                test_acc_sc, test_acc_gf = self.sess.run([self.acc_sc, self.acc_gf], feed_dict={self.img: batch_test_img,
                    self.label_sc: batch_test_label_sc[:, -1, :].reshape((-1, 1)),
                    self.label_gf: batch_test_label_gf[:, -1, :].reshape((-1, 2)),
                    self.keep_prob: 1})

                test_acc_sc_sum += test_acc_sc
                test_acc_gf_sum += test_acc_gf
                test_count += 1
            print("epoch %d, loss_sc %.4f, loss_gf %.4f, train_acc_sc %.4f, test_acc_sc %.4f, train_acc_gf %.4f, test_acc_gf %.4f" %
                    (e, train_loss_sum / count, train_gf_loss_sum / count, train_acc_sum / count, test_acc_sc_sum / test_count, train_gf_acc_sum / count,
                        test_acc_gf_sum / test_count))

        if checkpoint_dir is not None:
            self.save(checkpoint_dir)
            print("Model has been saved.")

    def save(self, checkpoint_dir):
        ckpt_path = os.path.join(checkpoint_dir, "model.ckpt")
        self.saver.save(self.sess, ckpt_path)

    def load(self, checkpoint_dir):
        ckpt_path = os.path.join(checkpoint_dir, "model.ckpt")
        if os.path.isfile(ckpt_path):
            self.saver.restore(self.sess, ckpt_path)


from DobotAPI.Dobot import Dobot
from Chessboard.Chessboard import Chessboard
from CheckersReinforcementLearning.Board import Board
from CheckersReinforcementLearning.AI import Alpha_beta
from DobotCamera import DobotCamera
from DobotCheckers import DobotCheckers
import numpy as np
import time
import cv2
import helpers

dobot = Dobot(homeX=64, homeY=142, homeZ=-6.5)
zUpValue = 5
zDownValue = -40
releaseX = 64
releaseY = 142

dobotCam = DobotCamera(
    xDobRefPoint=150.75,
    yDobRefPoint=-64.8,
    xCamRefPoint=56,
    yCamRefPoint=51,
    xScaleFactor=0.68,
    yScaleFactor=0.67
)

# img = helpers.captureBoard()
# img = cv2.imread("./img/checkers_scanner_pieces_2.jpg")

def AIMakeMove(img):
    chessboard = Chessboard()

    dobotCheckers = DobotCheckers(
        chessboard,
        dobot,
        dobotCam,
        zDownValue,
        zUpValue,
        releaseX,
        releaseY
    )

    boardDetected = chessboard.detectGameState(img)
    if boardDetected is False:
        return False

    chessboard.prettyPrintBoard()

    AIBoard = Board(chessboard.convertStateToAI(), the_player_turn=False)

    alphaBetaAI = Alpha_beta(False, 6)
    alphaBetaAI.set_board(AIBoard)

    # playerTurn = False
    # while not AIBoard.is_game_over():
    #     if not playerTurn:

    # Get moves from AI
    nextMove = alphaBetaAI.get_next_move()
    print("AI move: ", nextMove)

    captured = AIBoard.make_move(nextMove)

    dobotCheckers.grab(nextMove[0])

    dobotCheckers.release(nextMove[-1])

    if len(captured) > 0:
        dobotCheckers.capture(captured)

    # GoHome (away from camera)
    dobot.move(releaseX, releaseY)
    return True

while True:
    boardDetected = False
    while boardDetected is False:
        img = helpers.captureBoard()
        print("AI's turn...")
        boardDetected = AIMakeMove(img)
    print("AI finished it's turn.")
    humanFinishedMove = False
    handInsideBoard = False
    print("Human's turn...")
    while True:
        handDetected = helpers.detectHand(helpers.captureBoard())
        print("HandDetected: ", handDetected)
        if (handInsideBoard is True) and (handDetected is False):
            print("Human finished it's turn.")
            break
        else:
            handInsideBoard = handDetected