import RTIF.LowLevel.quaternion as quat
from RTIF.API import API
import numpy as np
import time

class HAPI(API):
    """
    Higher level API interface, including some feedback, robot state checking and so on
    """
    def __init__(self, IP_ADDRESS):
        super(HAPI, self).__init__(IP_ADDRESS)
        self.__base_t = np.asarray((0, 0, 0), dtype=np.float32)
        self.__base_q = np.asarray((1, 0, 0, 0), dtype=np.float32)

    def isLastMovementEnd(self, select=(1, 1, 1, 1, 1, 1)):
        """
        :param select: 1 if dimension is selected
        :return: boolean
        """
        select = np.asarray(select, dtype=np.float32)
        data = self.rtif.receive()
        tar_rad = np.asarray(data["Target Joint Positions"], dtype=np.float32)
        cur_rad = np.asarray(data["Actual Joint Positions"], dtype=np.float32)
        speed = np.asarray(data["Actual Joint Velocities"])
        return np.max(np.abs((tar_rad - cur_rad) * select)) < 1e-4 and np.max(np.abs(speed * select)) < 1e-1

    def RecordingActions(self):
        """
        start recording mode, input ' ' to capture and 'e' to exit
        :return: list of control points
        """
        def __cap_one_point():
            data = self.rtif.receive()
            cap = {"Tool Position": data["Actual Tool Coordinates"][:],
                   "Tool Force": data["Generalized Tool Force"][:],
                   "Joint Position": data["Actual Joint Positions"][:],
                   "Time Step": data["Time Step"]
                   }
            return cap
        caps = []

        self.TeachMode(18000)
        while True:
            i = raw_input("Input 'Enter' to capture, 'e+Enter' to exit, m+Enter to enter edit mode >>")
            if i == 'e':
                break
            elif i == '':
                caps.append(__cap_one_point())
                print ("captured %d points" % len(caps))
            elif i == 'm':
                current_id = 1
                while True:
                    for ind, c in enumerate(caps):
                        if ind == current_id:
                            print ("->%02d : End point (%04f, %04f, %04f, %04f, %04f, %04f)" % ((ind,) + c["Joint Position"]))
                        else:
                            print ("  %02d : End point (%04f, %04f, %04f, %04f, %04f, %04f)" % ((ind,) + c["Joint Position"]))
                    print ("Edit mode: 'e+Enter' to exit, 'u+Enter' to move cursor up, 'd+Enter' down,")
                    print (" 'm+Enter' to modify current, 't+Enter to delete, 'i' to insert before current, 'r+Enter to replay current'")
                    j = raw_input()
                    if j == 'e':
                        break
                    if j == 'u':
                        if current_id == len(caps)-1:
                            current_id = 0
                        else:
                            current_id += 1
                    elif j == 'd':
                        if current_id == 0:
                            current_id = len(caps)-1
                        else:
                            current_id -= 1
                    elif j == 'r':
                        self.MoveJointToRad(caps[current_id]["Joint Position"])
                        while self.is_direct_mode() and not self.isLastMovementEnd():
                            time.sleep(0.5)
                        self.TeachMode(18000)
                    elif j == 't':
                        caps.pop(current_id)
                        if current_id >= len(caps):
                            caps -= 1
                        if caps < 0:
                            break
                    elif j == 'm':
                        cap = __cap_one_point()
                        cap["Time Step"] = caps[current_id]["Time Step"]
                        caps[current_id] = cap
                    elif j == 'i':
                        cap = __cap_one_point()
                        if current_id == 0:
                            cap["Time Step"] = caps[current_id]["Time Step"] - 3.0
                        else:
                            cap["Time Step"] = (caps[current_id]["Time Step"] + caps[current_id - 1]["Time Step"]) / 2.0
                        caps.insert(current_id, cap)
                    else:
                        print ("Unknown input command %s", j)
                        break

            else:
                print ("Unknown input command %s", i)
        self.EndTeachMode()
        return caps

    def ReplayCapturedData(self, caps, with_time=False):
        """
        will replay captured data in fixed speed
        :param caps: list returned by 'RecordingActions'
        :return: None
        """
        if len(caps) <= 0:
            return
        if with_time:
            self.MoveJointToRad(caps[0]["Joint Position"], a=1.2, v=0.25)
            while self.is_direct_mode() and not self.isLastMovementEnd():
                time.sleep(0.5)
            for i in range(1, len(caps)):
                self.MoveJointToRad(caps[i]["Joint Position"], t=(caps[i]["Time Step"] - caps[i-1]["Time Step"]))
                while self.is_direct_mode() and not self.isLastMovementEnd():
                    time.sleep(0.5)
        else:
            for cap in caps:
                self.MoveJointToRad(cap["Joint Position"], a=1.2, v=0.25)
                while self.is_direct_mode() and not self.isLastMovementEnd():
                    time.sleep(0.5)

    def FineTuningPosition(self, interactive=False, add_value=((0.01, 0.01, 0.01), (1.0, 0.0, 0.0, 0.0)), factor=10.0):
        """
        :param interactive: if True, call curses to capture keyboard input, else add add_value to current position
        :param add_value: step size, for interactive mode, (dx, dy, dz, drx, dry, drz), else ((dx, dy, dz), (qw, qi, qj, qk)
        :param factor: if in interactive mode, means key_hold boosting factor, the max speed will achieve after holding
                        a key for some time
        :return: current position after fine tuning in format (x,y,z),(w,i,j,k)
        """
        if interactive:
            if len(add_value) == 6:
                import curses
                stdscr = curses.initscr()
                curses.noecho()
                curses.cbreak()
                stdscr.keypad(1)
                stdscr.nodelay(1)
                stdscr.addstr("Start Interactive Fine Tune Mode\n\r")
                stdscr.addstr("Press 'a','d' for moving along x axis, 'A', 'D' for rotate\n\r")
                stdscr.addstr("Press 's','w' for moving along y axis, 'S', 'W' for rotate\n\r")
                stdscr.addstr("Press 'q','e' for moving along z axis, 'Q', 'E' for rotate\n\r")
                stdscr.addstr("Press 'esc' to exit >>>\n\r")
                step = list(add_value)
                while 1:
                    char = stdscr.getch()
                    curses.flushinp()
                    p, q = self.GetCurrentEndPos()
                    stdscr.clrtoeol()
                    rot = self.from_q_to_rad_axis(q)
                    stdscr.addstr("\r(x=%f, y=%f, z=%f, rx=%f, ry=%f, rz=%f)" % (p[0], p[1], p[2], rot[0], rot[1], rot[2]))
                    if char == curses.KEY_LEFT or char == ord('a'):     # -X
                        p[0] -= step[0]
                        if step[0] < add_value[0] * factor:
                            step[0] += add_value[0] * factor / 10.0
                    elif char == curses.KEY_RIGHT or char == ord('d'):  # +X
                        p[0] += step[0]
                        if step[0] < add_value[0] * factor:
                            step[0] += add_value[0] * factor / 10.0
                    elif char == curses.KEY_UP or char == ord('w'):     # +Y
                        p[1] += step[1]
                        if step[1] < add_value[1] * factor:
                            step[1] += add_value[1] * factor / 10.0
                    elif char == curses.KEY_DOWN or char == ord('s'):   # -Y
                        p[1] -= step[1]
                        if step[1] < add_value[1] * factor:
                            step[1] += add_value[1] * factor / 10.0
                    elif char == curses.KEY_PPAGE or char == ord('e'):  # +Z
                        p[2] += step[2]
                        if step[2] < add_value[2] * factor:
                            step[2] += add_value[2] * factor / 10.0
                    elif char == curses.KEY_NPAGE or char == ord('q'):  # -Z
                        p[2] -= step[2]
                        if step[2] < add_value[2] * factor:
                            step[2] += add_value[2] * factor / 10.0
                    elif char == ord('A'):                              # -rX
                        q = quat.qmul(q, self.from_rad_axis_to_q((-step[3], 0, 0)))
                        if step[3] < add_value[3] * factor:
                            step[3] += add_value[3] * factor / 10.0
                    elif char == ord('D'):                              # +rX
                        q = quat.qmul(q, self.from_rad_axis_to_q((step[3], 0, 0)))
                        if step[3] < add_value[3] * factor:
                            step[3] += add_value[3] * factor / 10.0
                    elif char == ord('W'):                              # +rY
                        q = quat.qmul(q, self.from_rad_axis_to_q((0, step[4], 0)))
                        if step[4] < add_value[4] * factor:
                            step[4] += add_value[4] * factor / 10.0
                    elif char == ord('S'):                              # -rY
                        q = quat.qmul(q, self.from_rad_axis_to_q((0, -step[4], 0)))
                        if step[4] < add_value[4] * factor:
                            step[4] += add_value[4] * factor / 10.0
                    elif char == ord('E'):                              # +rZ
                        q = quat.qmul(q, self.from_rad_axis_to_q((0, 0, step[5])))
                        if step[5] < add_value[5] * factor:
                            step[5] += add_value[5] * factor / 10.0
                    elif char == ord('Q'):                              # -rZ
                        q = quat.qmul(q, self.from_rad_axis_to_q((0, 0, -step[5])))
                        if step[5] < add_value[5] * factor:
                            step[5] += add_value[5] * factor / 10.0
                    elif char == 27:
                        break
                    else:
                        step = list(add_value)
                        continue
                    self.MoveEndPointToPosition(pos=p, rotation=q, v=0.1, a=0.25)
                    time.sleep(0.1)
                curses.endwin()
                print ("Exit Interactive Fine Tune Mode")
            else:
                print ("Error add_value format, need (dx, dy, dz, drx, dry, drz) but get", add_value)
        else:
            if len(add_value) == 2 and len(add_value[0]) == 3 and len(add_value[1]) == 4:
                p, q = self.GetCurrentEndPos()
                p_next = p + np.asarray(add_value[0], dtype=np.float32)
                q_next = quat.qmul(add_value[1], q)
                self.MoveEndPointToPosition(pos=p_next, rotation=q_next)
                while not self.isLastMovementEnd():
                    time.sleep(0.5)
            else:
                print ("Error add_value format, need ((dx, dy, dz), (qw, qi, qj, qk)) but get", add_value)
        return self.GetCurrentEndPos()

    def set_coordinate_origin(self, ori=None):
        """
        Setting the Coordinate origin point. If ori is None, will automatically use teach mode.
        operating coordinate is: +X from robot center to end tool, +Z toward sky.
        tool coordinate is: the connector of tool towards +X and tool face -Z
        :param ori: 3d tuple (x, y, z) or ((x, y, z), (w ,i, j, k)), or None for teach mode
                    conversion will be automatically done
        :return: basic transform (x,y,z), (w,i,j,k)
        """
        if ori is None:
            self.switch_mode(True)
            self.TeachMode()
            print ("Please move robot arm to origin point")
            print ("Notice: the external sensor connector points to +X, and tool towards -Z")
            print ("And press Enter key >>>")
            if raw_input() == '':
                ori = self.GetCurrentEndPos()
                print ("New origin point is (%f, %f, %f)" % ori)
            else:
                print ("Cancel without changing coordinate")
                return

        if len(ori) != 2:
            if len(ori) != 3:
                print ("Error value! ori should be (x, y, z)")
                return
            # convert coordinate
            unix = ori / np.linalg.norm(ori[:2])
            unix[2] = 0
            uniz = np.asarray((0, 0, 1), dtype=np.float32)
            unio = ori

            q, t = quat.from_vector_to_q((1, 0, 0), unio + unix, (0, 0, 1), unio + uniz, (0, 0, 0), unio)
            ori = (t, q)

        elif len(ori[0]) != 3 or len(ori[1]) != 4:
            print ("Error value! ori should be (x, y, z), (w, i, j, k)")
            return

        self.__base_t = np.asarray(ori[0], dtype=np.float32)
        self.__base_q = np.asarray(ori[1], dtype=np.float32)

    def get_coordinate_origin(self):
        return self.__base_t, self.__base_q

    def GetCurrentEndForce(self):
        f = super(HAPI, self).GetCurrentEndForce()
        f, t = f[:3], f[3:]
        f = quat.qrotote_v(quat.qconj(self.__base_q), f)
        t = quat.qrotote_v(quat.qconj(self.__base_q), t)
        return np.asarray((f[0], f[1], f[2], t[0], t[1], t[2]), dtype=np.float32)

    def GetCurrentEndPos(self):
        p, q = super(HAPI, self).GetCurrentEndPos()
        p = quat.qrotote_v(quat.qconj(self.__base_q), p - self.__base_t)
        q = quat.qmul(quat.qconj(self.__base_q), q)
        return p, q

    def GetTargetEndPos(self, RAW=False):
        p, q = super(HAPI, self).GetTargetEndPos()
        p = quat.qrotote_v(quat.qconj(self.__base_q), p - self.__base_t)
        q = quat.qmul(quat.qconj(self.__base_q), q)
        if RAW:
            q = self.from_q_to_rad_axis(q)
            return np.asarray((p[0], p[1], p[2], q[0], q[1], q[2]), dtype=np.float32)
        else:
            return p, q

    def MoveEndPointToPosition(self, pos=None, rotation=None, a=1.2, v=0.25, t=None):
        if pos is not None:
            pos = quat.qrotote_v(self.__base_q, pos, self.__base_t)
        if rotation is not None:
            if len(rotation) == 4:
                rotation = quat.qmul(self.__base_q, rotation)
            else:
                rotation = quat.qmul(self.__base_q, self.from_rad_axis_to_q(rotation))
        super(HAPI, self).MoveEndPointToPosition(pos=pos, rotation=rotation, a=a, v=v, t=t)

    def ForceMode(self, taskframe=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
                  selection=(0, 0, 0, 0, 0, 0),
                  wrench=(0.1, 0.1, 0.1, 0.1, 0.1, 0.1),
                  limits=(0.2, 0.2, 0.2, 0.2, 0.2, 0.2), duration=180):
        p, q = taskframe[:3], self.from_rad_axis_to_q(taskframe[3:])
        p = quat.qrotote_v(self.__base_q, p)
        q = self.from_q_to_rad_axis(quat.qmul(self.__base_q, q))
        super(HAPI, self).ForceMode(taskframe=(p[0], p[1], p[2], q[0], q[1], q[2]),
                                    selection=selection, wrench=wrench, limits=limits, duration=duration)


import skimage.io as io
import skimage.transform as skt
import numpy as np
from PIL import Image
from src.models.class_patcher import patcher
from src.utils.imgproc import *

class patcher(patcher):
    def __init__(self, body='./body/body_lopolykon.png', **options):
        super().__init__('ロポリこん', body=body, pantie_position=[5, 662], **options)
        self.skin = Image.open('./material/lopolykon_skin.png')

    def convert(self, image):
        pantie = np.array(image)
        arrx = np.zeros(100)
        arrx[:50] = np.sin(np.linspace(0, np.pi, 50)) * -50
        arry = np.zeros(100)
        arry[12:-10] = np.sin(np.linspace(0, np.pi, 78)) * 350
        pantie = affine_transform_by_arr(pantie, arrx, arry)[:140,:300]
        pantie = np.uint8(resize(pantie, [0.84, 0.84])*255)
        return Image.fromarray(pantie)

    def patch(self, image, transparent=False):
        pantie = self.convert(image)
        if transparent:
            patched = Image.new("RGBA", self.body_size)
        else:
            patched = self.body.copy()
        patched = self.paste(patched, self.skin, [0, 0])
        patched = self.paste(patched, pantie, self.pantie_position)
        return patched


import os
import time
import cv2
import pickle
from numpy.random.mtrand import f
import yaml
import numpy as np
from glob import glob
from tqdm import tqdm
from PIL import ImageFile, Image
from plyfile import PlyData
from concurrent.futures import ProcessPoolExecutor
from argparse import ArgumentParser
try:
    from neupeak.utils.webcv2 import imshow, waitKey
except:
    from cv2 import imshow, waitKey

parser = ArgumentParser()
parser.add_argument(
    "--cls", type=str, default="ape",
    help="Target object from {ape, benchvise, cam, can, cat, driller, duck, \
    eggbox, glue, holepuncher, iron, lamp, phone} (default ape)"
)
parser.add_argument(
    '--fuse_num', type=int, default=10,
    help="Number of images you want to generate."
)
parser.add_argument(
    '--DEBUG', action="store_true",
    help="To show the generated images or not."
)
args = parser.parse_args()
DEBUG = args.DEBUG

Intrinsic_matrix = {
    'linemod': np.array([[572.4114, 0., 325.2611],
                          [0., 573.57043, 242.04899],
                          [0., 0., 1.]]),
    'blender': np.array([[700.,    0.,  320.],
                         [0.,  700.,  240.],
                         [0.,    0.,    1.]]),
    'Tless': np.array([[1075.65091572, 0.0, 224.06888344],
                       [0.0, 1073.90347929, 167.72159802],
                       [0.0, 0.0, 1.0]])
}
lm_obj_dict={
    'ape':1,
    'benchvise':2,
    'cam':4,
    'can':5,
    'cat':6,
    'driller':8,
    'duck':9,
    'eggbox':10,
    'glue':11,
    'holepuncher':12,
    'iron':13,
    'lamp':14,
    'phone':15,
}

file_dir = os.path.dirname(os.path.abspath(__file__))
linemod_dir = os.path.join(file_dir, '../Linemod_preprocessed')
SUNPAS_dir = os.path.join(file_dir, '../SUN2012pascalformat')
cache_dir = os.path.join(file_dir, '../cache')
output_dir = os.path.join(file_dir, '../output')

def ensure_dir(pth):
    if not os.path.exists(pth):
        os.system("mkdir -p {}".format(pth))

def read_lines(pth):
    with open(pth, 'r') as f:
        return [
            line.strip() for line in f.readlines()
        ]

def read_pickle(pkl_path):
    with open(pkl_path, 'rb') as f:
        return pickle.load(f)

def save_pickle(data, pkl_path):
    with open(pkl_path, 'wb') as f:
        pickle.dump(data, f)

def collect_train_info(cls_name):
    cls_id = lm_obj_dict[cls_name]
    cls_root = os.path.join(linemod_dir, 'data/{:02d}'.format(cls_id))
    tr_pth = os.path.join(
        cls_root, "train.txt"
    )
    train_fns = read_lines(tr_pth)

    return train_fns

def collect_linemod_set_info(linemod_cls_name):
    """
    """
    database = []
    if not os.path.exists(cache_dir):
        os.mkdir(cache_dir)
    if os.path.exists(
            os.path.join(cache_dir,'{}_info.pkl').format(linemod_cls_name)
    ):
        return read_pickle(
            os.path.join(cache_dir,'{}_info.pkl').format(linemod_cls_name)
        )

    train_fns = collect_train_info(linemod_cls_name)
    cls_id = lm_obj_dict[linemod_cls_name]
    cls_root = os.path.join(linemod_dir, 'data/{:02d}'.format(cls_id))
    meta_file = open(os.path.join(cls_root, 'gt.yml'), "r")
    meta_lst = yaml.load(meta_file)

    print('begin generate database {}'.format(linemod_cls_name))
    rgb_ptn = os.path.join(cls_root, "rgb/{}.png")
    msk_ptn = os.path.join(cls_root, "mask/{}.png")
    dpt_ptn = os.path.join(cls_root, "depth/{}.png")
    for item in train_fns:
        data={}
        data['rgb_pth'] = rgb_ptn.format(item)
        data['dpt_pth'] = dpt_ptn.format(item)
        data['msk_pth'] = msk_ptn.format(item)

        meta = meta_lst[int(item)]
        if cls_id == 2:
            for i in range(0, len(meta)):
                if meta[i]['obj_id'] == 2:
                    meta = meta[i]
                    break
        else:
            meta = meta[0]
        R = np.resize(np.array(meta['cam_R_m2c']), (3, 3))
        T = np.array(meta['cam_t_m2c']) / 1000.0
        RT = np.concatenate((R, T[:, None]), axis=1)
        data['RT'] = RT
        database.append(data)

    print(
        'successfully generate database {} len {}'.format(
            linemod_cls_name, len(database)
        )
    )
    save_pickle(
        database, os.path.join(cache_dir,'{}_info.pkl').format(linemod_cls_name)
    )
    return database

def randomly_read_background(background_dir,cache_dir):
    if os.path.exists(os.path.join(cache_dir,'background_SUN_info.pkl')):
        fns = read_pickle(os.path.join(cache_dir,'background_SUN_info.pkl'))
    else:
        fns = glob(os.path.join(background_dir,'*.jpg')) + \
            glob(os.path.join(background_dir,'*.png'))
        save_pickle(fns, os.path.join(cache_dir,'background_SUN_info.pkl'))

    return cv2.imread(fns[np.random.randint(0,len(fns))])[:, :, ::-1]

def fuse_regions(rgbs, masks, depths, begins, cls_ids, background, th, tw, cls):
    fuse_order = np.arange(len(rgbs))
    np.random.shuffle(fuse_order)
    fuse_img = background
    fuse_img = cv2.resize(fuse_img,(tw,th),interpolation=cv2.INTER_LINEAR)
    fuse_mask = np.zeros([fuse_img.shape[0],fuse_img.shape[1]],np.int32)
    INF = pow(2,15)
    fuse_depth = np.ones([fuse_img.shape[0], fuse_img.shape[1]], np.uint16) * INF
    t_cls_id = lm_obj_dict[cls]
    if len(background.shape) < 3:
        return None, None, None, None
    for idx in fuse_order:
        if len(rgbs[idx].shape) < 3:
            continue
        cls_id = cls_ids[idx]
        rh,rw = masks[idx].shape
        if cls_id == t_cls_id:
            bh, bw = begins[idx][0], begins[idx][1]
        else:
            bh = np.random.randint(0,fuse_img.shape[0]-rh)
            bw = np.random.randint(0,fuse_img.shape[1]-rw)

        silhouette = masks[idx]>0
        out_silhouette = np.logical_not(silhouette)
        fuse_depth_patch = fuse_depth[bh:bh+rh, bw:bw+rw].copy()
        cover = (depths[idx] < fuse_depth_patch) * silhouette
        not_cover = np.logical_not(cover)

        fuse_mask[bh:bh+rh,bw:bw+rw] *= not_cover.astype(fuse_mask.dtype)
        cover_msk = masks[idx] * cover.astype(masks[idx].dtype)
        fuse_mask[bh:bh+rh,bw:bw+rw] += cover_msk

        fuse_img[bh:bh+rh,bw:bw+rw] *= not_cover.astype(fuse_img.dtype)[:,:,None]
        cover_rgb = rgbs[idx] * cover.astype(rgbs[idx].dtype)[:,:,None]
        fuse_img[bh:bh+rh,bw:bw+rw] += cover_rgb

        fuse_depth[bh:bh+rh, bw:bw+rw] *= not_cover.astype(fuse_depth.dtype)
        cover_dpt = depths[idx] * cover.astype(depths[idx].dtype)
        fuse_depth[bh:bh+rh, bw:bw+rw] += cover_dpt.astype(fuse_depth.dtype)

        begins[idx][0] = -begins[idx][0]+bh
        begins[idx][1] = -begins[idx][1]+bw

    dp_bg = (fuse_depth == INF)
    dp_bg_filter = np.logical_not(dp_bg)
    fuse_depth *= dp_bg_filter.astype(fuse_depth.dtype)

    return fuse_img, fuse_mask, fuse_depth, begins

def randomly_sample_foreground(image_db):
    idx = np.random.randint(0,len(image_db))
    rgb_pth = image_db[idx]['rgb_pth']
    dpt_pth = image_db[idx]['dpt_pth']
    msk_pth = image_db[idx]['msk_pth']
    with Image.open(dpt_pth) as di:
        depth = np.array(di).astype(np.int16)
    with Image.open(msk_pth) as li:
        mask = np.array(li).astype(np.int16)
    with Image.open(rgb_pth) as ri:
        rgb = np.array(ri)[:, :, :3].astype(np.uint8)

    mask = np.sum(mask,2)>0
    mask = np.asarray(mask,np.int32)

    hs, ws = np.nonzero(mask)
    hmin, hmax = np.min(hs),np.max(hs)
    wmin, wmax = np.min(ws),np.max(ws)

    mask = mask[hmin:hmax,wmin:wmax]
    rgb = rgb[hmin:hmax,wmin:wmax]
    depth = depth[hmin:hmax, wmin:wmax]

    rgb *= mask.astype(np.uint8)[:,:,None]
    depth *= mask.astype(np.uint16)[:,:]
    begin = [hmin,wmin]
    pose = image_db[idx]['RT']

    return rgb, mask, depth, begin, pose

def save_fuse_data(
    output_dir, idx, fuse_img, fuse_mask, fuse_depth, fuse_begins, t_pose, cls
):
    cls_id = lm_obj_dict[cls]
    if (fuse_mask == cls_id).sum() < 20:
        return None
    os.makedirs(output_dir, exist_ok=True)
    fuse_mask = fuse_mask.astype(np.uint8)
    data = {}
    data['rgb'] = fuse_img
    data['mask'] = fuse_mask
    data['depth'] = fuse_depth.astype(np.float32) / 1000.0
    data['K'] = Intrinsic_matrix['linemod']
    data['RT'] = t_pose
    data['cls_typ'] = cls
    data['rnd_typ'] = 'fuse'
    data['begins'] = fuse_begins
    if DEBUG:
        imshow("rgb", fuse_img[:, :, ::-1])
        imshow("depth", (fuse_depth / fuse_depth.max() * 255).astype('uint8'))
        imshow("label", (fuse_mask / fuse_mask.max() * 255).astype("uint8"))
        waitKey(0)
    sv_pth = os.path.join(output_dir, "{}.pkl".format(idx))
    pickle.dump(data, open(sv_pth, 'wb'))
    sv_pth = os.path.abspath(sv_pth)
    return sv_pth

def prepare_dataset_single(
    output_dir, idx, linemod_dir, background_dir, cache_dir, seed, cls
):
    time_begin = time.time()
    np.random.seed(seed)
    rgbs, masks, depths, begins, poses, cls_ids = [], [], [], [], [], []
    image_dbs={}
    for cls_name in lm_obj_dict.keys():
        cls_id = lm_obj_dict[cls_name]
        image_dbs[cls_id] = collect_linemod_set_info(cls_name)

    for cls_name in lm_obj_dict.keys():
        cls_id = lm_obj_dict[cls_name]
        rgb, mask, depth, begin, pose = randomly_sample_foreground(
            image_dbs[cls_id]
        )
        if cls_name == cls:
            t_pose = pose
        mask *= cls_id
        rgbs.append(rgb)
        masks.append(mask)
        depths.append(depth)
        begins.append(begin)
        poses.append(pose)
        cls_ids.append(cls_id)

    background = randomly_read_background(background_dir, cache_dir)

    fuse_img, fuse_mask, fuse_depth, fuse_begins= fuse_regions(
        rgbs, masks, depths, begins, cls_ids, background, 480, 640, cls
    )

    if fuse_img is not None:
        sv_pth = save_fuse_data(
            output_dir, idx, fuse_img, fuse_mask, fuse_depth, fuse_begins,
            t_pose, cls
        )
        return sv_pth

def prepare_dataset_parallel(
        output_dir, linemod_dir, fuse_num, background_dir, cache_dir,
        worker_num=8, cls="ape"
):
    exector = ProcessPoolExecutor(max_workers=worker_num)
    futures = []

    for idx in np.arange(fuse_num):
        seed = np.random.randint(500000)
        futures.append(exector.submit(
            prepare_dataset_single, output_dir, idx, linemod_dir,
            background_dir, cache_dir, seed, cls
        ))

    pth_lst = []
    for f in tqdm(futures):
        res = f.result()
        if res is not None:
            pth_lst.append(res)
    f_lst_pth = os.path.join(output_dir, "file_list.txt")
    with open(f_lst_pth, "w") as f:
        for item in pth_lst:
            print(item, file=f)

if __name__=="__main__":
    cls = args.cls
    for cls_name in lm_obj_dict.keys():
        collect_linemod_set_info(cls_name)
    _ = randomly_read_background(
        os.path.join(SUNPAS_dir, 'JPEGImages'),
        cache_dir
    )

    ensure_dir(output_dir)
    fuse_num = args.fuse_num
    worker_num = 20
    prepare_dataset_parallel(
        os.path.join(output_dir, "fuse",  cls),
        linemod_dir, fuse_num,
        os.path.join(SUNPAS_dir, 'JPEGImages'),
        cache_dir, worker_num, cls=cls
    )


import numpy as np
import pygame

class Actor:
    def __init__(self, settings, spritePath, size):
        self.settings = settings
        self.type = "not set"
        self.collisionGroup = -1
        self.timeStep = -1
        self.active = True

        #Physical attributes
        self.pos = np.array([settings.gameRes / 2, settings.gameRes / 2])
        self.oldPos = self.pos.copy()
        self.size = int(settings.gameRes * size)
        # needed for collision checking
        self.vertices = []
        self.vertices.append(np.array([0, 0]))
        self.vertices.append(np.array([0, self.size]))
        self.vertices.append(np.array([self.size, 0]))
        self.vertices.append(np.array([self.size, self.size]))

        #Graphical attributes
        self.sprite = pygame.image.load(spritePath)
        self.sprite = pygame.transform.scale(self.sprite, (self.size, self.size))

    def onBoxCollision(self, other):
        self.active = False

    def onWallCollision(self):
        pass


import cv2
from keras.models import load_model
import numpy as np
from imutils import face_utils
import os
from random import randint
import dlib
import imutils

face_cascade = cv2.CascadeClassifier('frontal.xml')
detector = dlib.get_frontal_face_detector()
model = load_model('emojigram.h5')
frame_check = 20

def calcResult(pred_class, cpu):
    if pred_class == cpu:
        return 'ok'
    return 'nok'

def main():
    flag = 0
    cFlag = 0
    oflag = 0
    result = ''
    emojis, emoji_queue, queue = get_emojis()
    cap = cv2.VideoCapture(0)
    x, y, w, h = 300, 50, 350, 350

    while (cap.isOpened()):
        ret, img = cap.read()
        img = cv2.flip(img, 1)
        # img = imutils.resize(img, width=800)
        faces = detector(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), 0)
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        mask2 = cv2.inRange(hsv, np.array([2, 50, 60]), np.array([25, 150, 255]))
        # mask2 = cv2.inRange(hsv, np.array([110, 50, 50]), np.array([130, 255, 255]))
        res = cv2.bitwise_and(img, img, mask=mask2)
        gray = cv2.cvtColor(res, cv2.COLOR_BGR2GRAY)
        median = cv2.GaussianBlur(gray, (5, 5), 0)

        kernel_square = np.ones((5, 5), np.uint8)
        dilation = cv2.dilate(median, kernel_square, iterations=2)
        opening = cv2.morphologyEx(dilation, cv2.MORPH_CLOSE, kernel_square)
        ret, thresh = cv2.threshold(opening, 30, 255, cv2.THRESH_BINARY)

        thresh = thresh[y:y + h, x:x + w]
        contours = cv2.findContours(thresh.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)[1]
        if cFlag == 0:
            c1 = 0
            c2 = 0
            c3 = 255
            thick = 2
        if len(faces) > 0:
            for face in faces:
                (fx, fy, fw, fh) = face_utils.rect_to_bb(face)
                cv2.rectangle(img, (fx, fy), (fx + fw, fy + fh), (c1, c2, c3), thick)
                if len(contours) > 0:
                    contour = max(contours, key=cv2.contourArea)
                    if cv2.contourArea(contour) > 2500:
                        if len(queue) != 0:
                            cpu = queue[0]
                        x, y, w1, h1 = cv2.boundingRect(contour)
                        newImage = thresh[y:y + h1, x:x + w1]
                        newImage = cv2.resize(newImage, (50, 50))
                        pred_probab, pred_class = keras_predict(model, newImage)
                        # print(pred_class, pred_probab)
                        count = 0
                        for e in emoji_queue:
                            img = overlay(img, e, fx - 50 + count, 120, 40, 40)
                            count = count + 50
                        img = overlay(img, emojis[pred_class], 370, 50, 90, 90)
                        img = overlay(img, emojis[cpu], 530, 50, 90, 90)
                        result = calcResult(pred_class, cpu)

                elif len(contours) == 0:
                    flag = 0
                x, y, w, h = 300, 50, 350, 350
                cv2.putText(img, 'USER', (380, 40),
                            cv2.FONT_HERSHEY_COMPLEX, 0.7, (0, 0, 255), 2)
                cv2.putText(img, 'EMOJI', (550, 40),
                            cv2.FONT_HERSHEY_COMPLEX, 0.7, (0, 0, 255), 2)
                if len(emoji_queue) == 0:
                    cFlag = 1
                    c1 = 0
                    c2 = 255
                    c3 = 0
                    thick = 6
                    cv2.putText(img, 'COMPLETED', (420, 300),
                                cv2.FONT_HERSHEY_COMPLEX, 0.7, (0, 255, 0), 2)
                    cv2.putText(img, 'Getting new Emojis', (380, 350),
                                cv2.FONT_HERSHEY_COMPLEX, 0.7, (255, 0, 0), 2)
                    if oflag == 0:
                        img = overlay(img, emojis[9], fx, fy, fw, fh)
                    flag += 1
                    print(flag)
                    if flag >= 100:
                        emojis, emoji_queue, queue = get_emojis()
                        flag = 0
                        cFlag = 0

        if result == 'ok':
            cFlag = 1
            c1 = 0
            c2 = 255
            c3 = 0
            thick = 4
            flag += 1
            print(flag)
            if flag >= frame_check:
                result = 'nok'
                if len(emoji_queue) != 0:
                    emoji_queue.pop(0)
                    queue.pop(0)
                    flag = 0
                    cFlag = 0

        elif result == 'nok':
            cFlag = 0

        cv2.imshow("Frame", img)
        cv2.imshow("Contours", thresh)
        k = cv2.waitKey(10)
        if k == 27:
            break

def keras_predict(model, image):
    processed = keras_process_image(image)
    pred_probab = model.predict(processed)[0]
    pred_class = list(pred_probab).index(max(pred_probab))
    return max(pred_probab), pred_class

def keras_process_image(img):
    image_x = 50
    image_y = 50
    img = cv2.resize(img, (image_x, image_y))
    img = np.array(img, dtype=np.float32)
    img = np.reshape(img, (-1, image_x, image_y, 1))
    return img

def get_emojis():
    emojis_folder = 'Emojigram_emo/'
    emojis = []
    emoji_queue = []
    for emoji in range(len(os.listdir(emojis_folder))):
        print(emoji)
        emojis.append(cv2.imread(emojis_folder + str(emoji) + '.png', -1))
    x = [randint(1, 12) for p in range(0, 5)]
    for e in x:
        emoji_queue.append(cv2.imread(emojis_folder + str(e) + '.png', -1))
    return emojis,emoji_queue,x

def overlay(image, emoji, x, y, w, h):
    emoji = cv2.resize(emoji, (w, h))
    try:
        image[y:y + h, x:x + w] = blend_transparent(image[y:y + h, x:x + w], emoji)
    except:
        pass
    return image

def blend_transparent(face_img, overlay_t_img):
    # Split out the transparency mask from the colour info
    overlay_img = overlay_t_img[:, :, :3]  # Grab the BRG planes
    overlay_mask = overlay_t_img[:, :, 3:]  # And the alpha plane

    # Again calculate the inverse mask
    background_mask = 255 - overlay_mask

    # Turn the masks into three channel, so we can use them as weights
    overlay_mask = cv2.cvtColor(overlay_mask, cv2.COLOR_GRAY2BGR)
    background_mask = cv2.cvtColor(background_mask, cv2.COLOR_GRAY2BGR)

    # Create a masked out face image, and masked out overlay
    # We convert the images to floating point in range 0.0 - 1.0
    face_part = (face_img * (1 / 255.0)) * (background_mask * (1 / 255.0))
    overlay_part = (overlay_img * (1 / 255.0)) * (overlay_mask * (1 / 255.0))

    # And finally just add them together, and rescale it back to an 8bit integer image
    return np.uint8(cv2.addWeighted(face_part, 255.0, overlay_part, 255.0, 0.0))

keras_predict(model, np.zeros((50, 50, 1), dtype=np.uint8))
main()


import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd.function import Function
from torch.autograd import Variable

class CenterTripletLoss(nn.Module):
    """ Hetero-center-triplet-loss-for-VT-Re-ID
   "Parameters Sharing Exploration and Hetero-Center Triplet Loss for Visible-Thermal Person Re-Identification"
   [(arxiv)](https://arxiv.org/abs/2008.06223).

    Args:
    - margin (float): margin for triplet.
    """

    def __init__(self, batch_size, margin=0.3):
        super(CenterTripletLoss, self).__init__()
        self.margin = margin
        self.ranking_loss = nn.MarginRankingLoss(margin=margin)

    def forward(self, feats, labels):
        """
        Args:
        - inputs: feature matrix with shape (batch_size, feat_dim)
        - targets: ground truth labels with shape (num_classes)
        """
        label_uni = labels.unique()
        targets = torch.cat([label_uni,label_uni])
        label_num = len(label_uni)
        feat = feats.chunk(label_num*2, 0)
        center = []
        for i in range(label_num*2):
            center.append(torch.mean(feat[i], dim=0, keepdim=True))
        inputs = torch.cat(center)

        n = inputs.size(0)

        # Compute pairwise distance, replace by the official when merged
        dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n)
        dist = dist + dist.t()
        dist.addmm_(1, -2, inputs, inputs.t())
        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability

        # For each anchor, find the hardest positive and negative
        mask = targets.expand(n, n).eq(targets.expand(n, n).t())
        dist_ap, dist_an = [], []
        for i in range(n):
            dist_ap.append(dist[i][mask[i]].max().unsqueeze(0))
            dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0))
        dist_ap = torch.cat(dist_ap)
        dist_an = torch.cat(dist_an)

        # Compute ranking hinge loss
        y = torch.ones_like(dist_an)
        loss = self.ranking_loss(dist_an, dist_ap, y)

        # compute accuracy
        correct = torch.ge(dist_an, dist_ap).sum().item()
        return loss, correct

class CrossEntropyLabelSmooth(nn.Module):
    """Cross entropy loss with label smoothing regularizer.
    Reference:
    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
    Equation: y = (1 - epsilon) * y + epsilon / K.
    Args:
        num_classes (int): number of classes.
        epsilon (float): weight.
    """
    def __init__(self, num_classes, epsilon=0.1, use_gpu=True):
        super(CrossEntropyLabelSmooth, self).__init__()
        self.num_classes = num_classes
        self.epsilon = epsilon
        self.use_gpu = use_gpu
        self.logsoftmax = nn.LogSoftmax(dim=1)

    def forward(self, inputs, targets):
        """
        Args:
            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
            targets: ground truth labels with shape (num_classes)
        """
        log_probs = self.logsoftmax(inputs)
        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1)
        if self.use_gpu: targets = targets.cuda()
        targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
        loss = (- targets * log_probs).mean(0).sum()
        return loss

class OriTripletLoss(nn.Module):
    """Triplet loss with hard positive/negative mining.

    Reference:
    Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.
    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.

    Args:
    - margin (float): margin for triplet.
    """

    def __init__(self, batch_size, margin=0.3):
        super(OriTripletLoss, self).__init__()
        self.margin = margin
        self.ranking_loss = nn.MarginRankingLoss(margin=margin)

    def forward(self, inputs, targets):
        """
        Args:
        - inputs: feature matrix with shape (batch_size, feat_dim)
        - targets: ground truth labels with shape (num_classes)
        """
        n = inputs.size(0)

        # Compute pairwise distance, replace by the official when merged
        dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n)
        dist = dist + dist.t()
        dist.addmm_(1, -2, inputs, inputs.t())
        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability

        # For each anchor, find the hardest positive and negative
        mask = targets.expand(n, n).eq(targets.expand(n, n).t())
        dist_ap, dist_an = [], []
        for i in range(n):
            dist_ap.append(dist[i][mask[i]].max().unsqueeze(0))
            dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0))
        dist_ap = torch.cat(dist_ap)
        dist_an = torch.cat(dist_an)

        # Compute ranking hinge loss
        y = torch.ones_like(dist_an)
        loss = self.ranking_loss(dist_an, dist_ap, y)

        # compute accuracy
        correct = torch.ge(dist_an, dist_ap).sum().item()
        return loss, correct

# Adaptive weights
def softmax_weights(dist, mask):
    max_v = torch.max(dist * mask, dim=1, keepdim=True)[0]
    diff = dist - max_v
    Z = torch.sum(torch.exp(diff) * mask, dim=1, keepdim=True) + 1e-6 # avoid division by zero
    W = torch.exp(diff) * mask / Z
    return W

def normalize(x, axis=-1):
    """Normalizing to unit length along the specified dimension.
    Args:
      x: pytorch Variable
    Returns:
      x: pytorch Variable, same shape as input
    """
    x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
    return x

class TripletLoss_WRT(nn.Module):
    """Weighted Regularized Triplet'."""

    def __init__(self):
        super(TripletLoss_WRT, self).__init__()
        self.ranking_loss = nn.SoftMarginLoss()

    def forward(self, inputs, targets, normalize_feature=False):
        if normalize_feature:
            inputs = normalize(inputs, axis=-1)
        dist_mat = pdist_torch(inputs, inputs)

        N = dist_mat.size(0)
        # shape [N, N]
        is_pos = targets.expand(N, N).eq(targets.expand(N, N).t()).float()
        is_neg = targets.expand(N, N).ne(targets.expand(N, N).t()).float()

        # `dist_ap` means distance(anchor, positive)
        # both `dist_ap` and `relative_p_inds` with shape [N, 1]
        dist_ap = dist_mat * is_pos
        dist_an = dist_mat * is_neg

        weights_ap = softmax_weights(dist_ap, is_pos)
        weights_an = softmax_weights(-dist_an, is_neg)
        furthest_positive = torch.sum(dist_ap * weights_ap, dim=1)
        closest_negative = torch.sum(dist_an * weights_an, dim=1)

        y = furthest_positive.new().resize_as_(furthest_positive).fill_(1)
        loss = self.ranking_loss(closest_negative - furthest_positive, y)

        # compute accuracy
        correct = torch.ge(closest_negative, furthest_positive).sum().item()
        return loss, correct

def pdist_torch(emb1, emb2):
    '''
    compute the eucilidean distance matrix between embeddings1 and embeddings2
    using gpu
    '''
    m, n = emb1.shape[0], emb2.shape[0]
    emb1_pow = torch.pow(emb1, 2).sum(dim = 1, keepdim = True).expand(m, n)
    emb2_pow = torch.pow(emb2, 2).sum(dim = 1, keepdim = True).expand(n, m).t()
    dist_mtx = emb1_pow + emb2_pow
    dist_mtx = dist_mtx.addmm_(1, -2, emb1, emb2.t())
    # dist_mtx = dist_mtx.clamp(min = 1e-12)
    dist_mtx = dist_mtx.clamp(min = 1e-12).sqrt()
    return dist_mtx

def pdist_np(emb1, emb2):
    '''
    compute the eucilidean distance matrix between embeddings1 and embeddings2
    using cpu
    '''
    m, n = emb1.shape[0], emb2.shape[0]
    emb1_pow = np.square(emb1).sum(axis = 1)[..., np.newaxis]
    emb2_pow = np.square(emb2).sum(axis = 1)[np.newaxis, ...]
    dist_mtx = -2 * np.matmul(emb1, emb2.T) + emb1_pow + emb2_pow
    # dist_mtx = np.sqrt(dist_mtx.clip(min = 1e-12))
    return dist_mtx


#
# @lc app=leetcode.cn id=867 lang=python3
#
# [867] 转置矩阵
#
# https://leetcode-cn.com/problems/transpose-matrix/description/
#
# algorithms
# Easy (68.31%)
# Likes:    174
# Dislikes: 0
# Total Accepted:    64.1K
# Total Submissions: 93.9K
# Testcase Example:  '[[1,2,3],[4,5,6],[7,8,9]]'
#
# 给你一个二维整数数组 matrix， 返回 matrix 的 转置矩阵 。
#
# 矩阵的 转置 是指将矩阵的主对角线翻转，交换矩阵的行索引与列索引。
#
#
#
#
#
# 示例 1：
#
#
# 输入：matrix = [[1,2,3],[4,5,6],[7,8,9]]
# 输出：[[1,4,7],[2,5,8],[3,6,9]]
#
#
# 示例 2：
#
#
# 输入：matrix = [[1,2,3],[4,5,6]]
# 输出：[[1,4],[2,5],[3,6]]
#
#
#
#
# 提示：
#
#
# m == matrix.length
# n == matrix[i].length
# 1
# 1
# -10^9
#
#
#

# @lc code=start
class Solution:
    def transpose(self, matrix: List[List[int]]) -> List[List[int]]: ## 直接用zip转置，然后再换成list
        import numpy as np
        return np.transpose(matrix).tolist()
# @lc code=end

    # def transpose(self, matrix: List[List[int]]) -> List[List[int]]: ## 按照题意转置

    #     row = len(matrix)

    #     col = len(matrix[0])

    #     res = [[0]*row for _ in range(col)]

    #     for i in range(col):
    #         for j in range(row):
    #             res[i][j] = matrix[j][i]

    #     return res

    # def transpose(self, matrix: List[List[int]]) -> List[List[int]]: ## 直接用zip转置，然后再换成list
    #     tmatrix = list(zip(*matrix))

    #     row,col = len(tmatrix),len(matrix[0])
    #     res = []
    #     for i in range(row):
    #         res.append(list(tmatrix[i]))
    #     return res


# %% md   代码环境tf.1.14.0

'''
如何用显卡跑代码:
首先卸了tensorflow
安装tensorflow-gpu==对应版本号
然后运行代码就行了,他会自动有gpu了

'''

import warnings
warnings.filterwarnings("ignore") #忽略警告
# 基于自注意力机制的翻译系统

# %% md

## 1. 数据处理
# - 读取数据
# - 分别保存为inputs，outputs

# %%
#cmn 就是中英文翻译的数据集了.
with open('./cmn.txt', 'r', encoding='utf8') as f:
    data = f.readlines()

# %%

from tqdm import tqdm

inputs = []
outputs = []
for line in tqdm(data[:10000]):
    [en, ch] = line.strip('\n').split('\t')
    inputs.append(en.replace(',', ' ,')[:-1].lower())
    outputs.append(ch[:-1])

# %%

print(inputs[:10])

# %%

print(outputs[:10])

# %% md

### 1.1 英文分词
# 我们将英文用空格隔开即可，但是需要稍微修改一下，将大写字母全部用小写字母代替。在上文中使用
# `.lower
# `进行了替代。
#
# ```py
for line in tqdm(data):
    [en, ch] = line.strip('\n').split('\t')
    inputs.append(en[:-1].lower())
    outputs.append(ch[:-1])

# ```
# 此处我们只需要将英文用空格分开即可。

# %%

inputs = [en.split(' ') for en in inputs]

# %%

print(inputs[:10])

# %% md

### 1.2 中文分词
# - 中文分词选择结巴分词工具。
# ```py
import jieba

outputs = [[char for char in jieba.cut(''.join(line)) if char != ' '] for line in outputs]
# ```
# # - 也可以用hanlp。
# # ```py
# from pyhanlp import *
#
# outputs = [[term.word for term in HanLP.segment(line) if term.word != ' '] for line in outputs]
# # ```
# - 或者按字分词？

# %%
#
# import jieba
#
# jieba_outputs = [[char for char in jieba.cut(line) if char != ' '] for line in outputs[-10:]]
# print(jieba_outputs)

# %%

# outputs = [[char for char in jieba.cut(line) if char != ' '] for line in tqdm(outputs)]

# %% md

### 1.3 生成字典

# %%

def get_vocab(data, init=['<PAD>']):
    vocab = init
    for line in tqdm(data):
        for word in line:
            if word not in vocab:
                vocab.append(word)
    return vocab

SOURCE_CODES = ['<PAD>']
TARGET_CODES = ['<PAD>', '<GO>', '<EOS>'] #go是句子开始标志,eos是句子结束标志
encoder_vocab = get_vocab(inputs, init=SOURCE_CODES)
decoder_vocab = get_vocab(outputs, init=TARGET_CODES)

# %%

print(encoder_vocab[:10])
print(decoder_vocab[:10])

# %% md

### 1.4 数据生成器

# 翻译系统解码器端需要输入和输出，因此解码器需要拼凑出输入序列和输出序列。

# %%

#应该decoder_inputs,decoder_targets都加上go和eos
encoder_inputs = [[encoder_vocab.index(word) for word in line] for line in inputs]
decoder_inputs = [[decoder_vocab.index('<GO>')] + [decoder_vocab.index(word) for word in line] for line in outputs]
decoder_targets = [[decoder_vocab.index(word) for word in line] + [decoder_vocab.index('<EOS>')] for line in outputs]

# %%

print(decoder_inputs[:4])
print(decoder_targets[:4])

# %%

import numpy as np

def get_batch(encoder_inputs, decoder_inputs, decoder_targets, batch_size=4):
    batch_num = len(encoder_inputs) // batch_size
    for k in range(batch_num):
        begin = k * batch_size
        end = begin + batch_size
        en_input_batch = encoder_inputs[begin:end]
        de_input_batch = decoder_inputs[begin:end]
        de_target_batch = decoder_targets[begin:end]
        max_en_len = max([len(line) for line in en_input_batch])
        max_de_len = max([len(line) for line in de_input_batch])
        en_input_batch = np.array([line + [0] * (max_en_len - len(line)) for line in en_input_batch])
        de_input_batch = np.array([line + [0] * (max_de_len - len(line)) for line in de_input_batch])
        de_target_batch = np.array([line + [0] * (max_de_len - len(line)) for line in de_target_batch])
        yield en_input_batch, de_input_batch, de_target_batch

# %%

batch = get_batch(encoder_inputs, decoder_inputs, decoder_targets, batch_size=4)
next(batch)

# %% md

## 2. 构建模型

# %%

import tensorflow as tf

# %% md

# ### 2.1 构造建模组件
# 下面代码实现了图片结构中的各个功能组件。
# ![image.png](attachment: image.png)

#### layer norm层

# %%

def normalize(inputs,
              epsilon=1e-8,
              scope="ln",
              reuse=None):
    '''Applies layer normalization.

    Args:
      inputs: A tensor with 2 or more dimensions, where the first dimension has
        `batch_size`.
      epsilon: A floating number. A very small number for preventing ZeroDivision Error.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns:
      A tensor with the same shape and data dtype as `inputs`.
    '''
    with tf.variable_scope(scope, reuse=reuse):
        inputs_shape = inputs.get_shape()
        params_shape = inputs_shape[-1:]

        mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)
        beta = tf.Variable(tf.zeros(params_shape))
        gamma = tf.Variable(tf.ones(params_shape))
        normalized = (inputs - mean) / ((variance + epsilon) ** (.5))
        outputs = gamma * normalized + beta

    return outputs

# %% md

#### embedding层

# %%

def embedding(inputs,
              vocab_size,
              num_units,
              zero_pad=True,
              scale=True,
              scope="embedding",
              reuse=None):
    '''Embeds a given tensor.
    Args:
      inputs: A `Tensor` with type `int32` or `int64` containing the ids
         to be looked up in `lookup table`.
      vocab_size: An int. Vocabulary size.
      num_units: An int. Number of embedding hidden units.
      zero_pad: A boolean. If True, all the values of the fist row (id 0)
        should be constant zeros.
      scale: A boolean. If True. the outputs is multiplied by sqrt num_units.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.
    Returns:
      A `Tensor` with one more rank than inputs's. The last dimensionality
        should be `num_units`.

    For example,

    ```
    import tensorflow as tf

    inputs = tf.to_int32(tf.reshape(tf.range(2*3), (2, 3)))
    outputs = embedding(inputs, 6, 2, zero_pad=True)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        print sess.run(outputs)
    >>
    [[[ 0.          0.        ]
      [ 0.09754146  0.67385566]
      [ 0.37864095 -0.35689294]]
     [[-1.01329422 -1.09939694]
      [ 0.7521342   0.38203377]
      [-0.04973143 -0.06210355]]]
    ```

    ```
    import tensorflow as tf

    inputs = tf.to_int32(tf.reshape(tf.range(2*3), (2, 3)))
    outputs = embedding(inputs, 6, 2, zero_pad=False)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        print sess.run(outputs)
    >>
    [[[-0.19172323 -0.39159766]
      [-0.43212751 -0.66207761]
      [ 1.03452027 -0.26704335]]
     [[-0.11634696 -0.35983452]
      [ 0.50208133  0.53509563]
      [ 1.22204471 -0.96587461]]]
    ```
    '''
    with tf.variable_scope(scope, reuse=reuse):
        lookup_table = tf.get_variable('lookup_table',
                                       dtype=tf.float32,
                                       shape=[vocab_size, num_units],
                                       initializer=tf.contrib.layers.xavier_initializer())
        if zero_pad: #xavier_initializer 一种初始化方法
            lookup_table = tf.concat((tf.zeros(shape=[1, num_units]),
                                      lookup_table[1:, :]), 0)
        outputs = tf.nn.embedding_lookup(lookup_table, inputs)

        if scale:
            outputs = outputs * (num_units ** 0.5)

    return outputs

# %% md

#### multihead层
# 该层实现了下面功能：
# ![image.png](attachment: image.png)

# %%

def multihead_attention(key_emb,
                        que_emb,
                        queries,
                        keys,
                        num_units=None,
                        num_heads=8,
                        dropout_rate=0,
                        is_training=True,
                        causality=False,
                        scope="multihead_attention",
                        reuse=None):
    '''Applies multihead attention.

    Args:
      queries: A 3d tensor with shape of [N, T_q, C_q].
      keys: A 3d tensor with shape of [N, T_k, C_k].
      num_units: A scalar. Attention size.
      dropout_rate: A floating point number.
      is_training: Boolean. Controller of mechanism for dropout.
      causality: Boolean. If true, units that reference the future are masked.
      num_heads: An int. Number of heads.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns
      A 3d tensor with shape of (N, T_q, C)
    '''
    with tf.variable_scope(scope, reuse=reuse):
        # Set the fall back option for num_units
        if num_units is None:
            num_units = queries.get_shape().as_list[-1]

        # Linear projections
        Q = tf.layers.dense(queries, num_units, activation=tf.nn.relu)  # (N, T_q, C)
        K = tf.layers.dense(keys, num_units, activation=tf.nn.relu)  # (N, T_k, C)
        V = tf.layers.dense(keys, num_units, activation=tf.nn.relu)  # (N, T_k, C)

        # Split and concat
        Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0)  # (h*N, T_q, C/h)
        K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0)  # (h*N, T_k, C/h)
        V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0)  # (h*N, T_k, C/h)

        # Multiplication
        outputs = tf.matmul(Q_, tf.transpose(K_, [0, 2, 1]))  # (h*N, T_q, T_k)

        # Scale
        outputs = outputs / (K_.get_shape().as_list()[-1] ** 0.5)

        # Key Masking
        key_masks = tf.sign(tf.abs(tf.reduce_sum(key_emb, axis=-1)))  # (N, T_k)
        key_masks = tf.tile(key_masks, [num_heads, 1])  # (h*N, T_k)
        key_masks = tf.tile(tf.expand_dims(key_masks, 1), [1, tf.shape(queries)[1], 1])  # (h*N, T_q, T_k)

        paddings = tf.ones_like(outputs) * (-2 ** 32 + 1)
        outputs = tf.where(tf.equal(key_masks, 0), paddings, outputs)  # (h*N, T_q, T_k)

        # Causality = Future blinding
        if causality:
            diag_vals = tf.ones_like(outputs[0, :, :])  # (T_q, T_k)
            tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense()  # (T_q, T_k)
            masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(outputs)[0], 1, 1])  # (h*N, T_q, T_k)

            paddings = tf.ones_like(masks) * (-2 ** 32 + 1)
            outputs = tf.where(tf.equal(masks, 0), paddings, outputs)  # (h*N, T_q, T_k)

        # Activation
        outputs = tf.nn.softmax(outputs)  # (h*N, T_q, T_k)

        # Query Masking
        query_masks = tf.sign(tf.abs(tf.reduce_sum(que_emb, axis=-1)))  # (N, T_q)
        query_masks = tf.tile(query_masks, [num_heads, 1])  # (h*N, T_q)
        query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, tf.shape(keys)[1]])  # (h*N, T_q, T_k)
        outputs *= query_masks  # broadcasting. (N, T_q, C)

        # Dropouts
        outputs = tf.layers.dropout(outputs, rate=dropout_rate, training=tf.convert_to_tensor(is_training))

        # Weighted sum
        outputs = tf.matmul(outputs, V_)  # ( h*N, T_q, C/h)

        # Restore shape
        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2)  # (N, T_q, C)

        # Residual connection
        outputs += queries

        # Normalize
        outputs = normalize(outputs)  # (N, T_q, C)

    return outputs

# %% md

#### feedforward

# 两层全连接，用卷积模拟加速运算，也可以使用dense层。

# %%

def feedforward(inputs,
                num_units=[2048, 512],
                scope="multihead_attention",
                reuse=None):
    '''Point-wise feed forward net.

    Args:
      inputs: A 3d tensor with shape of [N, T, C].
      num_units: A list of two integers.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns:
      A 3d tensor with the same shape and dtype as inputs
    '''
    with tf.variable_scope(scope, reuse=reuse):
        # Inner layer
        params = {"inputs": inputs, "filters": num_units[0], "kernel_size": 1,
                  "activation": tf.nn.relu, "use_bias": True}
        outputs = tf.layers.conv1d(**params)

        # Readout layer
        params = {"inputs": outputs, "filters": num_units[1], "kernel_size": 1,
                  "activation": None, "use_bias": True}
        outputs = tf.layers.conv1d(**params)

        # Residual connection
        outputs += inputs

        # Normalize
        outputs = normalize(outputs)

    return outputs

# %% md

# #### label_smoothing.
# 对于训练有好处，将0变为接近零的小数，1
# 变为接近1的数，原文：
#
# During
# training, we
# employed
# label
# smoothing
# of
# value ls = 0.1[36].This
# hurts
# perplexity, as the
# model
# learns
# to
# be
# more
# unsure, but
# improves
# accuracy and BLEU
# score.
#
#
# # %%

def label_smoothing(inputs, epsilon=0.1):
    '''Applies label smoothing. See https://arxiv.org/abs/1512.00567.

    Args:
      inputs: A 3d tensor with shape of [N, T, V], where V is the number of vocabulary.
      epsilon: Smoothing rate.

    For example,

    ```
    import tensorflow as tf
    inputs = tf.convert_to_tensor([[[0, 0, 1],
       [0, 1, 0],
       [1, 0, 0]],
      [[1, 0, 0],
       [1, 0, 0],
       [0, 1, 0]]], tf.float32)

    outputs = label_smoothing(inputs)

    with tf.Session() as sess:
        print(sess.run([outputs]))

    >>
    [array([[[ 0.03333334,  0.03333334,  0.93333334],
        [ 0.03333334,  0.93333334,  0.03333334],
        [ 0.93333334,  0.03333334,  0.03333334]],
       [[ 0.93333334,  0.03333334,  0.03333334],
        [ 0.93333334,  0.03333334,  0.03333334],
        [ 0.03333334,  0.93333334,  0.03333334]]], dtype=float32)]
    ```
    '''
    K = inputs.get_shape().as_list()[-1]  # number of channels
    return ((1 - epsilon) * inputs) + (epsilon / K)

# %% md

### 2.2 搭建模型
# 模型实现下图结构：
# ![image.png](attachment: image.png)

# %%

class Graph():
    def __init__(self, is_training=True):
        tf.reset_default_graph()
        self.is_training = arg.is_training
        self.hidden_units = arg.hidden_units
        self.input_vocab_size = arg.input_vocab_size
        self.label_vocab_size = arg.label_vocab_size
        self.num_heads = arg.num_heads
        self.num_blocks = arg.num_blocks
        self.max_length = arg.max_length
        self.lr = arg.lr
        self.dropout_rate = arg.dropout_rate

        # input placeholder
        self.x = tf.placeholder(tf.int32, shape=(None, None))
        self.y = tf.placeholder(tf.int32, shape=(None, None))
        self.de_inp = tf.placeholder(tf.int32, shape=(None, None))

        # Encoder
        with tf.variable_scope("encoder"):
            # embedding
            self.en_emb = embedding(self.x, vocab_size=self.input_vocab_size, num_units=self.hidden_units, scale=True,
                                    scope="enc_embed")
            self.enc = self.en_emb + embedding(
                tf.tile(tf.expand_dims(tf.range(tf.shape(self.x)[1]), 0), [tf.shape(self.x)[0], 1]),
                vocab_size=self.max_length, num_units=self.hidden_units, zero_pad=False, scale=False, scope="enc_pe")
            # Dropout
            self.enc = tf.layers.dropout(self.enc,
                                         rate=self.dropout_rate,
                                         training=tf.convert_to_tensor(self.is_training))

            #Blocks
            for i in range(self.num_blocks):
                with tf.variable_scope("num_blocks_{}".format(i)):
                    # Multihead Attention
                    self.enc = multihead_attention(key_emb=self.en_emb,
                                                   que_emb=self.en_emb,
                                                   queries=self.enc,
                                                   keys=self.enc,
                                                   num_units=self.hidden_units,
                                                   num_heads=self.num_heads,
                                                   dropout_rate=self.dropout_rate,
                                                   is_training=self.is_training,
                                                   causality=False)

            # Feed Forward
            self.enc = feedforward(self.enc, num_units=[4 * self.hidden_units, self.hidden_units])

        # Decoder
        with tf.variable_scope("decoder"):
            # embedding
            self.de_emb = embedding(self.de_inp, vocab_size=self.label_vocab_size, num_units=self.hidden_units,
                                    scale=True, scope="dec_embed")
            self.dec = self.de_emb + embedding(
                tf.tile(tf.expand_dims(tf.range(tf.shape(self.de_inp)[1]), 0), [tf.shape(self.de_inp)[0], 1]),
                vocab_size=self.max_length, num_units=self.hidden_units, zero_pad=False, scale=False, scope="dec_pe")
            # Dropout
            self.dec = tf.layers.dropout(self.dec,
                                         rate=self.dropout_rate,
                                         training=tf.convert_to_tensor(self.is_training))

            # Multihead Attention ( self-attention)
            for i in range(self.num_blocks):
                with tf.variable_scope("num_blocks_{}".format(i)):
                    # Multihead Attention
                    self.dec = multihead_attention(key_emb=self.de_emb,
                                                   que_emb=self.de_emb,
                                                   queries=self.dec,
                                                   keys=self.dec,
                                                   num_units=self.hidden_units,
                                                   num_heads=self.num_heads,
                                                   dropout_rate=self.dropout_rate,
                                                   is_training=self.is_training,
                                                   causality=True,
                                                   scope='self_attention')

            #Multihead Attention ( vanilla attention)
            for i in range(self.num_blocks):
                with tf.variable_scope("num_blocks_{}".format(i)):
                    ### Multihead Attention
                    self.dec = multihead_attention(key_emb=self.en_emb,
                                                   que_emb=self.de_emb,
                                                   queries=self.dec,
                                                   keys=self.enc,
                                                   num_units=self.hidden_units,
                                                   num_heads=self.num_heads,
                                                   dropout_rate=self.dropout_rate,
                                                   is_training=self.is_training,
                                                   causality=True,
                                                   scope='vanilla_attention')

                    # Feed Forward
            self.outputs = feedforward(self.dec, num_units=[4 * self.hidden_units, self.hidden_units])

        # Final linear projection
        self.logits = tf.layers.dense(self.outputs, self.label_vocab_size)
        self.preds = tf.to_int32(tf.argmax(self.logits, axis=-1))
        self.istarget = tf.to_float(tf.not_equal(self.y, 0))
        self.acc = tf.reduce_sum(tf.to_float(tf.equal(self.preds, self.y)) * self.istarget) / (
            tf.reduce_sum(self.istarget))
        tf.summary.scalar('acc', self.acc)

        if is_training:
            # Loss
            self.y_smoothed = label_smoothing(tf.one_hot(self.y, depth=self.label_vocab_size))
            self.loss = tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.logits, labels=self.y_smoothed)
            self.mean_loss = tf.reduce_sum(self.loss * self.istarget) / (tf.reduce_sum(self.istarget))

            # Training Scheme
            self.global_step = tf.Variable(0, name='global_step', trainable=False)
            self.optimizer = tf.train.AdamOptimizer(learning_rate=self.lr, beta1=0.9, beta2=0.98, epsilon=1e-8)
            self.train_op = self.optimizer.minimize(self.mean_loss, global_step=self.global_step)

            # Summary
            tf.summary.scalar('mean_loss', self.mean_loss)
            self.merged = tf.summary.merge_all()

# %% md

## 3. 训练模型

# %% md

### 3.1 参数设定

# %%

def create_hparams():
    params = tf.contrib.training.HParams(
        num_heads=8,
        num_blocks=6,
        # vocab
        input_vocab_size=50,
        label_vocab_size=50,
        # embedding size
        max_length=100,
        hidden_units=512,
        dropout_rate=0.2,
        lr=0.0003,
        is_training=True)
    return params

arg = create_hparams()
arg.input_vocab_size = len(encoder_vocab)
arg.label_vocab_size = len(decoder_vocab)

# %% md

### 3.2 模型训练

# %%

import os

epochs = 0
batch_size = 64

g = Graph(arg)

saver = tf.train.Saver()
with tf.Session() as sess:
    merged = tf.summary.merge_all()
    sess.run(tf.global_variables_initializer())
    if os.path.exists('./logs/model2.meta'):
        print('已经有模型了,所以下面直接进行finetune')
        saver.restore(sess, './logs/model2')
    writer = tf.summary.FileWriter('tensorboard/lm', tf.get_default_graph())
    for k in range(epochs):
        total_loss = 0
        batch_num = len(encoder_inputs) // batch_size
        batch = get_batch(encoder_inputs, decoder_inputs, decoder_targets, batch_size)
        for i in tqdm(range(batch_num)):
            encoder_input, decoder_input, decoder_target = next(batch)
            feed = {g.x: encoder_input, g.y: decoder_target, g.de_inp: decoder_input}
            cost, _ = sess.run([g.mean_loss, g.train_op], feed_dict=feed)
            total_loss += cost
            if (k * batch_num + i) % 10 == 0:
                rs = sess.run(merged, feed_dict=feed)
                writer.add_summary(rs, k * batch_num + i)
        if (k + 1) % 5 == 0:
            print('epochs', k + 1, ': average loss = ', total_loss / batch_num)
    saver.save(sess, './logs/model2')
    writer.close()

# %% md

### 3.3 模型推断

# %%
import tensorflow as tf
arg.is_training = False

g = Graph(arg)

saver = tf.train.Saver()

with tf.Session() as sess:
    saver.restore(sess, './logs/model2')
    while True:
        line = input('输入测试英语句子: ')
        if line == 'exit': break
        line = line.lower().replace(',', ' ,').strip('\n').split(' ')
        tmp=[]
        for i in line:
            try:
                tmp.append(encoder_vocab.index(i))
            except:
                continue

        x=np.array(tmp)
        x = x.reshape(1, -1)
        de_inp = [[decoder_vocab.index('<GO>')]]
        while True:
            y = np.array(de_inp)
            preds = sess.run(g.preds, {g.x: x, g.de_inp: y})
            if preds[0][-1] == decoder_vocab.index('<EOS>'):
                break
            de_inp[0].append(preds[0][-1])#输入每一次,加进去预测出来的一个字.
        got = ''.join(decoder_vocab[idx] for idx in de_inp[0][1:])
        print(got)

# %%


"""
.. _tut_intro_pyton:

Introduction to Python
======================
"""

###############################################################################
# Python is a modern, general-purpose, object-oriented, high-level programming
# language. First make sure you have a working python environment and
# dependencies (see :ref:`install_python_and_mne_python`). If you are
# completely new to python, don't worry, it's just like any other programming
# language, only easier. Here are a few great resources to get you started:
#
# * `SciPy lectures <http://scipy-lectures.github.io>`_
# * `Learn X in Y minutes: Python <https://learnxinyminutes.com/docs/python/>`_
# * `NumPy for MATLAB users <https://docs.scipy.org/doc/numpy-dev/user/numpy-for-matlab-users.html>`_  # noqa
#
# We highly recommend watching the Scipy videos and reading through these
# sites to get a sense of how scientific computing is done in Python.
#
# Here are few bulletin points to familiarise yourself with python:
#
# Everything is dynamically typed. No need to declare simple data
# structures or variables separately.
a = 3
print(type(a))
b = [1, 2.5, 'This is a string']
print(type(b))
c = 'Hello world!'
print(type(c))

###############################################################################
# If you come from a background of matlab, remember that indexing in python
# starts from zero:
a = [1, 2, 3, 4]
print('This is the zeroth value in the list: {}'.format(a[0]))

###############################################################################
# No need to reinvent the wheel. Scipy and Numpy are battle field tested
# libraries that have a vast variety of functions for your needs. Consult the
# documentation and remember, you can always ask the interpreter for help with
# a question mark at the end of a function::
#
#    >>> import numpy as np
#    >>> np.arange?


"""Direct Plot

Educational library to directly plot single data points.

ATTENTION: This library is slow and not suited for production use!

It has been developed for educational purpose, especially to visualize
numerical algorithms e.g. for simulation or plotting measurement data.

It wraps matplotlib and pyplot commands in even simpler commands:

```
import math
import directplot as dp

dp.init()

for i in range(51):
    x = i*2*math.pi/50
    y = math.sin(x)
    dp.add(0, x, y)

dp.waitforclose()
```

The following functions are provided:

* `init()` Initializes and opens a Direct Plot window
* `add()` Adds a single point to a plot line
* `showMarker()` Shows or hides marker points on a plot line or on all plot lines
* `label()` Changes the label of a plot line used in the legend
* `title()` Changes the title of a sub-plot
* `xylabel()` Changes the axis lables of a sub-plot
* `refresh()` Refreshes the contents of the plot window
* `close()` Closes the Direct Plot window
* `clear()` Deletes the contents of the plot window
* `waitforclose()` Displays a new window title on the plot window and
                   blocks execution until user closes the window.
"""

__version__ = '0.4.1'
__author__ = 'Georg Braun'

import inspect as _inspect
from typing import Sequence as _Sequence
from .directplot import __DirectPlot
import platform as _platform
import matplotlib.pyplot as _plt

def init(titles: _Sequence[str] = ["Direct-Plot"], linesPerSubplot: int = 4, showMarker: bool = True, maxPoints: int = 10000) -> None:
    """Initializes and opens a Direct Plot window.

    Parameters:
    -----------
    * titles: A list or tuple containing 1 to 3 strings, resulting in 1 to 3 sub-plots on the plot window. Optional with a default title for a single sub-plot.
    * linesPerSubplot: Number of lines (data series) per sub-plot. Optional with default 4
    * showMarker: Determines if data points are emphasized with a little dot. Optional with default True
    * maxPoints: Maximum number of data points per line (data series). Optional with default 10000

    Returns:
    --------
    None

    Examples:
    ---------
    ```
    dp.init()
    dp.init(["Results"])
    dp.init(["Height", "Speed", "Forces"], 2, False)
    ```
    """

    global __dp
    if __dp is not None:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): YOU HAVE CALLED {_inspect.currentframe().f_code.co_name}() TOO OFTEN!")
    __dp = __DirectPlot(titles, linesPerSubplot, showMarker, maxPoints)

def close() -> None:
    """Closes the Direct Plot window.

    Example:
    --------
    ```
    dp.close()
    ```
    """
    global __dp
    try:
        __dp.close()
        del(__dp)
        __dp = None
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def clear() -> None:
    """Deletes the contents of the plot window.

    Keeps the number of sub-plots, the number of lines per sub-plot and the titles of the sub-plots. Everything else is reset/deleted.

    Example:
    --------
    ```
    dp.clear()
    ```
    """

    global __dp
    try:
        __dp.clear()
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def waitforclose(msg: str = None) -> None:
    """Displays a new window title on the plot window and blocks execution until user closes the window.

    Parameters:
    -----------
    * msg: A string to be shown on the window title. Optional with default None resulting in a standard title

    Returns:
    --------
    None

    Examples:
    ---------
    ```
    dp.waitforclose()
    dp.waitforclose("PLEASE CLOSE THIS WINDOW")
    ```
    """

    global __dp
    try:
        __dp.waitforclose(msg)
        del(__dp)
        __dp = None
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def add(id: int, x: float, y: float, refresh: bool = True) -> None:
    """Adds a single point to a plot line.

    Parameters:
    -----------
    * id: The id of the target plot line
    * x: x value
    * y: y value
    * refresh: Determines if the plot is refreshed immediately resulting in slower plotting speed. Optional with default True

    Returns:
    --------
    None

    Examples:
    ---------
    ```
    dp.add(0, 0.1, 2.7)
    dp.add(1, 1.1, 7.3, False)
    dp.add(1, 1.2, 7.2)
    ```
    """

    global __dp
    try:
        __dp.add(id, x, y, refresh)
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def refresh() -> None:
    """Refreshes the contents of the plot window.

    Mostly used in conjunction with add() and refresh=False.

    Example:
    --------
    ```
    dp.add(0, 0.1, 7.3, False)
    dp.add(0, 0.2, 6.9, False)
    dp.add(0, 0.3, 2.1, False)
    dp.refresh()
    ```
    """

    global __dp
    try:
        __dp.refresh()
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def showMarker(show: bool = True, id: int = None) -> None:
    """Shows or hides marker points on a plot line or on all plot lines.

    Parameters:
    -----------
    * show: Show or hide markes. Optional with default True
    * id: The id of the target plot line. Optional with default None resulting in a change of markers on all plot lines.

    Returns:
    --------
    None

    Examples:
    ---------
    ```
    dp.showMarker()
    dp.showMarker(False, 1)
    ```
    """

    global __dp
    try:
        __dp.showMarker(show, id)
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def label(id: int, label: str) -> None:
    """Changes the label of a plot line used in the legend.

    Parameters:
    -----------
    * id: The id of the target plot line
    * label: The new label text

    Returns:
    --------
    None

    Examples:
    ---------
    ```
    dp.label(0, "mass in kg")
    ```
    """

    global __dp
    try:
        __dp.label(id, label)
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def title(id: int, title: str) -> None:
    """Changes the title of a sub-plot

    Parameters:
    -----------
    * id: The id of the target plot line used to determine the corresponding sub-plot
    * title: The new title text

    Returns:
    --------
    None

    Examples:
    ---------
    ```
    dp.title(0, "Simulated Values")
    ```
    """

    global __dp
    try:
        __dp.title(id, title)
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def xylabel(id: int, xlabel: str, ylabel: str) -> None:
    """Changes the axis lables of a sub-plot

    Parameters:
    -----------
    * id: The id of the target plot line used to determine the corresponding sub-plot
    * xlabel: New label for the x axis
    * ylabel: New label for the y axis

    Returns:
    --------
    None

    Examples:
    ---------
    ```
    dp.xylabel(0, "time in s", "force in N")
    ```
    """

    global __dp
    try:
        __dp.xylabel(id, xlabel, ylabel)
    except AttributeError:
        raise Exception(f"ERROR in directplot.{_inspect.currentframe().f_code.co_name}(): NO MORE PLOT-WINDOW - DID YOU CLOSE IT ALREADY?")

def dp_selftest() -> None:
    """Runs some tests and quits."""

    import time

    print()
    print("Welcome to directplot. We will run some tests for you. Have fun...")
    print()

    print("test1:", _test1.__doc__)
    _test1()
    print("test2:", _test2.__doc__)
    _test2()
    print("test3:", _test3.__doc__)
    _test3()

    print()
    print("Done with testing. Bye bye ...")
    for i in range(3, 0, -1):
        print(i, end=" ", flush=True)
        time.sleep(1)
    print()

def _test1() -> None:
    """Plots a sine curve"""

    import math
    import time

    time.sleep(0.2)
    points=51
    init()
    for i in range(points):
        x = i*2*math.pi/(points-1)
        y = math.sin(x)
        add(0, x, y)
    time.sleep(2)
    close()

def _test2() -> None:
    """Combines sine and cosine to a circle in two sub-plots"""

    import math
    import time

    time.sleep(0.2)
    points=51
    init(["Sinus, Cosinus", "Circle"], 2, False)
    label(0, "Cosinus")
    label(1, "Sinus")
    for i in range(points):
        t = i*2*math.pi/(points-1)
        x = math.cos(t)
        y = math.sin(t)
        add(0, t, x, False)
        add(1, t, y, False)
        add(2, x, y, False)
        refresh()
    time.sleep(2)
    close()

def _test3() -> None:
    """Combines sine and cosine to a circle in two sub-plots and plots sinc() in a third one."""

    import math
    import time
    import numpy as _np

    time.sleep(0.2)
    points=51
    init(["Sinus, Cosinus", "Circle", "sinc()"], 2)
    label(0, "Cosinus")
    label(1, "Sinus")
    label(4, "sinc()")
    for i in range(points):
        t = i*2*math.pi/(points-1)
        x = math.cos(t)
        y = math.sin(t)
        add(0, t, x, False)
        add(1, t, y, False)
        add(2, x, y, False)
        add(4, t, _np.sinc(t-math.pi), False)
        refresh()
    time.sleep(3)
    close()

def onImport():
    print(f'directplot v{__version__} started with backend {_plt.get_backend()}')
    handleBackend()
    print()

def handleBackend():
    if _platform.system().lower() == 'darwin':
        targetBackend = 'TkAgg'
        print(f'Seems like MacOS. Trying to switch backend to {targetBackend}')
        try:
            _plt.switch_backend(targetBackend)
        except Exception as exceptionDetails:
            print(f'ERROR: Could not switch to {targetBackend}! Details:')
            print(f'{exceptionDetails}')
            print()
            print(f'Will try with current backend {_plt.get_backend()}')
        else:
            print(f'Successfully switched to backend {_plt.get_backend()}')

__dp = None
onImport()


import pyautogui
import time
from PIL import Image
import pytesseract
import os
import numpy as np
from skimage.color import rgb2gray
from skimage.feature import match_template
from skimage.io import imread
import pyautogui

startall = time.time()
# get screen resolution
screensize = pyautogui.size()
# screenshot directory
img_path = r'screen.png'
img_path_button = r'button.png'
img_path_template= r'template_match.png'
screenshot = "screencapture {}".format(img_path)
screenshot_button = "screencapture {}".format(img_path_button)
screenshot_template = "screencapture {}".format(img_path_template)

# CROPPING of screenshots
wood = (70, 180, 150, 210)
end_msg = (450, 300, 780, 360)
horizontal_button = (0,800,100,870)
vertical_button = (0,1400,80,1450)
map = (10,210,1240,690)

# POSITIONS
# speed button
speed = (552, 98)
# upgrade to lasers
upgrade = [(27,32),(0,-54),(29,-37),(0,-61)]
hero1 = (36, 183)
hero2 = (587, 136)
hero3 = (36, 138)
start = (337, 372)
# back button of vysor
backbutton = (102, 418) # horizontal
backbuttonV = (51, 708) # vertical
cancelbutton = (253, 314)
restartbutton =(361, 296)

def take2nd(elem):
    # for sorting 2nd position of a tuple within a list
    return elem[1]

def templatematching():
    # extract coordinates of all tower positions
    os.system(screenshot_template)
    # get screenshot & crop to reduce search area
    img = imread(img_path_template)
    img = rgb2gray(img)

    # find ratio of screen resolution with screenshot resolution
    ratio = screensize[0] / img.shape[1]

    cropped = img[0:650, 0:1230]
    # read in template
    sign = imread('sign.png')

    # Template matching, results gives a correlation coefficient for all pixels in image
    result = match_template(cropped, sign)
    # Flatten, sort, and top n coefficient values'
    d = result.flatten()
    e = np.sort(d)[::-1][:20]

    # extract all coordinates
    list = []
    for i in e:
        # locate indexes from top n values
        top = np.argwhere(d==i)[0][0]
        # obtain coordinates from index
        ij = np.unravel_index(top, result.shape)
        # swap position as it is in (y, x)
        x, y = ij[::-1]
        # centre the image based on template shape
        h, w = sign.shape
        x = int(x+h/2)
        y = int(y+w/2)
        list.append((x,y))

    # remove all duplicates
    list2 = list
    value = 5
    for i in list2:
        x = i[0]
        y = i[1]

        for i2 in list2:
            x2 = i2[0]
            y2 = i2[1]
            # if same coordinates do not remove as comparing same point
            if x-x2 == 0 and y-y2 == 0:
                continue
            # if coordinates are very similar, likely to be duplicates; remove
            elif -value <= x-x2 <= value and -value <= y-y2 <= value:
                list2.remove(i2)

    # change the coordinates of screenshot to actual computer screen coordinates
    list2 = [(one * ratio, two * ratio) for one, two in list2]
    # order it such that tower builds from bottom to top
    return sorted(list2, reverse=True, key=take2nd)

def objectrecognition():
    # screenshot to check for ending
    os.system(screenshot)
    img = Image.open(img_path)
    cropped = img.crop(end_msg)
    # OCR to text
    return pytesseract.image_to_string(cropped)

def colorrecognition(button):
    # screenshot to check for blue button position
    os.system(screenshot_button)
    img = Image.open(img_path_button)
    cropped = img.crop(button)
    # change tpositions into nparray so that can use np.histogram
    rbg = np.array(cropped.histogram())
    # get the first binned quantity
    return int(np.histogram(rbg, bins=200, range=(0,500))[0][0])

def upgrade1st():
    # 1st tower upgrade
    for i in tpositions:
        a = i
        pyautogui.click(a)
        time.sleep(0.5)
        z = upgrade[1]
        pyautogui.click(a[0]+z[0],a[1]+z[1])
        time.sleep(0.5)
        pyautogui.click(a[0]+z[0],a[1]+z[1])
        for aa in range(5):
            pyautogui.click(start)
            time.sleep(1)

        # OCR
        text = objectrecognition()
        if text == 'Emergency':
            global end
            end = 1
            print(1, 'Game Ended')
            # click cancel button when game ends
            pyautogui.click(cancelbutton)
            break
        else:
            pass

def upgrade2nd():
    # 2nd tower upgrade
    global end
    if end != 1:
        for i in tpositions:
            a = i
            pyautogui.click(a)
            time.sleep(0.5)
            z = upgrade[2]
            pyautogui.click(a[0]+z[0],a[1]+z[1])
            time.sleep(0.5)
            pyautogui.click(a[0]+z[0],a[1]+z[1])
            for aa in range(7):
                pyautogui.click(start)
                time.sleep(1)

            # OCR
            text = objectrecognition()
            if text == 'Emergency':
                end = 1
                print(2, 'Game Ended')
                # click cancel button when game ends
                pyautogui.click(cancelbutton)
                break
            else:
                pass

def finalupgrade():
    # Final tower upgrades
    global end
    if end != 1:
        for zz, i in enumerate(tpositions):
            # remove some tower upgrades so it won't kill the boss
            if zz != 2 and zz != 3 and zz != 4:
                a = i
                pyautogui.click(a)
                time.sleep(0.5)
                z = upgrade[3]
                pyautogui.click(a[0]+z[0],a[1]+z[1])
                time.sleep(0.5)
                pyautogui.click(a[0]+z[0],a[1]+z[1])
                for aa in range(5):
                    pyautogui.click(start)
                    time.sleep(1)

                # OCR
                text = objectrecognition()
                if text == 'Emergency':
                    end = 1
                    print(3, 'Game Ended')
                    # click cancel button when game ends
                    pyautogui.click(cancelbutton)
                    break
                else:
                    pass

def restart(num):
    # remove advert and restart game

    # begin restart if OCR finds game has ended
    if end != 1:
        for i in range(30):
            text = objectrecognition()
            if text == 'Emergency':
                print(4, 'Game Ended')
                # click cancel button when game ends
                pyautogui.click(cancelbutton)
                break
            else:
                time.sleep(5)
                pass

    for i in range(8):
        text = objectrecognition()
        time.sleep(3)
        if text == 'Defeat':
            # click restart button if it appears
            print('Restart Button Detected')
            pyautogui.click(restartbutton)
            break
        # account for vertical rotation when displaying ads
        # previous histogram testing shown the quantity of colors fall within 980-1050
        elif 980 < colorrecognition(vertical_button) < 1050:
            print ('Vertical Ad Detected')
            time.sleep(1)
            pyautogui.click(backbuttonV)
            time.sleep(3)
        elif 980 < colorrecognition(horizontal_button) < 1050:
            print ('Horizontal Ad Detected')
            time.sleep(1)
            pyautogui.click(backbutton)
            time.sleep(3)

    endt = time.time()
    if num == 1:
        print (num, 'round down,', 360*num, 'diamonds earned! after',round((endt-startt)/60,2), 'mins')
    else:
        print (num, 'rounds down,', 360*num, 'diamonds earned! after', round((endt-startt)/60,2), 'mins')

# get all tower coordinates
tpositions = templatematching()

# select vysor screen
pyautogui.click(tpositions[0])

# define number of rounds to loop
for num in range(1,10):
    # flag for ending
    global end
    end = 0
    startt = time.time()
    # shift heros
    pyautogui.click(hero1)
    pyautogui.click(141, 222)
    pyautogui.click(hero2)
    pyautogui.click(181, 272)
    pyautogui.click(hero3)
    pyautogui.click(435, 266)

    # speed up
    pyautogui.click(speed)
    # pyautogui.click(start) round
    pyautogui.click(start)
    time.sleep(2)

    # build towers
    for i in tpositions:
        a = i
        pyautogui.click(a)
        time.sleep(0.5)
        z = upgrade[0]
        pyautogui.click(a[0]+z[0],a[1]+z[1])
        time.sleep(0.5)
        pyautogui.click(a[0]+z[0],a[1]+z[1])
        time.sleep(0.15)
        pyautogui.click(start)

    # click for next round
    for aa in range(3):
        pyautogui.click(start)
        time.sleep(2)

    upgrade1st()

    # click for next round
    for aa in range(3):
        pyautogui.click(start)
        time.sleep(2)

    upgrade2nd()
    finalupgrade()
    restart(num)

endall = time.time()
print('Game Ended after {} mins'.format(round((endall-startall)/60,2)))
# print(pyautogui.position())