from typing import Iterable

import pygame
import math

from objects import Ball
from player import Player

SCREEN_X = 1920
SCREEN_Y = 1080
WALL_HEIGHT = 25000

TILE_SIZE = 60
QUALITY1 = TILE_SIZE / 16
QUALITY2 = TILE_SIZE / 256
FOV_MULTIP = 1

REACH = TILE_SIZE / 1

"""
class Player:
    rect: pygame.Rect

    def draw(self, screen: pygame.Surface):
        pygame.draw.rect(screen, (0, 0, 255), self.rect)

    def __init__(self, x, y):
        self.rect = pygame.Rect(0, 0, TILE_SIZE // 2, TILE_SIZE // 2)
        self.rect.center = (x, y)
"""

class GridList:
    tw: int = TILE_SIZE
    columns: int = 32
    rows: int = 0
    offset_x: int = SCREEN_X // 2
    grid: list

    def __init__(self):

        # MAP_1

        self.grid = [
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 3, 1, 3, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 3, 0, 0, 0, 3, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
            1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 3, 1, 1, 1, 1, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 3, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 2, 0, 0, 3, 0, 0, 1, 1, 1, 3, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1,
            1, 0, 0, 0, 3, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 3, 0, 0, 0, 9,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        ]
        self.rows = len(self.grid) // self.columns

    def get_spawn(self) -> tuple:
        index = 0
        for row in range(self.rows):
            for column in range(self.columns):
                if self.grid[index] == 2:
                    return column * self.tw + (self.tw // 2), row * self.tw + (self.tw // 2)
                index += 1

    def get_objects(self) -> Iterable[tuple]:
        index = 0
        for row in range(self.rows):
            for column in range(self.columns):
                if self.grid[index] == 4:
                    yield column * self.tw + (self.tw // 2), row * self.tw + (self.tw // 2)
                index += 1

    def get_tile(self, x, y):
        grid_x = int(math.floor(x) // TILE_SIZE)
        grid_y = int(math.floor(y) // TILE_SIZE)
        index = grid_x
        index += grid_y * SCREEN_X // TILE_SIZE
        return self.grid[index]

    def replace_tile(self, x, y, new_tile):
        grid_x = int(x // TILE_SIZE)
        grid_y = int(y // TILE_SIZE)
        index = grid_x
        index += grid_y * SCREEN_X // TILE_SIZE
        self.grid[index] = new_tile

    def draw(self, screen: pygame.Surface):
        index = 0
        for row in range(self.rows):
            y = row * self.tw
            for column in range(self.columns):
                x = column * self.tw

                if self.grid[index] == 1:
                    pygame.draw.rect(screen, (128, 128, 128), (x, y, self.tw, self.tw))

                elif self.grid[index] == 2:
                    pygame.draw.rect(screen, (0, 0, 0), (x, y, self.tw, self.tw))

                elif self.grid[index] == 3:
                    pygame.draw.rect(screen, (0, 192, 0), (x, y, self.tw, self.tw))

                elif self.grid[index] == 9:
                    pygame.draw.rect(screen, (255, 255, 0), (x, y, self.tw, self.tw))

                index += 1

def main():
    pygame.init()
    pygame.font.init()
    pygame.mixer.init()

    grid = GridList()
    screen = pygame.display.set_mode((SCREEN_X, SCREEN_Y), pygame.FULLSCREEN | pygame.HWACCEL | pygame.DOUBLEBUF)
    raster_screen = screen.copy()
    # raster_screen.set_alpha(50)
    map_screen = pygame.Surface((grid.columns * TILE_SIZE, grid.rows * TILE_SIZE))
    map_zone = pygame.Rect(0, 0, 1200, 1200)
    x_limit, y_limit = map_screen.get_size()
    px, py = grid.get_spawn()
    pdir = 0
    player = Player(px, py)

    pelf = 0

    objects = list()
    for x, y in grid.get_objects():
        obj = Ball(x, y)
        objects.append(obj)

    clock = pygame.time.Clock()
    pressed_keys = set()
    relmouse = 0, 0
    mousewheel = 0
    pygame.event.set_grab(True)
    pygame.mouse.set_visible(False)

    wall_red = 0
    wall_grn = 0
    wall_blu = 0

    running = True
    while running:
        frame_delay = clock.tick(60) / 1000
        running = not bool(pygame.event.get(pygame.QUIT))

        relmouse = 0, 0
        mousewheel = 0

        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                pressed_keys.add(event.key)
            elif event.type == pygame.KEYUP:
                pressed_keys.remove(event.key)
            if event.type == pygame.MOUSEMOTION:
                relmouse = event.rel
            if event.type == pygame.MOUSEWHEEL:
                mousewheel = event.y

        running = not pygame.K_BACKSPACE in pressed_keys and not pygame.K_ESCAPE in pressed_keys

        player.turn(relmouse[0], frame_delay)

        player.turn(
            15 if pygame.K_RIGHT in pressed_keys else
            -15 if pygame.K_LEFT in pressed_keys else
            0,
            frame_delay
        )

        player.change_turn_speed(mousewheel * 0.001)
        player.move(
            forward=pygame.K_w in pressed_keys,
            backward=pygame.K_s in pressed_keys,
            strafe_right=0,
            strafe_left=0
        )

        player.update(grid, frame_delay)

        if pygame.K_LSHIFT not in pressed_keys:
            player.move(
                forward=pygame.K_w in pressed_keys,
                backward=pygame.K_s in pressed_keys,
                strafe_right=0,
                strafe_left=0
            )

        player.update(grid, frame_delay)

        player.change_turn_speed(mousewheel * 0.001)
        player.move(
            forward=0,
            backward=0,
            strafe_right=pygame.K_d in pressed_keys,
            strafe_left=pygame.K_a in pressed_keys
        )

        player.update(grid, frame_delay)

        if pygame.K_LSHIFT not in pressed_keys:
            player.move(
                forward=0,
                backward=0,
                strafe_right=pygame.K_d in pressed_keys,
                strafe_left=pygame.K_a in pressed_keys
            )

            player.update(grid, frame_delay)

        current_x, current_y = player.get_location()
        pdir = player.get_dir()

        if pygame.K_e in pressed_keys:
            pelf = 1
            temp1 = current_x
            temp2 = current_y
            temp1 += math.cos(pdir) * REACH
            temp2 += math.sin(pdir) * REACH

            tile = grid.get_tile(current_x + math.cos(pdir) * REACH, current_y + math.sin(pdir) * REACH)

            if tile == 3:
                grid.replace_tile(temp1, temp2, 0)
            elif tile == 9:
                running = 0
                print("you win i guess")

        raster_screen.fill((0, 0, 0))
        line_x = -8
        angle: float = -0.60
        while angle <= 0.60:
            angle += 0.01
            line_x += 16
            rx = current_x
            ry = current_y
            dist = 0
            rdir: float = (angle * FOV_MULTIP + pdir)
            wall_red = 0
            ray_run = 1

            dx = math.cos(angle) * QUALITY1
            while ray_run == 1:
                rx += math.cos(rdir) * QUALITY1
                ry += math.sin(rdir) * QUALITY1
                dist += dx

                if (grid.get_tile(rx, ry)) == 1:
                    ray_run = 0
                elif (grid.get_tile(rx, ry)) == 3:
                    ray_run = 0
                elif (grid.get_tile(rx, ry)) == 9:
                    ray_run = 0

            ray_run = 1

            dx = math.cos(angle) * QUALITY2
            tile = (grid.get_tile(rx, ry))
            if tile == 1 or tile == 3 or tile == 9:
                while ray_run == 1:
                    rx -= math.cos(rdir) * QUALITY2
                    ry -= math.sin(rdir) * QUALITY2
                    dist -= dx

                    tile = (grid.get_tile(rx, ry))
                    if tile != 1 and tile != 3 and tile != 9:
                        ray_run = 0
                    else:
                        ray_run = 1

            rx += math.cos(rdir) * QUALITY2
            ry += math.sin(rdir) * QUALITY2
            dist += QUALITY2

            if (grid.get_tile(rx, ry)) == 1 or 3 or 9:
                if dist <= 0:
                    dist = 1
                center_y = SCREEN_Y / 2
                wall_height = (1 / dist) * WALL_HEIGHT

                if (grid.get_tile(rx, ry)) == 1:
                    wall_red = 0
                    wall_grn = 0
                    wall_blu = 255
                elif (grid.get_tile(rx, ry)) == 3:
                    wall_red = 0
                    wall_grn = 192
                    wall_blu = 0
                elif (grid.get_tile(rx, ry)) == 9:
                    wall_red = 255
                    wall_grn = 255
                    wall_blu = 0

                pygame.draw.line(
                    raster_screen,
                    (wall_red, wall_grn, wall_blu),
                    (line_x, center_y - wall_height),
                    (line_x, center_y + wall_height),
                    width=16
                )

        for obj in objects:
            obj.draw(raster_screen, current_x, current_y, 0 - pdir)

        map_screen.fill((0, 0, 0))
        screen.fill((0, 0, 0))
        grid.draw(map_screen)
        player.draw(map_screen)

        map_zone.center = player.get_location()
        if map_zone.left < 0:
            map_zone.left = 0
        if map_zone.top < 0:
            map_zone.top = 0
        if map_zone.right > x_limit:
            map_zone.right = x_limit
        if map_zone.bottom > y_limit:
            map_zone.bottom = y_limit

        if pygame.K_m in pressed_keys:
            _map_screen = pygame.transform.scale(map_screen.subsurface(map_zone), (600, 600))
        else:
            _map_screen = pygame.transform.scale(map_screen.subsurface(map_zone), (300, 300))

        screen.blit(raster_screen, (0, 0))
        screen.blit(_map_screen, (0, 0))
        pygame.display.flip()

if __name__ == '__main__':
    main()


import re

import nltk
from nltk import word_tokenize, TreebankWordDetokenizer
from nltk.corpus import stopwords

#nltk.download()

def extract_line_contains(data: str) -> str:
    return ""

def remove_unwanted_characters(text: str) -> str:
    cleanTxt = text.replace("_", "")
    cleanTxt = cleanTxt.replace("-", "")
    cleanTxt = cleanTxt.replace(".", "")
    cleanTxt = cleanTxt.replace(";", "")

    return cleanTxt

def process_text(text: str) -> str:
    cleaned_text = remove_unwanted_characters(text)
    filtered_text = remove_stopwords(cleaned_text, 'french')
    #correct_spell(cleaned_text)

    return filtered_text

def remove_stopwords(text: str, language: str):
    stop_words = set(stopwords.words(language))

    #regex will search for all words including \n which is needed to preserve the new lines in the detokenizer in line 36
    word_tokens = re.findall(r'\S+|\n', text)
    filtered_words =[]
    for word in word_tokens:
        if word not in stop_words or word.__eq__('\n'):
            filtered_words.append(word)

    filtered_text = TreebankWordDetokenizer().detokenize(filtered_words)
    return filtered_text

# def check_spell(text: str):
#     words = word_tokenize(text)
#     return misspelled

def correct_spell(text: str):
    return text


# Find sinh value of PI/2

# NumPy provides the ufuncs sinh(), cosh() and tanh() that take values in radians and produce the corresponding sinh, cosh and tanh values..
import numpy as np

x = np.sinh(np.pi/2)
print(x)

# Find cosh values for all of the values in arr

arr = np.array([np.pi/2, np.pi/3, np.pi/4, np.pi/5])
x = np.cosh(arr)
print(x)

# Find the angle of 1.0:

x = np.arcsinh(1.0) #vFinding Angles
print(x)

# Find the angle for all of the tanh values in array:

arr = np.array([0.1, 0.2, 0.5]) # Angles of Each Value in Arrays
x = np.arctanh(arr)
print(x)


from santaspkg.cost_function import reference_cost_function
from santaspkg.dataset import sample_submission

from pathlib import Path

submission_template = sample_submission.copy(deep=True)

submissions_dir = Path(__file__).parent.parent / 'submissions'
submissions_dir.mkdir(exist_ok=True)

def mk_submit(assignment):
    score = reference_cost_function(assignment)
    filename = str(submissions_dir / f'submission_{score}.csv')
    submission_template['assigned_day'] = assignment
    submission_template.to_csv(filename)


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

"""
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

#basico para api
import gzip
import os
import struct
import urllib2
import xmlrpclib

from fontesLegendas import FontesBase

class OpenSubtitles(FontesBase):

    _linkDownload = ""
    nomeLegenda = ""

    def getNomeFonte(self):
        return "OpenSubtitles"

    def __get_hash(self,arquivo):
            longlongformat = '<q'  # little-endian long long
            bytesize = struct.calcsize(longlongformat)

            f = open(arquivo, "rb")

            filesize = os.path.getsize(arquivo)
            hash = filesize

            if filesize < 65536 * 2:
                return "SizeError"

            for x in range(65536 / bytesize):
                buffer = f.read(bytesize)
                (l_value,) = struct.unpack(longlongformat, buffer)
                hash += l_value
                hash = hash & 0xFFFFFFFFFFFFFFFF  # to remain as 64bit number

            f.seek(max(0, filesize - 65536), 0)
            for x in range(65536 / bytesize):
                buffer = f.read(bytesize)
                (l_value,) = struct.unpack(longlongformat, buffer)
                hash += l_value
                hash = hash & 0xFFFFFFFFFFFFFFFF

            f.close()
            returnedhash = "%016x" % hash
            return returnedhash

    def procuraLegenda(self,arquivo):

        volta = False
        videoTamanho = os.path.getsize(arquivo)
        videoHash = self.__get_hash(arquivo)

        opensubtitles = xmlrpclib.ServerProxy("http://api.opensubtitles.org/xml-rpc")

        saidaLogin = opensubtitles.LogIn("","","pob","OSTestUserAgent")
        token = saidaLogin['token']

        stringProcura = []
        stringProcura.append({'sublanguageid':'pob', 'moviehash':videoHash, 'moviebytesize':str(videoTamanho)})

        listaResultado = opensubtitles.SearchSubtitles(token,stringProcura)
        if listaResultado['data']:
            for i in range(len(listaResultado['data'])):
                nomeLegenda =  listaResultado['data'][i]['SubFileName']
                nomeLegenda = nomeLegenda[0:len(nomeLegenda)-4]
                arquivo = arquivo[0:len(arquivo)-4]
                if nomeLegenda in arquivo:
                    linkDownload = listaResultado['data'][i]['SubDownloadLink']
                    self._linkDownload = linkDownload
                    __videoHash = videoHash
                    volta = True

        opensubtitles.LogOut(token)
        return volta

    def downloadLegenda(self, diretorio, arquivo):
        request = urllib2.Request(self._linkDownload)
        response = urllib2.urlopen(request)

        nomeLegenda = os.path.join(diretorio,arquivo)
        nomeLegenda = nomeLegenda[0:len(nomeLegenda)-4]

        with open("/tmp/algo.gz", "wb") as local_file:
            local_file.write(response.read())

        f = gzip.open("/tmp/algo.gz", "rb")
        file_content = f.read()
        f.close()
        os.remove("/tmp/algo.gz")

        legenda = open(nomeLegenda + ".pt-BR.srt", "w")
        legenda.write(file_content)
        self.nomeLegenda = nomeLegenda + ".pt-BR.srt"

    def getNomeLegenda(self):
        return self.nomeLegenda


import textwrap

from assistance._config import SIMPLER_OPENAI_MODEL
from assistance._keys import get_openai_api_key
from assistance._summarisation.thread import run_with_summary_fallback

OPEN_AI_API_KEY = get_openai_api_key()

MODEL_KWARGS = {
    "engine": SIMPLER_OPENAI_MODEL,
    "max_tokens": 256,
    "temperature": 0.7,
    "top_p": 1,
    "frequency_penalty": 0,
    "presence_penalty": 0,
}

PROMPT = textwrap.dedent(
    """
        # Get Correspondent

        Alex Carpenter is having an email conversation with a
        prospective student. The transcript of the email conversation
        is below. Please extract the first name of the student from the
        email transcript.

        If you are unable to extract the first name, please instead
        just write [Name not found in email transcript].

        ## Email transcript

        {transcript}

        ## Email address of the student

        {email_address}

        ## Response

        First Name:
    """
).strip()

async def get_first_name(scope: str, email_thread, their_email_address: str) -> str:
    prompt = PROMPT.replace("{email_address}", their_email_address)
    response, _ = await run_with_summary_fallback(
        scope=scope,
        prompt=prompt,
        email_thread=email_thread,
        api_key=OPEN_AI_API_KEY,
        **MODEL_KWARGS,
    )

    return response


'''
tree sum

Write a function, tree_sum, that takes in the root of a binary tree that contains number values. The function should return the total sum of all values in the tree.
test_00:

a = Node(3)
b = Node(11)
c = Node(4)
d = Node(4)
e = Node(-2)
f = Node(1)

a.left = b
a.right = c
b.left = d
b.right = e
c.right = f

#       3
#    /    \
#   11     4
#  / \      \
# 4   -2     1

tree_sum(a) # -> 21

test_01:

a = Node(1)
b = Node(6)
c = Node(0)
d = Node(3)
e = Node(-6)
f = Node(2)
g = Node(2)
h = Node(2)

a.left = b
a.right = c
b.left = d
b.right = e
c.right = f
e.left = g
f.right = h

#      1
#    /   \
#   6     0
#  / \     \
# 3   -6    2
#    /       \
#   2         2

tree_sum(a) # -> 10

test_02:

tree_sum(None) # -> 0
'''

# class Node:
#   def __init__(self, val):
#     self.val = val
#     self.left = None
#     self.right = None

# DFS recursive
def tree_sum(root):
  if root is None:
    return 0
  return root.val + tree_sum(root.left) + tree_sum(root.right)

# BFS
from collections import deque

def tree_sum(root):
  if root is None:
    return 0

  queue = deque([root])
  sum = 0
  while queue:
    current = queue.popleft()
    sum += current.val

    if current.left:
      queue.append(current.left)
    if current.right:
      queue.append(current.right)

  return sum


t=int(input())

'''
s[i]=chr(((ord(s[i]) - 65 + 26 - 13) % 26) + 65)

핵심은 A의 ord를 0으로 만든 다음에 이제 더해서 나머지!!
그걸로 해가지고 가져오는거.
'''

for i in range(t):
    a,b=map(int,input().split())
    s=list(input())
    for i in range(len(s)):
        s[i]=chr(((ord(s[i])-65)*a+b)%26+65)
    print(''.join(s))


import operator
import time
from typing import Dict, List, Optional

from acme import core
from acme.tf import losses
from acme.tf import networks
from acme.tf import savers as tf2_savers
from acme.tf import utils as tf2_utils
from acme.utils import counting
from acme.utils import loggers
import numpy as np
import reverb
import sonnet as snt
import tensorflow as tf
import tree

class RCRRLearner(core.Learner):
  """Recurrent CRR learner.

  This is the learning component of a RCRR agent. It takes a dataset as
  input and implements update functionality to learn from this dataset.
  """

  def __init__(self,
               policy_network: snt.RNNCore,
               critic_network: networks.CriticDeepRNN,
               target_policy_network: snt.RNNCore,
               target_critic_network: networks.CriticDeepRNN,
               dataset: tf.data.Dataset,
               accelerator_strategy: Optional[tf.distribute.Strategy] = None,
               behavior_network: Optional[snt.Module] = None,
               cwp_network: Optional[snt.Module] = None,
               policy_optimizer: Optional[snt.Optimizer] = None,
               critic_optimizer: Optional[snt.Optimizer] = None,
               discount: float = 0.99,
               target_update_period: int = 100,
               num_action_samples_td_learning: int = 1,
               num_action_samples_policy_weight: int = 4,
               baseline_reduce_function: str = 'mean',
               clipping: bool = True,
               policy_improvement_modes: str = 'exp',
               ratio_upper_bound: float = 20.,
               beta: float = 1.0,
               counter: Optional[counting.Counter] = None,
               logger: Optional[loggers.Logger] = None,
               checkpoint: bool = False):
    """Initializes the learner.

    Args:
      policy_network: the online (optimized) policy.
      critic_network: the online critic.
      target_policy_network: the target policy (which lags behind the online
        policy).
      target_critic_network: the target critic.
      dataset: dataset to learn from, whether fixed or from a replay buffer
        (see `acme.datasets.reverb.make_reverb_dataset` documentation).
      accelerator_strategy: the strategy used to distribute computation,
        whether on a single, or multiple, GPU or TPU; as supported by
        tf.distribute.
      behavior_network: The network to snapshot under `policy` name. If None,
        snapshots `policy_network` instead.
      cwp_network: CWP network to snapshot: samples actions
        from the policy and weighs them with the critic, then returns the action
        by sampling from the softmax distribution using critic values as logits.
        Used only for snapshotting, not training.
      policy_optimizer: the optimizer to be applied to the policy loss.
      critic_optimizer: the optimizer to be applied to the distributional
        Bellman loss.
      discount: discount to use for TD updates.
      target_update_period: number of learner steps to perform before updating
        the target networks.
      num_action_samples_td_learning: number of action samples to use to
        estimate expected value of the critic loss w.r.t. stochastic policy.
      num_action_samples_policy_weight: number of action samples to use to
        estimate the advantage function for the CRR weighting of the policy
        loss.
      baseline_reduce_function: one of 'mean', 'max', 'min'. Way of aggregating
        values from `num_action_samples` estimates of the value function.
      clipping: whether to clip gradients by global norm.
      policy_improvement_modes: one of 'exp', 'binary', 'all'. CRR mode which
        determines how the advantage function is processed before being
        multiplied by the policy loss.
      ratio_upper_bound: if policy_improvement_modes is 'exp', determines
        the upper bound of the weight (i.e. the weight is
          min(exp(advantage / beta), upper_bound)
        ).
      beta: if policy_improvement_modes is 'exp', determines the beta (see
        above).
      counter: counter object used to keep track of steps.
      logger: logger object to be used by learner.
      checkpoint: boolean indicating whether to checkpoint the learner.
    """

    if accelerator_strategy is None:
      accelerator_strategy = snt.distribute.Replicator()
    self._accelerator_strategy = accelerator_strategy
    self._policy_improvement_modes = policy_improvement_modes
    self._ratio_upper_bound = ratio_upper_bound
    self._num_action_samples_td_learning = num_action_samples_td_learning
    self._num_action_samples_policy_weight = num_action_samples_policy_weight
    self._baseline_reduce_function = baseline_reduce_function
    self._beta = beta

    # When running on TPUs we have to know the amount of memory required (and
    # thus the sequence length) at the graph compilation stage. At the moment,
    # the only way to get it is to sample from the dataset, since the dataset
    # does not have any metadata, see b/160672927 to track this upcoming
    # feature.
    sample = next(dataset.as_numpy_iterator())
    self._sequence_length = sample.action.shape[1]

    self._counter = counter or counting.Counter()
    self._logger = logger or loggers.TerminalLogger('learner', time_delta=1.)
    self._discount = discount
    self._clipping = clipping

    self._target_update_period = target_update_period

    with self._accelerator_strategy.scope():
      # Necessary to track when to update target networks.
      self._num_steps = tf.Variable(0, dtype=tf.int32)

      # (Maybe) distributing the dataset across multiple accelerators.
      distributed_dataset = self._accelerator_strategy.experimental_distribute_dataset(
          dataset)
      self._iterator = iter(distributed_dataset)

      # Create the optimizers.
      self._critic_optimizer = critic_optimizer or snt.optimizers.Adam(1e-4)
      self._policy_optimizer = policy_optimizer or snt.optimizers.Adam(1e-4)

    # Store online and target networks.
    self._policy_network = policy_network
    self._critic_network = critic_network
    self._target_policy_network = target_policy_network
    self._target_critic_network = target_critic_network

    # Expose the variables.
    self._variables = {
        'critic': self._target_critic_network.variables,
        'policy': self._target_policy_network.variables,
    }

    # Create a checkpointer object.
    self._checkpointer = None
    self._snapshotter = None
    if checkpoint:
      self._checkpointer = tf2_savers.Checkpointer(
          objects_to_save={
              'counter': self._counter,
              'policy': self._policy_network,
              'critic': self._critic_network,
              'target_policy': self._target_policy_network,
              'target_critic': self._target_critic_network,
              'policy_optimizer': self._policy_optimizer,
              'critic_optimizer': self._critic_optimizer,
              'num_steps': self._num_steps,
          },
          time_delta_minutes=30.)

      raw_policy = snt.DeepRNN(
          [policy_network, networks.StochasticSamplingHead()])
      critic_mean = networks.CriticDeepRNN(
          [critic_network, networks.StochasticMeanHead()])
      objects_to_save = {
          'raw_policy': raw_policy,
          'critic': critic_mean,
      }
      if behavior_network is not None:
        objects_to_save['policy'] = behavior_network
      if cwp_network is not None:
        objects_to_save['cwp_policy'] = cwp_network
      self._snapshotter = tf2_savers.Snapshotter(
          objects_to_save=objects_to_save, time_delta_minutes=30)
    # Timestamp to keep track of the wall time.
    self._walltime_timestamp = time.time()

  def _step(self, sample: reverb.ReplaySample) -> Dict[str, tf.Tensor]:
    # Transpose batch and sequence axes, i.e. [B, T, ...] to [T, B, ...].
    sample = tf2_utils.batch_to_sequence(sample)
    observations = sample.observation
    actions = sample.action
    rewards = sample.reward
    discounts = sample.discount

    dtype = rewards.dtype

    # Cast the additional discount to match the environment discount dtype.
    discount = tf.cast(self._discount, dtype=discounts.dtype)

    # Loss cumulants across time. These cannot be python mutable objects.
    critic_loss = 0.
    policy_loss = 0.

    # Each transition induces a policy loss, which we then weight using
    # the `policy_loss_coef_t`; shape [B], see https://arxiv.org/abs/2006.15134.
    # `policy_loss_coef` is a scalar average of these coefficients across
    # the batch and sequence length dimensions.
    policy_loss_coef = 0.

    per_device_batch_size = actions.shape[1]

    # Initialize recurrent states.
    critic_state = self._critic_network.initial_state(per_device_batch_size)
    target_critic_state = critic_state
    policy_state = self._policy_network.initial_state(per_device_batch_size)
    target_policy_state = policy_state

    with tf.GradientTape(persistent=True) as tape:
      for t in range(1, self._sequence_length):
        o_tm1 = tree.map_structure(operator.itemgetter(t - 1), observations)
        a_tm1 = tree.map_structure(operator.itemgetter(t - 1), actions)
        r_t = tree.map_structure(operator.itemgetter(t - 1), rewards)
        d_t = tree.map_structure(operator.itemgetter(t - 1), discounts)
        o_t = tree.map_structure(operator.itemgetter(t), observations)

        if t != 1:
          # By only updating the target critic state here we are forcing
          # the target critic to ignore observations[0]. Otherwise, the
          # target_critic will be unrolled for one more timestep than critic.
          # The smaller the sequence length, the more problematic this is: if
          # you use RNN on sequences of length 2, you would expect the code to
          # never use recurrent connections. But if you don't skip updating the
          # target_critic_state on observation[0] here, it won't be the case.
          _, target_critic_state = self._target_critic_network(
              o_tm1, a_tm1, target_critic_state)

        # ========================= Critic learning ============================
        q_tm1, next_critic_state = self._critic_network(o_tm1, a_tm1,
                                                        critic_state)
        target_action_distribution, target_policy_state = self._target_policy_network(
            o_t, target_policy_state)

        sampled_actions_t = target_action_distribution.sample(
            self._num_action_samples_td_learning)
        # [N, B, ...]
        tiled_o_t = tf2_utils.tile_nested(
            o_t, self._num_action_samples_td_learning)
        tiled_target_critic_state = tf2_utils.tile_nested(
            target_critic_state, self._num_action_samples_td_learning)

        # Compute the target critic's Q-value of the sampled actions.
        sampled_q_t, _ = snt.BatchApply(self._target_critic_network)(
            tiled_o_t, sampled_actions_t, tiled_target_critic_state)

        # Compute average logits by first reshaping them to [N, B, A] and then
        # normalizing them across atoms.
        new_shape = [self._num_action_samples_td_learning, r_t.shape[0], -1]
        sampled_logits = tf.reshape(sampled_q_t.logits, new_shape)
        sampled_logprobs = tf.math.log_softmax(sampled_logits, axis=-1)
        averaged_logits = tf.reduce_logsumexp(sampled_logprobs, axis=0)

        # Construct the expected distributional value for bootstrapping.
        q_t = networks.DiscreteValuedDistribution(
            values=sampled_q_t.values, logits=averaged_logits)
        critic_loss_t = losses.categorical(q_tm1, r_t, discount * d_t, q_t)
        critic_loss_t = tf.reduce_mean(critic_loss_t)

        # ========================= Actor learning =============================
        action_distribution_tm1, policy_state = self._policy_network(
            o_tm1, policy_state)
        q_tm1_mean = q_tm1.mean()

        # Compute the estimate of the value function based on
        # self._num_action_samples_policy_weight samples from the policy.
        tiled_o_tm1 = tf2_utils.tile_nested(
            o_tm1, self._num_action_samples_policy_weight)
        tiled_critic_state = tf2_utils.tile_nested(
            critic_state, self._num_action_samples_policy_weight)
        action_tm1 = action_distribution_tm1.sample(
            self._num_action_samples_policy_weight)
        tiled_z_tm1, _ = snt.BatchApply(self._critic_network)(
            tiled_o_tm1, action_tm1, tiled_critic_state)
        tiled_v_tm1 = tf.reshape(tiled_z_tm1.mean(),
                                 [self._num_action_samples_policy_weight, -1])

        # Use mean, min, or max to aggregate Q(s, a_i), a_i ~ pi(s) into the
        # final estimate of the value function.
        if self._baseline_reduce_function == 'mean':
          v_tm1_estimate = tf.reduce_mean(tiled_v_tm1, axis=0)
        elif self._baseline_reduce_function == 'max':
          v_tm1_estimate = tf.reduce_max(tiled_v_tm1, axis=0)
        elif self._baseline_reduce_function == 'min':
          v_tm1_estimate = tf.reduce_min(tiled_v_tm1, axis=0)

        # Assert that action_distribution_tm1 is a batch of multivariate
        # distributions (in contrast to e.g. a [batch, action_size] collection
        # of 1d distributions).
        assert len(action_distribution_tm1.batch_shape) == 1
        policy_loss_batch = -action_distribution_tm1.log_prob(a_tm1)

        advantage = q_tm1_mean - v_tm1_estimate
        if self._policy_improvement_modes == 'exp':
          policy_loss_coef_t = tf.math.minimum(
              tf.math.exp(advantage / self._beta), self._ratio_upper_bound)
        elif self._policy_improvement_modes == 'binary':
          policy_loss_coef_t = tf.cast(advantage > 0, dtype=dtype)
        elif self._policy_improvement_modes == 'all':
          # Regress against all actions (effectively pure BC).
          policy_loss_coef_t = 1.
        policy_loss_coef_t = tf.stop_gradient(policy_loss_coef_t)

        policy_loss_batch *= policy_loss_coef_t
        policy_loss_t = tf.reduce_mean(policy_loss_batch)

        critic_state = next_critic_state

        critic_loss += critic_loss_t
        policy_loss += policy_loss_t
        policy_loss_coef += tf.reduce_mean(policy_loss_coef_t)  # For logging.

      # Divide by sequence length to get mean losses.
      critic_loss /= tf.cast(self._sequence_length, dtype=dtype)
      policy_loss /= tf.cast(self._sequence_length, dtype=dtype)
      policy_loss_coef /= tf.cast(self._sequence_length, dtype=dtype)

    # Compute gradients.
    critic_gradients = tape.gradient(critic_loss,
                                     self._critic_network.trainable_variables)
    policy_gradients = tape.gradient(policy_loss,
                                     self._policy_network.trainable_variables)

    # Delete the tape manually because of the persistent=True flag.
    del tape

    # Sync gradients across GPUs or TPUs.
    ctx = tf.distribute.get_replica_context()
    critic_gradients = ctx.all_reduce('mean', critic_gradients)
    policy_gradients = ctx.all_reduce('mean', policy_gradients)

    # Maybe clip gradients.
    if self._clipping:
      policy_gradients = tf.clip_by_global_norm(policy_gradients, 40.)[0]
      critic_gradients = tf.clip_by_global_norm(critic_gradients, 40.)[0]

    # Apply gradients.
    self._critic_optimizer.apply(critic_gradients,
                                 self._critic_network.trainable_variables)
    self._policy_optimizer.apply(policy_gradients,
                                 self._policy_network.trainable_variables)

    source_variables = (
        self._critic_network.variables + self._policy_network.variables)
    target_variables = (
        self._target_critic_network.variables +
        self._target_policy_network.variables)

    # Make online -> target network update ops.
    if tf.math.mod(self._num_steps, self._target_update_period) == 0:
      for src, dest in zip(source_variables, target_variables):
        dest.assign(src)
    self._num_steps.assign_add(1)

    return {
        'critic_loss': critic_loss,
        'policy_loss': policy_loss,
        'policy_loss_coef': policy_loss_coef,
    }

  @tf.function
  def _replicated_step(self) -> Dict[str, tf.Tensor]:
    sample = next(self._iterator)
    fetches = self._accelerator_strategy.run(self._step, args=(sample,))
    mean = tf.distribute.ReduceOp.MEAN
    return {
        k: self._accelerator_strategy.reduce(mean, fetches[k], axis=None)
        for k in fetches
    }

  def step(self):
    # Run the learning step.
    with self._accelerator_strategy.scope():
      fetches = self._replicated_step()

    # Update our counts and record it.
    new_timestamp = time.time()
    time_passed = new_timestamp - self._walltime_timestamp
    self._walltime_timestamp = new_timestamp
    counts = self._counter.increment(steps=1, wall_time=time_passed)
    fetches.update(counts)

    # Checkpoint and attempt to write the logs.
    if self._checkpointer is not None:
      self._checkpointer.save()
      self._snapshotter.save()
    self._logger.write(fetches)

  def get_variables(self, names: List[str]) -> List[List[np.ndarray]]:
    return [tf2_utils.to_numpy(self._variables[name]) for name in names]


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
lines_bars_and_markers/using_span_where.py

Matplotlib > Gallery > Lines, bars and markers > Using span_where
https://matplotlib.org/gallery/lines_bars_and_markers/span_regions.html#using-span-where
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections as collections

t  = np.arange(0.0, 2, 0.01)
s1 = np.sin(2*np.pi*t)
# This line exists in the example, but it's not used.
#s2 = 1.2*np.sin(4*np.pi*t)

fig, ax = plt.subplots()
ax.set_title('using span_where')
ax.plot(t, s1, color='black')       # Plot s1
ax.axhline(0, color='black', lw=2)  # Draw a horizontal line for y=0

# Green shady area is applied where s1>0
collection = collections.BrokenBarHCollection.span_where(
    t, ymin=0, ymax=1, where=s1 > 0, facecolor='green', alpha=0.5)
ax.add_collection(collection)

# Green shady area is applied where s1<0
collection = collections.BrokenBarHCollection.span_where(
    t, ymin=-1, ymax=0, where=s1 < 0, facecolor='red', alpha=0.5)
ax.add_collection(collection)

plt.show()


import numpy as np
# https://github.com/cgohlke/transformations/blob/master/transformations/transformations.py
def quaternion_about_axis(angle, axis, _EPS=1e-6):
    """Return quaternion for rotation about axis.

    >>> q = quaternion_about_axis(0.123, [1, 0, 0])
    >>> numpy.allclose(q, [0.99810947, 0.06146124, 0, 0])
    True

    """
    axis = np.asarray(axis)
    q = np.zeros((*axis.shape[:-1], axis.shape[-1]+1))
    q[..., 1:] = axis
    qlen = np.linalg.norm(q, axis=-1)
    if qlen > _EPS:
        q *= np.sin(angle / 2.0) / qlen
    q[..., 0] = np.cos(angle / 2.0)
    return q

def axangle_from_quat(quat):
    """Return axis-angle for rotation from quaternion

    >>> q = quaternion_about_axis(0.123, [1, 0, 0])
    >>> axis, angle = axangle_from_quat(q)
    >>> numpy.allclose(axis, [1, 0, 0])
    True
    >>> numpy.allclose(angle, 0.123)
    True

    """
    quat /= np.linalg.norm(quat, axis=-1)
    sin_theta_by_2 = np.linalg.norm(quat[..., 1:4], axis=-1)
    cos_theta_by_2 = quat[..., 0]
    theta = 2 * np.arctan2(sin_theta_by_2, cos_theta_by_2)
    axis = quat[..., 1:5] / sin_theta_by_2
    return axis, theta

def rotmat(theta):
    return np.array([[np.cos(theta), -np.sin(theta)],
                     [np.sin(theta), np.cos(theta)]])


import pytest
from sqlite_utils.utils import column_affinity

EXAMPLES = [
    # Examples from https://www.sqlite.org/datatype3.html#affinity_name_examples
    ("INT", int),
    ("INTEGER", int),
    ("TINYINT", int),
    ("SMALLINT", int),
    ("MEDIUMINT", int),
    ("BIGINT", int),
    ("UNSIGNED BIG INT", int),
    ("INT2", int),
    ("INT8", int),
    ("CHARACTER(20)", str),
    ("VARCHAR(255)", str),
    ("VARYING CHARACTER(255)", str),
    ("NCHAR(55)", str),
    ("NATIVE CHARACTER(70)", str),
    ("NVARCHAR(100)", str),
    ("TEXT", str),
    ("CLOB", str),
    ("BLOB", bytes),
    ("REAL", float),
    ("DOUBLE", float),
    ("DOUBLE PRECISION", float),
    ("FLOAT", float),
    # Numeric, treated as float:
    ("NUMERIC", float),
    ("DECIMAL(10,5)", float),
    ("BOOLEAN", float),
    ("DATE", float),
    ("DATETIME", float),
]

@pytest.mark.parametrize("column_def,expected_type", EXAMPLES)
def test_column_affinity(column_def, expected_type):
    assert expected_type is column_affinity(column_def)

@pytest.mark.parametrize("column_def,expected_type", EXAMPLES)
def test_columns_dict(fresh_db, column_def, expected_type):
    fresh_db.execute("create table foo (col {})".format(column_def))
    assert {"col": expected_type} == fresh_db["foo"].columns_dict


import open3d as op3d, numpy as np, logging

from scipy.io import loadmat
from scipy.spatial import KDTree
from typing import Union, Tuple, List
from pathlib import Path
from frontalization.mesh import Mesh
from utils_mesh import MeshOperator, create_sphere_mesh

path_bfm_model = Path('../BFM/01_MorphableModel.mat')
path_additional_bfm = Path('../BFM/facemodel_info.mat')
path_landmarks_bfm_indeces = Path('../BFM/landmarks68_BFM.anl')
path_bfm_prepared = Path('../BFM/BFM_model_front.mat')

# https://ibug.doc.ic.ac.uk/resources/300-W/
class LandmarksAnnotator():

    def get_left_eye_landmarks_indeces(self):
        return [37, 38, 39, 40, 41, 42]

    def get_right_eye_landmarks_indeces(self):
        return [43, 44, 45, 46, 47, 48]

    def get_nose_vertical_landmarks(self):
        return [28, 29, 30, 31]

    def get_nose_gorizontal_landmarks(self):
        return [32, 33, 34, 35, 36]

    def get_upper_lip_landmarks(self):
        return [49, 50, 51, 52, 53, 54, 55, 61, 62, 63, 64, 65]

    def get_bottom_lip_landmarks(self):
        return [56, 57, 58, 59, 60, 66, 67, 68]

    def get_left_border_of_the_face(self):
        return [1, 2, 3, 4, 5, 6, 7, 8, 9]

    def get_right_border_of_the_face(self):
        return [10, 11, 12, 13, 14, 15, 16, 17]

class BaselFaceModel2009():

    _default_pathes = {"path_bfm_model" : Path('../BFM/01_MorphableModel.mat'),
                       "path_additional_bfm": Path('../BFM/facemodel_info.mat'),
                       "path_landmarks_bfm_indeces" : Path('../BFM/landmarks68_BFM.anl'),
                       "path_bfm_prepared" : Path('../BFM/BFM_model_front.mat')}

    def __init__(self, path_bfm_model, path_bfm_triangulation=None, path_bfm_prepared_model=None,
                 path_landmark_indeces=None, scale_basis:bool=True, create_mean_face_mesh:bool=False):
        self._path_bfm_model = self.check_path(path_bfm_model)
        assert self._path_bfm_model.suffix == '.mat', 'Incorrect format BFM 2009!'
        if path_bfm_triangulation:
            self._path_bfm_additional_info = self.check_path(path_bfm_triangulation)
        if path_landmark_indeces:
            self._path_landmark_label = self.check_path(path_landmark_indeces)
        self._num_sh_basis = 199
        self._num_tex_basis = 199

        # load model
        if path_bfm_prepared_model:
            path_bfm_prepared_model = self.check_path(path_bfm_prepared_model)
            self._init_bfm_model(path_bfm_prepared_model, scale_basis=scale_basis)
        else:
            # without this filepath
            self._init_bfm_model()

        if create_mean_face_mesh:
            self._path_mean_face_shape_obj = self._save_mean_face_mesh_as_obj()

    def check_path(self, path, type='file'):
        if isinstance(path, str):
            path = Path(path)
        if type=='file':
            assert path.is_file(), f'The file {path} not found!'
        else:
            assert path.is_dir(), f'The directory {path} not found!'
        return path

    @property
    def shape_basis(self):
        return self._shape_basis

    def _init_bfm_model(self, path_bfm_prepare_model=None, scale_basis:bool=True):
        self._bfm_model = loadmat(str(self._path_bfm_model))
        self._shape_basis = self._bfm_model['shapePC'].astype(np.float64).reshape((-1, 3, self._num_sh_basis))/1e5  # shape basis
        self._shape_sigmas = self._bfm_model['shapeEV'].astype(np.float64).reshape((-1, self._num_sh_basis))  # corresponding eigen values
        self._texture_basis = self._bfm_model['texPC'].astype(np.float64).reshape((-1, 3, self._num_sh_basis))  # texture basis
        self._texture_sigmas = self._bfm_model['texEV'].astype(np.float64).reshape((-1, self._num_sh_basis))  # eigen value
        if scale_basis:
            self._shape_basis = self._shape_sigmas * self._shape_basis.reshape(-1, self._num_sh_basis)
            self._shape_basis = self.shape_basis.reshape((-1, 3, self._num_sh_basis))
            self._texture_basis = self._texture_sigmas * self._texture_basis.reshape(-1, self._num_sh_basis)
            self._texture_basis = self._texture_basis.reshape((-1, 3, self._num_sh_basis))

        self._tex_mean = self._bfm_model['texMU'].astype(np.float64).reshape((-1, 3))  # mean texture
        self._shape_mean = self._bfm_model['shapeMU'].astype(np.float64).reshape((-1, 3))/1e5 # mean face
        if path_bfm_prepare_model:
            if isinstance(path_bfm_prepare_model, str):
                path_bfm_prepare_model = Path(path_bfm_prepare_model)
            assert path_bfm_prepare_model.is_file(), f'The file {path_bfm_prepare_model} not found!'
            bfm_prepared = loadmat(str(path_bfm_prepare_model)) # load triangulation
            self._triangles = bfm_prepared['tri'].astype(np.int64)-1
            logging.info(f'Extract triangles from prepared BFM model: {path_bfm_prepare_model}!')
        else:
            #self._triangles = bfm_additional_info['tri'].astype(np.int64).reshape((-1, 3))-1
            self._triangles = self._bfm_model['tl'].astype(np.int64) - 1
            self._triangles = self._triangles[:, [1, 0, 2]]

        if hasattr(self, '_path_landmark_label'):
            keypoints_indeces = []
            with self._path_landmark_label.open() as reader:
                for idx_line, line in enumerate(reader):
                    if idx_line == 0:
                        continue
                    index_keypts = int(line.rstrip(' \n'))
                    keypoints_indeces.append(index_keypts)
            self._keypts_indeces = keypoints_indeces
        logging.info(f'The basel face model is load!')

    def get_face_colors(self, tex_coefs:np.ndarray, normalize=True):
        num_basis = tex_coefs.shape[0]
        basis = self._texture_basis.reshape(-1, self._texture_basis.shape[-1])
        basis = basis[:, :num_basis]
        face_colors = basis @ tex_coefs[:, None] + self.mean_texture.reshape(-1, 1)
        if normalize:
            face_colors = face_colors / 255.
        return face_colors.reshape(-1, 3)

    def get_face_shape(self, id_coeff:np.ndarray):
        num_basis = id_coeff.shape[0]
        basis = self._shape_basis.reshape(-1, self._shape_basis.shape[-1])
        basis = basis[:, :num_basis]
        face_shape = basis @ id_coeff[:, None] + self.mean_shape.reshape(-1, 1)
        return face_shape.reshape(-1, 3)

    def get_right_glob_idx_by_left(self, glob_l_idx:int):
        if not hasattr(self, 'left_idx_to_right_idx_sym') or not hasattr(self, 'right_idx_to_left_idx_sym'):
            assert False
        if glob_l_idx not in self.left_idx_to_right_idx_sym:
            return None
        return self.left_idx_to_right_idx_sym[glob_l_idx]

    def get_left_glob_idx_by_right(self, glob_r_idx: int):
        if not hasattr(self, 'left_idx_to_right_idx_sym') or not hasattr(self, 'right_idx_to_left_idx_sym'):
            assert False
        if glob_r_idx not in self.right_idx_to_left_idx_sym:
            return None
        return self.right_idx_to_left_idx_sym[glob_r_idx]

    def get_left_symmetry_index_pt_test(self, glob_left_idx:int):
        if not hasattr(self, 'left_idx_to_right_idx_sym') or not hasattr(self, 'right_idx_to_left_idx_sym'):
            assert False
        if glob_left_idx not in self.left_idx_to_right_idx_sym:
            if glob_left_idx in self.right_idx_to_left_idx_sym:
                print(f'Glob index {glob_left_idx} correspond to right_idx_to_left_idx_sym.')
                return self.right_idx_to_left_idx_sym[glob_left_idx], 'right'
            else:
                print(f'Glob index {glob_left_idx} does not match any point.')
                return None
        return self.left_idx_to_right_idx_sym[glob_left_idx], 'left'

    def get_right_symmetry_index_pt_test(self, glob_right_idx: int):
        if not hasattr(self, 'right_idx_to_left_idx_sym') or not hasattr(self, 'left_idx_to_right_idx_sym'):
            assert False
        if glob_right_idx not in self.right_idx_to_left_idx_sym:
            if glob_right_idx in self.left_idx_to_right_idx_sym:
                print(f'Glob index {glob_right_idx} correspond to left_idx_to_right_idx_sym.')
                return self.left_idx_to_right_idx_sym[glob_right_idx], 'left'
            else:
                print(f'Glob index {glob_right_idx} does not match any point.')
                return None

        return self.right_idx_to_left_idx_sym[glob_right_idx], 'right'

    def get_symmetry_points(self, left_idx:int):
        result = self.get_left_symmetry_index_pt_test(left_idx)
        if result is None:
            print(f'Incorrect global left index: {left_idx}.')
            return None
        right_idx, flag = result
        # need swap
        if flag == 'right':
            print('Need swap left <-> right.')
            temp = left_idx
            left_idx = right_idx
            right_idx = temp

        left_pt, right_pt = self._shape_mean[left_idx], self._shape_mean[right_idx]
        symm_left_right_pts = np.vstack([left_pt[None, :], right_pt[None, :]])
        return symm_left_right_pts

    #TODO: IN COMPLEX CASE NEED CORRECT
    def build_vert_symmetry_pt_correspondence(self, complex:bool=False) -> dict:
        mean_shape = self._shape_mean
        left_part_of_face_mask = (mean_shape[:, 0] < 0).astype(np.bool)
        right_part_of_face_mask = (mean_shape[:, 0] >= 0).astype(np.bool)
        left_part_of_face = mean_shape[left_part_of_face_mask]
        right_part_of_face = mean_shape[right_part_of_face_mask]
        kdtree_r_part_yz_proj = KDTree(data=right_part_of_face[:, 1:], copy_data=True)
        left_idx_to_glob_idx = np.arange(self._shape_mean.shape[0])[left_part_of_face_mask]
        right_idx_to_glob_idx = np.arange(self._shape_mean.shape[0])[right_part_of_face_mask]
        left_idx_to_right_idx_sym = {}
        right_idx_to_left_idx_sym = {}
        for idx_pt_left, pt_left in enumerate(left_part_of_face):
            pt_left_yz = pt_left[1:]
            dist, indeces = kdtree_r_part_yz_proj.query(pt_left_yz, k = np.arange(1, 6))
            # TODO: need check also np.abs(x)
            if complex:
                pts_right_candidates = right_part_of_face[indeces]
                diffs_dists_x0_plane = np.abs(np.abs(pts_right_candidates[:, 0]) - np.abs(pt_left[0]))
                min_right_idx = np.argmin(diffs_dists_x0_plane)
                right_closest_idx = indeces[min_right_idx]
            else:
                right_closest_idx = indeces[0]
            glob_left_idx = left_idx_to_glob_idx[idx_pt_left]
            glob_right_idx = right_idx_to_glob_idx[right_closest_idx]

            left_idx_to_right_idx_sym[glob_left_idx] = glob_right_idx
            right_idx_to_left_idx_sym[glob_right_idx] = glob_left_idx

        self.left_idx_to_right_idx_sym = left_idx_to_right_idx_sym
        self.right_idx_to_left_idx_sym = right_idx_to_left_idx_sym

        return {'left_to_right': left_idx_to_right_idx_sym, 'right_to_left': right_idx_to_left_idx_sym}

    def draw_face_mesh_with_pts(self, added_pts:np.ndarray):
        face_mean_mesh = MeshOperator.get_mesh_by(vertices=self._shape_mean, triangles=self._triangles)
        face_mean_mesh.compute_vertex_normals()
        list_circle_meshs = []
        for pt in added_pts:
            mesh_circle = create_sphere_mesh(pt, radius=0.01, color=[0, 1., 0.])
            list_circle_meshs.append(mesh_circle)
        list_circle_meshs += [face_mean_mesh]
        op3d.visualization.draw_geometries(list_circle_meshs, width=800, height=600)

    @property
    def keypoints_indeces(self):
        return self._keypts_indeces

    @property
    def get_mean_face_mesh_custom(self):
        return Mesh(vertices=self._shape_mean, triangles=self._triangles, textures=self._tex_mean)

    @property
    def mean_shape(self):
        return self._shape_mean

    @property
    def mean_texture(self):
        return self._tex_mean

    @property
    def triangles(self):
        return self._triangles

    @property
    def mean_face_pcd(self):
        return MeshOperator.get_point_cloud_by(self._shape_mean, self._tex_mean)

    @property
    def mean_face_mesh_open3d(self):
        return MeshOperator.get_mesh_by(self._shape_mean, self._triangles, self._tex_mean)

    def get_mean_face_mesh_open3d_with_kpts(self, color_kpts=(0., 1., 0.),
                                                  radius:float=0.005) -> List[op3d.geometry.TriangleMesh]:
        if not hasattr(self, "_keypts_indeces"):
            assert False, "Keypoints indeces not defined!"
        keypts_vertices = self._shape_mean[self._keypts_indeces]
        meshs_keypts = []
        for idx, key_pts in enumerate(keypts_vertices):
            sphere_pt = self._create_sphere_mesh(center=key_pts, radius=radius, color=color_kpts)
            meshs_keypts += [sphere_pt]

        mean_face_mesh = self.mean_face_mesh_open3d
        return [mean_face_mesh] + meshs_keypts

    def _create_sphere_mesh(self, center: np.ndarray, radius: float, color:Union[Tuple, List]):
        mesh_sphere = op3d.geometry.TriangleMesh.create_sphere(radius=radius)
        mesh_sphere.paint_uniform_color(color)
        return mesh_sphere.translate(center)

    def _save_mean_face_mesh_as_obj(self, filepath='/tmp/mean_face_mesh_basel.obj', save_texture_file:bool=True):
        mean_face_mesh = self.mean_face_mesh_open3d
        if isinstance(filepath, str):
            filepath = Path(filepath)
        output = op3d.io.write_triangle_mesh(str(filepath), mesh=mean_face_mesh, write_triangle_uvs=save_texture_file,
                                    print_progress=True)
        logging.info(f'Save mean face shape as {filepath}, Success: {output}.')
        return filepath

if __name__ == '__main__':
    bfm = BaselFaceModel2009(path_bfm_model)
    bfm.build_vert_symmetry_pt_correspondence(complex=True)


import logging
import time
import json
import os

from arena_wrapper.arena_orchestrator import ArenaOrchestrator
from arena_wrapper.enums.object_output_wrapper import ObjectOutputType
from modeling.inference.model_executors.placeholder_model_executor import ArenaPHModel

BASE_DIR = os.environ['SUBMISSIONS_DIR_PATH']

class ArenaHelper:
    def __init__(self):
        self.arena_orchestrator = ArenaOrchestrator()
        self.model_handler = ArenaPHModel(object_output_type="OBJECT_MASK", data_path=None)
        self.cv_model = self.model_handler.load_default_cv_model()
        self.logger = logging.getLogger("ArenaHelper")
        self.logger.setLevel(logging.DEBUG)
        self.color_images = None
        self.predicted_actions_list = None

    def launch_game(self, cdf_data):
        try:
            self.predicted_actions_list = []
            self.arena_orchestrator.init_game(cdf_data)
            return_val, error_code = self.execute_dummy_action()
            if not return_val:
                self.logger.error("Error in executing initial dummy action %s" % error_code)
                return False
            else:
                self.logger.info("Game is initialized successfully")
            self.color_images = self.arena_orchestrator.get_images_from_metadata("colorImage")
            return True
        except Exception as ex:
            self.logger.error("Error in initializing the game: %s" % ex)
        return False

    def execute_dummy_action(self):
        dummy_action = [{
            "id": "1",
            "type": "Rotate",
            "rotation": {
                "direction": "Right",
                "magnitude": 0,
            }
        }]
        return_val, error_code = False, None
        action_counter = 0
        while not return_val and action_counter <= 10:
            return_val, error_code = self.arena_orchestrator.execute_action(dummy_action, ObjectOutputType.OBJECT_MASK, None)
            time.sleep(1)
            action_counter += 1
            self.logger.info(f"Error code: {error_code}")
        return return_val, error_code

    def execute_instruction(self, utterance):
        try:
            actions = self.model_handler.predict(utterance, self.color_images, instruction_history='', cv_model=self.cv_model)
            self.predicted_actions_list.extend(actions)
            # execute actions on arena
            ret_val, error_code = self.arena_orchestrator.execute_action(actions, ObjectOutputType.OBJECT_MASK, None)
            if not ret_val:
                self.logger.error("Action could not be executed successfully.")
            self.color_images = self.arena_orchestrator.get_images_from_metadata("colorImage")
            return self.arena_orchestrator.get_scene_data()
        except Exception as ex:
            self.logger.error("Error in initializing the game: %s" % ex)
        return None

    def get_predicted_actions(self):
        return self.predicted_actions_list

    def get_latest_game_state(self):
        exclude_keys = ["sceneMetadata", "colorImage", "depthImage", "normalsImage", "instanceSegmentationImage", "objects"]
        return {key: self.arena_orchestrator.response[key] for key in self.arena_orchestrator.response if key not in exclude_keys}

def run_instructions(instructions, arena_helper):
    for data in instructions:
        command = data["instruction"]
        arena_helper.execute_instruction(command)
    response = {
        "predicted_actions": arena_helper.get_predicted_actions(),
        "last_game_state": arena_helper.get_latest_game_state()
    }
    return response

def generate_submission_file(test_case_responses, test_case_id, index):
    if not os.path.exists(BASE_DIR):
        os.makedirs(BASE_DIR)
        print("The new directory is created!")
    file_path = BASE_DIR + test_case_id + "_" + str(index) + ".json"
    with open(file_path, "w") as actions_file:
        actions_file.write(json.dumps(test_case_responses))

def main():
    arena_helper = ArenaHelper()
    with open('sample_test_data.json', 'r') as openfile:
        test_data = json.load(openfile)
    test_case_ids = test_data.keys()
    for test_case_id in test_case_ids:
        test_case = test_data[test_case_id]
        human_annotations = test_case["human_annotations"]
        cdf_data = test_case["CDF"]
        for index, human_annotation in enumerate(human_annotations):
            if arena_helper.launch_game(cdf_data):
                result = run_instructions(human_annotation["instructions"], arena_helper)
                generate_submission_file(result, test_case_id, index)
            else:
                print("Error while launching the game")

if __name__ == '__main__':
    main()


import requests
import json

# id = input("Enter student id: ")
# URL = "http://127.0.0.1:8000/stuinfo/"+id

# r = requests.get(url = URL)

# print(r.json())
URL = "http://127.0.0.1:8000/stucreate/"
data = {
	'name': 'kfjlsadkjf',
	'roll': 1983,
	'city': 'jamalpur'
}
json_data = json.dumps(data)
r = requests.post(url = URL, data = json_data)

data = r.json()
print(data)


from astropy.nddata import support_nddata

from astropy import log
import numpy as np
from astropy.wcs import wcs
import craspy.axes as axes

from .index import _slab
from .homogen import synthesize_bubbles
import matplotlib.pyplot as plt

def _draw_spectra(data, wcs=None, unit=None,velocities=False):
    try:
        freq_axis = axes._get_axis(wcs, "FREQ")
    except ValueError:
        log.warning("Data does not have a spectral dimension")
        return
    ll = list(range(wcs.naxis))
    ll.remove(freq_axis)
    yvals = np.nansum(data, axis=tuple(ll))
    plt.ylabel(unit)
    if velocities == True:
        xvals = axes.spectral_velocities(data, wcs=wcs)
        plt.xlabel("VEL [Km/s]")
    else:
        dim = wcs.wcs.spec
        fqis = np.arange(data.shape[data.ndim - dim - 1])
        idx = np.zeros((fqis.size, data.ndim))
        idx[:, dim] = fqis
        vals = wcs.all_pix2world(idx, 0)
        xvals = vals[:, dim]
        plt.xlabel("FREQ [Hz]")
    plt.plot(xvals, yvals)

@support_nddata
def visualize_spectra(data, wcs=None, unit=None,velocities=False):
    if wcs is None:
        if data.ndim != 1:
            log.info("Only 1D data can be shown without WCS")
            return
        else:
            visualize_plot(data,wcs,unit)
            return
    _draw_spectra(data,wcs,unit,velocities)
    plt.show()

@support_nddata
def visualize(data,wcs=None,unit=None):
    """
    Generic function to visualize data, line-plot for 1D and image for 2D.

    Parameters
    ----------
    data : numpy.ndarray or astropy.nddata.NDData or astropy.nddata.NDDataRef
        Astronomical image

    wcs : astropy.wcs.WCS
        World Coordinate System from the image (not needed if contained in NDData)

    unit : astropy.unit
        Image units (not needed if contained in NDData)
    """

    if data.ndim == 1:
        return visualize_plot(data,wcs,unit)
    elif data.ndim == 2:
        return visualize_image(data,wcs,unit)
    elif data.ndim == 3:
        if axes._is_spectra(data,wcs):
            visualize_spectra(data,wcs,unit)
        else:
            #TODO: Moment 0
            visualize_image(data,wcs,unit)
    else:
        log.error("Data dimensions must be between 1 and 3")

@support_nddata
def visualize_plot(data,wcs=None,unit=None):
    """
    Plot 1D data for astronomical data.

    Parameters
    ----------
    data : numpy.ndarray or astropy.nddata.NDData
        Astronomical image

    wcs : astropy.wcs.WCS
        World Coordinate System from the image (not needed if contained in NDData)

    unit : astropy.unit
        Image units (not needed if contained in NDData)
    """
    if wcs is None:
        plt.plot(data)
        plt.ylabel(unit)
    else:
        #TODO: Implement x vector, but check why the wcs cannot be onedimensional!
        plt.plot(data)
        plt.ylabel(unit)
        plt.xlabel(wcs.axis_type_names[0])
    plt.show()

def _draw_image(data,gax,wcs,unit,cmap=plt.cm.gist_heat):
    plt.imshow(data, origin='lower', cmap=cmap)
    g0 = gax.coords[0]
    g1 = gax.coords[1]
    g0.set_axislabel(wcs.axis_type_names[0])
    g1.set_axislabel(wcs.axis_type_names[1])
    cb = plt.colorbar()
    cb.ax.set_ylabel(unit)

@support_nddata
def show_subcube(data,wcs=None,meta=None,mask=None,unit=None,lower=None,upper=None):
    plt.figure(figsize=(10,5))
    mslab=_slab(data, lower, upper)
    #scube=reduction.cut(data,wcs=wcs,meta=meta,mask=mask,unit=unit,lower=lower,upper=upper)
    if data.ndim == 3:
        freq_axis = axes._get_axis(wcs, "FREQ")
        wcs=wcs.slice(mslab, numpy_order=True)
        wcsp=wcs.dropaxis(wcs.wcs.spec)
        im1=121
        #data2d= np.nansum(data, axis=(freq_axis))
        #scube2d=np.nansum(scube.data, axis=(freq_axis))
    else:
        wcs=wcs.slice(mslab, numpy_order=True)
        wcsp=wcs
        im1=111
        #im2=122
        #data2d=data
        #scube2d=scube.data
    gax=plt.subplot(im1,projection=wcsp)
    if data.ndim == 3:
        _draw_image(np.nansum(data[mslab],axis=freq_axis),gax,wcsp,unit,cmap=plt.cm.seismic)
    else:
        _draw_image(data[mslab], gax, wcsp, unit, cmap=plt.cm.seismic)
    #gax=plt.subplot(im2, projection=swcs2d)
    #_draw_image(scube2d, gax, swcs2d, unit, cmap=plt.cm.seismic)
    if (data.ndim == 3):
        #plt.subplot(223)
        #_draw_spectra(data,wcs=wcs,unit=unit)
        plt.subplot(122)
        _draw_spectra(data[mslab], wcs=wcs, unit=unit)
    plt.show()
    #return(scube)

def show_clumps(hrtree,node=0,print_numbers=True):
    fig = plt.figure(figsize=(10, 10))
    shape = hrtree.synth.data.shape
    parent=hrtree.find_parent(node,hrtree.tree)
    n_colors = hrtree.visible_clumps(parent,node)
    gs = gridspec.GridSpec(2, 2,width_ratios=[shape[2] / shape[0], 1], height_ratios=[shape[1] / shape[0], 1])
    tempXY = hrtree.synth.data.sum(axis=(0))
    nax = plt.subplot(gs[0])
    nax.imshow(tempXY, origin='lower', cmap=plt.cm.gray_r)
    nax.get_xaxis().set_visible(False)
    nax.get_yaxis().set_visible(False)

    tempXZ = hrtree.synth.data.sum(axis=(1))
    nax2 = plt.subplot(gs[2])
    nax2.imshow(tempXZ, origin='lower', cmap=plt.cm.gray_r)
    nax2.get_xaxis().set_visible(False)
    nax2.get_yaxis().set_visible(False)

    tempYZ = hrtree.synth.data.sum(axis=(2)).T
    nax3 = plt.subplot(gs[1])
    nax3.imshow(tempYZ, origin='lower', cmap=plt.cm.gray_r)
    nax3.get_xaxis().set_visible(False)
    nax3.get_yaxis().set_visible(False)

    color = plt.cm.rainbow(np.linspace(0, 1, n_colors))

    hrtree.t_n = 0
    def plotme(key):
        newSyn = np.zeros(shape)
        newSyn = synthesize_bubbles(newSyn, hrtree.clumps[key], hrtree.kernel, hrtree.noise, hrtree.delta)
        ct = nax.contour(newSyn.sum(axis=(0)), levels=[0.0], alpha=1.0, colors=[color[hrtree.t_n]])
        if print_numbers:
            plt.clabel(ct, fontsize=10, inline=1, fmt={0.0: str(key)})
        ct = nax2.contour(newSyn.sum(axis=(1)), levels=[0.0], alpha=1.0, colors=[color[hrtree.t_n]])
        if print_numbers:
            plt.clabel(ct, fontsize=10, inline=1, fmt={0.0: str(key)})
        ct = nax3.contour(newSyn.sum(axis=(2)).T, levels=[0.0], alpha=1.0, colors=[color[hrtree.t_n]])
        if print_numbers:
            plt.clabel(ct, fontsize=10, inline=1, fmt={0.0: str(key)})
        hrtree.t_n += 1

    def recu(tree):
        for (key, val) in tree.items():
            if hrtree.enabled[key]:
                if hrtree.display[key]:
                   plotme(key)
                if isinstance(val, dict):
                    recu(val)
    if hrtree.enabled[node]:
        if hrtree.display[node]:
            plotme(node)
    if isinstance(parent[node], dict):
        recu(parent[node])

    #i = 0
    #for b in bco:
    #    imask = (labels == b)
    #    npos = sol[imask]
    #    newSyn = np.zeros(shape)
    #    newSyn = synthesize_bubbles(newSyn, npos, mould, noise, delta)
    #    nax.contour(newSyn.sum(axis=(0)), levels=[0.0], alpha=1.0, colors=[color[i]])
    #    nax2.contour(newSyn.sum(axis=(1)), levels=[0.0], alpha=1.0, colors=[color[i]])
    #    nax3.contour(newSyn.sum(axis=(2)).T, levels=[0.0], alpha=1.0, colors=[color[i]])
    #    i += 1
    #plt.tight_layout()
    plt.show()
    #plt.figure()

    # nax4=fig.add_subplot(1,3,3)
    #nax4 = plt.subplot(111)
    #i = 0
    #for b in bco:
    #    imask = (labels == b)
    #    npos = sol[imask]
    #    newSyn = np.zeros(shape)
    #    newSyn = synthesize_bubbles(newSyn, npos, mould, noise, delta)
    #    nax.contour(newSyn.sum(axis=(0)), levels=[0.0], alpha=1.0, colors=[color[i]])
    #    nax2.contour(newSyn.sum(axis=(1)), levels=[0.0], alpha=1.0, colors=[color[i]])
    #    nax3.contour(newSyn.sum(axis=(2)).T, levels=[0.0], alpha=1.0, colors=[color[i]])
    #    nax4.plot(newSyn.sum(axis=(1, 2)), color=color[i])
    #    i += 1
    #plt.tight_layout()
    #plt.show()
    #return (bco)

@support_nddata
def visualize_image(data,wcs=None,unit=None,contour=False,cmap=None):
    """
    Plot 2D astronomical data.

    Parameters
    ------------
    data : numpy.ndarray or astropy.nddata.NDData or astropy.nddata.NDDataRef
        Astronomical image

    wcs : astropy.wcs.WCS
        World Coordinate System from the image (not needed if contained in NDData)

    unit : astropy.unit
        Image units (not needed if contained in NDData)

    contour : numpy.ndarray
        For plotting Contourns
    """
    if cmap is None:
        cmap=plt.cm.gray_r
    if wcs is None:
        if data.ndim != 2:
            log.info("Cannot visualize image data with no WCS and dimension != 2")
            return
        plt.imshow(data, origin='lower', cmap=cmap)
        cb = plt.colorbar()
        cb.ax.set_ylabel(unit)
    else:
        # Only work for spectral cubes... what else?
        if data.ndim == 3:
            try:
                freq_axis = axes._get_axis(wcs, "FREQ")
            except ValueError:
                log.warning("Data does not have a FREQ dimension to stack, aborting visualization")
                return
            wcs = wcs.dropaxis(wcs.naxis - freq_axis - 1)
            data = np.nansum(data, axis=(freq_axis))
        gax = plt.subplot(111, projection=wcs)
        _draw_image(data, gax, wcs, unit,cmap=cmap)
        g0 = gax.coords[0]
        g1 = gax.coords[1]
        g0.grid(color='yellow', alpha=0.5, linestyle='solid')
        g1.grid(color='yellow', alpha=0.5, linestyle='solid')
        if contour:
            arms=rms(data)
            dmax=data.max()
            crs=np.arange(1,dmax/arms)
            plt.contour(data,levels=arms*crs,alpha=0.5)
    plt.show()

# Specific plotting function
def plot_snr_estimation(target,snr_results):
    fig = plt.figure(figsize=(10,4))
    # Unpack results
    (slimit,rms,x,y,v,n,p)=snr_results
    # Plot
    ax = fig.add_subplot(1,1,1)
    ax.plot(x,y,color='b')
    ax.plot(x,x,color='b',linestyle="--")
    ax.axvline(x=slimit,color='r',label="$"+str(slimit)+" \sigma$")
    ax.legend(loc=7)
    ax.set_title(target)
    ax.set_xlabel("SNR ($ \\tau / \sigma$)")
    ax.set_ylabel('RMS', color='b')
    for tl in ax.get_yticklabels():
        tl.set_color('b')
    axp = ax.twinx()
    axp.plot(x[:-1],v,color='grey')
    axp.set_ylabel('$\Delta$RMS', color='grey')
    for tl in axp.get_yticklabels():
        tl.set_color('grey')

def visualize_rgb(rdata,gdata,bdata):
    wcs=rdata.wcs
    gax = plt.subplot(111, projection=wcs)
    sh = rdata.data.shape
    data = np.zeros((sh[0],sh[1],3))
    #plt.imshow(gdata.data)
    #plt.show()
    (rdata,_,_)=flux.standarize(rdata)
    (gdata,_,_)=flux.standarize(gdata)
    (bdata,_,_)=flux.standarize(bdata)
    data[:,:,0] = rdata.data/np.nanmax(rdata.data)
    data[:,:,1] = gdata.data/np.nanmax(gdata.data)
    data[:,:,2] = bdata.data/np.nanmax(bdata.data)
    plt.imshow(data, origin='lower')
    g0 = gax.coords[0]
    g1 = gax.coords[1]
    g0.set_axislabel(wcs.axis_type_names[0])
    g1.set_axislabel(wcs.axis_type_names[1])
    #cb = plt.colorbar()
    #cb.ax.set_ylabel(unit)
    g0 = gax.coords[0]
    g1 = gax.coords[1]
    g0.grid(color='yellow', alpha=0.5, linestyle='solid')
    g1.grid(color='yellow', alpha=0.5, linestyle='solid')
    plt.show()

from matplotlib import gridspec

# TODO: vmax can be computed by default
def show_image_grid(img_list,side,vmax,cmap=None):
    if cmap == None:
        cmap="gray_r"
    origin="lower"

    plt.figure(figsize=(10,10))
    gs = gridspec.GridSpec(side, side,wspace=0.0, hspace=0.0)
    for i in range(side*side):
        ax  = plt.subplot(gs[i])
        ax.imshow(img_list[i],origin=origin,vmax=vmax,vmin=0,cmap=cmap)
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
    plt.show()


import random
import chaos
import rune

weapon_enchantment_dict = {
    '1': 'Additional Damage',
    '2': 'Characteristic Gain',
    '3': 'Characteristic Drain',
    '4': 'Bane Weapon',
    '5': 'Flame Attack',
    '6': 'Poison Attack',
    '7': 'Degeneration Attack',
    '8': 'Freeze Attack',
    '9': 'Warp Attack',
    '10': 'Sleep',
    '11': 'Flight',
    '12': 'Breather Underwater',
    '13': 'Confusion',
    '14': 'Fear ',
    '15': 'Instability',
    '16': 'Protection',
    '17': 'Resist Fire',
    '18': 'Animated',
    '19': 'Invisibility',
    '20': 'Magic Damper',
    '21': 'Berserk',
    '22': 'Repel Undead',
    '23': 'Repel Demons',
    '24': 'Destroy Magical Weapon',
    '25': 'Spell Absorption',
    '26': 'Mighty Strike'
}

melee_weapon_enchantment_list = ['1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1',
                                 '1', '1', '1', '1', '1', '1', '1', '1', '2', '2', '2', '2', '2', '2', '2', '2', '2',
                                 '2', '2', '2', '2', '2', '3', '3', '3', '3', '3', '3', '3', '3', '3', '4', '4', '4',
                                 '4', '4', '4', '4', '4', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14',
                                 '14', '15', '15', '16', '16', '17', '18', '19', '20', '21', '21', '22', '22', '23',
                                 '23', '24', '25', '25', '26', '26']

range_weapon_enchantment_list = ['1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1', '1',
                                 '1', '1', '1', '1', '1', '1', '1', '1', '2', '2', '2', '2', '2', '2', '2', '2', '2',
                                 '2', '2', '2', '2', '2', '3', '3', '3', '3', '3', '3', '3', '3', '3', '4', '4', '4',
                                 '4', '4', '4', '4', '4', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '15',
                                 '15', '16', '16', '17', '19', '20', '22', '22', '23', '23', '25', '25', '26', '26']

melee_weapon_list = ['1', '2', '2', '2', '3', '4', '5', '6', '7', '7', '8', '8', '9', '10', '11', '12', '13', '14',
                     '15', '15', '15', '16', '17', '18', '18', '18', '18', '18', '19', '19', '19', '19', '19', '20',
                     '20', '20', '21', '21']

range_weapon_list = ['1', '2', '2', '3', '4', '5', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16',
                     '17']

melee_weapon_dict = {
    '1': ['Buckler', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '2': ['Dagger', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '3': ['Cavalry Spear', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '4': ['Flail', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '5': ['Foil', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '6': ['Knuckle-duster ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '7': ['Great Sword', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '8': ['Great Axe', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '9': ['Halberd', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '10': ['Lance', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '11': ['Main Gauche ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '12': ['Morning Star ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '13': ['Quarterstaff', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '14': ['Rapier ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '15': ['Shield ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '16': ['Spear ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '17': ['Sword-breaker ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '18': ['Hand Axe', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '19': ['Long Sword', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '20': ['One-Handed Hammer', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '21': ['Great Hammer', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
}

range_weapon_dict = {
    '1': ['Blunderbuss', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '2': ['Pistol', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '3': ['Repeater Pistol', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '4': ['Bola', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '5': ['Crossbow', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '6': ['Crossbow Pistol ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '7': ['Javelin ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '8': ['Lasso', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '9': ['Net', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '10': ['Repeater Crossbow', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '11': ['Sling ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '12': ['Staff Sling', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '13': ['Throwing Axe', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '14': ['Throwing Hammer ', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '15': ['Throwing Dagger', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '16': ['Throwing Star', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188'],
    '17': ['Whip', 'Enchanted weapon. See manual for a description of its properties.', 'CORE1', '188']
}

# this function picks a random object
# from the list, then takes the information from
# the corresponding dictionary
# entry.

def pick_weapon():
    type_of_magic_weapon = random.randint(1, 100)
    if type_of_magic_weapon > 84 and type_of_magic_weapon < 96:
        return rune.pick_weapon()
    if type_of_magic_weapon > 95:
        weapon = chaos.pick_weapon_chaos()
        return weapon

    amount_properties_roll = random.randint(1, 100)
    if amount_properties_roll <= 20:
        amount_properties = 0
    elif amount_properties_roll <= 70:
        amount_properties = 1
    elif amount_properties_roll <= 90:
        amount_properties = 2
    elif amount_properties_roll <= 95:
        amount_properties = 3
    elif amount_properties_roll <= 99:
        amount_properties = 4
    elif amount_properties_roll == 100:
        amount_properties = 5
    weapon_type_roll = random.randint(1, 100)
    if weapon_type_roll <= 70:
        choice = str(random.choice(melee_weapon_list))
        chosen_object = melee_weapon_dict[choice]
        enchantment_list = melee_weapon_enchantment_list
    elif weapon_type_roll <= 100:
        choice = str(random.choice(range_weapon_list))
        chosen_object = range_weapon_dict[choice]
        enchantment_list = range_weapon_enchantment_list
    object_id = choice
    name = chosen_object[0]

    if amount_properties == 0:
        name += ' (Magical, '
    elif amount_properties != 0:
        name += ' ('
    description = chosen_object[1]
    source = chosen_object[2]
    page = chosen_object[3]

    object_enchantments = []
    while len(object_enchantments) < amount_properties:
        pick = random.choice(enchantment_list)
        if pick not in object_enchantments:
            object_enchantments.append(pick)
    for enchantment in object_enchantments:
        if enchantment == '1':
            damage_roll = random.randint(1, 100)
            name += 'Additional Damage ('
            if damage_roll <= 50:
                name += '+1), '
            elif damage_roll <= 75:
                name += '+2), '
            elif damage_roll <= 90:
                name += '+3), '
            elif damage_roll <= 95:
                name += '+4), '
            elif damage_roll <= 98:
                name += 'double normal damage), '
            elif damage_roll <= 100:
                name += 'double normal damage, +4), '
        elif enchantment == '2':
            char_gain_roll = random.randint(1, 89)
            name += 'Characteristic Gain ('
            if char_gain_roll <= 9:
                name += 'WS +10), '
            elif char_gain_roll <= 18:
                name += 'S +10), '
            elif char_gain_roll <= 27:
                name += 'T +10), '
            elif char_gain_roll <= 36:
                name += 'Agi +10), '
            elif char_gain_roll <= 45:
                name += 'A +1), '
            elif char_gain_roll <= 54:
                name += 'Fel +10), '
            elif char_gain_roll <= 63:
                name += 'Int +10), '
            elif char_gain_roll <= 72:
                name += 'WP +10), '
            elif char_gain_roll <= 75:
                name += 'WS +'+str(random.randint(1, 3)*10)+'), '
            elif char_gain_roll <= 77:
                name += 'S +'+str(random.randint(1, 3)*10)+'), '
            elif char_gain_roll <= 79:
                name += 'T +'+str(random.randint(1, 3)*10)+'), '
            elif char_gain_roll <= 81:
                name += 'Agi +'+str(random.randint(1, 3)*10)+'), '
            elif char_gain_roll <= 83:
                name += 'A +'+str(random.randint(1, 3))+'), '
            elif char_gain_roll <= 85:
                name += 'Fel +'+str(random.randint(1, 3)*10)+'), '
            elif char_gain_roll <= 87:
                name += 'Int +'+str(random.randint(1, 3)*10)+'), '
            elif char_gain_roll <= 89:
                name += 'WP +'+str(random.randint(1, 3)*10)+'), '
        elif enchantment == '3':
            char_drain_roll = random.randint(1, 80)
            name += 'Characteristic Drain ('
            if char_drain_roll <= 15:
                name += 'WS -10), '
            elif char_drain_roll <= 30:
                name += 'S -10), '
            elif char_drain_roll <= 40:
                name += 'T -10), '
            elif char_drain_roll <= 50:
                name += 'Agi -10), '
            elif char_drain_roll <= 60:
                name += 'Fel -10), '
            elif char_drain_roll <= 70:
                name += 'Int -10), '
            elif char_drain_roll <= 80:
                name += 'WP -10), '
        elif enchantment == '4':
            name += 'Bane ('
            bane_roll = random.randint(1, 100)
            if bane_roll <= 5:
                name += 'Goblins and Snotlings'
            elif bane_roll <= 10:
                name += 'Hobgoblins'
            elif bane_roll <= 20:
                name += 'Orcs'
            elif bane_roll <= 25:
                name += 'All Goblinoids'
            elif bane_roll <= 27:
                name += 'Elementals'
            elif bane_roll <= 30:
                name += 'Daemons'
            elif bane_roll <= 35:
                name += 'Undead'
            elif bane_roll <= 45:
                name += 'Non-Daemonic creatures and followers of Chaos'
            elif bane_roll <= 50:
                name += 'Dragons'
            elif bane_roll <= 55:
                name += 'All reptiles'
            elif bane_roll <= 60:
                name += 'Dwarfs, Gnomes and Halflings'
            elif bane_roll <= 65:
                name += 'Elves'
            elif bane_roll <= 70:
                name += 'Fimir'
            elif bane_roll <= 75:
                name += 'Giant Animals'
            elif bane_roll <= 80:
                name += 'Skaven'
            elif bane_roll <= 85:
                name += 'Giants'
            elif bane_roll <= 90:
                name += 'Ogres and Trolls'
            elif bane_roll <= 95:
                name += 'Werecreatures'
            elif bane_roll <= 100:
                name += 'Vampires'
            name += '), '
        elif enchantment == '16':
            name += 'Protection ('
            protection_roll = random.randint(1, 100)
            if protection_roll <= 75:
                name += '+1 AP in all areas'
            elif protection_roll <= 90:
                name += '+2 AP in all areas'
            elif protection_roll <= 95:
                name += '+3 AP in all areas'
            elif protection_roll <= 98:
                name += '1 Attack per round automatically parried'
            elif protection_roll <= 100:
                name += '2 Attacks per round automatically parried'
            name += '), '
        else:
            name += weapon_enchantment_dict[enchantment]+', '
    name = name[:-2]
    name += ')'
    final_object = [object_id, name, description, source, page]
    return final_object


from reportlab.lib.units import inch, cm, pica

class Length:
    def __init__(self, size, unit):
        self.size = size
        if unit in ['inch', 'in']:
            self.unit = 'in'
            self._scale = inch
        elif unit in ['cm', 'mm']:
            self.unit = 'cm'
            self._scale = cm
            if unit == 'mm':
                self.size = self.size / 10
        elif unit in ['pica']:
            self.unit = 'pica'
            self._scale = pica
        else:
            self.unit = 'pt'
            self._scale = 1

    def __call__(self):
        return self.size * self._scale

    def __repr__(self):
        return '{} {}'.format(self.size, self.unit)

class Page:
    def __init__(self, width, height, unit='pt'):
        self.width = Length(width, unit)
        self.height = Length(height, unit)

    @property
    def size(self):
        return (self.width(), self.height())

    def flip(self):
        self.width, self.height = self.height, self.width

def get_paper(size_name, orientation='landscape'):
    if size_name in ['US Letter']:
        paper = Page(11, 8.5, 'in')
    elif size_name in ['US Legal']:
        paper = Page(14, 8.5, 'in')
    elif size_name in ['US Tabloid']:
        paper = Page(17, 11, 'in')
    if orientation == 'portrait':
        paper.flip()
    return paper


import os
import sys
import random
import itertools
import numpy as np
import datetime

import tensorflow as tf
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

def main():
    BASE_MODEL_NAMES = ["SpCas9-NG", "SpCas9", "VRQR variant", "SpCas9-NG", "SpCas9-NRRH", "SpCas9-NRTH", "SpCas9-NRCH", "SpG", "SpRY", "Sc++"]
    base_model_names = ["NG", "SpCas9", "VRQR", "NG", "NRRH", "NRTH", "NRCH", "SpG", "SpRY", "Sc++"]
    EDITOR_NAMES = ["YE1-BE4max", "SsAPOBEC3B", "ABE8e(V106W)", "ABE8.17m+V106W", "CGBE1", "miniCGBE1", "APOBEC-nCas9-Ung"]
    editor_names = ["YE1", "Ss", "8e", "17m", "CGBE1", "mini", "Bi"]
    starts = [7, 6, 6, 6, 7, 7, 7]
    lengths = [5, 7, 8, 8, 5, 5, 5]

    path = sys.argv[1]

    FILE = open(os.path.join(path, "MODEL.txt"), "r")
    MODEL = int(FILE.readlines()[0][0])
    FILE.close()
    if MODEL == 1:
        EDITOR_NAMES = EDITOR_NAMES[2:4]
        editor_names = editor_names[2:4]
        starts = starts[2:4]
        lengths = lengths[2:4]
    elif MODEL == 2:
        EDITOR_NAMES = EDITOR_NAMES[:2]
        editor_names = editor_names[:2]
        starts = starts[:2]
        lengths = lengths[:2]
    else:
        EDITOR_NAMES = EDITOR_NAMES[4:]
        editor_names = editor_names[4:]
        starts = starts[4:]
        lengths = lengths[4:]

    with open(os.path.join(path, "input.txt"), "r") as FILE:
        lines = FILE.readlines()

    TIME = open(os.path.join(path, "TIME.txt"), "w")
    TIME_START = datetime.datetime.now()

    time_start = datetime.datetime.now()
    seq, SEQ_PAM = preprocess_seq_pam(lines[3:])

    SEQS = {}
    for editor_i, (editor_name, start, length) in enumerate(zip(editor_names, starts, lengths)):
        SEQS[editor_name] = preprocess_seq(lines[3:], start, length)
    time_end = datetime.datetime.now()
    TIME.write("\t".join(["%30s" % 'preprocess', "%30s" % 'elapsed time:', str(time_end - time_start)]) + "\n")

    MODELS = {}
    time_start = datetime.datetime.now()
    for base_i, base_model_name in enumerate(base_model_names):
        if base_i != 0:
            MODELS['PAM_variant_%s_model.h5' % base_model_name] = \
                tf.keras.models.load_model('bin/DeepCas9variants/PAM_variant_%s_model.h5' % base_model_name, compile=False)

    for editor_i, (editor_name, start, length) in enumerate(zip(editor_names, starts, lengths)):
        MODELS['DeepNG-BE_%s_model' % editor_name] = \
            tf.keras.models.load_model('bin/DeepBE/DeepNG-BE_%s_model' % editor_name, compile=False)
        MODELS['DeepBE_%s_model' % editor_name] = \
            tf.keras.models.load_model('bin/DeepBE/DeepBE_%s_model' % editor_name, compile=False)
    time_end = datetime.datetime.now()
    TIME.write("\t".join(["%30s" % "model_load", "%30s" % 'elapsed time:', str(time_end - time_start)]) + "\n")

    scores = []
    for base_i, base_model_name in enumerate(base_model_names):
        time_start = datetime.datetime.now()
        if base_i != 0:
            base_score = MODELS['PAM_variant_%s_model.h5' % base_model_name].predict(SEQ_PAM, batch_size=4096)
        time_end = datetime.datetime.now()
        TIME.write("\t".join(["%30s" % (base_model_name+"-pred"), "%30s" % 'elapsed time:', str(time_end - time_start)]) + "\n")

        for editor_i, (editor_name, start, length) in enumerate(zip(editor_names, starts, lengths)):
            time_start = datetime.datetime.now()
            if base_i == 0:
                score = MODELS['DeepNG-BE_%s_model' % editor_name].predict([SEQ_PAM, SEQ_PAM], batch_size=4096)
            else:
                score = MODELS['DeepBE_%s_model' % editor_name].predict([[SEQS[editor_name], base_score], SEQ_PAM], batch_size=4096)
            time_end = datetime.datetime.now()
            TIME.write("\t".join(["%30s" % (editor_name + "-pred"), "%30s" % 'elapsed time:', str(time_end - time_start)]) + "\n")

            time_start = datetime.datetime.now()
            alist = []
            for x in range(length):
                b = []
                for i in range(1, 2**x + 1):
                    for k in range(1, 2**(length-1 - x) + 1):
                        b.append(2**(length - x) * i - k - 1)
                alist.append(b)
            time_end = datetime.datetime.now()
            TIME.write("\t".join(["%30s" % "create-alist", "%30s" % 'elapsed time:', str(time_end - time_start)]) + "\n")

            time_start = datetime.datetime.now()
            for i in range(len(score)):
                eff = score[i].sum()
                motif = seq[i][start:start+length]
                non = []
                for k in range(length):
                    if ((editor_name in ["YE1", "Ss"] and motif[k] in 'Cc') or
                        (editor_name in ["8e", "17m"] and motif[k] in 'Aa') or
                        (editor_name in ["CGBE1", "mini", "Bi"] and motif[k] in 'Cc')):
                        non.append(2**(length-1 - k))
                on = []
                for j in range(len(non)):
                    iter = list(itertools.combinations(non, j+1))
                    for k in iter:
                        on.append(sum(k))
                on = [x - 1 for x in on]
                off = list(set(list(range(2**(length)-1))) - set(on))

                # for m in off:
                #     score[score == score[i, m]] = 0
                score[i, off] = 0

                patt = score[i].sum()
                score[i] = (score[i] / patt) * eff * 100
            time_end = datetime.datetime.now()
            TIME.write("\t".join(["%30s" % "set-score", "%30s" % 'elapsed time:', str(time_end - time_start)]) + "\n")

            time_start = datetime.datetime.now()
            if base_i == 0:
                output = "DeepNG-BE_%s_%s_Results.txt" % (BASE_MODEL_NAMES[base_i], EDITOR_NAMES[editor_i])
            else:
                output = "DeepBE_%s_%s_Results.txt" % (BASE_MODEL_NAMES[base_i], EDITOR_NAMES[editor_i])

            with open(os.path.join(path, output), "w") as OUTPUT:
                OUTPUT.write("\t".join(["Model", "", "", "", "", "", "", "DeepNG-BE" if base_i == 0 else "DeepBE"]) + "\n")
                OUTPUT.write("\t".join(["Nuclease or nickase", "", "", "", "", "", "", BASE_MODEL_NAMES[base_i]]) + "\n")
                OUTPUT.write("\t".join(["ID", "Location", "Input Sequence (30bp)", "Target Sequence (30bp)", "Guide RNA Sequence (20bp)",
                                        "Outcome Sequence (30bp)", "GC Contents (%)", EDITOR_NAMES[editor_i]]) + "\n")

                scores.append([])
                for i in range(len(seq)):
                    strings = []

                    data = lines[i+3].strip().split("\t")
                    score_sum = 0
                    for j in range(2**(length)-1):
                        strings.append(data[:5] + [seq[i], data[5], "%.3f" % score[i, j] if data[6+len(scores)-1] == "." else "0"])
                        if float(strings[-1][-1]) > 0:
                            score_sum += float(strings[-1][-1])
                    scores[-1].append(score_sum)

                    for k in range(len(alist)):
                        for j in alist[k]:
                            if float(strings[j][-1]) > 0:
                                if editor_name in ["YE1", "Ss"]:
                                    strings[j][5] = strings[j][5][:k+start] + "t" + strings[j][5][k+1+start:]
                                elif editor_name in ["8e", "17m"]:
                                    strings[j][5] = strings[j][5][:k+start] + "g" + strings[j][5][k+1+start:]
                                elif editor_name in ["CGBE1", "mini", "Bi"]:
                                    strings[j][5] = strings[j][5][:k+start] + "g" + strings[j][5][k+1+start:]

                    strings_dict = {}
                    for string in strings:
                        if float(string[-1]) > 0:
                            v = float(string[-1]) * -1
                            while 1:
                                if v not in strings_dict:
                                    strings_dict[v] = string
                                    break
                                else:
                                    v += random.random() * 0.0000001

                    for _, string in sorted(strings_dict.items()):
                        OUTPUT.write("\t".join(string) + "\n")

            time_end = datetime.datetime.now()
            TIME.write("\t".join(["%30s" % "print", "%30s" % 'elapsed time:', str(time_end - time_start)]) + "\n")

    time_start = datetime.datetime.now()
    with open(os.path.join(path, "DeepBE_Results.txt"), "w") as OUTPUT:
        for l in range(len(lines)):
            if l < 3:
                OUTPUT.write(lines[l])
            else:
                data = lines[l].strip().split("\t")
                for i in range(len(data[6:])):
                    if data[6+i] == ".":
                        data[6+i] = "%.3f" % scores[i][l-3]# if scores[i][l-3] > 0.1 else "No activity"
                OUTPUT.write("\t".join(data) + "\n")
    time_end = datetime.datetime.now()
    TIME.write("\t".join(["%30s" % 'print-all', "%30s" % 'elapsed time:', str(time_end - time_start)]) + "\n")

    TIME_END = datetime.datetime.now()
    TIME.write("\n" + "\t".join(["%30s" % 'TOTAL:', "%30s" % 'elapsed time:', str(TIME_END - TIME_START)]) + "\n")
    TIME.close()

    return

def preprocess_seq(lines, start, length):
    DATA_X = np.zeros((len(lines), length + 2, 4), dtype=int)
    for l in range(len(lines)):
        seq = lines[l].split("\t")[3]
        for i in range(length + 2):
            if seq[i + start - 1] in "Aa":    DATA_X[l, i, 0] = 1
            elif seq[i + start - 1] in "Cc":  DATA_X[l, i, 1] = 1
            elif seq[i + start - 1] in "Gg":  DATA_X[l, i, 2] = 1
            elif seq[i + start - 1] in "Tt":  DATA_X[l, i, 3] = 1

    return DATA_X

def preprocess_seq_pam(lines, seq_length=30):
    seqs = []
    DATA_X_PAM = np.zeros((len(lines), seq_length, 4), dtype=int)
    for l in range(len(lines)):
        seq = lines[l].split("\t")[3]
        seqs.append(seq)

        for i in range(seq_length):
            if seq[i]in "Aa":    DATA_X_PAM[l, i, 0] = 1
            elif seq[i] in "Cc": DATA_X_PAM[l, i, 1] = 1
            elif seq[i] in "Gg": DATA_X_PAM[l, i, 2] = 1
            elif seq[i] in "Tt": DATA_X_PAM[l, i, 3] = 1

    return seqs, DATA_X_PAM

if __name__ == '__main__':
    main()