# Exercise: Convolution and ReLU

# Introduction
# In this exercise, you'll work on building some intuition around feature extraction.
# irst, we'll walk through the example we did in the tutorial again, but this time, with a kernel you choose yourself.
# We've mostly been working with images in this course, but what's behind all of the operations we're learning about
# is mathematics. So, we'll also take a look at how these feature maps can be represented instead as arrays of numbers
# and what effect convolution with a kernel will have on them.

# Run the cell below to get started!
# Setup feedback system
from learntools.core import binder
binder.bind(globals())
from learntools.computer_vision.ex2 import *

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

plt.rc('figure', autolayout=True)
plt.rc('axes', labelweight='bold', labelsize='large',
       titleweight='bold', titlesize=18, titlepad=10)
plt.rc('image', cmap='magma')

tf.config.run_functions_eagerly(True)

# Apply Transformations
# The next few exercises walk through feature extraction just like the example in the tutorial.
# Run the following cell to load an image we'll use for the next few exercises.
image_path = '../input/computer-vision-resources/car_illus.jpg'
image = tf.io.read_file(image_path)
image = tf.io.decode_jpeg(image, channels=1)
image = tf.image.resize(image, size=[400, 400])

img = tf.squeeze(image).numpy()
plt.figure(figsize=(6, 6))
plt.imshow(img, cmap='gray')
plt.axis('off')
plt.show();
# You can run this cell to see some standard kernels used in image processing.
import learntools.computer_vision.visiontools as visiontools
from learntools.computer_vision.visiontools import edge, bottom_sobel, emboss, sharpen

kernels = [edge, bottom_sobel, emboss, sharpen]
names = ["Edge Detect", "Bottom Sobel", "Emboss", "Sharpen"]

plt.figure(figsize=(12, 12))
for i, (kernel, name) in enumerate(zip(kernels, names)):
    plt.subplot(1, 4, i+1)
    visiontools.show_kernel(kernel)
    plt.title(name)
plt.tight_layout()

# 1) Define Kernel
# Use the next code cell to define a kernel. You have your choice of what kind of kernel to apply.
# One thing to keep in mind is that the sum of the numbers in the kernel determines how bright the final image is.
# Generally, you should try to keep the sum of the numbers between 0 and 1
# (though that's not required for a correct answer).

# In general, a kernel can have any number of rows and columns. For this exercise, let's use a  3×3
# kernel, which often gives the best results. Define a kernel with tf.constant.
# YOUR CODE HERE: Define a kernel with 3 rows and 3 columns.
kernel = tf.constant([
    [-1, -1, -1],
    [-1,  8, -1],
    [-1, -1, -1],
])
# Uncomment to view kernel
visiontools.show_kernel(kernel)

# Check your answer
q_1.check()

# Now we'll do the first step of feature extraction, the filtering step.
# First run this cell to do some reformatting for TensorFlow.
# Reformat for batch compatibility.
image = tf.image.convert_image_dtype(image, dtype=tf.float32)
image = tf.expand_dims(image, axis=0)
kernel = tf.reshape(kernel, [*kernel.shape, 1, 1])
kernel = tf.cast(kernel, dtype=tf.float32)

# 2) Apply Convolution
# Now we'll apply the kernel to the image by a convolution. The layer in Keras that does this is layers.Conv2D.
# What is the backend function in TensorFlow that performs the same operation?
# YOUR CODE HERE: Give the TensorFlow convolution function (without arguments)
conv_fn = tf.nn.conv2d

# Check your answer
q_2.check()

# Once you've got the correct answer, run this next cell to execute the convolution and see the result!
image_filter = conv_fn(
    input=image,
    filters=kernel,
    strides=1, # or (1, 1)
    padding='SAME',
)

plt.imshow(
    # Reformat for plotting
    tf.squeeze(image_filter)
)
plt.axis('off')
plt.show();
# Can you see how the kernel you chose relates to the feature map it produced?

# 3) Apply ReLU
# Now detect the feature with the ReLU function. In Keras, you'll usually use this as the activation function
# in a Conv2D layer. What is the backend function in TensorFlow that does the same thing?
# YOUR CODE HERE: Give the TensorFlow ReLU function (without arguments)
relu_fn = tf.nn.relu

# Check your answer
q_3.check()

# Once you've got the solution, run this cell to detect the feature with ReLU and see the result!

# The image you see below is the feature map produced by the kernel you chose.
# If you like, experiment with some of the other suggested kernels above, or, try to invent one that will extract
# a certain kind of feature.
image_detect = relu_fn(image_filter)

plt.imshow(
    # Reformat for plotting
    tf.squeeze(image_detect)
)
plt.axis('off')
plt.show();

# In the tutorial, our discussion of kernels and feature maps was mainly visual.
# We saw the effect of Conv2D and ReLU by observing how they transformed some example images.

# But the operations in a convolutional network (like in all neural networks) are usually defined through mathematical
# functions, through a computation on numbers. In the next exercise, we'll take a moment to explore this point of view.

# Let's start by defining a simple array to act as an image, and another array to act as the kernel.
# Run the following cell to see these arrays.
# Sympy is a python library for symbolic mathematics. It has a nice
# pretty printer for matrices, which is all we'll use it for.
import sympy
sympy.init_printing()
from IPython.display import display

image = np.array([
    [0, 1, 0, 0, 0, 0],
    [0, 1, 0, 0, 0, 0],
    [0, 1, 0, 0, 0, 0],
    [0, 1, 0, 0, 0, 0],
    [0, 1, 0, 1, 1, 1],
    [0, 1, 0, 0, 0, 0],
])

kernel = np.array([
    [1, -1],
    [1, -1],
])

display(sympy.Matrix(image))
display(sympy.Matrix(kernel))
# Reformat for Tensorflow
image = tf.cast(image, dtype=tf.float32)
image = tf.reshape(image, [1, *image.shape, 1])
kernel = tf.reshape(kernel, [*kernel.shape, 1, 1])
kernel = tf.cast(kernel, dtype=tf.float32)

# 4) Observe Convolution on a Numerical Matrix
# What do you see? The image is simply a long vertical line on the left and a short horizontal line on the lower right.
# What about the kernel? What effect do you think it will have on this image? After you've thought about it,
# run the next cell for the answer.

# In the tutorial, we talked about how the pattern of positive numbers will tell you the kind of features the kernel
# will extract. This kernel has a vertical column of 1's, and so we would expect it to return features of vertical
# lines.

# Now let's try it out. Run the next cell to apply convolution and ReLU to the image and display the result.
image_filter = tf.nn.conv2d(
    input=image,
    filters=kernel,
    strides=1,
    padding='VALID',
)
image_detect = tf.nn.relu(image_filter)

# The first matrix is the image after convolution, and the second is
# the image after ReLU.
display(sympy.Matrix(tf.squeeze(image_filter).numpy()))
display(sympy.Matrix(tf.squeeze(image_detect).numpy()))
# Is the result what you expected?

# Conclusion
# In this lesson, you learned about the first two operations a convolutional classifier uses for feature extraction:
# filtering an image with a convolution and detecting the feature with the rectified linear unit.


from sqlalchemy import Column, ForeignKey, Integer, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship
from sqlalchemy import create_engine
from flask_login import UserMixin
from flask_dance.consumer.backend.sqla import OAuthConsumerMixin

Base = declarative_base()

class User(UserMixin, Base):
    __tablename__ = 'user'

    id = Column(Integer, primary_key=True)
    username = Column(String(250), unique=True)

class OAuth(OAuthConsumerMixin, Base):
    __tablename__ = 'oauth'

    user_id = Column(Integer, ForeignKey(User.id))
    user = relationship(User)

class Category(Base):
    __tablename__ = 'category'

    id = Column(Integer, primary_key=True)
    name = Column(String(250), nullable=False)

    @property
    def serialize(self):
        return {
            'name': self.name,
            'id': self.id,
        }

class Item(Base):
    __tablename__ = 'item'

    id = Column(Integer, primary_key=True)
    name = Column(String(80), nullable=False)
    description = Column(String(250))
    category_id = Column(Integer, ForeignKey('category.id'))
    category = relationship(Category)

    @property
    def serialize(self):
        return {
            'name': self.name,
            'description': self.description,
            'id': self.id,
            'category_id': self.category_id
        }

engine = create_engine('sqlite:///catalog.db')
Base.metadata.create_all(engine)


# Problem 54 - Poker Hands
#The file, poker.txt, contains one-thousand random hands dealt to two players.
# Each line of the file contains ten cards (separated by a single space):
#  the first five are Player 1's cards and the last five are Player 2's cards.
#  You can assume that all hands are valid (no invalid characters or repeated cards),
# each player's hand is in no specific order, and in each hand there is a clear winner.

# Open the file and put it into a string

# This solution cannot decide 2 of the hands yet, and one of these player 1 wins (full house not implemented, and hands with same two pairs)

from collections import defaultdict

# dict for rank of cards
card_dict = {"2": 2, "3": 3, "4": 4, "5": 5, "6": 6, "7": 7, "8": 8, "9": 9, "T": 10, "J": 11, "Q": 12, "K": 13, "A": 14}
player1wins = 0
player2wins = 0

def hasPair(value, playercards, cards):
    removeCharacters = ["S", "C", "D", "H", " "]
    for character in removeCharacters:
        playercards = playercards.replace(character,"")
    rankcount = defaultdict(lambda:0)
    card = [x[0] for x in playercards]
    for y in card:
        rankcount[y] += 1
    for k,v in rankcount.items():
        if(v==2):
            print(k, "key")
            return 2, str(k)
    return value, cards

def hasFourOfKind(value, playercards, cards):
    removeCharacters = ["S", "C", "D", "H", " "]
    for character in removeCharacters:
        playercards = playercards.replace(character,"")
    rankcount = defaultdict(lambda:0)
    card = [x[0] for x in playercards]
    for y in card:
        rankcount[y] += 1
    for x in set(rankcount.values()):
        if(x==4):
            return 8, playercards
    return value, cards

def hasFullHouse(value, playercards, cards):
    return value, cards

def hasThreeOfKind(value, playercards, cards):
    removeCharacters = ["S", "C", "D", "H", " "]
    for character in removeCharacters:
        playercards = playercards.replace(character,"")
    rankcount = defaultdict(lambda:0)
    card = [x[0] for x in playercards]
    for y in card:
        rankcount[y] += 1

    for x in set(rankcount.values()):
        if(x==3):
            return 4, playercards

    return value, cards

def hasTwoPairs(value, playercards, cards):
    removeCharacters = ["S", "C", "D", "H", " "]
    for character in removeCharacters:
        playercards = playercards.replace(character,"")
    rankcount = defaultdict(lambda:0)
    card = [x[0] for x in playercards]
    for y in card:
        rankcount[y] += 1
    count = 0
    for x in list(rankcount.values()):
        if(x==2):
            count +=1
            if(count == 2):
                return 3, playercards

    return value, cards

def hasRoyalFlush(value, playercards, cards):
    if('TH' in playercards and 'JH' in playercards and 'QH' in playercards and 'KH' in playercards and 'AH' in playercards):
        return 10, playercards
    if('TD' in playercards and 'JD' in playercards and 'QD' in playercards and 'KD' in playercards and 'AD' in playercards):
        return 10, playercards
    if('TC' in playercards and 'JC' in playercards and 'QC' in playercards and 'KC' in playercards and 'AC' in playercards):
        return 10,playercards
    if('TS' in playercards and 'JS' in playercards and 'QS' in playercards and 'KS' in playercards and 'AS' in playercards):
        return 10,playercards
    return value, cards

def hasStraight(value, playercards, cards):
    removeCharacters = ["S", "C", "D", "H", " "]
    for character in removeCharacters:
        playercards = playercards.replace(character,"")
    #todo check if playercards is straight from 5 ranks close to each other
    rankcount = defaultdict(lambda:0)
    card = [x[0] for x in playercards]
    for y in card:
        rankcount[y] += 1
    cardranks = [card_dict[x] for x in card]
    if(max(cardranks)-min(cardranks) ==4 and len(set(rankcount.values()))==1):
        return 5, playercards
    return value, cards

def hasFlush(value, playercards, cards):
    countH, countC, countD, countS = 0,0,0,0
    for x in playercards:
        if(x == 'H'):
            countH += 1
        elif(x == 'S'):
            countS += 1
        elif(x == 'C'):
            countC += 1
        elif(x == 'D'):
            countD += 1
    if((countH or countC or countD or countS) == 5):
        return 6, playercards
    return value, cards

def hasStraighFlush(value, playercards, cards):
    #check if containing a straight flush by checking if it has both a straight and flush
    if(hasStraight(value, playercards, cards)[0] ==5 and hasFlush(value, playercards, cards)[0]==6):
        return 9, playercards
    return value, cards

def checkHighestCard(cards1, cards2):
    if not cards1 or not cards2:
        return 0
    removeCharacters = ["S", "C", "D", "H", " "]
    for character in removeCharacters:
        cards1 = cards1.replace(character,"")
        cards2 = cards2.replace(character,"")
    card1 = [x[0] for x in cards1]
    card2 = [x[0] for x in cards2]
    cardranks1 = [card_dict[x] for x in card1]
    cardranks2 = [card_dict[x] for x in card2]
    print(cardranks1,cardranks2, "here")
    if(max(cardranks1)>max(cardranks2)):
        return 1
    elif(max(cardranks2)>max(cardranks1)):
        return 2
    else:
        mcards1 = cards1.replace(str(max(cardranks1)), "")
        mcards2 = cards2.replace(str(max(cardranks2)), "")
        print(mcards1, mcards2, "new cards")
        checkHighestCard(mcards1, mcards2)

def onlyHighCardCheck(player1cards, player2cards, valuePlayer1, valuePlayer2):
    global player1wins, player2wins
    if(valuePlayer1 == 0 and valuePlayer2 == 0):
        checkHighest = checkHighestCard(player1cards, player2cards)
        if(checkHighest==1):
            print("Player1 wins with high card", x)
            player1wins += 1
        elif(checkHighest==2):
            print("Player2 wins with high card", x)
            player2wins += 1

#assigns a value to the cards from 0-10 depending on the rules
def valueOfCards(playercards):
    value, cards = 0, playercards
    value, cards = hasPair(value, playercards, cards)
    value, cards = hasTwoPairs(value, playercards, cards)
    value, cards = hasThreeOfKind(value, playercards, cards)
    value, cards = hasStraighFlush(value, playercards, cards)
    value, cards = hasStraight(value, playercards, cards)
    value, cards = hasFlush(value, playercards, cards)
    value, cards = hasFullHouse(value, playercards, cards)
    value, cards = hasFourOfKind(value, playercards, cards)
    value, cards = hasRoyalFlush(value, playercards, cards)
    # todo add additional rules / functions

    return value, cards

#save the data of poker hands to a string
datastring = ""
datafile = open('p054_poker.txt',"r")
if(datafile.mode == 'r'):
    datastring = datafile.read()
game = datastring.split('\n')
game = game[0:len(game)-1]
removeCharacters = [" "]

count = 0
for x in game:
    count += 1
    print("Game:", count)
    player1cards = x[:14]
    player2cards = x[15:]

    #for character in removeCharacters:
    #    player1cards = player1cards.replace(character,"")
    #    player2cards = player2cards.replace(character,"")

    # check value of cards
    valuePlayer1, cards1 = valueOfCards(player1cards)
    valuePlayer2, cards2 = valueOfCards(player2cards)

    print("CARDS OF PLAYER 1:", player1cards, "VALUE:", valuePlayer1)
    print("CARDS OF PLAYER 2:", player2cards, "VALUE:", valuePlayer2)
    # no better than high cards, check highest card, TODO recursive check next highest card
    onlyHighCardCheck(player1cards, player2cards, valuePlayer1, valuePlayer2)

    # if value is the same, check who has best rank cards
    if(valuePlayer1 == valuePlayer2 and (valuePlayer1 != 0 and valuePlayer2 != 0)):
        print(cards1, cards2)
        checkHighest = checkHighestCard(cards1, cards2)
        if(checkHighest==0):
            onlyHighCardCheck(player1cards, player2cards, valuePlayer1, valuePlayer2)
        if(checkHighest==1):
            print("Player1 wins with high card", x)
            player1wins += 1
        elif(checkHighest==2):
            print("Player2 wins with high card", x)
            player2wins += 1

    if(valuePlayer1>valuePlayer2):
        print("Player1 wins", x)
        player1wins += 1
    elif(valuePlayer2>valuePlayer1):
        print("Player2 wins", x)
        player2wins += 1
print("player1:", player1wins, "player2:", player2wins, "decided games:", player1wins+player2wins)


from typing import Callable, List, TypeVar

StringFunc = Callable[[], str]
T = TypeVar('T')
Constructor = Callable[[], T]

def get(of: Constructor[T], *,
        stdin: StringFunc = input,
        strip: bool = True) -> T:
    """Get value of type 'of' from 'stdin'."""

    raw_input = stdin().strip() if strip else stdin()
    return of(raw_input)

def get_list(of: Constructor[T], *,
             stdin: StringFunc = input,
             strip: bool = True,
             delimiter: str = ' ') -> List[T]:
    """Get list of values of type 'of' from 'stdin'."""

    raw_input = stdin().strip() if strip else stdin()
    delimited_raw_input = raw_input.split(delimiter)
    return list(map(of, delimited_raw_input))


# mod parser 3d-news for TorrentPier-II 2.1.5 alpha
# ver 1.0 by dr_brown

from urllib.request import Request, urlopen
from urllib.error import HTTPError
from urllib.parse import quote
from bs4 import BeautifulSoup
from time import sleep
import pymysql.cursors
import urllib
import time
import re

'''If use re-uploading images to your image hosting, set True (only chevereto support)'''
image_api = False
key_api = ''  # api key
url_api = ''  # url image hosting

network_news_forum_id = 2  # id forum news
bot_uid = -746
user_agent = 'Mozilla/5.0 (Windows NT 6.3; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) ' \
             'Chrome/78.0.3904.70 Safari/537.36'

global_url = 'https://3dnews.ru/'
url = 'https://3dnews.ru/news'

'''Connect to Database'''
db_connect = dict(host=' ', user=' ', password=' ', db=' ', charset='utf8mb4',
                  cursorclass=pymysql.cursors.DictCursor)

connection = pymysql.connect(**db_connect)

def run():
    try:
        with connection.cursor() as cursor:
            page = parse_html(url)
            if page is None:
                return False
            links = parse_links(page, cursor)
            if links is None:
                return False
            for link in links:
                html = parse_html(link)
                if html is None:
                    return False
                data = parse_inner_page(html)
                save_row(data, cursor)
                sleep(1)
    finally:
        connection.close()

def parse_html(parse_url):
    req = Request(parse_url, headers={'User-Agent': user_agent})
    try:
        page = BeautifulSoup(urlopen(req).read().decode('UTF-8'), 'html.parser')
    except HTTPError:
        return None
    return page

def parse_links(page, cursor):
    name_id = []
    try:
        for wrap in page.find_all('div', {'class': 'article-entry article-infeed marker_sw nImp0 nIcat10 cat_10 nIaft '
                                                   'newsAllFeedHideItem'}):
            name_id.append(wrap.get('id'))
    except AttributeError:
        return None
    links = []
    for ids in name_id:
        sql = "SELECT * FROM bb_news_grab WHERE import_id = %s"
        cursor.execute(sql, (ids,))
        result = cursor.fetchone()
        if result is not None:
            continue
        import_sql = "INSERT INTO bb_news_grab (import_id) VALUES (%s)"
        cursor.execute(import_sql, (ids,))
        connection.commit()
        links.append(global_url + ids)
    return links

def parse_inner_page(html):
    try:
        name_text = html.find('div', {'class': 'js-mediator-article'})
    except AttributeError:
        name_text = ''
    else:
        name_text = clear_url(name_text)
    try:
        name_title = html.find('h1', {'itemprop': 'headline'}).get_text('h1')
    except AttributeError:
        name_title = ''
    result = dict(text=name_text, title=name_title)
    return result

def clear_url(name_text):
    if image_api:
        images = name_text.find_all('img')
        upload_image = []
        pic_image = []
        for image in images:
            try:
                image_clr = urllib.parse.quote(image['src'], safe='-._~:/?#[]@!$&\'()*+,;=')
                pic_req = Request(url_api.format(key_api, image_clr))
                pic = urlopen(pic_req).read().decode('UTF-8')
                if pic == 'Invalid file source':
                    continue
                elif pic == 'File too big - max 2 MB':
                    continue
            except HTTPError:
                continue
            else:
                upload_image.append(image['src'])
                pic_image.append(pic)
    regex = r'<a.*?>([\s\S]*?)<\/a>'
    name_text = str(name_text)
    clear_text = re.sub(regex, r'\1', name_text)
    clear_text = re.sub(r'<br\s*?>', r' ', clear_text)
    clear_text = re.sub(r'\s+', r' ', clear_text)
    if image_api:
        for (a, b) in zip(upload_image, pic_image):
            clear_text = re.sub(a, b, clear_text)
    return clear_text

def save_row(data, cursor):
    topic_sql = "INSERT INTO bb_topics (topic_title, topic_poster, topic_time, forum_id, topic_status, " \
                "topic_type, topic_dl_type, topic_vote, topic_last_post_time) VALUES (%s, %s, %s, %s, " \
                "%s, %s, %s, %s, %s)"
    cursor.execute(topic_sql, (data['title'], bot_uid, time.time(), network_news_forum_id, '0', '0',
                               '0', '0', time.time(),))
    connection.commit()
    topic_last_id = "SELECT LAST_INSERT_ID() as last"
    cursor.execute(topic_last_id)
    topic_id = cursor.fetchone()
    post_sql = "INSERT INTO bb_posts (topic_id, forum_id, poster_id, post_username, post_time, " \
               "poster_ip) VALUES (%s, %s, %s, %s, %s, %s)"
    cursor.execute(post_sql, (topic_id['last'], network_news_forum_id, bot_uid, 'Bot', time.time(),
                              '127.0.0.1',))
    connection.commit()
    post_last_id = "SELECT LAST_INSERT_ID() as last"
    cursor.execute(post_last_id)
    post_id = cursor.fetchone()
    update_topic = "UPDATE bb_topics SET topic_first_post_id = %s, topic_last_post_id = %s WHERE " \
                   "topic_id = %s"
    cursor.execute(update_topic, (post_id['last'], post_id['last'], topic_id['last'],))
    post_text_sql = "INSERT INTO bb_posts_text (post_id, post_text) VALUES (%s, %s)"
    cursor.execute(post_text_sql, (post_id['last'], data['text'],))
    update_forum = "UPDATE bb_forums SET forum_posts = forum_posts + 1, forum_last_post_id = %s, " \
                   "forum_topics = forum_topics + 1 WHERE forum_id = %s"
    cursor.execute(update_forum, (post_id['last'], network_news_forum_id,))
    connection.commit()

run()


'''Mine Sweeper Game'''

import os
import sys
import Tkinter as tk

modulePath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "mine_sweeper")
sys.path.append(modulePath)
import Grid

TOP = tk.Tk()

GRID_WIDTH = 20
GRID_HEIGHT = 20

CELL_SIZE = 30
# Code to add widgets will go here...
CANVAS = tk.Canvas(TOP, bg="blue", height=GRID_HEIGHT * CELL_SIZE, width=GRID_WIDTH * CELL_SIZE)

grid = Grid.Grid(CANVAS, GRID_WIDTH, GRID_HEIGHT, CELL_SIZE)

CANVAS.pack()

TOP.mainloop()


import logging
import os
import datetime
import socket
import copy

from stellar_base.keypair import Keypair

from ..design import Node

log = logging.getLogger(__name__)

def subparser(subparser):
    parser = subparser.add_parser(
        'fix-design',
        help='fix the incorrect design',
    )
    parser.set_defaults(command='fix_design')

    return

def run(parser, args):
    nodes = copy.deepcopy(args.design.design_yaml['nodes'])
    for k, v in nodes.items():
        if v is None:
            v = Node.get_defaults_design()
        elif 'secret_seed' not in v:
            v['secret_seed'] = Keypair.random().seed().decode()

        args.design.design_yaml['nodes'][k] = v

    print('''%(sep)s
# generated at %(date)s from '%(user)s@%(hostname)s'
%(content)s''' % dict(
        sep='#' * 80,
        date=datetime.datetime.now().isoformat(),
        user=os.environ.get('USER'),
        hostname=socket.gethostname(),
        content=args.design.serialize(),
    ))

    return 0


import random

class Card:
    def __init__(self, suit, value):
        self.suit = suit
        self.value = value

    def __repr__(self):
        return f'{self.value} of {self.suit}'

class Deck:
    def __init__(self):
        cardsInDeck = []
        suits = ['Spades', 'Hearts', 'Diamonds', 'Clubs']
        values = ['A', 'K', 'Q', 'J', '10', '9', '8', '7', '6', '5',
                  '4', '3', '2']
        for suit in suits:
            for value in values:
                cardsInDeck.append(Card(suit, value))

        self.cards = cardsInDeck

    def shuffle(self):
        if len(self.cards) > 1:
            random.shuffle(self.cards)

    def deal(self):
        if len(self.cards) > 1:
            return self.cards.pop(0)

class Hand:
    def __init__(self, dealer=False):
        self.dealer = dealer
        self.cards = []
        self.value = 0

    def add_card(self, card):
        self.cards.append(card)

    def calculate_value(self):
        self.value = 0
        has_ace = False
        for card in self.cards:
            if card.value.isnumeric():
                self.value += int(card.value)
            else:
                if card.value == 'A':
                    has_ace = True
                    self.value += 11
                else:
                    self.value += 10
        if has_ace and self.value > 21:
            self.value -= 10

    def get_value(self):
        self.calculate_value()
        return self.value

    def display(self, player_turn=True):
        if self.dealer and player_turn is True:
            print('Dealer\'s Hand is:\nHidden')
            print(self.cards[1])
        elif self.dealer and player_turn is False:
            print('Dealer\'s Hand is:')
            for card in self.cards:
                print(card)
            print(f'Value: {self.get_value()}')
        else:
            print('Your hand is:')
            for card in self.cards:
                print(card)
            print(f'Value: {self.get_value()}')

class Game:
    def __init__(self):
        pass

    def play(self, player_score, dealer_score):
        divider = '*' * 20
        playing = True

        while playing:
            self.player_score = player_score
            self.dealer_score = dealer_score
            self.deck = Deck()
            self.deck.shuffle()
            self.player_hand = Hand()
            self.dealer_hand = Hand(dealer=True)
            for i in range(2):
                self.player_hand.add_card(self.deck.deal())
                self.dealer_hand.add_card(self.deck.deal())

            self.player_hand.display()
            # print(divider)
            self.dealer_hand.display()
            if self.dealer_hand.get_value() == 21:
                print('Game over dealer has 21.')
                self.dealer_score += 1
                menu(self.player_score, self.dealer_score)
            else:
                while self.player_hand.get_value() < 21:
                    hit_me = input('Do you want another card? ')
                    if 'y' in hit_me.lower():
                        self.player_hand.add_card(self.deck.deal())
                        self.player_hand.display()
                        if self.player_hand.get_value() > 21:
                            print('Player busted, dealer wins')
                            self.dealer_score += 1
                            menu(self.player_score, self.dealer_score)
                    else:
                        while self.dealer_hand.get_value() < 21:
                            if self.dealer_hand.get_value() >= self.player_hand.get_value():
                                self.dealer_hand.display(player_turn=False)
                                print('Dealer has better score than player, dealer wins')
                                self.dealer_score += 1
                                menu(self.player_score, self.dealer_score)
                            else:
                                self.dealer_hand.add_card(self.deck.deal())
                                self.dealer_hand.display(player_turn=False)
                                if self.dealer_hand.get_value() > 21:
                                    print('Dealer busted, player wins')
                                    self.player_score += 1
                                    menu(self.player_score, self.dealer_score)
                                elif self.dealer_hand.get_value() >= self.player_hand.get_value():
                                    print('Dealer has better score than player, dealer wins')
                                    self.dealer_score += 1
                                    menu(self.player_score, self.dealer_score)
                                else:
                                    continue

def menu(player_score, dealer_score):
    print('Scoreboard')
    print(f'{player_name.title()}: {player_score}')
    print(f'Dealer: {dealer_score}')
    print('*' * 25)
    play = input('Do you want to play? ')
    if 'y' in play.lower():
        game = Game()
        game.play(player_score, dealer_score)

if __name__ == "__main__":
    print('Welcome to Blackjack. Good Luck!')
    player_name = input('Player Name: ')
    player_score = 0
    dealer_score = 0
    menu(player_score, dealer_score)


import numpy as np
from sklearn.metrics import jaccard_similarity_score
import os
from os import listdir
from os.path import isfile, join
import cv2

def load_image(img_org,INPUT_SIZE):
    h, w, c = img_org.shape
    img = cv2.resize(img_org, (INPUT_SIZE, INPUT_SIZE))
    img = np.array(img, dtype=np.float32)
    x = np.expand_dims(img, axis=0)
    x = preprocess_input(x)
    return x

def load_label(img_org,INPUT_SIZE,nb_classes):
    img = cv2.resize(img_org,(INPUT_SIZE, INPUT_SIZE))
    img = np.array(img, dtype=np.uint8)
    img[img==255] = 0
    y = np.zeros((1, img.shape[0], img.shape[1], nb_classes), dtype=np.float32)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            y[0, i, j, img[i][j]] = 1
    return y

def get_IOU(prediction_dir, gt_dir):
    input_size = 321
    nb_classes = 21
    all_imgs= [f for f in listdir(prediction_dir) if isfile(join(prediction_dir, f))]
    predictions, gt = [], []
    for im_name in all_imgs:
        pred_img_org = cv2.imread(os.path.join(prediction_dir, im_name))
        #print(pred_img_org)
        gt_img_org = cv2.imread(gt_dir + im_name[:-9] + ".png")
        #gt_img_org = cv2.imread(os.path.join(gt_dir, im_name))
        #print(gt_dir + im_name[:-9] + ".png")
        predictions.append(load_label(pred_img_org, input_size, nb_classes))
        gt.append(load_label(gt_img_org, input_size, nb_classes))

    iou = jaccard_similarity_score(np.array(gt).flatten(), np.array(predictions).flatten())
    print(iou)

if __name__ == "__main__":
    get_IOU("output/prediction/", "/om2/user/cfmata/voc12/train/VOCdevkit/VOC2012/SegmentationClass_1D")
    #get_IOU("output/deeplab_resnet_ckpt_prediction_test/", "output/deeplab_resnet_ckpt_prediction_test/") #debug


import sqlite3
from som_logg import (LoginCreate, Prof, helper)
conn = sqlite3.connect('sample.db')
cur = conn.cursor()
cur.execute('''CREATE TABLE IF NOT EXISTS users(login TEXT, pasw TEXT, alternative TEXT)''')

conn.commit()

cmd = input('Введите команду, которую хотите использовать: \n')

command_dict = {'/help': helper,
                '/new_user': LoginCreate,
                '/selector': Prof}

def income_func():
    if cmd == '/help':
        command_dict['/help']()
    if cmd == '/new_user':
        command_dict['/new_user']('sample.db').create_user()
    if cmd == '/selector':
        command_dict['/selector'](cur, 'users').get_info_user()

income_func()


import os
import hashlib
import argparse

# Grab command line arguments for directory and signature file
parser = argparse.ArgumentParser()
parser.add_argument('--directory', '-d', help="The directory to scan for viruses in", required=True)
parser.add_argument('--signatures', '-s', help="The signature file", required=True)
args = parser.parse_args()

print("[*] Scanning directory: %s"%args.directory)
print("[*] Using signature file: %s"%args.signatures)

blocksize=65536

for root, subFolders, files in os.walk(args.directory):
    # Take MD5 Of the file

    for fi in files:
        with open(os.path.join(root, fi), 'rb') as fin:
            hash_md5 = hashlib.md5()
            for block in iter(lambda: fin.read(blocksize), b""):
                hash_md5.update(block)
            digest = hash_md5.hexdigest()
            with open(args.signatures) as signatures_file:
                for sig in signatures_file:
                    if sig.strip() == str(digest):
                        print("[!] File %s is a virus"%(os.path.join(root, fi)))


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

import serial.tools.list_ports as ports
import subprocess
import argparse
import psutil
import yaml
import os

def get_com_by_sn(serial = ''):
    com_ports = list(ports.comports())
    for i in com_ports:
        if serial == '' or serial == i.serial_number:
            return i.device

def get_com_by_id():
    # Arduino Micro
    pid = '2341'
    hid = '8037'
    com_ports = list(ports.comports())
    for i in com_ports:
        if pid and hid in i.hwid:
            return i.device

# Retrieve CLI arguments
parser = argparse.ArgumentParser(description='Updates Deej config and reloads the app.')
parser.add_argument('--com_port', type=str, help='Manual COM port (optional).')
parser.add_argument('--cfg_path', type=str, help='Manual config path (optional).')
parser.add_argument('--ard_sn', type=str, help='Arduino serial number (optional).')
args = parser.parse_args()

# Setup paths
if args.cfg_path != None:
    deej_dir = args.cfg_path
    print('Manual config path ' + deej_dir + '.')
else:
    deej_dir = os.path.abspath(os.path.dirname(__file__))
    print('Config path ' + deej_dir + '.')

deej_exe = deej_dir + '\deej.exe'
deej_cfg = deej_dir + '\config.yaml'

if not (os.path.exists(deej_exe) and os.path.exists(deej_cfg)):
    print('Deej files not found at the specified path!')
    quit()

# Get com device
if args.com_port != None:
    deej_com = args.com_port
elif args.ard_sn != None:
    deej_com = get_com_by_sn(args.ard_sn)
else:
    deej_com = get_com_by_id()

if deej_com == None:
    print('No device found!')
    quit()
else:
    if args.com_port != None:
        print('Manual device port ' + deej_com + '.')
    elif args.ard_sn != None:
        print('Device s/n ' + args.ard_sn + ' found on port ' + deej_com + '.')
    else:
        print('Device found on port ' + deej_com + '.')

# Exit Deej if runnning
for proc in psutil.process_iter():
    if proc.name() == deej_exe.split('\\')[-1]:
        proc.kill()

# Modify Deej config
with open(deej_cfg) as cfg:
    elm_list = yaml.safe_load(cfg)
    if elm_list['com_port'] != deej_com:
        elm_list['com_port'] = deej_com
        print('Deej config updated.')

        with open(deej_cfg, 'w') as cfg:
            yaml.dump(elm_list, cfg)

# Start Deej
# Specify working dir for Deej to load it's config from
process = subprocess.Popen(deej_exe, cwd = deej_dir)
print('Deej started')


from __future__ import annotations

import asyncio
import contextlib
from concurrent.futures.thread import ThreadPoolExecutor
from functools import partial
from typing import Optional, AsyncGenerator, Union
from typing_extensions import Literal

from serial import Serial, serial_for_url  # type: ignore[import]

TimeoutProperties = Union[Literal["write_timeout"], Literal["timeout"]]

class AsyncSerial:
    """Async wrapper around Serial."""

    @classmethod
    async def create(
        cls,
        port: str,
        baud_rate: int,
        timeout: Optional[float] = None,
        write_timeout: Optional[float] = None,
        loop: Optional[asyncio.AbstractEventLoop] = None,
        reset_buffer_before_write: bool = False,
    ) -> AsyncSerial:
        """
        Create an AsyncSerial instance.

        Args:
            port: url or port name
            baud_rate: the baud rate
            timeout: optional timeout in seconds
            write_timeout: optional write timeout in seconds
            loop: optional event loop. if None get_running_loop will be used
            reset_buffer_before_write: reset the serial input buffer before
             writing to it
        """
        loop = loop or asyncio.get_running_loop()
        executor = ThreadPoolExecutor(max_workers=1)
        serial = await loop.run_in_executor(
            executor=executor,
            func=partial(
                serial_for_url,
                url=port,
                baudrate=baud_rate,
                timeout=timeout,
                write_timeout=write_timeout,
            ),
        )
        return cls(
            serial=serial,
            executor=executor,
            loop=loop,
            reset_buffer_before_write=reset_buffer_before_write,
        )

    def __init__(
        self,
        serial: Serial,
        executor: ThreadPoolExecutor,
        loop: asyncio.AbstractEventLoop,
        reset_buffer_before_write: bool,
    ) -> None:
        """
        Constructor

        Args:
            serial: connected Serial object
            executor: a thread pool executor
            loop: event loop
        """
        self._serial = serial
        self._executor = executor
        self._loop = loop
        self._reset_buffer_before_write = reset_buffer_before_write

    async def read_until(self, match: bytes) -> bytes:
        """
        Read data until match.

        Args:
            match: a sequence of bytes to match
            timeout: optional timeout in seconds. this is a temporary override
                of parameter supplied to `create`

        Returns:
            read data.
        """
        return await self._loop.run_in_executor(
            executor=self._executor,
            func=partial(self._serial.read_until, expected=match),
        )

    async def write(self, data: bytes) -> None:
        """
        Write data

        Args:
            data: data to write.

        Returns:
            None
        """
        await self._loop.run_in_executor(
            executor=self._executor,
            func=lambda: self._sync_write(data=data),
        )

    def _sync_write(self, data: bytes) -> None:
        """
        The synchronous write function

        Args:
            data: data to write.

        Returns:
            None
        """
        if self._reset_buffer_before_write:
            self._serial.reset_input_buffer()
        self._serial.write(data=data)
        self._serial.flush()

    async def open(self) -> None:
        """
        Open the connection.

        Returns: None
        """
        return await self._loop.run_in_executor(
            executor=self._executor, func=self._serial.open
        )

    async def close(self) -> None:
        """
        Close the connection

        Returns: None
        """
        return await self._loop.run_in_executor(
            executor=self._executor, func=self._serial.close
        )

    async def is_open(self) -> bool:
        """
        Check if connection is open.

        Returns: boolean
        """
        return self._serial.is_open is True

    def reset_input_buffer(self) -> None:
        """Reset the input buffer"""
        self._serial.reset_input_buffer()

    @contextlib.asynccontextmanager
    async def timeout_override(
        self, timeout_property: TimeoutProperties, timeout: Optional[float]
    ) -> AsyncGenerator[None, None]:
        """Context manager that will temporarily override the default timeout."""
        default_timeout = getattr(self._serial, timeout_property)
        override = timeout is not None and default_timeout != timeout
        try:
            if override:
                await self._loop.run_in_executor(
                    executor=self._executor,
                    func=lambda: setattr(self._serial, timeout_property, timeout),
                )
            yield
        finally:
            if override:
                await self._loop.run_in_executor(
                    executor=self._executor,
                    func=lambda: setattr(
                        self._serial, timeout_property, default_timeout
                    ),
                )


import sys
sys.stdin = open('0827_트리.txt', 'r')

V, E = map(int, input().split())
arr = list(map(int, input().split()))

L = [0] * (V + 1) # 왼쪽 노드 저장
R = [0] * (V + 1) # 오른쪽 노드 저장
P = [0] * (V + 1) # 부모 노드 저장

for i in range(0, E*2, 2):
    p, c = arr[i], arr[i+1]
    if L[p] == 0:
        L[p] = c
    else:
        R[p] = c
    P[c] = p

def inorder(v):
    if v == 0: return
    inorder(L[v])
    print(v, end=' ')
    inorder(R[v])

inorder(1)


import random
import os
import glob

import pele_platform.constants.constants as cs
from pele_platform.Errors import custom_errors
from pele_platform.Utilities.Parameters.SimulationParams.MSMParams import msm_params
from pele_platform.Utilities.Parameters.SimulationParams.GlideParams import glide_params
from pele_platform.Utilities.Parameters.SimulationParams.BiasParams import bias_params
from pele_platform.Utilities.Parameters.SimulationParams.InOutParams import inout_params
from pele_platform.Adaptive.interaction_restrictions import (
    InteractionRestrictionsBuilder,
)
from pele_platform.Utilities.Parameters.SimulationParams.PCA import pca
from pele_platform.Utilities.Parameters.SimulationParams.site_finder import site_finder
from pele_platform.Utilities.Parameters.SimulationParams.PPI import ppi
import pele_platform.Utilities.Helpers.helpers as hp

LOGFILE = '\n"simulationLogPath" : "$OUTPUT_PATH/logFile.txt",'

class SimulationParams(
    msm_params.MSMParams,
    glide_params.GlideParams,
    bias_params.BiasParams,
    inout_params.InOutParams,
    pca.PCAParams,
    site_finder.SiteFinderParams,
    ppi.PPIParams,
):
    def __init__(self, args):
        self.simulation_type(args)
        self.main_pele_params(args)
        self.singularity_params(args)
        self.main_adaptive_params(args)
        self.optative_params(args)
        self.system_preparation_params(args)
        self.ligand_params(args)
        self.anm_params(args)
        self.water_params(args)
        self.box_params(args)
        self.output_params(args)
        self.analysis_params(args)
        self.constraints_params(args)
        self.interaction_restrictions_params(args)
        self.covalent_docking_params(args)
        self.check_flags(args)

        # Create all simulation types (could be more efficient --> chnage in future)
        super().generate_msm_params(args)
        super().generate_glide_params(args)
        super().generate_bias_params(args)
        super().generate_inout_params(args)
        super().generate_pca_params(args)
        super().generate_site_finder_params(args)
        super().generate_ppi_params(args)
        # rna.RNAParams.__init__(self, args)

    def simulation_type(self, args):
        self.adaptive = (
            True if args.package in ["site_finder", "adaptive", "PPI"] else None
        )
        self.frag_pele = True if args.package == "frag" else None
        # Trick to let frag handle control fodler parameters --> Improve
        self.complexes = "$PDB" if self.frag_pele else "$COMPLEXES"
        self.frag_pele_steps = "$STEPS" if self.frag_pele else "$PELE_STEPS"
        self.output_path = "$RESULTS_PATH" if self.frag_pele else "$OUTPUT_PATH"

    def main_pele_params(self, args):
        if "*" in args.system:
            self.system = glob.glob(args.system)[0]
            self.input = glob.glob(args.system) if not args.input else args.input
        elif args.input and not args.system:
            self.system = args.input[0]
            self.input = args.input
        else:
            self.system = args.system
            self.input = args.input

        self.residue = args.residue
        self.chain = args.chain
        if self.adaptive:
            assert (
                    self.system and self.residue and self.chain
            ), "User must define input, residue and chain"
        self.debug = args.debug if args.debug else False
        self.pele_steps = (
            args.pele_steps
            if args.pele_steps
            else self.simulation_params.get("pele_steps", 8)
        )
        self.logfile = LOGFILE if args.log else ""
        self.license = (
            args.pele_license
            if args.pele_license
            else cs.DEFAULT_PELE_LICENSE
        )
        self.anm_freq = (
            args.anm_freq
            if args.anm_freq is not None
            else self.simulation_params.get("anm_freq", 4)
        )
        self.sidechain_freq = (
            args.sidechain_freq
            if args.sidechain_freq is not None
            else self.simulation_params.get("sidechain_freq", 2)
        )
        self.min_freq = (
            args.min_freq
            if args.min_freq is not None
            else self.simulation_params.get("min_freq", 1)
        )
        self.water_freq = (
            args.water_freq
            if args.water_freq is not None
            else self.simulation_params.get("water_freq", 1)
        )
        conformation_freq = (
            args.conformation_freq
            if args.conformation_freq is not None
            else self.simulation_params.get("conformation_freq", 4)
        )

        if args.ligand_conformations:
            self.conformation_freq = cs.CONFORMATION_FREQUENCY.format(conformation_freq)
        else:
            self.conformation_freq = ""

        self.temperature = (
            args.temperature
            if args.temperature
            else self.simulation_params.get("temperature", 1500)
        )
        self.sidechain_resolution = (
            args.sidechain_resolution
            if args.sidechain_resolution
            else self.simulation_params.get("sidechain_resolution", 10)
        )
        self.proximityDetection = (
            "false"
            if args.proximityDetection is False
            else self.simulation_params.get("proximityDetection", "true")
        )
        self.steric_trials = (
            args.steric_trials
            if args.steric_trials
            else self.simulation_params.get("steric_trials", 250)
        )
        self.overlap_factor = (
            args.overlap_factor
            if args.overlap_factor
            else self.simulation_params.get("overlap_factor", 0.65)
        )
        self.overlap_factor_conformation = (
            args.overlap_factor_conformation
            if args.overlap_factor_conformation
            else self.simulation_params.get("overlap_factor_conformation", 0.65)
        )
        self.steering = (
            args.steering
            if args.steering
            else self.simulation_params.get("steering", 0)
        )
        self.perturbation = (
            ""
            if args.perturbation is False
            else self.simulation_params.get("perturbation", cs.PERTURBATION)
        )
        self.perturbation_params(args)
        self.com = (
            args.com
            if args.com
            else self.simulation_params.get("COMligandConstraint", 0)
        )
        self.minimum_steps = cs.MINIMUMSTEPS if args.minimum_steps else ""
        self.conformation_perturbation = (
            ""
            if not args.ligand_conformations
            else self.simulation_params.get("conformation_perturbation", cs.CONFORMATION_PERTURBATION)
        )

    def anm_params(self, args):
        self.anm_displacement = (
            args.anm_displacement
            if args.anm_displacement
            else self.simulation_params.get("anm_displacement", 0.75)
        )
        self.anm_modes_change = (
            args.anm_modes_change
            if args.anm_modes_change
            else self.simulation_params.get("anm_modes_change", 4)
        )
        self.anm_direction = (
            args.anm_direction
            if args.anm_direction
            else self.simulation_params.get("anm_direction", "random")
        )
        self.anm_mix_modes = (
            args.anm_mix_modes
            if args.anm_mix_modes
            else self.simulation_params.get(
                "anm_mix_modes", "mixMainModeWithOthersModes"
            )
        )
        self.anm_picking_mode = (
            args.anm_picking_mode
            if args.anm_picking_mode
            else self.simulation_params.get("anm_picking_mode", "RANDOM_MODE")
        )
        self.anm_num_of_modes = (
            args.anm_num_of_modes
            if args.anm_num_of_modes
            else self.simulation_params.get("anm_num_of_modes", 6)
        )
        self.anm_relaxation_constr = (
            args.anm_relaxation_constr
            if args.anm_relaxation_constr
            else self.simulation_params.get("anm_relaxation_constr", 0.5)
        )
        self.remove_constraints = (
            args.remove_constraints
            if args.remove_constraints is not None
            else self.simulation_params.get("remove_constraints", False)
        )

    def perturbation_params(self, args):
        if self.perturbation:
            self.selection_to_perturb = (
                args.selection_to_perturb
                if args.selection_to_perturb
                else self.simulation_params.get(
                    "selection_to_perturb", cs.SELECTION_TO_PERTURB
                )
            )
            self.parameters = (
                args.parameters
                if args.parameters
                else self.simulation_params.get("params", True)
            )
            self.ligand = cs.LIGAND if self.perturbation else ""
            self.binding_energy = (
                args.binding_energy
                if args.binding_energy
                else self.simulation_params.get("binding_energy", cs.BE)
            )
            self.sasa = (
                args.sasa if args.sasa else self.simulation_params.get("sasa", cs.SASA)
            )
        else:
            self.selection_to_perturb = ""
            self.parameters = ""
            self.ligand = ""
            self.binding_energy = ""
            self.sasa = ""

        self.sidechain_radius = (
            args.sidechain_radius
            if args.sidechain_radius is not None
            else self.simulation_params.get("sidechain_radius", 6)
        )

    def main_adaptive_params(self, args):
        self.spawning = (
            args.spawning
            if args.spawning
            else self.simulation_params.get("spawning_type", "independent")
        )
        self.spawning_condition = (
            args.spawning_condition
            if args.spawning_condition
            else self.simulation_params.get("spawning_condition", None)
        )
        self.spawning_condition = (
            '"condition": "{}",'.format(self.spawning_condition)
            if self.spawning_condition
            else ""
        )
        self.density = (
            args.density
            if args.density
            else self.simulation_params.get("density", "null")
        )
        self.simulation_type = (
            args.simulation_type
            if args.simulation_type
            else self.simulation_params.get("simulation_type", "pele")
        )
        iterations = 1 if self.spawning == "independent" else 30
        self.iterations = (
            args.iterations
            if args.iterations
            else self.simulation_params.get("iterations", iterations)
        )
        self.clust_type = (
            args.clust_type
            if args.clust_type
            else self.simulation_params.get("clust_type", "rmsd")
        )
        self.cluster_values = (
            args.cluster_values
            if args.cluster_values
            else self.simulation_params.get("cluster_values", "[1.75, 2.5, 4, 6]")
        )
        self.cluster_conditions = (
            args.cluster_conditions
            if args.cluster_conditions
            else self.simulation_params.get("cluster_conditions", "[1, 0.6, 0.4, 0.0]")
        )
        self.seed = args.seed if args.seed else random.randrange(1, 70000)
        self.templates = os.path.abspath(
            os.path.join(os.path.dirname(os.path.dirname(__file__)), "PeleTemplates")
        )
        self.usesrun = "true" if args.usesrun else "false"

        if args.usesrun:
            mpi_params_name = "srunParameters"
        else:
            mpi_params_name = "mpiParameters"

        # Adaptive needs a json-like string for pele_mpi_params
        # Frag needs the direct parameter string for pele_mpi_params
        if not self.frag_pele:
            if (args.pele_mpi_params is not None and
                    args.pele_mpi_params != ""):
                self.pele_mpi_params = \
                    f'"{mpi_params_name}": "{args.pele_mpi_params}",'
            else:
                self.pele_mpi_params = ""
        else:
            if (args.pele_mpi_params is not None and
                    args.pele_mpi_params != ""):
                self.pele_mpi_params = args.pele_mpi_params
            else:
                self.pele_mpi_params = ""

    def optative_params(self, args):
        if args.forcefield is None:
            self.forcefield = self.simulation_params.get("forcefield",
                                                         "OPLS2005")
        else:
            self.forcefield = args.forcefield

        # Keep in mind the the default solvent is "VDGBNP" when using OPLS2005
        # and "OBC" when using any OpenFF force field
        if args.solvent is None:
            if 'openff' in self.forcefield.lower():
                self.solvent = 'OBC'
            else:
                self.solvent = 'VDGBNP'
        else:
            self.solvent = args.solvent

        self.verbose = (
            "true" if args.verbose else self.simulation_params.get("verbose", "false")
        )

        self.cpus = args.cpus = (
            args.cpus if args.cpus else self.simulation_params.get("cpus", 60)
        )

        self.cpus_per_mutation = args.cpus_per_mutation if args.cpus_per_mutation else 0

        self.restart = (
            args.restart
            if args.restart is not None
            else False
        )
        self.test = args.test
        # +1 to avoid being 0
        self.equil_steps = (
            args.eq_steps
            if args.eq_steps is not None
            else self.simulation_params.get("equilibration_steps", 1)
        )
        self.equilibration = (
            "true"
            if args.equilibration
            else self.simulation_params.get("equilibration", "false")
        )
        self.equilibration_mode = (
            args.equilibration_mode
            if args.equilibration_mode
            else self.simulation_params.get("equilibration_mode", "equilibrationSelect")
        )
        self.adaptive_restart = args.adaptive_restart
        self.poses = (
            args.poses
            if args.poses
            else self.simulation_params.get("poses", self.cpus - 1)
        )
        self.pele_exec = (
            args.pele_exec if args.pele_exec else cs.DEFAULT_PELE_EXEC
        )
        self.pele_data = (
            args.pele_data if args.pele_data else cs.DEFAULT_PELE_DATA
        )
        self.pele_documents = (
            args.pele_documents
            if args.pele_documents
            else cs.DEFAULT_PELE_DOCUMENTS
        )
        self.polarize_metals = args.polarize_metals if args.polarize_metals else False
        self.polarization_factor = (
            args.polarization_factor if args.polarization_factor else 2.0
        )
        self.skip_refinement = args.skip_refinement if args.skip_refinement else False
        self.bandwidth = (
            args.bandwidth
            if args.bandwidth
            else self.simulation_params.get("bandwidth", "auto")
        )
        self.clustering_method = (
            args.clustering_method
            if args.clustering_method
            else self.simulation_params.get("clustering_method", "meanshift")
        )

    def system_preparation_params(self, args):
        self.skip_prep = (
            args.skip_prep
            if args.skip_prep
            else self.simulation_params.get("skip_prep", False)
        )
        self.nonstandard = (
            args.nonstandard
            if args.nonstandard
            else self.simulation_params.get("nonstandard", [])
        )
        self.constraints = None
        self.external_constraints = (
            hp.retrieve_constraints_for_pele(args.external_constraints, self.system)
            if args.external_constraints
            else []
        )
        self.permissive_metal_constr = (
            args.permissive_metal_constr if args.permissive_metal_constr else []
        )
        self.constrain_core = (
            args.constrain_core
            if args.constrain_core
            else self.simulation_params.get("constrain_core", None)
        )
        self.constrain_core_spring = (
            args.constrain_core_spring if args.constrain_core_spring else 50.0
        )
        self.no_ppp = (
            args.no_ppp if args.no_ppp else self.simulation_params.get("no_ppp", False)
        )

    def ligand_params(self, args):
        self.spython = os.path.join(cs.SCHRODINGER, "utilities/python")
        if not os.path.exists(self.spython):
            self.spython = os.path.join(cs.SCHRODINGER, "run")
        self.mae_lig = (
            args.mae_lig
            if args.mae_lig
            else self.simulation_params.get("mae_lig", None)
        )
        self.external_template = (
            args.template
            if args.template
            else self.simulation_params.get("template", [])
        )
        self.external_rotamers = (
            args.rotamers
            if args.rotamers
            else self.simulation_params.get("rotamers", [])
        )
        self.core = args.core
        self.mtor = args.mtor
        self.n = args.n
        self.forcefield = args.forcefield if args.forcefield is not None else "OPLS2005"
        self.lig = self.mae_lig if self.mae_lig else "{}.mae".format(self.residue)
        self.gridres = args.gridres if args.gridres else self.simulation_params.get("gridres", 10)
        self.use_peleffy = args.use_peleffy if args.use_peleffy is not None else False

        # Take into account that the defaults for the parameterization method
        # are the following:
        #  - OPLS2005 when OPLS2005 force field is used
        #  - am1bcc when any OpenFF force field is used
        self.charge_parametrization_method = args.charge_parametrization_method

        self.exclude_terminal_rotamers = (
            args.exclude_terminal_rotamers
            if args.exclude_terminal_rotamers is not None
            else True
        )
        self.skip_ligand_prep = args.skip_ligand_prep if args.skip_ligand_prep else []
        self.solvent_template = args.solvent_template

        if not self.use_peleffy and self.core is None:  # plop
            self.core = -1

        if not self.use_peleffy and self.forcefield.upper() != "OPLS2005":  # plop
            raise custom_errors.IncompatibleForcefield(f"PlopRotTemp is incompatible with {self.forcefield}. Set "
                                                       f"'use_peleffy: true' in input.yaml.")
        self.ligand_conformations = args.ligand_conformations if args.ligand_conformations else []

    def water_params(self, args):
        self.water_temp = (
            args.water_temp
            if args.water_temp
            else self.simulation_params.get("water_temp", 5000)
        )
        self.water_overlap = (
            args.water_overlap
            if args.water_overlap
            else self.simulation_params.get("water_overlap", 0.78)
        )
        self.water_constr = (
            args.water_constr
            if args.water_constr
            else self.simulation_params.get("water_constr", 0)
        )
        self.water_trials = (
            args.water_trials
            if args.water_trials
            else self.simulation_params.get("water_trials", 10000)
        )

        self.allow_empty_selectors = (
            '"allowEmptyWaterSelectors": true,' if args.water_empty_selector else ""
        )
        self.water_arg = (
            hp.retrieve_all_waters(self.system)
            if args.waters == "all_waters"
            else args.waters
        )  # IDs of waters
        self.n_waters = (
            args.n_waters
            if args.n_waters
            else self.simulation_params.get("n_waters", 0)
        )
        self.waters = (
            args.waters if args.waters else self.simulation_params.get("waters", [])
        )
        self.water_radius = args.water_radius if args.water_radius else 6
        self.water_center = None
        self.water = ""
        self.water_energy = None
        self.water_ids_to_track = []

    def box_params(self, args):

        self.box = args.box if args.box is not None else True

        self.box_radius = (
            args.box_radius
            if args.box_radius
            else self.simulation_params.get("box_radius", None)
        )
        if args.box_center:
            if ":" in args.box_center:
                self.box_center = [
                    str(coord)
                    for coord in hp.get_coords_from_residue(
                        self.system, args.box_center
                    )
                ]
                self.box_center = "[" + ", ".join(self.box_center) + "]"
            else:
                self.box_center = [str(x) for x in args.box_center]
                self.box_center = (
                        "[" + ",".join([str(coord) for coord in self.box_center]) + "]"
                )
        else:
            self.box_center = self.simulation_params.get("box_center", None)

    def output_params(self, args):
        self.folder = args.folder
        self.output = args.output if args.output is not None else "output"
        self.report_name = args.report_name if args.report_name else "report"
        self.traj_name = args.traj_name if args.traj_name else "trajectory.pdb"
        self.xtc = self.traj_name.endswith(".xtc")
        self.pdb = self.traj_name.endswith(".pdb")
        self.inter_step_logger = cs.INTERSTEPLOGGER if args.inter_step_logger else ""

    def analysis_params(self, args):
        self.analyse = args.analyse if args.analyse is not None else True
        self.mae = args.mae if args.mae else False
        self.only_analysis = args.only_analysis
        self.analysis_nclust = args.analysis_nclust
        self.overwrite = args.overwrite
        self.be_column = args.be_column
        self.te_column = args.te_column
        self.limit_column = args.limit_column
        self.kde = args.kde if args.kde is not None else False
        self.kde_structs = args.kde_structs if args.kde_structs else 1000
        self.min_population = (
            args.min_population if args.min_population is not None else 0.01
        )
        self.max_top_clusters = (
            args.max_top_clusters if args.max_top_clusters is not None else 8
        )
        self.max_top_poses = (
            args.max_top_poses if args.max_top_poses is not None else 100
        )
        self.top_clusters_criterion = (
            args.top_clusters_criterion
            if args.top_clusters_criterion is not None
            else "interaction_25_percentile"
        )
        self.cluster_representatives_criterion = (
            args.cluster_representatives_criterion
            if args.cluster_representatives_criterion is not None
            else "interaction_min"
        )
        self.clustering_filtering_threshold = args.clustering_filtering_threshold if args.clustering_filtering_threshold is not None else 0.25
        self.plot_filtering_threshold = args.plot_filtering_threshold if args.plot_filtering_threshold is not None else 0.02
        self.clustering_coverage = args.clustering_coverage if args.clustering_coverage is not None else 0.75

    def constraints_params(self, args):
        """
        Sets backbone constraints parameters, such as interval and spring constants. It prioritises the user-defined
        values, then checks default values for specified constraints level (if any), then checks simulation_params for
        a specific package to finally fall back to level 1 default.
        Args:
            args (EnviroBuilder): arguments defined by the user in input.yaml

        Returns:
            Backbone CA and terminal CA spring constants, backbone CA interval.
        """
        constraints_flags = [
            "ca_interval",
            "ca_constr",
            "terminal_constr",
        ]  # list of relevant flags to check

        # Checks, if the user specified one of the built-in constraint levels.
        self.constraint_level = (
            args.constraint_level if args.constraint_level is not None else None
        )
        self.level_params = cs.constraint_levels.get(self.constraint_level, {})

        # Dictionary with default constraints params, i.e. level 1
        defaults = cs.constraint_levels.get(1, {})

        for flag in constraints_flags:
            flag_value = getattr(args, flag, None)
            if not flag_value:
                flag_value = self.level_params.get(
                    flag, self.simulation_params.get(flag, defaults[flag])
                )
            setattr(self, flag, flag_value)

    def interaction_restrictions_params(self, args):
        """
        Sets parameters for interaction restrictions.
        Fills the pele_params INTERACTION_RESTRICTIONS template with an additional section of parameters change.
        """
        if args.interaction_restrictions:
            restrictions = InteractionRestrictionsBuilder()
            restrictions.parse_interaction_restrictions(
                self.system, args.interaction_restrictions
            )
            self.met_interaction_restrictions = restrictions.metrics_to_json()
            self.interaction_restrictions = restrictions.conditions_to_json()
            self.parameters = restrictions.fill_template(self.parameters)
        else:
            self.met_interaction_restrictions = ""
            self.interaction_restrictions = ""

    def singularity_params(self, args):
        """
        Sets parameters for singularity containers.
        """
        if args.pele_mpi_params is None:
            args.pele_mpi_params = cs.DEFAULT_PELE_MPI_PARAMS

        if args.singularity_exec:
            if args.pele_mpi_params != "":
                args.pele_mpi_params = args.pele_mpi_params + ' ' + args.singularity_exec
            else:
                args.pele_mpi_params = args.singularity_exec

            args.pele_exec = "Pele_mpi" if not self.frag_pele else args.pele_mpi_params + " Pele_mpi"

    def covalent_docking_params(self, args):
        """
        Sets covalent docking parameters.
        """
        self.covalent_residue = args.covalent_residue if args.covalent_residue else None
        self.nonbonding_radius = args.nonbonding_radius if args.nonbonding_radius is not None else 20.0
        self.perturbation_trials = args.perturbation_trials if args.perturbation_trials is not None else self.simulation_params.get(
            "perturbation_trials", 10)
        self.max_trials_for_one = self.perturbation_trials * 2

        if self.covalent_residue:
            # Refinement distance should be empty for the general simulation (handled in CovalentDocking runner).
            self.refinement_angle = args.refinement_angle if args.refinement_angle is not None else self.simulation_params.get(
                "refinement_angle", 10)
            self.refinement_angle = cs.refinement_angle.format(self.refinement_angle)
            self.sidechain_perturbation = cs.SIDECHAIN_PERTURBATION
            self.covalent_sasa = cs.SASA_COVALENT.format(self.covalent_residue)
            self.residue_type = args.residue_type
        else:
            self.sidechain_perturbation = ""
            self.covalent_sasa = ""
            self.refinement_angle = ""

    @staticmethod
    def check_flags(args):
        """
        Recognises any flags incompatible with each other.
        """
        if args.mae_lig and args.template:
            mae_lig_residue = os.path.basename(os.path.splitext(args.mae_lig)[0]).upper()

            if isinstance(args.template, str):
                args.template = [args.template]

            template_residues = [os.path.basename(os.path.splitext(file)[0]).upper().rstrip("Z")
                                 for file in args.template]

            if mae_lig_residue in template_residues:
                raise custom_errors.IncompatibleYamlFlags(
                    "Argument mae_lig cannot be used in conjunction with ligand template,")


import argparse
import colorsys
import glob
import json
import os
import shutil
import sys
import tqdm

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import pycocotools.mask as rletools

from PIL import Image
from multiprocessing import Pool
# from mots_common.io import load_sequences, load_seqmap
from functools import partial
from subprocess import call

from tools.visualize import generate_colors, colormap
from tools.video_utils import frames2video

# =======================================================
# Global variables
# =======================================================
"""
known_tao_ids: set of tao ids that can be mapped exactly to coco ids.
neighbor_classes: tao classes that are similar to coco_classes.
unknown_tao_ids: all_tao_ids that exclude known_tao_ids and neighbor_classes..
"""

IOU_THRESHOLD = 0.5
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = '/'.join(ROOT_DIR.split('/')[:-1])
all_ids = set([i for i in range(1, 1231)])

# Category IDs in TAO that are known (appeared in COCO)
with open(ROOT_DIR + "/datasets/tao/coco_id2tao_id.json") as f:
    coco_id2tao_id = json.load(f)
known_tao_ids = set([v for k, v in coco_id2tao_id.items()])
# Category IDs in TAO that are unknown (comparing to COCO)
unknown_tao_ids = all_ids.difference(known_tao_ids)

# neighbor classes
with open(ROOT_DIR + "/datasets/tao/neighbor_classes.json") as f:
    coco2neighbor_classes = json.load(f)
# Gather tao_ids that can be categorized in the neighbor_classes
neighbor_classes = set()
for coco_id, neighbor_ids in coco2neighbor_classes.items():
    neighbor_classes = neighbor_classes.union(set(neighbor_ids))

# Exclude neighbor classes from unknown_tao_ids
unknown_tao_ids = unknown_tao_ids.difference(neighbor_classes)

# =======================================================
# =======================================================

class SegmentedObject:
    def __init__(self, bbox, mask, score, class_id, track_id):
        self.bbox = bbox
        self.mask = mask
        self.score = score
        self.class_id = class_id
        self.track_id = track_id

def box_IoU_xywh(boxA, boxB):
    """input box: [x,y,w,h]"""
    # convert [x,y,w,h] to [x1,y1,x2,y2]
    xA, yA, wA, hA = boxA
    boxA = [xA, yA, xA + wA, yA + hA]
    xB, yB, wB, hB = boxB
    boxB = [xB, yB, xB + wB, yB + hB]

    # determine the (x, y)-coordinates of the intersection rectangle
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    # compute the area of intersection rectangle
    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
    # compute the area of both the prediction and ground-truth
    # rectangles
    boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = interArea / float(boxAArea + boxBArea - interArea)
    # return the intersection over union value
    return iou

# from https://github.com/matterport/Mask_RCNN/blob/master/mrcnn/visualize.py
def apply_mask(image, mask, color, alpha=0.5):
    """Apply the given mask to the image.
    """
    for c in range(3):
        image[:, :, c] = np.where(mask == 1,
                                  image[:, :, c] * (1 - alpha) + alpha * color[c],
                                  image[:, :, c])
    return image

def load_txt(path):
    objects_per_frame = {}
    track_ids_per_frame = {}  # To check that no frame contains two objects with same id
    combined_mask_per_frame = {}  # To check that no frame contains overlapping masks
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            fields = line.split(",")

            frame = int(fields[0])
            if frame not in objects_per_frame:
                objects_per_frame[frame] = []
            if frame not in track_ids_per_frame:
                track_ids_per_frame[frame] = set()
            if int(fields[1]) in track_ids_per_frame[frame]:
                assert False, "Multiple objects with track id " + fields[1] + " in frame " + fields[0]
            else:
                track_ids_per_frame[frame].add(int(fields[1]))

            # class_id = int(fields[2])
            class_id = 1
            if not (class_id == 1 or class_id == 2 or class_id == 10):
                assert False, "Unknown object class " + fields[2]

            fields[12] = fields[12].strip()
            score = float(fields[6])
            mask = {'size': [int(fields[10]), int(fields[11])], 'counts': fields[12].encode(encoding='UTF-8')}
            bbox = [float(fields[2]), float(fields[3]), float(fields[4]), float(fields[5])]  # [x, y, w, h]
            # if frame not in combined_mask_per_frame:
            #   combined_mask_per_frame[frame] = mask
            # elif rletools.area(rletools.merge([combined_mask_per_frame[frame], mask], intersect=True)) > 0.0:
            #   assert False, "Objects with overlapping masks in frame " + fields[0]
            # else:
            #   combined_mask_per_frame[frame] = rletools.merge([combined_mask_per_frame[frame], mask], intersect=False)

            objects_per_frame[frame].append(SegmentedObject(
                bbox,
                mask,
                score,
                class_id,
                int(fields[1])
            ))

    return objects_per_frame

def load_sequences(seq_paths):
    objects_per_frame_per_sequence = {}
    print("Loading Sequences")
    for seq_path_txt in tqdm.tqdm(seq_paths):
        seq = seq_path_txt.split("/")[-1].replace(".txt", "")
        if os.path.exists(seq_path_txt):
            objects_per_frame_per_sequence[seq] = load_txt(seq_path_txt)
            # Sort proposals in each frame by score
            for frame_id, props in objects_per_frame_per_sequence[seq].items():
                props.sort(key=lambda p: p.score, reverse=True)
        else:
            assert False, "Can't find data in directory " + seq_path_txt

    return objects_per_frame_per_sequence

def process_sequence(seq_fpaths, tracks_folder, img_folder, output_folder, max_frames, all_frames_dict,
                     annot_frames_dict,
                     topN_proposals, gt_frame2anns, tao_id2name, only_annotated, draw_boxes=True, create_video=True):
    folder_name = tracks_folder.split("/")[-1]
    # print("Processing sequence", seq_name)
    os.makedirs(output_folder, exist_ok=True)
    # tracks = load_sequences(seq_fpaths)
    for seq_fpath in seq_fpaths:
        tracks = load_sequences([seq_fpath])
        seq_id = seq_fpath.split('/')[-1].replace(".txt", "")
        max_frames_seq = max_frames[seq_id]
        all_frames = all_frames_dict[seq_id]
        if seq_id in annot_frames_dict.keys():
            annot_frames = annot_frames_dict[seq_id]

            visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, all_frames, annot_frames, output_folder,
                                topN_proposals, gt_frame2anns, tao_id2name, 'unknown', only_annotated, draw_boxes,
                                create_video)
            visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, all_frames, annot_frames, output_folder,
                                topN_proposals, gt_frame2anns, tao_id2name, 'known', only_annotated, draw_boxes,
                                create_video)
            visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, all_frames, annot_frames, output_folder,
                                topN_proposals, gt_frame2anns, tao_id2name, 'neighbor', only_annotated, draw_boxes,
                                create_video)
        else:
            visualize_all_sequences(seq_id, tracks, max_frames_seq, all_frames, output_folder, topN_proposals, draw_boxes=False,
                                    create_video=True)

def process_sequence_coco(track_result_map, img_id2name, datasrc, img_folder, output_folder, max_frames,
                          annot_frames_dict,
                          topN_proposals, gt_frame2anns, tao_id2name, draw_boxes=True, create_video=True):
    """
    Args:
        track_result_map: Dict { img_id: List[detections] }. tracking result
        img_id2name: a Dict, that maps img_id to img_name (datasrc/video_name/frame_name)
        img_folder:
        output_folder:
        max_frames:
        annot_frames_dict:
        draw_boxes:
        create_video:

    Returns:

    """

    # print("Processing sequence", seq_name)
    os.makedirs(output_folder, exist_ok=True)

    split = img_folder.split('/')
    if "train" in split:
        gt_path = os.path.join(BASE_DIR, "data/TAO/annotations/train.json")
    elif "val" in split:
        gt_path = os.path.join(BASE_DIR, "data/TAO/annotations/validation.json")
    else:
        raise Exception("invalid img_folder")

    with open(gt_path, 'r') as f:
        gt_dict = json.load(f)
    gt_images = gt_dict["images"]
    img_name2id = dict()
    for img in gt_images:
        img_name = img["file_name"]
        if img["id"] not in img_id2name.keys():
            img_name2id[img_name] = -1
        img_name2id[img_name] = img["id"]

    # tracks: List of
    #   { video_name: {frame_id: List[SegmentedObject]}, sorted in the order of frame }
    tracks = dict()
    for video, frames in annot_frames_dict.items():
        frames.sort()
        if video not in tracks.keys():
            tracks[video] = dict()
        for idx, frame in enumerate(frames):
            frame_id = idx + 1
            tracks[video][frame_id] = list()
            frame_name = '/'.join(frame.split('/')[7:])
            img_id = img_name2id[frame_name]
            if img_id not in track_result_map.keys():
                # There are not detections in the current frame
                dets = list()
            else:
                dets = track_result_map[img_id]
            # Convert each dets to SegmentedObject
            dets.sort(key=lambda k: k['score'])
            for det in dets:
                tracks[video][frame_id].append(SegmentedObject(bbox=det["bbox"],
                                                               class_id=det["category_id"],
                                                               track_id=det["track_id"]))
            # only keep detections with top 10 highest scores.
            tracks[video][frame_id] = tracks[video][frame_id][:10]

    # Now we have tracks
    # Get the seq_id, which is just the video name
    # Get the max_frames

    for seq_id in annot_frames_dict.keys():
        annot_frames = annot_frames_dict[seq_id]
        max_frames_seq = len(annot_frames)
        visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, annot_frames, output_folder,
                            topN_proposals, draw_boxes, create_video)

def visualize_all_sequences(seq_id, tracks, max_frames_seq, all_frames, output_folder, topN_proposals, draw_boxes=False,
                            create_video=True):
    colors = colormap()
    dpi = 100.0
    frames_with_annotations = [frame.split('/')[-1] for frame in all_frames]
    all_frame_names = [frame.split('/')[-1] for frame in all_frames]
    for t in range(max_frames_seq):
        print("Processing frame", all_frames[t])
        filename_t = all_frames[t]
        img = np.array(Image.open(filename_t), dtype="float32") / 255
        img_sizes = img.shape
        fig = plt.figure()
        fig.set_size_inches(img_sizes[1] / dpi, img_sizes[0] / dpi, forward=True)
        fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
        ax = fig.subplots()
        ax.set_axis_off()

        if t + 1 in tracks[seq_id].keys():
            for obj in tracks[seq_id][t + 1][:topN_proposals]:
                color = colors[obj.track_id % len(colors)]
                if obj.class_id == 1:
                    category_name = ""
                elif obj.class_id == 2:
                    category_name = "Pedestrian"
                else:
                    category_name = "Ignore"
                    color = (0.7, 0.7, 0.7)

                if obj.class_id == 1 or obj.class_id == 2:  # Don't show boxes or ids for ignore regions
                    x, y, w, h = rletools.toBbox(obj.mask)
                    if draw_boxes:
                        import matplotlib.patches as patches
                        rect = patches.Rectangle((x, y), w, h, linewidth=1,
                                                 edgecolor=color, facecolor='none', alpha=1.0)
                        ax.add_patch(rect)
                    category_name += ":" + str(obj.track_id)
                    # ax.annotate(category_name, (x + 0.5 * w, y + 0.5 * h), color=color, weight='bold',
                    #             fontsize=7, ha='center', va='center', alpha=1.0)
                binary_mask = rletools.decode(obj.mask)
                apply_mask(img, binary_mask, color)

        ax.imshow(img)
        if not os.path.exists(os.path.join(output_folder + "/" + seq_id)):
            os.makedirs(os.path.join(output_folder + "/" + seq_id))
        fig.savefig(output_folder + "/" + seq_id + "/" + all_frame_names[t])
        plt.close(fig)
    if create_video:
        fps = 10
        frames2video(pathIn=output_folder + "/" + seq_id + '/',
                     pathOut=output_folder + "/" + seq_id + ".mp4", fps=fps)

        # Delete the frames
        shutil.rmtree(output_folder + "/" + seq_id)

def visualize_sequences(seq_id, tracks, max_frames_seq, img_folder, all_frames, annot_frames, output_folder,
                        topN_proposals, gt_frame2anns, tao_id2name, split, only_annotated=True, draw_boxes=True,
                        create_video=True):
    # colors = generate_colors(min(60, topN_proposals))
    colors = colormap()
    dpi = 100.0
    # frames_with_annotations = [frame for frame in tracks.keys() if len(tracks[frame]) > 0]
    # img_sizes = next(iter(tracks[frames_with_annotations[0]])).mask["size"]
    frames_with_annotations = [frame.split('/')[-1] for frame in annot_frames]
    all_frame_names = [frame.split('/')[-1] for frame in all_frames]
    # img_sizes = next(iter(tracks[seq_id])).bbox
    if only_annotated:
        frame_idx_with_annotations = list()
        for frame_name in frames_with_annotations:
            frame_idx = all_frame_names.index(frame_name)
            frame_idx_with_annotations.append(frame_idx)
        annot_frame_idx = 0
    for t in range(max_frames_seq):
        if only_annotated:
            if t != frame_idx_with_annotations[annot_frame_idx]:
                continue
            else:
                annot_frame_idx += 1
        print("Processing frame", all_frames[t])
        # filename_t = img_folder + "/" + seq_id + "/" + frames_with_annotations[t]
        filename_t = all_frames[t]
        img = np.array(Image.open(filename_t), dtype="float32") / 255
        img_sizes = img.shape
        fig = plt.figure()
        fig.set_size_inches(img_sizes[1] / dpi, img_sizes[0] / dpi, forward=True)
        fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
        ax = fig.subplots()
        ax.set_axis_off()

        # GT bbox and masks
        valid_cat_ids = list()
        if split == 'known':
            valid_cat_ids = known_tao_ids
        elif split == "neighbor":
            valid_cat_ids = neighbor_classes
        elif split == "unknown":
            valid_cat_ids = unknown_tao_ids
        # Filter out gt-annotations that don't belong to current split
        gt_anns_filtered = list()
        k = '/'.join(filename_t.split('/')[-2:]).replace(".jpg", "")
        for ann in gt_frame2anns[k]:
            if ann["category_id"] in valid_cat_ids:
                gt_anns_filtered.append(ann)

        # Apply bbox and maks from Proposals to image
        if t + 1 in tracks[seq_id].keys():
            for obj in tracks[seq_id][t + 1][:topN_proposals]:
                # Compare obj.bbox with every GT-bbox
                # bbox from GT and proposal is in the form of [x,y,w,h]
                matched_gt = None
                max_IoU_score = 0
                for ann in gt_anns_filtered:
                    iou = box_IoU_xywh(obj.bbox, ann['bbox'])
                    if iou >= IOU_THRESHOLD and iou > max_IoU_score:
                        matched_gt = ann
                        max_IoU_score = max(max_IoU_score, iou)
                if not matched_gt:
                    continue

                color = colors[obj.track_id % len(colors)]
                if obj.class_id == 1:
                    category_name = ""
                elif obj.class_id == 2:
                    category_name = "Pedestrian"
                else:
                    category_name = "Ignore"
                    color = (0.7, 0.7, 0.7)

                if obj.class_id == 1 or obj.class_id == 2:  # Don't show boxes or ids for ignore regions
                    x, y, w, h = obj.bbox
                    if draw_boxes:
                        rect = patches.Rectangle((x, y), w, h, linewidth=1,
                                                 edgecolor=color, facecolor='none', alpha=1.0)
                        ax.add_patch(rect)
                    category_name += ":" + str(obj.track_id)
                    # ax.annotate(category_name, (x + 0.5 * w, y + 0.5 * h), color=color, weight='bold',
                    #             fontsize=7, ha='center', va='center', alpha=1.0)
                binary_mask = rletools.decode(obj.mask)
                apply_mask(img, binary_mask, color)
            # Draw GT bbox
            color_gt = (1.0, 0.0, 0.0)
            for ann in gt_anns_filtered:
                x, y, w, h = ann['bbox']
                rect = patches.Rectangle((x, y), w, h, linewidth=5, linestyle='-.',
                                         edgecolor=color_gt, facecolor='none', alpha=0.5)
                ax.add_patch(rect)
                category_name = tao_id2name[ann["category_id"]]
                ax.annotate(category_name, (x + 0.5 * w, y + 0.5 * h), color=color_gt, weight='bold',
                            fontsize=7, ha='center', va='center', alpha=1.0)

        ax.imshow(img)
        if not os.path.exists(os.path.join(output_folder + "/" + seq_id)):
            os.makedirs(os.path.join(output_folder + "/" + seq_id))
        fig.savefig(output_folder + "/" + seq_id + "/" + all_frame_names[t])
        plt.close(fig)

    if create_video:
        fps = 1 if only_annotated else 10
        frames2video(pathIn=output_folder + "/" + seq_id + '/',
                     pathOut=output_folder + '/' + split + "/" + seq_id + ".mp4", fps=fps)

        # Delete the frames
        shutil.rmtree(output_folder + "/" + seq_id)

        # os.chdir(output_folder + "/" + seq_id)
        # call(["ffmpeg", "-framerate", "10", "-y", "-i", "%06d.jpg", "-c:v", "libx264", "-profile:v", "high", "-crf",
        #       "20",
        #       "-pix_fmt", "yuv420p", "-vf", "pad=\'width=ceil(iw/2)*2:height=ceil(ih/2)*2\'", "output.mp4"])

def load_and_preprocessing_gt(gt_path, datasrc):
    with open(gt_path, 'r') as f:
        gt = json.load(f)

    # Get map: image_id --> frame_name
    img_id2frame_name = dict()
    frame_name2ima_id = dict()
    for image in gt["images"]:
        img_id2frame_name[image["id"]] = image["file_name"]
        frame_name2ima_id[image["file_name"]] = image["id"]

    # Get map: frame_name --> annotations
    frame_name2anns = dict()
    for ann in gt["annotations"]:
        frame_path = img_id2frame_name[ann["image_id"]]
        curr_datasrc = frame_path.split('/')[1]
        if curr_datasrc == datasrc:
            video_name = frame_path.split('/')[-2]
            frame_name = frame_path.split('/')[-1].replace(".jpg", "")
            k = video_name + '/' + frame_name
            if k not in frame_name2anns.keys():
                frame_name2anns[k] = list()
            frame_name2anns[k].append(ann)

    tao_id2name = dict()
    for cat in gt["categories"]:
        tao_id2name[cat["id"]] = cat["name"]

    # There are soem frames although listed as annotated frames, but have no annotations
    with open(ROOT_DIR + '/datasets/tao/val_annotated_{}.txt'.format(datasrc), 'r') as f:
        content = f.readlines()
    for frame in content:
        frame = frame.strip()
        frame_name = '/'.join(frame.split('/')[2:]).replace(".jpg", "")
        if frame_name not in frame_name2anns.keys():
            # print(datasrc + '/' + frame_name)
            frame_name2anns[frame_name] = list()

    return frame_name2anns, tao_id2name

def main():
    # if len(sys.argv) != 5:
    #     print("Usage: python visualize_mots.py tracks_folder(gt or tracker results) img_folder output_folder seqmap")
    #     sys.exit(1)

    # tracks_folder = sys.argv[1]
    # img_folder = sys.argv[2]
    # output_folder = sys.argv[3]
    # seqmap_filename = sys.argv[4]

    parser = argparse.ArgumentParser(description='Visualization script for tracking result')
    parser.add_argument("--tracks_folder", type=str, default="/home/kloping/OpenSet_MOT/Tracking/unovost_RAFT_noReID/")
    parser.add_argument("--gt_path", type=str, default="/home/kloping/OpenSet_MOT/data/TAO/annotations/validation.json")
    parser.add_argument("--img_folder", type=str, default="/home/kloping/OpenSet_MOT/data/TAO/frames/val/")
    parser.add_argument("--datasrc", type=str, default='LaSOT')
    parser.add_argument("--phase", default="objectness", help="objectness, score or one_minus_bg_score", type=str)
    parser.add_argument("--only_annotated", action="store_true")
    parser.add_argument("--tao_subset", action="store_true", help="if only process fixed tao-subsets")
    parser.add_argument("--topN_proposals", default="1000",
                        help="for each frame, only display top N proposals (according to their scores)", type=int)
    parser.add_argument("--track_format", help="The file format of the tracking result", type=str, default='mot')
    args = parser.parse_args()

    if args.track_format == "mot":
        """
        The tracking result is stored in the mot-format.
            https://motchallenge.net/instructions/
        """
        curr_data_src = args.datasrc

        tracks_folder = args.tracks_folder + '/_' + args.phase + '/' + curr_data_src
        img_folder = args.img_folder + '/' + curr_data_src
        output_folder = args.tracks_folder + "/viz_" + args.phase + str(args.topN_proposals) + '/' + curr_data_src
        topN_proposals = args.topN_proposals
        print("For each frame, display top {} proposals".format(topN_proposals))

        if args.tao_subset:
            with open('../datasets/tao/tao_val_subset.txt', 'r') as f:
                content = f.readlines()
            content = [c.strip() for c in content]
            seqmap_filenames = [
                os.path.join(args.tracks_folder, '_' + args.phase, curr_data_src, c.split('/')[-1] + '.txt')
                for c in content if c.split('/')[0] == curr_data_src]
        else:
            seqmap_filenames = sorted(glob.glob(tracks_folder + '/*' + '.txt'))

        # Pre-processing GT
        gt_frame_name2anns, tao_id2name = load_and_preprocessing_gt(args.gt_path, curr_data_src)

        # Image path in all sequences
        all_frames = dict()  # {seq_name: List[frame_paths]}
        sequences = [fn.split('/')[-1] for fn in sorted(glob.glob(os.path.join(img_folder, '*')))]
        for video in sequences:
            fpath = os.path.join(img_folder, video)
            frames = sorted(glob.glob(fpath + '/*' + '.jpg'))
            all_frames[video] = frames

        annot_frames = dict()  # annotated frames for each sequence
        if args.only_annotated:
            # Get the annotated frames in the current sequence.
            txt_fname = ROOT_DIR + "/datasets/tao/val_annotated_{}.txt".format(curr_data_src)
            with open(txt_fname) as f:
                content = f.readlines()
            content = ['/'.join(c.split('/')[2:]) for c in content]
            annot_seq_paths = [os.path.join(img_folder, x.strip()) for x in content]

            for s in annot_seq_paths:
                seq_name = s.split('/')[-2]
                if seq_name not in annot_frames.keys():
                    annot_frames[seq_name] = []
                annot_frames[seq_name].append(s)

        max_frames = dict()
        for seq_name, frames in annot_frames.items():
            max_frames[seq_name] = len(frames)

        for seq_fpath in seqmap_filenames:
            seq_name = seq_fpath.split('/')[-1][:-4]
            max_frames[seq_name] = len(all_frames[seq_name])

        process_sequence_part = partial(process_sequence, max_frames=max_frames, all_frames_dict=all_frames,
                                        annot_frames_dict=annot_frames,
                                        tracks_folder=tracks_folder, img_folder=img_folder,
                                        gt_frame2anns=gt_frame_name2anns, tao_id2name=tao_id2name,
                                        only_annotated=args.only_annotated,
                                        output_folder=output_folder, topN_proposals=topN_proposals)
        process_sequence_part(seqmap_filenames)

    elif args.track_format == "coco":
        """
        The tracking result is stored in the coco-format.
            https://github.com/TAO-Dataset/tao/blob/master/docs/evaluation.md
        """
        max_frames = dict()  # dummy variable here that is not used.
        datasrc = "Charades"  # ["ArgoVerse", "BDD", "Charades", "LaSOT", "YFCC100M"]

        tracks_folder = "/home/kloping/OpenSet_MOT/Tracking/tao_track_results/SORT_val/sort_tao_val_results.json"
        img_folder = "/home/kloping/OpenSet_MOT/data/TAO/frames/val/"
        output_folder = "/home/kloping/OpenSet_MOT/Tracking/tao_track_results/viz/SORT_val/" + datasrc

        annot_frames = dict()  # annotated frames for each sequence
        # Get the annotated frames in the current sequence.
        txt_fname = "../data/tao/val_annotated_{}.txt".format(datasrc)
        with open(txt_fname) as f:
            content = f.readlines()
        content = ['/'.join(c.split('/')[1:]) for c in content]
        annot_seq_paths = [os.path.join(img_folder, datasrc, x.strip()) for x in content]

        for s in annot_seq_paths:
            seq_name = s.split('/')[-2]
            if seq_name not in annot_frames.keys():
                annot_frames[seq_name] = []
            annot_frames[seq_name].append(s)
            annot_frames[seq_name].sort()

        # Load tracking result json file
        with open(tracks_folder, 'r') as f:
            track_results = json.load(f)

        # Get the (train/val) GT annotation json from img_folder
        #   expected img_folder: /home/kloping/OpenSet_MOT/data/TAO/frames/val/
        split = img_folder.split('/')
        if "train" in split:
            gt_path = os.path.join(BASE_DIR, "data/TAO/annotations/train.json")
        elif "val" in split:
            gt_path = os.path.join(BASE_DIR, "data/TAO/annotations/validation.json")
        else:
            raise Exception("invalid img_folder")

        with open(gt_path, 'r') as f:
            gt_dict = json.load(f)
        gt_images = gt_dict["images"]
        img_id2name = dict()
        for img in gt_images:
            img_name = img["file_name"][4:-4]
            if img["id"] not in img_id2name.keys():
                img_id2name[img["id"]] = ''
            img_id2name[img["id"]] = img_name

        # Filter, let the images in the current datasrc remain.
        track_result_map = dict()  # {img_id: List[detections]}
        for det in track_results:
            img_name = img_id2name[det["image_id"]]
            video_name = '/'.join(img_name.split('/')[:2])  # datasrc/video_name
            curr_d = img_name.split('/')[0]
            if curr_d == datasrc:
                if det["image_id"] not in track_result_map.keys():
                    track_result_map[det["image_id"]] = list()
                track_result_map[det["image_id"]].append(det)

        process_sequence_part = partial(process_sequence_coco, max_frames=max_frames, datasrc=datasrc,
                                        annot_frames_dict=annot_frames,
                                        img_id2name=img_id2name, img_folder=img_folder, output_folder=output_folder)
        process_sequence_part(track_result_map)

    # with Pool(10) as pool:
    #     pool.map(process_sequence_part, seqmap)
    # for seq in seqmap:
    #  process_sequence_part(seq)

if __name__ == "__main__":
    main()