# -*- coding:utf-8 -*-

import sys
import joystick
import serial
import threading
import numpy as np
from rx import Observable
from rx.subjects import Subject
from rx.concurrency import NewThreadScheduler

FEEDRATE = 400 # mm / minute
MAX_REACH = 0.2 # mm / feedrate
MIN_NORM = 0.2

device_fn = sys.argv[1] if len(sys.argv)>1 else joystick.Joystick.available_sticks()[0]
print("Using input device: {0}".format(device_fn))

in_stick = joystick.Joystick(device_fn)
tty = serial.Serial('/dev/ttyUSB0',250000)

def send_gcode(s):
  tty.flushInput()
  tty.write(s + b'\n')
  return tty.readline().strip()

def poll_joystick(_):
  poll_result = in_stick.poll()
  x, y = None, None
  if poll_result.changed_axis is not None:
    axis = poll_result.changed_axis
    if axis.number == 0:
        x = axis.fvalue
    elif axis.number == 1:
        y = axis.fvalue
  return x, y

def update_position(accu, update):
  return tuple( (current if current is not None else previous) for current, previous in zip(update, accu) )

joystick_positions = Observable.interval(1).map(poll_joystick).scan(update_position, (0,0))

cv = threading.Condition()
new_position = None

def new_position_available():
  global new_position
  return new_position is not None

def execute_move():
  global cv, new_position
  while True:
    with cv:
      cv.wait_for(new_position_available)
      dx, dy = new_position
      new_position = None

    # Relative positioning
    send_gcode(b'G91')
    # Rapid move
    send_gcode('G1 X{:.3f} Y{:.3f} F{}'.format(dx, -dy, FEEDRATE).encode('ascii'))

consumer_thread = threading.Thread(target=execute_move)
consumer_thread.daemon = True
consumer_thread.start()

def move_printer(delta):
  global cv, new_position
  if np.linalg.norm(delta) > MIN_NORM:
    dx, dy = np.array(delta) * MAX_REACH
    print(dx, dy)
    with cv:
      new_position = dx, dy
      cv.notify()

joystick_positions \
  .filter(lambda pos: all(val is not None for val in pos)) \
  .combine_latest(Observable.interval(20), lambda a,b: a) \
  .observe_on(NewThreadScheduler()) \
  .subscribe(on_next=move_printer)

input("Press any key to exit\n")


from facenet_pytorch import MTCNN, InceptionResnetV1
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
import numpy as np
import pandas as pd
import os
import cv2

def create_embeddings():

    workers = 0 if os.name == 'nt' else 4

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('Running on device: {}'.format(device))

    mtcnn = MTCNN(
        image_size=160, margin=0, min_face_size=20,
        thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
        device=device
    )

    resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)

    pictures=[]

    directory=os.listdir('images_cropped')
    for i in range(len(directory)):
        count = 0
        flag = False
        while(flag==False):
            path='images_cropped'+'/'+str(directory[i])+'/'+str(count)+'.jpg'
            pic=cv2.imread(path)
            print(count)
            if pic is not None:
                tens=mtcnn(pic)
                if tens is not None:
                    pictures.append(tens)
                    flag=True
                    break
                    count=0
                    print(directory[i])
                    print('here')
                else:
                    count=count+1
            else:
                print(directory[i])
                count=count+1

    print(pictures)

    embeddings=[]

    for x in pictures:
        embeddings.append(resnet(x.unsqueeze(0)).detach().cpu())

    dists = [[(e1 - e2).norm().item() for e2 in embeddings] for e1 in embeddings]
    print(pd.DataFrame(dists, columns=directory, index=directory))
    torch.save(embeddings, 'embeddings.pt')  # Of course, you can also save it in a file.
    torch.save(directory, 'names.pt')


# ********** modules ********** #
# chainer
import chainer
from chainer import cuda
from chainer.training import extensions

# others
import numpy as np
import argparse
import glob, os
import random

# network, which named "Looking to Listen at the Cocktail Party"
from network import Audio_Visual_Net

# ********** setup ********** #
np.random.seed(0)
INDEX = 0
DATA_DIR_SPEC = ""
DATA_DIR_VISUAL = ""

# ********** main ********** #
def main():
    # ===== Argparse ===== #
    parser = argparse.ArgumentParser()
    parser.add_argument("--gpu", "-g", type=int, default=-1,
                        help="specify GPU")
    parser.add_argument("--iteration", "-i", type=int, default=5000,
                        help="# of iterations")
    parser.add_argument("--batch_size", "-b", type=int, default=6,
                        help="batch size")
    parser.add_argument("--units", "-u", type=int,
                        default=5000, help="# of FC units")
    parser.add_argument("--data_visual", type=str, default=DATA_DIR_VISUAL,
                        help="Visual data directory, which has csv files")
    parser.add_argument("--data_speech", type=str, default=DATA_DIR_SPEC,
                        help="Spectrogram data directory, which has npz files")
    parser.add_argument("--result_dir", "-r", type=str,
                        default="result-{}/".format(INDEX))
    parser.add_argument("--resume", default="",
                        help="Resume the training from snapshot")
    args = parser.parse_args()

    # ===== GPU or CPU ===== #
    if args.gpu >= 0:
        xp = cuda.cupy
        cuda.get_device(args.gpu).use()
    else:
        xp = np
    chainer.backends.cuda.get_device_from_id(args.gpu).use()

    # ===== Load model ===== #
    print("loading model...")
    model = Audio_Visual_Net(spec_len=49, face_len=12, num_fusion_units=args.units, gpu=args.gpu)
    if args.gpu >= 0:
        model.to_gpu(args.gpu)
    optimizer = chainer.optimizers.Adam(alpha=3*1e-5)
    optimizer.setup(model)

    # ===== Set data ===== #
    print("loading data...")
    spec_input = sorted(glob.glob(os.path.join(args.data_speech, "*.npz")))
    vis_input = sorted(glob.glob(os.path.join(args.data_visual, "*")))
    assert len(spec_input)==len(vis_input), "# of files are different between faces and audios."
    all_nums = range(len(spec_input))
    threshold = int(len(all_nums) * 0.99)
    all_nums_train = all_nums[:threshold]
    all_nums_test = all_nums[threshold:]
    train = [(i) for i in all_nums_train]
    test = [(i) for i in all_nums_test]
    train_iter = chainer.iterators.SerialIterator(dataset=train, batch_size=args.batch_size, shuffle=True, repeat=True)
    test_iter = chainer.iterators.SerialIterator(dataset=test, batch_size=args.batch_size, shuffle=False, repeat=False)

    # ===== Define trainer ===== #
    print("setting trainer...")
    updater = chainer.training.updaters.StandardUpdater(train_iter, optimizer, device=args.gpu)
    trainer = chainer.training.Trainer(updater, (args.iteration, "iteration"), out=args.result_dir)

    iter_trigger = 10
    trainer.extend(extensions.Evaluator(test_iter, model, device=args.gpu), trigger=(iter_trigger, "iteration"))
    trainer.extend(extensions.LogReport(trigger=(iter_trigger, "iteration")), trigger=(iter_trigger, "iteration"))
    trainer.extend(extensions.ProgressBar(update_interval=2))
    trainer.extend(extensions.PlotReport(["main/loss", "validation/main/loss"], "iteration", file_name="loss.png", trigger=(10, "iteration")))
    trainer.extend(extensions.PrintReport(["epoch", "iteration", "main/loss", "validation/main/loss", "elapsed_time"]), trigger=(iter_trigger, "iteration"))
    trainer.extend(extensions.snapshot(), trigger=(int(iter_trigger*10), "iteration"))

    if args.resume:
        chainer.serializers.load_npz(args.resume, trainer)

    # ===== Training ===== #
    print("start training...")
    trainer.run()

    # ===== Save model ===== #
    print("saving model...")
    model.to_cpu()
    chainer.serializers.save_npz(os.path.join(args.result_dir, "model-{}.npz".format(INDEX)), model)
    chainer.serializers.save_npz(os.path.join(args.result_dir, "optimizer-{}.npz".format(INDEX)), optimizer)

    print("done!!")

if __name__ == "__main__":
    main()


import numpy as np

import chainer
from chainer.backends import cuda
import chainer.functions as F

from chainercv.links.model.faster_rcnn.utils.anchor_target_creator import\
    AnchorTargetCreator
from chainercv.links.model.faster_rcnn.utils.proposal_target_creator import\
    ProposalTargetCreator

class FasterRCNNTrainChain(chainer.Chain):

    """Calculate losses for Faster R-CNN and report them.

    This is used to train Faster R-CNN in the joint training scheme
    [#FRCNN]_.

    The losses include:

    * :obj:`rpn_loc_loss`: The localization loss for \
        Region Proposal Network (RPN).
    * :obj:`rpn_cls_loss`: The classification loss for RPN.
    * :obj:`roi_loc_loss`: The localization loss for the head module.
    * :obj:`roi_cls_loss`: The classification loss for the head module.

    .. [#FRCNN] Shaoqing Ren, Kaiming He, Ross Girshick, Jian Sun. \
    Faster R-CNN: Towards Real-Time Object Detection with \
    Region Proposal Networks. NIPS 2015.

    Args:
        faster_rcnn (~chainercv.links.model.faster_rcnn.FasterRCNN):
            A Faster R-CNN model that is going to be trained.
        rpn_sigma (float): Sigma parameter for the localization loss
            of Region Proposal Network (RPN). The default value is 3,
            which is the value used in [#FRCNN]_.
        roi_sigma (float): Sigma paramter for the localization loss of
            the head. The default value is 1, which is the value used
            in [#FRCNN]_.
        anchor_target_creator: An instantiation of
            :class:`~chainercv.links.model.faster_rcnn.AnchorTargetCreator`.
        proposal_target_creator: An instantiation of
            :class:`~chainercv.links.model.faster_rcnn.ProposalTargetCreator`.

    """

    def __init__(self, faster_rcnn, rpn_sigma=3., roi_sigma=1.,
                 anchor_target_creator=AnchorTargetCreator(),
                 proposal_target_creator=ProposalTargetCreator()):
        super(FasterRCNNTrainChain, self).__init__()
        with self.init_scope():
            self.faster_rcnn = faster_rcnn
        self.rpn_sigma = rpn_sigma
        self.roi_sigma = roi_sigma

        self.anchor_target_creator = anchor_target_creator
        self.proposal_target_creator = proposal_target_creator

        self.loc_normalize_mean = faster_rcnn.loc_normalize_mean
        self.loc_normalize_std = faster_rcnn.loc_normalize_std

    def __call__(self, imgs, bboxes, labels, scales):
        """Forward Faster R-CNN and calculate losses.

        Here are notations used.

        * :math:`N` is the batch size.
        * :math:`R` is the number of bounding boxes per image.

        Currently, only :math:`N=1` is supported.

        Args:
            imgs (~chainer.Variable): A variable with a batch of images.
            bboxes (~chainer.Variable): A batch of bounding boxes.
                Its shape is :math:`(N, R, 4)`.
            labels (~chainer.Variable): A batch of labels.
                Its shape is :math:`(N, R)`. The background is excluded from
                the definition, which means that the range of the value
                is :math:`[0, L - 1]`. :math:`L` is the number of foreground
                classes.
            scales (~chainer.Variable): Amount of scaling applied to
                each input image during preprocessing.

        Returns:
            chainer.Variable:
            Scalar loss variable.
            This is the sum of losses for Region Proposal Network and
            the head module.

        """
        if isinstance(bboxes, chainer.Variable):
            bboxes = bboxes.array
        if isinstance(labels, chainer.Variable):
            labels = labels.array
        if isinstance(scales, chainer.Variable):
            scales = scales.array

        n = bboxes.shape[0]
        if n != 1:
            raise ValueError('Currently only batch size 1 is supported.')
        scales = cuda.to_cpu(scales).tolist()

        _, _, H, W = imgs.shape
        img_size = (H, W)

        features = self.faster_rcnn.extractor(imgs)
        rpn_locs, rpn_scores, rois, roi_indices, anchor = self.faster_rcnn.rpn(
            features, img_size, scales)

        # Since batch size is one, convert variables to singular form
        bbox = bboxes[0]
        label = labels[0]
        rpn_score = rpn_scores[0]
        rpn_loc = rpn_locs[0]
        roi = rois

        # Sample RoIs and forward
        sample_roi, gt_roi_loc, gt_roi_label = self.proposal_target_creator(
            roi, bbox, label,
            self.loc_normalize_mean, self.loc_normalize_std)
        sample_roi_index = self.xp.zeros((len(sample_roi),), dtype=np.int32)
        roi_cls_loc, roi_score = self.faster_rcnn.head(
            features, sample_roi, sample_roi_index)

        # RPN losses
        gt_rpn_loc, gt_rpn_label = self.anchor_target_creator(
            bbox, anchor, img_size)
        rpn_loc_loss = _fast_rcnn_loc_loss(
            rpn_loc, gt_rpn_loc, gt_rpn_label, self.rpn_sigma)
        rpn_cls_loss = F.softmax_cross_entropy(rpn_score, gt_rpn_label)

        # Losses for outputs of the head.
        n_sample = roi_cls_loc.shape[0]
        roi_cls_loc = roi_cls_loc.reshape((n_sample, -1, 4))
        roi_loc = roi_cls_loc[self.xp.arange(n_sample), gt_roi_label]
        roi_loc_loss = _fast_rcnn_loc_loss(
            roi_loc, gt_roi_loc, gt_roi_label, self.roi_sigma)
        roi_cls_loss = F.softmax_cross_entropy(roi_score, gt_roi_label)

        loss = rpn_loc_loss + rpn_cls_loss + roi_loc_loss + roi_cls_loss
        chainer.reporter.report({'rpn_loc_loss': rpn_loc_loss,
                                 'rpn_cls_loss': rpn_cls_loss,
                                 'roi_loc_loss': roi_loc_loss,
                                 'roi_cls_loss': roi_cls_loss,
                                 'loss': loss},
                                self)
        return loss

def _smooth_l1_loss(x, t, in_weight, sigma):
    sigma2 = sigma ** 2
    diff = in_weight * (x - t)
    abs_diff = F.absolute(diff)
    flag = (abs_diff.array < (1. / sigma2)).astype(np.float32)

    y = (flag * (sigma2 / 2.) * F.square(diff) +
         (1 - flag) * (abs_diff - 0.5 / sigma2))

    return F.sum(y)

def _fast_rcnn_loc_loss(pred_loc, gt_loc, gt_label, sigma):
    xp = chainer.backends.cuda.get_array_module(pred_loc)

    in_weight = xp.zeros_like(gt_loc)
    # Localization loss is calculated only for positive rois.
    in_weight[gt_label > 0] = 1
    loc_loss = _smooth_l1_loss(pred_loc, gt_loc, in_weight, sigma)
    # Normalize by total number of negtive and positive rois.
    loc_loss /= xp.sum(gt_label >= 0)
    return loc_loss


# -*- coding: utf-8 -*-

import tensorflow as tf
import numpy as np

from . import actor_critic as ac
from .. import config_info as ci
from . import common as rl_common
from ..module import common as ml_common

class PolicyUpdater:
    def policy_loss(self, pi, pi_old, dist, dist_old, adv, num_actions):
        raise NotImplementedError('PolicyUpdater::policy_loss is not implemented!')

    def updating_ctrl(self):
        raise NotImplementedError('PolicyUpdater::updating_ctrl is not implemented!')

    def after_actor_update(self):
        raise NotImplementedError('PolicyUpdater::after_actor_update is not implemented!')

class KLPenaltyPolicyUpdater(PolicyUpdater):
    def __init__(self, lam, KL_target, xi, alpha, beta_high, beta_low, lam_lower, lam_higher, num_e):
        self.lam=tf.Variable(lam, trainable=False)
        self.KL_target=KL_target
        self.xi=xi
        self.alpha=alpha
        self.beta_high=beta_high
        self.beta_low=beta_low
        self.lam_lower=lam_lower
        self.lam_higher=lam_higher
        self.num_e=num_e

    def policy_loss(self, pi, pi_old, dist, dist_old, adv, continuous_dist):
        self.ratio=pi/(pi_old+self.num_e)
        self.surr=tf.multiply(self.ratio, adv)
        if continuous_dist:
            self.kl=tf.distributions.kl_divergence(dist_old, dist)
        else:
            self.kl=tf.reduce_sum(tf.multiply(dist_old, tf.log(dist_old/(dist+self.num_e)+self.num_e)), axis=-1, keep_dims=True)
        self.kl_mean=tf.reduce_mean(self.kl)
        loss=self.lam*self.kl+self.xi*tf.square(tf.maximum(0.0, self.kl-2*self.KL_target))-self.surr
        return loss

    def updating_ctrl(self):
        return tf.less_equal(self.kl_mean, 4*self.KL_target)

    def after_actor_update(self):
        return tf.cond(tf.greater(self.kl_mean, self.beta_high*self.KL_target)
                     , lambda : tf.assign(self.lam, tf.clip_by_value(self.lam*self.alpha, self.lam_lower, self.lam_higher))
                     , lambda : tf.cond(tf.less(self.kl_mean, self.beta_low*self.KL_target)
                                      , lambda : tf.assign(self.lam, tf.clip_by_value(self.lam/self.alpha, self.lam_lower, self.lam_higher))
                                      , lambda : tf.assign(self.lam, self.lam)))

class ClipPolicyUpdater(PolicyUpdater):
    def __init__(self, epsilon, c, num_e):
        self.epsilon=epsilon
        self.c=c
        self.num_e=num_e

    def policy_loss(self, pi, pi_old, dist, dist_old, adv, continuous_dist):
        self.ratio=pi/(pi_old+self.num_e)
        self.surr=tf.multiply(self.ratio, adv)
        loss=-tf.minimum(self.surr, tf.clip_by_value(self.ratio, 1-self.epsilon, 1+self.epsilon)*adv)
        if self.c>0:
            if continuous_dist:
                loss-=self.c*dist.entropy()
            else:
                loss+=self.c*tf.reduce_sum(tf.multiply(dist, tf.log(dist+self.num_e)), axis=-1, keep_dims=True)
        return loss

    def updating_ctrl(self):
        return tf.constant(True, dtype=tf.bool)

    def after_actor_update(self):
        return tf.no_op()

class ProximalPolicyOptimizer(rl_common.RL):
    def __init__(self, running_type
               , actor_net_generator, state, sequence_length
               , discrete_action_num_list, action_lower, action_upper, sigma_delta
               , critic_net_generator, action, state_, reward
               , use_double, use_dueling, add_actor_action_to_critic
               , critic_tau, replace_critic_target_span, estimate_val_instead
               , policy_updater_list, actor_update_num, critic_update_num, recursive_reward_target
               , actor_learning_rate, critic_learning_rate, gamma
               , max_grad_value, max_grad_norm, optimizer_getter
               , loss_mask=None
               , extra_actor_eval_vars=[], extra_actor_target_vars=[]
               , extra_critic_eval_vars=[], extra_critic_target_vars=[]):
        self.state=state
        self.sequence_length=sequence_length
        self.action=action
        self.state_=state_
        self.reward=reward
        self.running_type=running_type
        self.discrete_action_num_list=discrete_action_num_list
        self.use_double=use_double
        self.replace_critic_target_span=replace_critic_target_span
        self.actor_learning_rate=actor_learning_rate
        self.critic_learning_rate=critic_learning_rate
        self.gamma=gamma
        self.policy_updater_list=policy_updater_list
        self.actor_update_num=actor_update_num
        self.critic_update_num=critic_update_num
        if loss_mask is not None:
            self.loss_mask=tf.expand_dims(loss_mask, axis=-1)
            self.mask_count=tf.maximum(tf.reduce_sum(loss_mask), 1)
        else:
            self.loss_mask=None
            self.mask_count=None

        with tf.variable_scope('actor'):
            self.actor=ac.Actor(running_type
                              , actor_net_generator, self.state, discrete_action_num_list
                              , True, 1
                              , action_lower, action_upper, sigma_delta, False
                              , self.state
                              , extra_actor_eval_vars, extra_actor_target_vars)
        self.discrete_action_op, self.continuous_action_op=self.actor.get_eval_op()
        if self.continuous_action_op is None:
            self.continuous_action_op=[]
        else:
            self.continuous_action_op=[self.continuous_action_op[1]]

        if running_type==ci.RunningType.train or running_type==ci.RunningType.valid or running_type==ci.RunningType.test:
            discrete_action_info_eval, continuous_action_info_eval=self.actor.get_eval_op()
            discrete_action_info_target, continuous_action_info_target=self.actor.get_target_op()
            discrete_eval_action_vec_list=[tf.one_hot(act, num_actions) for act, num_actions in zip(self.action[:len(discrete_action_num_list)], discrete_action_num_list)]

            with tf.variable_scope('critic_input'):
                if add_actor_action_to_critic:
                    discrete_action_info_target, continuous_action_info_target=self.actor.get_target_op()
                    critic_input=(self.state, discrete_eval_action_vec_list, None if continuous_action_info_eval is None else self.action[len(discrete_action_num_list)])
                    critic_input_=(self.state_, discrete_action_info_target, None if continuous_action_info_target is None else continuous_action_info_target[1])
                    #critic_input=tf.concat([self.state]+discrete_eval_action_vec_list+([] if continuous_action_info_eval is None else [self.action[len(discrete_action_num_list)]]), axis=-1)
                    #critic_input_=tf.concat([self.state_]+discrete_action_info_target+([] if continuous_action_info_target is None else [continuous_action_info_target[1]]), axis=-1)
                else:
                    critic_input=self.state
                    critic_input_=self.state_

            with tf.variable_scope('critic'):
                self.critic=ac.Critic(running_type
                                    , critic_net_generator, critic_input, critic_input_, None, discrete_action_num_list, 0 if action_lower is None else len(action_lower)
                                    , use_dueling, critic_tau if use_double else 1, estimate_val_instead
                                    , extra_critic_eval_vars, extra_critic_target_vars)

            if len(discrete_action_info_eval)!=len(self.critic.get_eval_op()[0]):
                raise ValueError('ProximalPolicyOptimizer::__init__: Discrete count mismatch in actor ({}) and critic ({})!'.format(len(discrete_action_info_eval), len(self.critic.get_eval_op()[0])))
            if continuous_action_info_eval is not None and continuous_action_info_eval[1].shape[-1].value!=self.critic.get_eval_op()[1][0].shape[-1].value:
                raise ValueError('ProximalPolicyOptimizer::__init__: Continuous count mismatch in actor ({}) and critic ({})!'.format(continuous_action_info_eval[1].shape[-1].value, self.critic.get_eval_op()[1][0].shape[-1].value))

            with tf.variable_scope('critic_loss'):
                discrete_q_info_eval, continuous_q_info_eval=self.critic.get_eval_op()
                discrete_q_info_target, continuous_q_info_target=self.critic.get_target_op()
                loss_param_list=[]
                for idx, (action_vec, q_info_eval, q_info_target) in enumerate(zip(discrete_eval_action_vec_list, discrete_q_info_eval, discrete_q_info_target)):
                    if estimate_val_instead:
                        the_q_eval=q_info_eval[1]
                        target_op=q_info_target[1]
                    else:
                        the_q_eval=tf.reduce_sum(tf.multiply(action_vec, q_info_eval[0]), axis=-1, keep_dims=True)
                        target_op=tf.reduce_sum(tf.multiply(action_vec, q_info_target[0]), axis=-1, keep_dims=True)
                    loss_param_list.append(('discrete_{}'.format(idx), the_q_eval, target_op))
                if continuous_q_info_eval is not None:
                    loss_param_list.append(('continuous', continuous_q_info_eval, continuous_q_info_target))
                q_target_list=[]
                critic_loss_list=[]
                for (scope_name, the_q_eval, target_op) in loss_param_list:
                    with tf.variable_scope('q_target'):
                        with tf.variable_scope(scope_name):
                            if recursive_reward_target:
                                def update_q_target(i, last_target, the_q_target_array):
                                    last_target=self.reward[:, i]+self.gamma*last_target
                                    the_q_target_array=the_q_target_array.write(i, last_target)
                                    return [i-1, last_target, the_q_target_array]
                                _, _, the_q_target_array=tf.while_loop(lambda i, last_target, the_q_target_array : i>=0
                                                                     , update_q_target
                                                                     , [self.sequence_length-1, target_op[:,-1], tf.TensorArray(dtype=target_op.dtype, size=self.sequence_length)]
                                                                     , parallel_iterations=1)
                                the_q_target=tf.transpose(the_q_target_array.stack(), perm=[1, 0]+[i for i in range(2, len(target_op.shape))])
                            else:
                                the_q_target=self.reward+self.gamma*target_op
                        q_target_list.append(the_q_target)
                        critic_loss_list.append(tf.squared_difference(the_q_eval, the_q_target))
                self.critic_loss=tf.concat(critic_loss_list, axis=-1, name='critic_loss')
                if self.loss_mask is not None:
                    self.critic_cost=tf.truediv(tf.reduce_sum(self.critic_loss*self.loss_mask), self.mask_count, name='critic_cost')
                else:
                    self.critic_cost=tf.reduce_mean(tf.reduce_sum(self.critic_loss, axis=-1), name='critic_cost')

            discrete_critic_info_eval, continuous_critic_info_eval=self.critic.get_eval_op()
            discrete_critic_info_target, continuous_critic_info_target=self.critic.get_target_op()
            with tf.variable_scope('actor_loss'):
                discrete_adv_list=[]
                for q_target, (q_eval, val_eval, adv_eval) in zip(q_target_list[:len(discrete_action_num_list)], discrete_critic_info_eval):
                    if use_dueling:
                        discrete_adv_list.append(adv_eval)
                    elif estimate_val_instead:
                        discrete_adv_list.append(q_target-val_eval)
                    else:
                        discrete_adv_list.append(q_target-tf.reduce_sum(tf.multiply(action_vec, q_eval), axis=-1, keep_dims=True))
                actor_loss_list=[]
                actor_updating_ctrl_op_list=[]
                self.after_actor_update_op=[]
                for action_vec, adv, action_prob_eval, action_prob_target, policy_updater in zip(discrete_eval_action_vec_list, discrete_adv_list, discrete_action_info_eval, discrete_action_info_target, self.policy_updater_list[:len(discrete_action_num_list)]):
                    pi=tf.reduce_sum(tf.multiply(action_vec, action_prob_eval), axis=-1, keep_dims=True)
                    pi_old=tf.reduce_sum(tf.multiply(action_vec, action_prob_target), axis=-1, keep_dims=True)
                    dist=action_prob_eval
                    dist_old=action_prob_target
                    actor_loss_list.append(policy_updater.policy_loss(pi, pi_old, dist, dist_old, adv, False))
                    actor_updating_ctrl_op_list.append(policy_updater.updating_ctrl())
                    self.after_actor_update_op.append(policy_updater.after_actor_update())
                if continuous_critic_info_target is not None:
                    adv=q_target_list[len(discrete_action_num_list)]-continuous_critic_info_eval
                    act=self.action[len(discrete_action_num_list)]
                    normal_dist_eval, _=continuous_action_info_eval
                    normal_dist_target, _=continuous_action_info_target
                    policy_updater=self.policy_updater_list[len(discrete_action_num_list)]
                    pi=normal_dist_eval.prob(act)
                    pi_old=normal_dist_target.prob(act)
                    dist=normal_dist_eval
                    dist_old=normal_dist_target
                    actor_loss_list.append(policy_updater.policy_loss(pi, pi_old, dist, dist_old, adv, True))
                    actor_updating_ctrl_op_list.append(policy_updater.updating_ctrl())
                    self.after_actor_update_op.append(policy_updater.after_actor_update())
                self.actor_loss=tf.concat(actor_loss_list, axis=-1, name='actor_loss')
                if self.loss_mask is not None:
                    self.actor_cost=tf.truediv(tf.reduce_sum(self.actor_loss*self.loss_mask), self.mask_count, name='actor_cost')
                else:
                    self.actor_cost=tf.reduce_mean(tf.reduce_sum(self.actor_loss, axis=-1), name='actor_cost')
                self.actor_updating_ctrl_op=tf.reduce_all(tf.stack(actor_updating_ctrl_op_list))

            if running_type==ci.RunningType.train:
                with tf.variable_scope('training_actor'):
                    self.actor_lr, self.actor_optimizer, self.actor_grads, self.actor_train_op, self.actor_global_step=ml_common.add_train_op(cost=self.actor_cost, learning_rate=self.actor_learning_rate, optimizer_getter=optimizer_getter, max_grad_value=max_grad_value, max_grad_norm=max_grad_norm, global_step=None, var_list=self.actor.get_trainable_variables())
                with tf.variable_scope('training_critic'):
                    self.critic_lr, self.critic_optimizer, self.critic_grads, self.critic_train_op, self.critic_global_step=ml_common.add_train_op(cost=self.critic_cost, learning_rate=self.critic_learning_rate, optimizer_getter=optimizer_getter, max_grad_value=max_grad_value, max_grad_norm=max_grad_norm, global_step=None, var_list=self.critic.get_trainable_variables())
                self.replace_actor_target_op=self.actor.get_replace_target_op()
                self.replace_critic_target_op=self.critic.get_replace_target_op()

    def get_continuous_lower_op(self):
        return self.actor.action_lower

    def get_continuous_upper_op(self):
        return self.actor.action_upper

    def replace_target(self, sess, iters):
        if self.running_type==ci.RunningType.train:
            sess.run(self.replace_actor_target_op)
            if not self.use_double or iters%self.replace_critic_target_span==0:
                sess.run(self.replace_critic_target_op)

    def run_sess_and_cost(self, sess, for_training, iters):
        self.replace_target(sess, iters)
        cost=0
        if for_training:
            for i in range(self.actor_update_num):
                _, actor_cost, updating_ctrl=sess.run([self.actor_train_op, self.actor_cost, self.actor_updating_ctrl_op])
                cost+=actor_cost
                if updating_ctrl==False:
                    i+=1
                    break
            sess.run(self.after_actor_update_op)
            cost/=i
            cost_critic=0
            for i in range(self.critic_update_num):
                _, critic_cost=sess.run([self.critic_train_op, self.critic_cost])
                cost_critic+=critic_cost
            cost+=cost_critic/self.critic_update_num
        else:
            actor_cost, critic_cost=sess.run([self.actor_cost, self.critic_cost])
            cost+=actor_cost+critic_cost
        return cost

    def run_sess_and_cost_for_placeholder(self, sess, for_training, iters, feed_dict):
        self.replace_target(sess, iters)
        cost=0
        if for_training:
            for i in range(self.actor_update_num):
                _, actor_cost, updating_ctrl=sess.run([self.actor_train_op, self.actor_cost, self.actor_updating_ctrl_op], feed_dict=feed_dict)
                cost+=actor_cost
                if updating_ctrl==False:
                    break
            sess.run(self.after_actor_update_op, feed_dict=feed_dict)
            cost/=i+1
            cost_critic=0
            for i in range(self.critic_update_num):
                _, critic_cost=sess.run([self.critic_train_op, self.critic_cost], feed_dict=feed_dict)
                cost_critic+=critic_cost
            cost+=cost_critic/self.critic_update_num
        else:
            actor_cost, critic_cost=sess.run([self.actor_cost, self.critic_cost], feed_dict=feed_dict)
            cost+=actor_cost+critic_cost
        return cost


"""main server script
will sit onboard host and operate as Nebula --- its dynamic soul"""

# --------------------------------------------------
#
# Embodied AI Engine Prototype v0.10
# 2021/01/25
#
# © Craig Vear 2020
# cvear@dmu.ac.uk
#
# Dedicated to Fabrizio Poltronieri
#
# --------------------------------------------------

import trio
from random import randrange, random
from time import time
from tensorflow.keras.models import load_model
import numpy as np
import sys
import pickle

# --------------------------------------------------
#
# instantiate an object for each neural net
#
# --------------------------------------------------

# v4 models were trained with 1st batch of Blue Haze datasets

class MoveRNN:
    def __init__(self):
        print('MoveRNN initialization')
        self.move_rnn = load_model('models/EMR-3_RNN_skeleton_data.nose.x.h5')

    def predict(self, in_val):
        # predictions and input with localval
        self.pred = self.move_rnn.predict(in_val)
        return self.pred

class AffectRNN:
    def __init__(self):
        print('AffectRNN initialization')
        self.affect_rnn = load_model('models/EMR-3_RNN_bitalino.h5')

    def predict(self, in_val):
        # predictions and input with localval
        self.pred = self.affect_rnn.predict(in_val)
        return self.pred

class MoveAffectCONV2:
    def __init__(self):
        print('MoveAffectCONV2 initialization')
        self.move_affect_conv2 = load_model('models/EMR-3_conv2D_move-affect.h5')

    def predict(self, in_val):
        # predictions and input with localval
        self.pred = self.move_affect_conv2.predict(in_val)
        return self.pred

class AffectMoveCONV2:
    def __init__(self):
        print('AffectMoveCONV2 initialization')
        self.affect_move_conv2 = load_model('models/EMR-3_conv2D_affect-move.h5')

    def predict(self, in_val):
        # predictions and input with localval
        self.pred = self.affect_move_conv2.predict(in_val)
        return self.pred

# --------------------------------------------------
#
# controls all thought-trains and affect responses
#
# --------------------------------------------------

class AiDataEngine():
    def __init__(self, speed=1):
        print('building engine server')
        self.interrupt_bang = False
        self.running = False
        # self.PORT = 8000
        # self.IP_ADDR = "127.0.0.1"
        self.global_speed = speed
        self.rnd_stream = 0

        # make a default dict for the engine
        self.datadict = {'move_rnn': 0,
                         'affect_rnn': 0,
                         'move_affect_conv2': 0,
                         'affect_move_conv2': 0,
                         'master_output': 0,
                         'user_in': 0,
                         'rnd_poetry': 0,
                         'rhythm_rnn': 0,
                         'affect_net': 0,
                         'self_awareness': 0,
                         'affect_decision': 0,
                         'rhythm_rate': 0.1}

        # name list for nets
        self.netnames = ['move_rnn',
                         'affect_rnn',
                         'move_affect_conv2',
                         'affect_move_conv2',
                         'self_awareness',  # Net name for self-awareness
                         'master_output']  # input for self-awareness

        # names for affect listening
        self.affectnames = ['user_in',
                            'rnd_poetry',
                            'affect_net',
                            'self_awareness']

        self.rhythm_rate = 0.1
        self.affect_listen = 0

        # fill with random values
        self.dict_fill()
        print(self.datadict)

        # instantiate nets as objects and make  models
        self.move_net = MoveRNN()
        self.affect_net = AffectRNN()
        self.move_affect_net = MoveAffectCONV2()
        self.affect_move_net = AffectMoveCONV2()
        self.affect_perception = MoveAffectCONV2()

        # logging on/off switches
        self.net_logging = False
        self.master_logging = True
        self.streaming_logging = False
        self.affect_logging = False

    # --------------------------------------------------
    #
    # prediction and rnd num gen zone
    #
    # --------------------------------------------------

    # makes a prediction for a given net and defined input var
    async def make_data(self):
        while self.running:

            # calc rhythmic intensity based on self-awareness factor & global speed
            intensity = self.datadict.get('self_awareness')
            self.rhythm_rate = (self.rhythm_rate * intensity) * self.global_speed
            self.datadict['rhythm_rate'] = self.rhythm_rate

            # get input vars from dict (NB not always self)
            in_val1 = self.get_in_val(0) # move RNN as input
            in_val2 = self.get_in_val(1) # affect RNN as input
            in_val3 = self.get_in_val(2) # move - affect as input
            in_val4 = self.get_in_val(1) # affect RNN as input

            # send in vals to net object for prediction
            pred1 = self.move_net.predict(in_val1)
            pred2 = self.affect_net.predict(in_val2)
            pred3 = self.move_affect_net.predict(in_val3)
            pred4 = self.affect_move_net.predict(in_val4)

            # special case for self awareness stream
            self_aware_input = self.get_in_val(5) # main movement as input
            self_aware_pred = self.affect_perception.predict(self_aware_input)

            if self.net_logging:
                print(f"  'move_rnn' in: {in_val1} predicted {pred1}")
                print(f"  'affect_rnn' in: {in_val2} predicted {pred2}")
                print(f"  move_affect_conv2' in: {in_val3} predicted {pred3}")
                print(f"  'affect_move_conv2' in: {in_val4} predicted {pred4}")
                print(f"  'self_awareness' in: {self_aware_input} predicted {self_aware_pred}")

            # put predictions back into the dicts and master
            self.put_pred(0, pred1)
            self.put_pred(1, pred2)
            self.put_pred(2, pred3)
            self.put_pred(3, pred4)
            self.put_pred(4, self_aware_pred)

            # outputs a stream of random poetry
            rnd_poetry = random()
            self.datadict['rnd_poetry'] = random()
            if self.streaming_logging:
                print(f'random poetry = {rnd_poetry}')

            await trio.sleep(self.rhythm_rate)

    # function to get input value for net prediction from dictionary
    def get_in_val(self, which_dict):
        # get the current value and reshape ready for input for prediction
        input_val = self.datadict.get(self.netnames[which_dict])
        input_val = np.reshape(input_val, (1, 1, 1))
        return input_val

    # function to put prediction value from net into dictionary
    def put_pred(self, which_dict, pred):
        # randomly chooses one of te 4 predicted outputs
        out_pred_val = pred[0][randrange(4)]
        if self.master_logging:
            print(f"out pred val == {out_pred_val},   master move output == {self.datadict['master_output']}")
        # save to data dict and master move out ONLY 1st data
        self.datadict[self.netnames[which_dict]] = out_pred_val
        self.datadict['master_output'] = out_pred_val

    # fills the dictionary with rnd values for each key of data dictionary
    def dict_fill(self):
        for key in self.datadict.keys():
            rnd = random()
            self.datadict[key] = rnd

    # --------------------------------------------------
    #
    # affect and streaming methods
    #
    # --------------------------------------------------

    # define which feed to listen to, and duration
    # and a course of affect response
    async def affect(self):
        # daddy cycle = is the master running on?
        while self.running:
            if self.affect_logging:
                print('\t\t\t\t\t\t\t\t=========HIYA - DADDY cycle===========')

            # flag for breaking on big affect signal
            self.interrupt_bang = True

            # calc master cycle before a change
            master_cycle = randrange(6, 26) * self.global_speed
            loop_dur = time() + master_cycle
            if self.affect_logging:
                print(f"                 interrupt_listener: started! sleeping now for {loop_dur}...")

            # refill the dicts?????
            self.dict_fill()

            # child cycle - waiting for interrupt  from master clock
            while time() < loop_dur:
                if self.affect_logging:
                        print('\t\t\t\t\t\t\t\t=========Hello - child cycle 1 ===========')

                # if a major break out then go to Daddy cycle
                if not self.interrupt_bang:
                    break

                # randomly pick an input stream for this cycle
                rnd = randrange(4)
                self.rnd_stream = self.affectnames[rnd]
                self.datadict['affect_decision'] = rnd
                print(self.rnd_stream)
                if self.affect_logging:
                    print(self.rnd_stream)

                # hold this stream for 1-4 secs, unless interrupt bang
                end_time = time() + (randrange(1000, 4000) / 1000)
                if self.affect_logging:
                    print('end time = ', end_time)

                # baby cycle 2 - own time loops
                while time() < end_time:
                    if self.affect_logging:
                        print('\t\t\t\t\t\t\t\t=========Hello - baby cycle 2 ===========')

                    # go get the current value from dict
                    affect_listen = self.datadict[self.rnd_stream]
                    if self.affect_logging:
                        print('current value =', affect_listen)

                    # make the master output the current value of the stream
                    self.datadict['master_output'] = affect_listen
                    if self.master_logging:
                        print(f'\t\t ==============  master move output = {affect_listen}')

                    # calc affect on behaviour
                    # if input stream is LOUD then smash a random fill and break out to Daddy cycle...
                    if affect_listen > 0.60:
                        if self.affect_logging:
                            print('interrupt > HIGH !!!!!!!!!')

                        # A - refill dict with random
                        self.dict_fill()

                        # B - jumps out of this loop into daddy
                        self.interrupt_bang = False
                        if self.affect_logging:
                            print('interrupt bang = ', self.interrupt_bang)

                        # C break out of this loop, and next (cos of flag)
                        break

                    # if middle loud fill dict with random, all processes norm
                    elif 0.30 < affect_listen < 0.59:
                        if self.affect_logging:
                            print('interrupt MIDDLE -----------')
                            print('interrupt bang = ', self.interrupt_bang)

                        # refill dict with random
                        self.dict_fill()

                    elif affect_listen <= 0.30:
                        if self.affect_logging:
                            print('interrupt LOW_______________')
                            print('interrupt bang = ', self.interrupt_bang)

                    # and wait for a cycle
                    await trio.sleep(self.rhythm_rate)

    # --------------------------------------------------
    #
    # parent threading solution
    #
    # --------------------------------------------------

    async def flywheel(self):
        print("parent: started!")
        while self.running:
            # print("parent: connecting to 127.0.0.1:{}".format(self.PORT))
            # client_stream = await trio.open_tcp_stream(self.IP_ADDR, self.PORT)
            # async with client_stream:
            #     # self.interrupt_bang = True
            async with trio.open_nursery() as nursery:
                # spawning all the data making
                print("parent: spawning making data ...")
                nursery.start_soon(self.make_data)

                # spawning affect listener and master clocks
                print("parent: spawning affect listener and clocks ...")
                nursery.start_soon(self.affect)

                # # spawning listening port for user input
                # print("parent: spawning receiver...")
                # nursery.start_soon(self.receiver, client_stream)
                #
                # # spawning sending port for output data
                # print("parent: spawning sender...")
                # nursery.start_soon(self.sender, client_stream)

    # --------------------------------------------------
    #
    # user accessible methods
    #
    # --------------------------------------------------

    # returns the live output from the class to user
    async def sender(self, client_stream):
        print("sender: started!")
        while self.running:
            # data = {'master_AI_output': self.datadict.get('master_output'),
            #         'intensity/rhythm': self.datadict.get('self_awareness'),
            #         'affect stream': self.rnd_stream,
            #         'individual NN outs':
            #             {'move RNN': self.datadict.get('move_rnn'),
            #              'affect RNN': self.datadict.get('affect_rnn'),
            #              'move_affect_conv2': self.datadict.get('move_affect_conv2'),
            #              'affect_move_conv2': self.datadict.get('affect_move_conv2'),
            #              'self_awareness': self.datadict.get('self_awareness')
            #              }
            #         }

            # serialise dict into json for Tx
            serial_data = pickle.dumps(self.datadict, -1)
            if self.streaming_logging:
                print("sender: sending {!r}".format(serial_data))
            #  & send
            await client_stream.send_all(serial_data)
            await trio.sleep(self.rhythm_rate)

    # receives user data from client (typically live audio input)
    async def receiver(self, client_stream):
        print("receiver: started!")
        while self.running:
            async for data in client_stream:
                load_data = pickle.loads(data)
                self.parse_got_dict(load_data)

                # self.datadict['user_in'] = load_data['mic_level']
                if self.streaming_logging:
                    print("receiver: got data {!r}".format(load_data))
            print("receiver: connection closed")
        sys.exit()

    def parse_got_dict(self, got_dict):
        self.datadict['user_in'] = got_dict['mic_level']

        # user change the overall speed of the engine
        self.global_speed = got_dict['speed']

        # user change tempo of outputs and parsing
        self.rhythm_rate = got_dict['tempo']

    # stop start methods
    def go(self):
        self.running = True
        trio.run(self.flywheel)
        print('I got here daddy')

    def quit(self):
        self.running = False

if __name__ == '__main__':
    engine = AiDataEngine()
    engine.go()


import numpy as np
from Game import Game
import random
from MCTS import MCTS
from Player import Player
import config
from ResNet import ResNet_Training
import os
import ResNet
import torch
import ray
import copy
from tqdm import tqdm
import json

ray.init(num_gpus=4)

def main():
    nn_budget = 100
    improved = 0
    best_nn = load_or_create_neural_net()
    best_nn.to(torch.device("cuda"))
    elos = []
    last_elo = 30
    data = []

    for i in range(1000):
        new_data = self_play(best_nn, nn_budget + improved)

        previous_best = copy.deepcopy(best_nn)
        improve_model_resnet(best_nn, new_data, improved)

        print('Compering the new NN vs the old one')
        winrate = test_player(best_nn, nn_budget, previous_best, nn_budget, 10)
        print("Winrate: " + str(winrate))
        if winrate < 0.52:
            best_nn = previous_best
        else:
            improved += 1
            #  last_elo = check_elo(best_nn, nn_budget, last_elo)
            #  elos.append(last_elo)
            torch.save(best_nn.state_dict(), './best_model_resnet.pth')
            with open('elos.txt', 'w') as filehandle:
                json.dump(elos, filehandle)

    print(elos)

def test_player(main_player_nn, main_player_budget, compare_player_nn, compare_player_budget, loops):
    games_count = 0
    main_player_wins = 0
    main_player_nn_id = ray.put(main_player_nn)
    compare_player_nn_id = ray.put(compare_player_nn)
    for i in range(loops):
        result_ids = []
        for _ in range(config.CPUS):
            if i < loops // 2:  # check if making a copy is enough
                result_ids.append(
                    [play_game.remote(main_player_nn_id, main_player_budget, compare_player_nn_id,
                                      compare_player_budget, False), 1])
            else:
                result_ids.append(
                    [play_game.remote(compare_player_nn_id, compare_player_budget, main_player_nn_id,
                                      main_player_budget, False), -1])

        for result_id, player_number in result_ids:
            winner, data = ray.get(result_id)
            if winner == player_number:
                main_player_wins += 1
            if winner != 0:
                games_count += 1

    del main_player_nn_id
    del compare_player_nn_id

    winrate = main_player_wins / games_count

    return winrate

def check_elo(nn, budget, start_elo):
    mcts_budget = start_elo
    print('Playing vs MCTS')
    for _ in range(10):
        winrate = test_player(nn, budget, False, mcts_budget, 2)
        if winrate > 0.5:
            mcts_budget += 10
        else:
            mcts_budget = max(0, mcts_budget - 10)
            if mcts_budget == 0:
                break

    print("Elo: " + str(mcts_budget))
    return mcts_budget

def self_play(nn, budget):
    print('Self playing')

    cpus = config.CPUS
    if not nn:
        cpus = 96
    winners = []
    all_data = []
    nn_id = ray.put(nn)
    for _ in tqdm(range(30)):
        result_ids = []
        for i in range(cpus):
            result_ids.append(play_game.remote(nn_id, budget, nn_id, budget, True))

        for result_id in result_ids:
            winner, data = ray.get(result_id)
            winners.append(winner)
            if len(all_data) == 0:
                all_data = data
            else:
                all_data = np.vstack((all_data, data))
    del nn_id

    return all_data

@ray.remote(num_gpus=0.2)
def play_game_with_gpu(player_one_nn, player_one_budget, player_two_nn, player_two_budget, is_selfplay):
    return play_game(player_one_nn, player_one_budget, player_two_nn, player_two_budget, is_selfplay)

@ray.remote
def play_game_no_gpu(player_one_nn, player_one_budget, player_two_nn, player_two_budget, is_selfplay):
    return play_game(player_one_nn, player_one_budget, player_two_nn, player_two_budget, is_selfplay)

@ray.remote(num_gpus=0.2)
def play_game(player_one_nn, player_one_budget, player_two_nn, player_two_budget, is_selfplay):
    player_one = Player(player_one_nn, player_one_budget, is_selfplay)
    player_two = Player(player_two_nn, player_two_budget, is_selfplay)
    gameover, winner, turn = False, 0, 0
    game = Game()
    new_data_for_the_game = np.zeros((3 * config.L * config.H + config.L + 1))

    while not gameover:
        turn += 1
        if turn % 2 == 1:
            move, turn_data = player_one.get_move(game.state, turn)
        else:
            move, turn_data = player_two.get_move(game.state, turn)

        if is_selfplay:
            new_data_for_the_game = np.vstack((new_data_for_the_game, turn_data))

        game.takestep(move)
        gameover, winner = game.gameover()

    if is_selfplay:
        new_data_for_the_game = np.delete(new_data_for_the_game, 0, 0)
        history_size = new_data_for_the_game.shape[0]

        z = 0
        if winner == 1:
            z = 1 - config.long_game_factor * history_size / 42  # the reward is bigger for shorter games
        elif winner == -1:
            z = -1 + config.long_game_factor * history_size / 42  # the reward is less negative for long games

        z_vec = np.zeros(history_size)

        for i in range(history_size):
            z_vec[i] = ((-1) ** i) * z

        new_data_for_the_game[:, -1] = z_vec
        new_data_for_the_game = extend_data(new_data_for_the_game)

    return winner, new_data_for_the_game

def extend_data(data):
    board_size = config.L * config.H
    extended_data = np.zeros((data.shape[0], data.shape[1]))

    for i in range(extended_data.shape[0]):
        board = np.copy(data[i, 0:3 * board_size]).reshape((3, config.H, config.L))
        yellowboard = board[0]
        redboard = board[1]
        player_turn = board[2]

        # parity operation on array for both yellow and red boards
        flip_yellow = np.fliplr(yellowboard)
        flip_red = np.fliplr(redboard)

        extended_data[i, 0:board_size] = flip_yellow.flatten()
        extended_data[i, board_size:2 * board_size] = flip_red.flatten()
        extended_data[i, 2 * board_size: 3 * board_size] = player_turn.flatten()

        # parity operation on the Pi's
        pi_s = np.copy(data[i, 3 * board_size:3 * board_size + config.L])
        flip_pi = np.flip(pi_s, axis=0)

        extended_data[i, 3 * board_size:3 * board_size + config.L] = flip_pi
        extended_data[i, -1] = np.copy(data[i, -1])

    return np.vstack((data, extended_data))

def load_or_create_neural_net():
    file_path = './best_model_resnet.pth'
    if os.path.exists(file_path):
        print('loading already trained model')
        best_player_so_far = ResNet.resnet18()
        best_player_so_far.load_state_dict(torch.load(file_path))

    else:
        print('Trained model doesnt exist. Starting from scratch.')
        best_player_so_far = ResNet.resnet18()

    best_player_so_far.eval()
    return best_player_so_far

# ---------------------------------------------------------------------------- #
# Neural Net training
def improve_model_resnet(nn, data, i):
    # here i is the number of times NN has improved : it will be used for learning rate annealing

    min_data = config.MINIBATCH * config.MINBATCHNUMBER
    max_data = config.MINIBATCH * config.MAXBATCHNUMBER
    size_training_set = data.shape[0]

    if size_training_set >= min_data:
        print('size of train set', size_training_set)

        # random sample from past experiences :
        if size_training_set >= max_data:
            X = data[np.random.choice(data.shape[0], max_data, replace=False)]
            # X = data[-size_training_set:-1,:]
        else:
            X = data

        # learning rate annealing: decrease by config.lrdecay the learning rate every config.annealing succesfull improvements of the model
        j = i // config.annealing

        if config.optim == 'sgd':
            lr_decay = config.sgd_lr / ((config.lrdecay) ** j)
        elif config.optim == 'adam':
            lr_decay = config.adam_lr / ((config.lrdecay) ** j)

        k = (i - 1) // config.annealing
        if j == k + 1:
            print('learning rate is now = ', lr_decay)

        training = ResNet_Training(nn, config.MINIBATCH, config.EPOCHS, lr_decay, X, X, 1)
        training.trainNet()

    else:
        print('Not enough training data. Please increase number of self play games or use previous data = True')
        print('train set size', data.shape[0])
        raise ValueError

def load_or_create_neural_net():
    file_path = './best_model_resnet.pth'
    best_player_so_far = ResNet.resnet18()
    if os.path.exists(file_path):
        print('loading already trained model')
        best_player_so_far.load_state_dict(torch.load(file_path))

    best_player_so_far.eval()

    return best_player_so_far

if __name__ == "__main__":
    main()


"""
Implementation of BIMODAL to generate SMILES
"""

import numpy as np
import torch
import torch.nn as nn
from bidir_lstm import BiDirLSTM

torch.manual_seed(1)
np.random.seed(5)

class BIMODAL():

    def __init__(self, molecule_size=7, encoding_dim=55, lr=.01, hidden_units=128):

        self._molecule_size = molecule_size
        self._input_dim = encoding_dim
        self._output_dim = encoding_dim
        self._layer = 2
        self._hidden_units = hidden_units

        # Learning rate
        self._lr = lr

        # Build new model
        self._lstm = BiDirLSTM(self._input_dim, self._hidden_units, self._layer)

        # Check availability of GPUs
        self._gpu = torch.cuda.is_available()
        self._device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        if torch.cuda.is_available():
            self._lstm = self._lstm.cuda()
            print('GPU available')

        # Adam optimizer
        self._optimizer = torch.optim.Adam(self._lstm.parameters(), lr=self._lr, betas=(0.9, 0.999))
        # Cross entropy loss
        self._loss = nn.CrossEntropyLoss(reduction='mean')

    def build(self, name=None):
        """Build new model or load model by name
        :param name:    model name
        """

        if (name is None):
            self._lstm = BiDirLSTM(self._input_dim, self._hidden_units, self._layer)

        else:
            self._lstm = torch.load(name + '.dat', map_location=self._device)

        if torch.cuda.is_available():
            self._lstm = self._lstm.cuda()

        self._optimizer = torch.optim.Adam(self._lstm.parameters(), lr=self._lr, betas=(0.9, 0.999))

    def print_model(self):
        '''Print name and shape of all tensors'''
        for name, p in self._lstm.state_dict().items():
            print(name)
            print(p.shape)

    def train(self, data, label, epochs=1, batch_size=1):
        '''Train the model
        :param  data:   data array (n_samples, molecule_size, encoding_length)
        :param  label:  label array (n_samples, molecule_size)
        :param  epochs: number of epochs for the training
        :param  batch_size: batch size for the training
        :return statistic:  array storing computed losses (epochs, batch size)
        '''

        # Number of samples
        n_samples = data.shape[0]

        # Change axes from (n_samples, molecule_size, encoding_dim) to (molecule_size, n_samples, encoding_dim)
        data = np.swapaxes(data, 0, 1)

        # Create tensor from label
        label = torch.from_numpy(label).to(self._device, dtype=torch.long)

        # Calculate number of batches per epoch
        if (n_samples % batch_size) is 0:
            n_iter = n_samples // batch_size
        else:
            n_iter = n_samples // batch_size + 1

        # To store losses
        statistic = np.zeros((epochs, n_iter))

        # Prepare model for training
        self._lstm.train()

        # Iteration over epochs
        for i in range(epochs):

            # Iteration over batches
            for n in range(n_iter):

                # Set gradient to zero for batch
                self._optimizer.zero_grad()

                # Store losses in list
                losses = []

                # Compute indices used as batch
                batch_start = n * batch_size
                batch_end = min((n + 1) * batch_size, n_samples)

                # Reset model with correct batch size
                self._lstm.new_sequence(batch_end - batch_start, self._device)

                # Current batch
                batch_data = torch.from_numpy(data[:, batch_start:batch_end, :].astype('float32')).to(self._device)

                # Initialize start and end position of sequence read by the model
                start = self._molecule_size // 2
                end = start + 1

                for j in range(self._molecule_size - 1):
                    self._lstm.new_sequence(batch_end - batch_start, self._device)

                    # Select direction for next prediction
                    if j % 2 == 0:
                        dir = 'right'
                    else:
                        dir = 'left'

                    # Predict next token
                    pred = self._lstm(batch_data[start:end], dir, self._device)

                    # Compute loss and extend sequence read by the model
                    if j % 2 == 0:
                        loss = self._loss(pred, label[batch_start:batch_end, end])
                        end += 1

                    else:
                        loss = self._loss(pred, label[batch_start:batch_end, start - 1])
                        start -= 1

                    # Append loss of current position
                    losses.append(loss.item())

                    # Accumulate gradients
                    # (NOTE: This is more memory-efficient than summing the loss and computing the final gradient for the sum)
                    loss.backward()

                # Store statistics: loss per token (middle token not included)
                statistic[i, n] = np.sum(losses) / (self._molecule_size - 1)

                # Perform optimization step
                self._optimizer.step()

        return statistic

    def validate(self, data, label, batch_size=128):
        ''' Validation of model and compute error
        :param data:    test data (n_samples, molecule_size, encoding_size)
        :param label:   label data (n_samples_molecules_size)
        :param batch_size:  batch size for validation
        :return:            mean loss over test data
        '''

        # Use train mode to get loss consistent with training
        self._lstm.train()

        # Gradient is not computed to reduce memory requirements
        with torch.no_grad():
            # Compute tensor of labels
            label = torch.from_numpy(label).to(self._device, dtype=torch.long)

            # Number of samples
            n_samples = data.shape[0]

            # Change axes from (n_samples, molecule_size, encoding_dim) to (molecule_size , n_samples, encoding_dim)
            data = np.swapaxes(data, 0, 1).astype('float32')

            # Initialize loss for complete validation set
            tot_loss = 0

            # Calculate number of batches per epoch
            if (n_samples % batch_size) is 0:
                n_iter = n_samples // batch_size
            else:
                n_iter = n_samples // batch_size + 1

            for n in range(n_iter):

                # Compute indices used as batch
                batch_start = n * batch_size
                batch_end = min((n + 1) * batch_size, n_samples)

                # Data used in this batch
                batch_data = torch.from_numpy(data[:, batch_start:batch_end, :].astype('float32')).to(self._device)

                # Initialize loss for molecule
                molecule_loss = 0

                # Reset model with correct batch size and device
                self._lstm.new_sequence(batch_end - batch_start, self._device)

                start = self._molecule_size // 2
                end = start + 1

                for j in range(self._molecule_size - 1):
                    self._lstm.new_sequence(batch_end - batch_start, self._device)

                    # Select direction for next prediction
                    if j % 2 == 0:
                        dir = 'right'
                    if j % 2 == 1:
                        dir = 'left'

                    # Predict next token
                    pred = self._lstm(batch_data[start:end], dir, self._device)

                    # Extend reading of the sequence
                    if j % 2 == 0:
                        loss = self._loss(pred, label[batch_start:batch_end, end])
                        end += 1

                    if j % 2 == 1:
                        loss = self._loss(pred, label[batch_start:batch_end, start - 1])
                        start -= 1
                    # Sum loss over molecule
                    molecule_loss += loss.item()

                # Add loss per token to total loss (start token and end token not counted)
                tot_loss += molecule_loss / (self._molecule_size - 1)

            return tot_loss / n_iter

    def sample(self, middle_token, T=1):
        '''Generate new molecule
        :param middle_token:    starting sequence
        :param T:               sampling temperature
        :return molecule:       newly generated molecule (molecule_length, encoding_length)
        '''

        # Prepare model
        self._lstm.eval()

        # Gradient is not compute to reduce memory requirements
        with torch.no_grad():
            # Output array with merged forward and backward directions

            # New sequence
            seq = np.zeros((self._molecule_size, 1, self._output_dim))
            seq[self._molecule_size // 2, 0] = middle_token

            # Create tensor for data and select correct device
            seq = torch.from_numpy(seq.astype('float32')).to(self._device)

            # Define start/end values for reading
            start = self._molecule_size // 2
            end = start + 1

            for j in range(self._molecule_size - 1):
                self._lstm.new_sequence(1, self._device)

                # Select direction for next prediction
                if j % 2 == 0:
                    dir = 'right'
                if j % 2 == 1:
                    dir = 'left'

                pred = self._lstm(seq[start:end], dir, self._device)

                # Compute new token
                token = self.sample_token(np.squeeze(pred.cpu().detach().numpy()), T)

                # Set new token within sequence
                if j % 2 == 0:
                    seq[end, 0, token] = 1.0
                    end += 1

                if j % 2 == 1:
                    seq[start - 1, 0, token] = 1.0
                    start -= 1

        return seq.cpu().numpy().reshape(1, self._molecule_size, self._output_dim)

    def sample_token(self, out, T=1.0):
        ''' Sample token
        :param out: output values from model
        :param T:   sampling temperature
        :return:    index of predicted token
        '''
        # Explicit conversion to float64 avoiding truncation errors
        out = out.astype('float64')

        # Compute probabilities with specific temperature
        out_T = out / T
        p = np.exp(out_T) / np.sum(np.exp(out_T))

        # Generate new token at random
        char = np.random.multinomial(1, p, size=1)
        return np.argmax(char)

    def save(self, name='test_model'):
        torch.save(self._lstm, name + '.dat')


# -*- coding: utf-8 -*-
import os
import numpy as np
from PIL import Image, ImageGrab
import time
import matplotlib.pyplot as plt
import pytesseract
import win32api
import win32gui
import win32ui
import win32con
from ctypes import windll
import cv2
############################鼠标点击使用常数########################################
X1 = float(1032)                                         #STEP 1
Y1 = float(578)
################################获取信息的常数################
X6 = float(1032)
tem = float(578)
Y6A = float(331/tem)
Y6B = float(487/tem)
Y6C = float(87/tem)
Y6D = float(243/tem)

CARDS_POSITONS = [(float(45/X6), Y6A, float(165/X6), Y6B), (float(250/X6), Y6A, float(370/X6), Y6B),\
                  (float(455/X6), Y6A, float(575/X6), Y6B), (float(660/X6), Y6A, float(780/X6), Y6B),\
                  (float(865/X6), Y6A, float(985/X6), Y6B), (float(275/X6), Y6C, float(395/X6), Y6D),\
                  (float(460/X6), Y6C, float(580/X6), Y6D), (float(645/X6), Y6C, float(765/X6), Y6D)]

BATTLE = (float(850/1025) , float(433/581), float(975/1025) , float(550/581))
ROUND = [float(690/X1), float(6/Y1), float(746/X1), float(34/Y1)]          #第几回合
# END = (float(45/X1), float(340/Y1), float(989/X1), float(442/Y1))          #结束界面判定
END = (float(800/X1), float(520/Y1), float(990/X1), float(570/Y1))          #结束界面判定

CHOICE_CHECK_BOX = (float(880/X1), float(516/Y1), 1, 1)   #主界面
CHECK_APPLE_BOX = (float(258/X1), float(96/Y1), float(346/X1), float(193/Y1))    #确认苹果界面

BORDER = [40, 55, 0, 0]

BOXS = {
    "HOME":  CHOICE_CHECK_BOX,
    "APPLE": CHECK_APPLE_BOX,
    "BATTLE":BATTLE,
#    "END":   END,
    "ROUND": ROUND,
    "CARDS": CARDS_POSITONS
}
################################对比的常数#########################
path = os.path.join(os.path.abspath(os.path.dirname(os.path.realpath('__file__'))), "data")
CHECK_POINT = Image.open(os.path.join(path,"check.png"))    #选关卡图片  用来对比
CHECK_APPLE = Image.open(os.path.join(path,"apple.png"))          #吃苹果图片
BATTLE_SCREEM = Image.open(os.path.join(path,"battle.png"))       #人物头像，用来确认是否在技能界面
CHECH_SUPPORT = Image.open(os.path.join(path,"teasupport.png"))
QP_SUPPORT = Image.open(os.path.join(path,"qpsupport.png"))

# END_SCREEM = Image.open(os.path.join(path,"end.png"))
# STATE = ["HOME", "APPLE", "BATTLE", "END"]
STATE = ["HOME", "APPLE", "BATTLE"]
# IMAGES = {"HOME":CHECK_POINT, "APPLE":CHECK_APPLE, "BATTLE":BATTLE_SCREEM, "END": END_SCREEM}
IMAGES = {"HOME":CHECK_POINT, "APPLE":CHECK_APPLE, "BATTLE":BATTLE_SCREEM}

##################################################################外部方法
def True_Type(input):
    if input:
        return False
    else:
        return True

def show_plot(image):
    plt.figure("Image")
    plt.imshow(image)
    plt.show()

def Absolute_Position(position, size):
    """return Absolute position in image
    """
    rp = position*np.array(size*2)
    return rp.astype(np.int32)
'''
def Crop_Border(image):
    """crop windows border to get game image
    """
    size = image.size
    box = [6, 46, 1286 ,766]
    # box = [BORDER[2], BORDER[0], size[0] - BORDER[3], size[1] - BORDER[1]]
    image = image.crop(box)
    return image
'''
def Crop(image, box):
    """crop image using box Relative_Box get image to confirm state
    """
    true_box = Absolute_Position(box, list(image.size))
    image = image.crop(true_box)
    return image
########################################################################################
def is_end(image):
    image = np.array(image)
    size = np.size(image)
    energy = np.sum(image)/size
    dark = np.sum(image < 30)/size
    color = np.sum(image > 235)/size
    if dark > 0.8 and color > 0.025:
        print(dark, color)
        return True
    else:
        False

def getcolor(image):
    COLOR = ["R", "G", "B"]
    image = np.array(image.convert("RGB"))
    c_min = np.min(image, axis=2)
    array = image - np.expand_dims(c_min, axis=2)
    c_max = [np.max(array[:,:,i]) for i in range(3)]
    c_sum = [np.sum(array[:,:,i]) for i in range(3)]
    if (c_max > np.array([230, 230, 230])).any():
        idx = np.argmax(c_sum)
        color = COLOR[idx]
        return color
    else:
        return None

def Compare(img1, img2, size = (64, 64)):
    img1 = np.array(img1.resize(size).convert("RGB"))/256
    img2 = np.array(img2.resize(size).convert("RGB"))/256
    loss = np.max(np.abs(img1 - img2), axis=2)
    loss = np.sum(loss < 0.2)/(64*64)
    if loss > 0.75:
        return True
    else:
        return False
######################################################################
def getcard(image):
    # image = Crop_Border(image)
    cards = [Crop(image, box) for box in BOXS["CARDS"]]
    cards = [getcolor(card) for card in cards]
    return cards

def getstate(self, image, s = None):
    # image = Crop_Border(image)

    '''
    b = image.getpixel((923, 56))
    print(b)
    print(b==(123,124,124))
    '''

    # image.save("test.png")
    if s:
        search_state = s
    else:
        search_state = STATE

    for i in search_state:
        img = Crop(image, BOXS[i])
        # img.save(i+".png")
        if IMAGES[i]:
            if Compare(img, IMAGES[i]):
                if i == "APPLE":
                    self.click(49)
                    self.click(50)
                    time.sleep(1)
                    while True:
                        image = capture(self.hwnd)
                        self.image=image
                        if is_support(self.image):
                            break
                        time.sleep(1)
                if i == "HOME":
                    self.click(48)
                    time.sleep(1)
                return i
    # print(image.getpixel((1006,682)))
    # if image.getpixel((949, 675))==(0xB4,0xAF,0x9E) and image.getpixel((1006,682)) == (0xD2,0xD2,0xD2) and image.getpixel((1059,693))==(0x5A,0x66,0x75) and image.getpixel((1155,685))==(0x2C,0x33,0x68):
    if image.getpixel((912, 59))==(0x8E,0x8E,0x8E) and image.getpixel((923,57)) == (0x8D,0x8D,0x8D) and image.getpixel((911,29))==(0x77,0x77,0x77) and image.getpixel((913,53))==(0x83,0x83,0x83):
        return "END"
    elif image.getpixel((273,158))==(0xc3,0x92,0x70) and image.getpixel((282,159))==(0x47,0x42,0x46)\
        and image.getpixel((282,162))==(0x55,0xe9,0xd3)and image.getpixel((316,162))==(0xfb,0xc6,0x88):
        print("REBOOTING")
        while True:
            windll.user32.SetCursorPos(298,213)    #鼠标移动到
            win32api.mouse_event(win32con.MOUSEEVENTF_LEFTDOWN, 0, 0, 0, 0)    #左键按下
            win32api.mouse_event(win32con.MOUSEEVENTF_LEFTUP, 0, 0, 0, 0)
            time.sleep(3)
            image = capture(self.hwnd)
            if is_comfirm(image):
                print("REBOOT SUCCESS")
                self.click(50)
                break
            elif is_cha(image):
                print("click cha")
                self.click(55)
                break

        self.click(90)

    elif image.getpixel((832,557))==(0x7c,0x7c,0x7c) and image.getpixel((776,557))==(0x04,0x04,0x04)\
        and image.getpixel((776,559))==(0x22,0x22,0x22) and image.getpixel((776,560))==(8,8,8)\
        and image.getpixel((776,567))==(3,3,3) and image.getpixel((842,557))==(0x8d,0x8d,0x8d):
        print("network error")
        self.click(50)
        time.sleep(2)
    elif (i!="APPLE"):
        return False
    image = capture(self.hwnd)
    self.image=image
    s=getstate(self,image)
    return s

def is_support(image):
    # image = Crop_Border(image)
    # print(image.getpixel((859,275)))
    if image.getpixel((843,258))==(21,106,1) and image.getpixel((928,258))==(25,59,116)\
        and image.getpixel((1013,258))==(123,17,8):
        return True
    return False

def is_pcards(image):
    # image = Crop_Border(image)
    # print(image.getpixel((859,275)))
    if image.getpixel((1059,59))==(0xff,0xff,0xff) and image.getpixel((1147,680))==(0x0f,0xac,0xd6) and image.getpixel((1180,681))==(0xbd,0xdc,0xfa)and image.getpixel((1212,681))==(0x0d,0x3f,0x64):
        return True
    return False

# 师匠红茶学妹
'''
require =[
    (51, 326, 96, 126, 172), (129,340, 203, 159, 176), (162,351, 255, 241, 233), (191, 359, 248, 251, 179),
    (57, 288, 216, 158, 249), (80, 266, 69, 31, 67), (143, 260, 167, 77, 77), (190, 286, 195, 222, 118),
    (1138, 343, 0x7E, 0xB2, 0x37), (1160, 332, 0xDF, 0xFF, 0xAD)]
'''
'''
require =[
    (51, 326, 103, 134, 173), (129,340, 231, 182, 206), (162,351, 255, 241, 230), (191, 359, 248, 251, 179),
    (57, 288, 216, 158, 249), (80, 266, 62, 22, 71), (143, 260, 156, 59, 59), (190, 286, 195, 222, 118),
    (1138, 343, 0x7E, 0xB2, 0x37), (1160, 332, 0xDF, 0xFF, 0xAD)]
'''
require =[
    (51, 326, 0x5e, 0x82, 0xb8), (129,340, 0xc6, 0x94, 0xd7), (162,351, 255, 241, 233), (191, 359, 248, 251, 179),
    (57, 288, 216, 158, 249), (80, 266, 0x3d, 0x1a, 0x45), (143, 260, 0xbd, 0x31, 0x38), (190, 286, 195, 222, 118),
    (1138, 343, 0x85, 0xb7, 0x33), (1160, 332, 0xe1, 0xfe, 0xaf)]

'''
# 蒙娜丽莎
require =[(61,335,73,96,192),(110,353,22,47,72),(130,364,255,189,148),(190,358,254,254,168),
    (1138, 343, 0x7E, 0xB2, 0x37), (1160, 332, 0xDF, 0xFF, 0xAD)]
'''
def Compare_support(img1, img2, size = (64, 64)):
    img1 = np.array(img1.resize(size).convert("RGB"))/256
    img2 = np.array(img2.resize(size).convert("RGB"))/256
    loss = np.max(np.abs(img1 - img2), axis=2)
    loss = np.sum(loss < 0.2)/(64*64)
    if loss > 0.95:
        return True
    else:
        return False

def is_comfirm(image):
    # image = Crop_Border(image)
    if image.getpixel((736,376))==(0xff,0,0) and image.getpixel((746,555))==(0xdc,0xdd,0xdd)\
        and image.getpixel((770,560))==(5,5,5)and image.getpixel((898,575))==(0x33,0x33,0x34):
        return True
    return False

def is_cha(image):
    if image.getpixel((1210,0))==(0x5c,0x5d,0x5d) and image.getpixel((1211,0))==(0xe5,0xe5,0xe5)\
        and image.getpixel((1213,2))==(0x99,0x99,0x99)and image.getpixel((1213,3))==(0x32,0x32,0x32)\
        and image.getpixel((1214,2))==(0x7f,0x7f,0x7f):
        return True
    return False

def is_pool(image):
    if image.getpixel((454,686))==(0x12,0x7a,0xb3) and image.getpixel((458,686))==(0xaf,0xb9,0xc5)\
        and image.getpixel((475,686))==(0xc7,0xda,0xe0)and image.getpixel((501,686))==(0xd1,0xd5,0xd9)\
        and image.getpixel((449,361))==(0xf8,0x88,0x39):
        return True
    return False

def if_pool_re(image):
    print(image.getpixel((200,426)),image.getpixel((200,440)),image.getpixel((200,463)),image.getpixel((203,485)))

    if image.getpixel((200,426))==(0x09,0x50,0x6f) and image.getpixel((200,440))==(0x07,0x5a,0x75)\
        and image.getpixel((200,463))==(0x05,0x5f,0x78)and image.getpixel((203,485))==(0xff,0xff,0xff):
        return True
    return False

def is_requiresupport(image):
    # return -1
    # image = Crop_Border(image)
    '''
    if image.getpixel((191,359))==(248,251,179) and image.getpixel((190,286))==(195,222,118)\
        and image.getpixel((1138,343))==(0x7e,0xb2,0x37)and image.getpixel((1160,332))==(0xdf,0xff,0xad):
        box = [51, 326, 208, 370]
        image1 = image.crop(box)
        if Compare_support(image1, QP_SUPPORT):
            return 1
    if image.getpixel((191,559))==(248,251,179) and image.getpixel((190,486))==(195,222,118)\
        and image.getpixel((1138,543))==(0x7e,0xb2,0x37)and image.getpixel((1160,532))==(0xdf,0xff,0xad):
        box = [51, 526, 208, 570]
        image2 = image.crop(box)
        if Compare_support(image2, QP_SUPPORT):
            return 2
    '''
    '''
    if image.getpixel((191,359))==(248,251,179) and image.getpixel((190,286))==(195,222,118)\
        and image.getpixel((1138,343))==(0x85,0xb7,0x33) and image.getpixel((1160,332))==(0xe1,0xfe,0xaf):
        # and image.getpixel((1138,343))==(0x7e,0xb2,0x37)and image.getpixel((1160,332))==(0xdf,0xff,0xad):

        box = [51, 240, 208, 370]
        image1 = image.crop(box)
        if Compare_support(image1, CHECH_SUPPORT):
            return 1
    if image.getpixel((191,559))==(248,251,179) and image.getpixel((190,486))==(195,222,118)\
        and image.getpixel((1138,343))==(0x85,0xb7,0x33) and image.getpixel((1160,332))==(0xe1,0xfe,0xaf):
        # and image.getpixel((1138,543))==(0x7e,0xb2,0x37)and image.getpixel((1160,532))==(0xdf,0xff,0xad):
        box = [51, 440, 208, 570]
        image2 = image.crop(box)
        if Compare_support(image2, CHECH_SUPPORT):
            return 2
    '''

    # 刷池子cba助战
    if image.getpixel((53,376))==(0xe6,0xbc,0x50) and image.getpixel((62,376))==(0xef,0xcd,0x66)\
        and image.getpixel((70,288))==(0xd7,0x9e,0xf8) and image.getpixel((114,277))==(0xff,0xff,0xe0)\
        and image.getpixel((163,275))==(0x80,0x3e,0x70) and image.getpixel((190,286))==(195,222,118) \
        and image.getpixel((1138,343))==(0x85,0xb7,0x33) and image.getpixel((1160,332))==(0xe1,0xfe,0xaf)\
        and image.getpixel((191,359))==(248,251,179):
        return 1
    if image.getpixel((53,576))==(0xe6,0xbc,0x50) and image.getpixel((62,576))==(0xef,0xcd,0x66)\
        and image.getpixel((70,488))==(0xd7,0x9e,0xf8) and image.getpixel((114,477))==(0xff,0xff,0xe0)\
        and image.getpixel((163,475))==(0x80,0x3e,0x70) and image.getpixel((190,486))==(195,222,118)\
        and image.getpixel((1138,543))==(0x85,0xb7,0x33) and image.getpixel((1160,532))==(0xe1,0xfe,0xaf)\
        and image.getpixel((191,559))==(248,251,179):
        return 2

    return 0
    '''
    for i in require:
        if image.getpixel((i[0], i[1])) == (i[2], i[3], i[4]):
            print(True)
        else: print(False)

    for i in require:
        if image.getpixel((i[0], i[1])) != (i[2], i[3], i[4]):
            flag=2
            break
    if flag == 1:
        return 1
    for i in require:
        if image.getpixel((i[0], i[1]+200)) != (i[2], i[3], i[4]):
            flag=0
            break
    return flag
    '''

'''
def is_xuemei(image):
    image = Crop_Border(image)
    if image.getpixel((51, 326))==(172,126,96) and image.getpixel((129,340)) == (176,158,203) and image.getpixel((162,351))==(233,241,255) and image.getpixel((191,359))==(179,251,248):
        return True
    return False

def is_shijiang(image):
    image = Crop_Border(image)
    if image.getpixel((57, 488))==(249,158,216) and image.getpixel((80,466)) == (67,31,69) and image.getpixel((143,460))==(77,77,167) and image.getpixel((190,486))==(118,222,195):
        return True
    return False
'''
def get_round(image):
    # config = "--psm 10 digits"
    # image = Crop_Border(image)

    # image = Crop(image, BOXS["ROUND"])
    # image = image.convert("L")
    # image.save("round.png")
    # code = pytesseract.image_to_string(image, config=config)[0]

    # print(pytesseract.image_to_string(image, config=config))

    # round 3 876,29 (255,255,255) 870,16 (250,250,250) 875,18 (255,255,255)
    if image.getpixel((873, 31))==(255,255,255) and image.getpixel((873,15)) == (92,92,92) and image.getpixel((871,17))==(234,234,234) and image.getpixel((867,19))==(200,200,200):
        return 0
    elif image.getpixel((866, 17))==(113,113,113) and image.getpixel((873,16)) == (254,254,254) and image.getpixel((876,22))==(235,235,235) and image.getpixel((873,33))==(214,214,214):
        return 1
    elif image.getpixel((867, 16))==(118,118,118) and image.getpixel((875,16)) == (194,194,194) and image.getpixel((875,27))==(198,198,198) and image.getpixel((867,33))==(174,174,174):
        return 2
    elif image.getpixel((273,158))==(0xc3,0x92,0x70) and image.getpixel((282,159))==(0x47,0x42,0x46)\
        and image.getpixel((282,162))==(0x55,0xe9,0xd3)and image.getpixel((316,162))==(0xfb,0xc6,0x88):
        return -1
    else:
        return 3

def capture(hwnd):
    """capture applications screem"""
    # left, top, right, bot = win32gui.GetWindowRect(hwnd)
    # w = (right - left)
    # h = (bot - top)
    # 104,37   106,67
    # print(hwnd)
    # win32gui.ShowWindow(hwnd,win32con.SW_RESTORE)
    # win32gui.SetForegroundWindow(hwnd)
    game_rect=win32gui.GetWindowRect(hwnd)
    src_image=ImageGrab.grab(game_rect)
    # src_image.save("123466.png")
    '''
    hwndDC = win32gui.GetWindowDC(hwnd)
    mfcDC  = win32ui.CreateDCFromHandle(hwndDC)
    saveDC = mfcDC.CreateCompatibleDC()
    saveBitMap = win32ui.CreateBitmap()

    saveBitMap.CreateCompatibleBitmap(mfcDC, w, h)
    saveDC.SelectObject(saveBitMap)

    result = windll.user32.PrintWindow(hwnd, saveDC.GetSafeHdc(), 0)
    bmpinfo = saveBitMap.GetInfo()
    bmpstr = saveBitMap.GetBitmapBits(True)
    im = Image.frombuffer('RGB',(bmpinfo['bmWidth'], bmpinfo['bmHeight']),
                           bmpstr, 'raw', 'BGRX', 0, 1)

    im.save("21234.png")
    win32gui.DeleteObject(saveBitMap.GetHandle())
    saveDC.DeleteDC()
    mfcDC.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwndDC)
    '''
    # if result == 1:
        # PrintWindow Succeeded
        # im.save("test.png")
        # im.show()

    '''
    if w <= 200 or h <= 200:
        raise "captrue error"
    '''
    return src_image

if __name__ == '__main__':
    scrpt = None


import numpy as np
import random

class ExperienceReplay(object):

    def __init__(self, maxlen, dimensions, minibatch_size, stack_num):
        self.maxlen = maxlen
        self.dimensions = dimensions
        self.minibatch_size = minibatch_size
        self.stack_num = stack_num

        self.screens = np.empty((maxlen,) + self.dimensions, dtype=np.uint8)
        self.actions = np.empty(maxlen, dtype=np.uint8)
        self.rewards = np.empty(maxlen, dtype=np.integer)
        self.terminals = np.empty(maxlen, dtype=np.bool)
        self.count = 0
        self.current = 0

        self.current_state = np.empty((1, self.stack_num) + self.dimensions, dtype=np.uint8)
        self.prestates = np.empty((self.minibatch_size, self.stack_num) + self.dimensions, dtype=np.uint8)
        self.poststates = np.empty((self.minibatch_size, self.stack_num) + self.dimensions, dtype=np.uint8)

    def add(self, memory):
        screen, action, reward, terminal = memory
        self.screens[self.current, ...] = screen
        self.actions[self.current] = action
        self.rewards[self.current] = reward
        self.terminals[self.current] = terminal
        self.count = max(self.count, self.current + 1)
        self.current = (self.current + 1) % self.maxlen
        return

    def getCurrentState(self):
        self.current_state[0] = self.getStackedState(self.current - 1)
        return self.current_state

    def getStackedState(self, index):
        index = index % self.count

        if index >= self.stack_num - 1:
            return self.screens[(index - (self.stack_num - 1)):(index + 1), ...]
        else:
            indices = [(index - 1) % self.count for i in reversed(range(self.stack_num))]
            return self.screens[indices, ...]

    def sample(self):
        indices = []
        while len(indices) < self.minibatch_size:
            while True:
                index = random.randint(self.stack_num, self.maxlen - 1)
                if index >= self.current and (index - self.stack_num) < self.current:
                    continue
                if self.terminals[(index - self.stack_num):index].any():
                    continue
                break
            self.prestates[len(indices), ...] = self.getStackedState(index - 1)
            self.poststates[len(indices), ...] = self.getStackedState(index)
            indices.append(index)

        actions = self.actions[indices]
        rewards = self.rewards[indices]
        terminals = self.terminals[indices]

        return self.prestates, actions, rewards, self.poststates, terminals


import numpy as np
import h5py as h5
import astropy.units as u
import os.path

ids = 10
loops = 10
n_AGN = 5000
files = 0
base_filename = "BASE"

for i in range(ids):
    for j in range(loops):
        if os.path.isfile(base_filename + "_id{}_loop{}.h5".format(i, j)):
            files += 1

with h5.File(base_filename + ".h5", "w") as main:
    main["m_1"] = np.zeros(n_AGN * files) * u.Msun
    main["m_2"] = np.zeros(n_AGN * files) * u.Msun
    main["a_enc"] = np.zeros(n_AGN * files) * u.AU
    main["e_enc"] = np.zeros(n_AGN * files)
    main["e_LISA"] = np.zeros(n_AGN * files)
    main["f_orb_LISA"] = np.zeros(n_AGN * files) * u.Hz
    main["dist"] = np.zeros(n_AGN * files) * u.Gpc
    main["snr"] = np.zeros(n_AGN * files)
    main["max_snr_harmonic"] = np.zeros(n_AGN * files).astype(int)
    main["age"] = np.zeros(n_AGN * files) * u.yr
    main["t_e2m"] = np.zeros(n_AGN * files) * u.yr
    main["t_se"] = np.zeros(n_AGN * files) * u.yr
    main["m_oligarch_final"] = np.zeros(n_AGN * files) * u.Msun

    k = 0
    for i in range(ids):
        for j in range(loops):
            filename = base_filename + "_id{}_loop{}.h5".format(i, j)
            print(filename)
            if os.path.isfile(filename):
                with h5.File(filename, "r") as file:
                    main["m_1"][k * n_AGN: (k + 1) * n_AGN] = file["m_1"]
                    main["m_2"][k * n_AGN: (k + 1) * n_AGN] = file["m_2"]
                    main["a_enc"][k * n_AGN: (k + 1) * n_AGN] = file["a_enc"]
                    main["e_enc"][k * n_AGN: (k + 1) * n_AGN] = file["e_enc"]
                    main["e_LISA"][k * n_AGN: (k + 1) * n_AGN] = file["e_LISA"]
                    main["f_orb_LISA"][k * n_AGN: (k + 1) * n_AGN] = file["f_orb_LISA"]
                    main["dist"][k * n_AGN: (k + 1) * n_AGN] = file["dist"]
                    main["snr"][k * n_AGN: (k + 1) * n_AGN] = file["snr"]
                    main["max_snr_harmonic"][k * n_AGN: (k + 1) * n_AGN] = file["max_snr_harmonic"]
                    main["age"][k * n_AGN: (k + 1) * n_AGN] = file["age"]
                    main["t_e2m"][k * n_AGN: (k + 1) * n_AGN] = file["t_e2m"]
                    main["t_se"][k * n_AGN: (k + 1) * n_AGN] = file["t_se"]
                    main["m_oligarch_final"][k * n_AGN: (k + 1) * n_AGN] = file["m_oligarch_final"]
                k += 1
                print(k)