import gc
import logging
import os
import glob
import pandas as pd

import sys
# sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')

import time

from collections import defaultdict

import torch
import torch.nn as nn
import torch.optim as optim
from math import exp

import numpy as np
torch.backends.cudnn.benchmark = True

from matplotlib import pyplot as plt
import matplotlib as mpl
import matplotlib.patches as patches
from matplotlib import pyplot as plt
from argoverse.map_representation.map_api import ArgoverseMap
from argoverse.data_loading.argoverse_forecasting_loader import ArgoverseForecastingLoader
from argoverse.visualization.visualize_sequences import viz_sequence
avm = ArgoverseMap()

num = 10

data_path="/datasets/argoverse/val/data"
infer_path="../../inn"

import os
import sys
sys.path.append("../ddn/")
sys.path.append("./")
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader
#from bernstein import bernstesin_coeff_order10_new

from argoverse.map_representation.map_api import ArgoverseMap
from argoverse.data_loading.argoverse_forecasting_loader import ArgoverseForecastingLoader
from argoverse.visualization.visualize_sequences import viz_sequence
avm = ArgoverseMap()

def denoise(gt_x, gt_y, w = 7):
    # denoising
    gt_x_t = []
    gt_y_t = []
    for iq in range(len(gt_x)):
        if iq >= w and iq + w <= len(gt_x):
            gt_x_t.append(np.mean(gt_x[iq: iq + w]))
            gt_y_t.append(np.mean(gt_y[iq: iq + w]))
        elif iq < w:
            okx = np.mean(gt_x[w: w + w])
            gt_x_t.append(gt_x[0] + (okx - gt_x[0]) * (iq) / w)
            oky = np.mean(gt_y[w: w + w])
            gt_y_t.append(gt_y[0] + (oky - gt_y[0]) * (iq) / w)
        else:
            okx = np.mean(gt_x[len(gt_x) - w:len(gt_x) - w  + w])
            oky = np.mean(gt_y[len(gt_x) - w: len(gt_x) - w + w])
            gt_x_t.append(okx + (gt_x[-1] - okx) * (w - (len(gt_x) - iq)) / w)
            gt_y_t.append(oky + (gt_y[-1] - oky) * (w - (len(gt_y) - iq)) / w)

    gt_x = gt_x_t
    gt_y = gt_y_t
    return gt_x, gt_y

from shapely.geometry.polygon import Polygon, Point

output_dir="../results/"
t_obs=20
dt=0.3
t_obs=20
pred=False
pred_array=None
batch_size = 512
dpi=100
w,h=512,512
res=0.5
paths = glob.glob(os.path.join(data_path, "*.csv"))
color = {
    'polygon': '#e6cf93',
    'polygon-outline': '#e6cf93',
    'centerline': '#fceec7',
    'agent': 'blue',
    'av': 'grey',
    'other': 'grey',
    'outline': 'black'
}
color = {
    'polygon': 'white',
    'polygon-outline': 'white',
    'centerline': 'white',
    'agent': 'white',
    'av': 'white',
    'other': 'white',
    'outline': 'black'
}

from tqdm import tqdm
for idx in tqdm(range(len(paths))):
    if idx < 19:
        continue
    path = paths[idx]
    dff = pd.read_csv(path)

    city = dff['CITY_NAME'].values[0]

    agent_df = dff[dff['OBJECT_TYPE'] == 'AGENT']
    x_a = agent_df['X'].values
    y_a = agent_df['Y'].values
    x_a, y_a = denoise(x_a, y_a)

    av_df = dff[dff['OBJECT_TYPE'] == 'AV']
    x_av = av_df['X'].values
    y_av = av_df['Y'].values
    x_av, y_av = denoise(x_av, y_av)

    others_df = dff[dff['OBJECT_TYPE'] == 'OTHERS']
    others_dfs = np.array([v for k, v in others_df.groupby('TRACK_ID')], dtype=object)
    x_o = {}
    y_o = {}

    for other_df in others_dfs:
        x_other, y_other = other_df['X'].values, other_df['Y'].values
        x_other, y_other = denoise(x_other, y_other)
        x_o[other_df['TRACK_ID'].values[0]] = x_other
        y_o[other_df['TRACK_ID'].values[0]] = other_df['Y'].values

    # group by timestamp
    dfs = [x for _, x in dff.groupby('TIMESTAMP')]

    grids_lanes = np.zeros((20, h, w))
    grids_obstacles = np.zeros((20, h, w))
    grids_centerlines = np.zeros((20, h, w))
    grids_agent = np.zeros((20, h, w))

    total_successors = []
    current = []
    das_polygons = []
    das_polygons_mp = []
    das_ids = []

    agent_polygons = []

    others_polygons = []

    for indd in range(0, 20):
        lane_id = avm.get_nearest_centerline(np.array([x_a[indd],y_a[indd]]), city_name=city)[0].id
        current.append(lane_id)
        successors = avm.get_lane_segment_successor_ids(lane_id, city)
        if successors == None:
            continue
        for successor in successors:
            total_successors.append(successor)
            successors_2d = avm.get_lane_segment_successor_ids(successor, city)
            for successorr in successors_2d:
                if successors_2d == None:
                    continue
                total_successors.append(successorr)
        polygons = [ avm.get_lane_segment_polygon(successor, city) for successor in successors]
    current = np.unique(np.array(current))
    total_successors = np.unique(np.array(total_successors))
    for curr in current:
        current_polygon = avm.get_lane_segment_polygon(curr, city)
        das_polygons.append(current_polygon)
        das_polygons_mp.append(avm.get_lane_segment_polygon(curr, city))
        das_ids.append(curr)
#         plt.fill(current_polygon[:, 0], current_polygon[:, 1], color='white', zorder=4)
    for successor in total_successors :
        polygon = avm.get_lane_segment_polygon(successor, city)
        das_polygons.append(polygon)
        das_polygons_mp.append(avm.get_lane_segment_polygon(successor, city))
        das_ids.append(successor)
#         plt.fill(polygon[:, 0], polygon[:, 1], color='white', zorder=4)
    das_polygons_mp = np.array(das_polygons_mp)
    x_off = 75
    y_off = 75
    points = np.array([[x_a[20] - x_off, y_a[20] + y_off],[x_a[20] + x_off, y_a[20] + y_off], [x_a[20] + x_off, y_a[20] - y_off],[x_a[20] - x_off, y_a[20] - y_off],[x_a[20] - x_off, y_a[20] + y_off]])

    for ind, df in enumerate(dfs):
        agent_df = df[df['OBJECT_TYPE'] == 'AGENT']
        others_df = df[df['OBJECT_TYPE'] == 'OTHERS']
        others_dfs = [x for _, x in others_df.groupby('TRACK_ID')]
        av_df = df[df['OBJECT_TYPE'] == 'AV']

        # agent
        x_traj = agent_df['X'].values
        y_traj = agent_df['Y'].values
        offsets = [x_a[0], y_a[0]] # offsets for other agents
        others_polyon = []
        if ind < len(dfs) - 1:
            x_off = 2 #0.75
            y_off = 2.25 #1.25
            points = np.array([[x_traj[0] - x_off, y_traj + y_off],[x_traj[0] + x_off, y_traj + y_off], [x_traj[0] + x_off, y_traj - y_off],[x_traj[0] - x_off, y_traj - y_off],[x_traj[0] - x_off, y_traj + y_off]])
            theta = np.arctan2((y_a[ind + 1] - y_a[ind]) , (x_a[ind + 1] - x_a[ind])) - np.pi/2
            ww = np.zeros(points.shape)
            A = np.matrix([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
            points = points - np.array([x_traj[0], y_traj[0]])
            for i,v in enumerate(points): ww[i] = A @ points[i]
            ww[:, 0] += x_traj[0]
            ww[:, 1] += y_traj[0]
            try:
                agent_polygons.append(Polygon(ww))
            except:
                print("AGENT problem")

        for indoo, other in enumerate(others_dfs):
            x_traj = other['X'].values
            y_traj = other['Y'].values
            indo = other['TRACK_ID'].values[0]
            if ind < len(dfs) - 1 and ind < len(x_o[indo]) - 1 and ind < len(y_o[indo]) - 1:
                x_off = 2
                y_off = 2.25
                points = np.array([[x_traj[0] - x_off, y_traj + y_off],[x_traj[0] + x_off, y_traj + y_off], [x_traj[0] + x_off, y_traj - y_off],[x_traj[0] - x_off, y_traj - y_off],[x_traj[0] - x_off, y_traj + y_off]])

                theta = np.arctan2((y_o[indo][ind + 1] - y_o[indo][ind]) , (x_o[indo][ind + 1] - x_o[indo][ind])) - np.pi/2
                ww = np.zeros(points.shape)
                A = np.matrix([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
                points = points - np.array([x_traj[0], y_traj[0]])
                for i,v in enumerate(points): ww[i] = A @ points[i]
                ww[:, 0] += x_traj[0]
                ww[:, 1] += y_traj[0]
                try:
                    others_polyon.append(Polygon(ww))
                except:
                    print("OTHERS")
        others_polygons.append(others_polyon)

    sample = np.zeros((h, w))
    lx = x_a[20] - res*(h/2)
    ly = y_a[20] - res*(w/2)

#     seq_lane_props = avm.city_lane_centerlines_dict[city]
#     for lane_id, lane_props in seq_lane_props.items():
#         lane_cl = lane_props.centerline
#         if (np.min(lane_cl[:, 0]) < x_max and np.min(lane_cl[:, 1]) < y_max and np.max(lane_cl[:, 0]) > x_min and np.max(lane_cl[:, 1]) > y_min):
#             lane_centerlines.append(lane_cl)

    for i in tqdm(range(h)):
        for j in range(w):
            px = lx + i * res
            py = ly + j * res
            point_xy = Point(px, py)
            flag = 0
            for k in range(len(das_polygons)):
                if Polygon(das_polygons[k]).contains(point_xy):
                    flag = 1
            sample[j,i] = flag

            for k in range(20):
                # get obstacle polygon
                for l in range(len(others_polygons[k])):
                    if others_polygons[k][l].contains(point_xy):
                        grids_obstacles[k, j, i] = 1

                # get agent polygon
                if agent_polygons[k].contains(point_xy):
                    grids_agent[k, j, i] = 1

    print("DONE")

    print(grids_agent.shape)
    for i in range(20): grids_lanes[i] = sample
    print(str(infer_path) + "/das/{}.npy".format(idx))
    np.save(str(infer_path) + "/das/{}.npy".format(idx), grids_lanes)
    np.save(str(infer_path) + "/agents/{}.npy".format(idx), grids_agent)
    np.save(str(infer_path) + "/others/{}.npy".format(idx), grids_obstacles)


from .dbtest import (
    DbTest,
    dbconnect
)

import os
from psycopg2.extras import (
    RealDictCursor,
    RealDictRow
)

PATH_TO_SQL_DIR = os.path.abspath(
    os.path.join(
        os.path.dirname(__file__),
        "..",
        "sql"
    )
)

class TestExample(DbTest):
    @dbconnect
    def test_select_organizations(self, conn):
        self.load_fixtures(
            conn,
            os.path.join(PATH_TO_SQL_DIR, "organizations.sql")
        )

        sql = """
        SELECT * FROM organizations;
        """
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(sql)
            organizations = cur.fetchall()

            assert len(organizations) == 7

    @dbconnect
    def test_count_the_number_of_subordinates(self, conn):
        self.load_fixtures(
            conn,
            os.path.join(PATH_TO_SQL_DIR, "organizations.sql")
        )

        sql = """
        SELECT COUNT(enterprise_sales_enterprise_customers.sales_organization_id) as subordinates_count, organizations."id" from organizations
        LEFT JOIN enterprise_sales_enterprise_customers ON organizations.id=enterprise_sales_enterprise_customers.sales_organization_id
        GROUP BY enterprise_sales_enterprise_customers.sales_organization_id, organizations."id" ORDER BY organizations."id";
        """
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(sql)
            actual = cur.fetchall()
            assert len(actual) == 7
            assert actual == [
                RealDictRow(**{
                    "subordinates_count": 0,
                    "id": 1,
                })
                , RealDictRow(**{
                    "subordinates_count": 4,
                    "id": 2,
                })
                , RealDictRow(**{
                    "subordinates_count": 0,
                    "id": 3,
                })
                , RealDictRow(**{
                    "subordinates_count": 0,
                    "id": 4,
                })
                , RealDictRow(**{
                    "subordinates_count": 0,
                    "id": 5,
                })
                , RealDictRow(**{
                    "subordinates_count": 1,
                    "id": 6,
                })
                , RealDictRow(**{
                    "subordinates_count": 0,
                    "id": 7,
                })
            ]

    @dbconnect
    def test_calculate_center_of_each_segment(self, conn):
        self.load_fixtures(
            conn,
            os.path.join(PATH_TO_SQL_DIR, "japan_segments.sql")
        )

        sql = """
        SELECT sub_query.id, ST_X(sub_query.bounds_center) as longitude, ST_Y(sub_query.bounds_center) as latitude
        FROM (SELECT japan_segments.id as id, st_centroid(bounds) as bounds_center FROM japan_segments) as sub_query;
        """
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(sql)
            actual = cur.fetchall()
            assert len(actual) == 10
            assert actual == [
                RealDictRow(**{
                    "id": "KAGOSHIMA_1",
                    "longitude": 130.642228315775,
                    "latitude": 30.7045454545455,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_2",
                    "longitude": 130.694183864916,
                    "latitude": 30.7045454545455,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_3",
                    "longitude": 130.746139414057,
                    "latitude": 30.7045454545455,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_4",
                    "longitude": 129.707028431231,
                    "latitude": 30.75,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_5",
                    "longitude": 129.758983980373,
                    "latitude": 30.75,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_6",
                    "longitude": 129.810939529514,
                    "latitude": 30.75,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_7",
                    "longitude": 129.862895078655,
                    "latitude": 30.75,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_8",
                    "longitude": 129.914850627797,
                    "latitude": 30.75,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_9",
                    "longitude": 129.966806176937,
                    "latitude": 30.75,
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_10",
                    "longitude": 130.018761726079,
                    "latitude": 30.75,
                })
            ]

    @dbconnect
    def test_segments_using_geojson_boundary(self, conn):
        self.load_fixtures(
            conn,
            os.path.join(PATH_TO_SQL_DIR, "japan_segments.sql")
        )

        sql = """
        SELECT sub.id from (SELECT * from japan_segments, (SELECT ST_GeomFromEWKT('SRID=4326;POLYGON((130.27313232421875 30.519681272749402,131.02020263671875 30.519681272749402,
        131.02020263671875 30.80909017893796,130.27313232421875 30.80909017893796,130.27313232421875 30.519681272749402))') as boundary) as sub_query) as sub where ST_Contains(sub.boundary, sub.bounds)
        """
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(sql)
            actual = cur.fetchall()
            print(actual)
            assert len(actual) == 3
            assert actual == [
                RealDictRow(**{
                    "id": "KAGOSHIMA_1",
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_2",
                })
                , RealDictRow(**{
                    "id": "KAGOSHIMA_3",
                })
            ]


import json
import boto3
from botocore.exceptions import ClientError
import os

region = os.environ['AWS_REGION']
sess = boto3.session.Session(region_name=region)

# def get_bucket_name():
#     ssmClient = sess.client('ssm')
#     response = ssmClient.get_parameter(
#             Name = 'ProserveProject_S3BucketName',
#             WithDecryption = True)
#     return response['Parameter']['Value']

def lambda_handler(event, context):

    s3Client = sess.client('s3')

    # try:
    #     bucketName = get_bucket_name()
    # except ClientError as e:
    #     print(e)
    #     return {
    #         'statusCode': 500,
    #         'body': json.dumps("An error occurred")
    #     }
    bucketName = os.environ['BUCKET_NAME']

    objectKey = json.loads(event['body'])["objectKey"].strip()

    response = s3Client.delete_object(
        Bucket = bucketName,
        Key = objectKey,
        VersionId = "null",
    )

    return {
        'statusCode': 204,
        'headers': {
            'Access-Control-Allow-Headers': '*',
            'Access-Control-Allow-Origin': '*',
            'Access-Control-Allow-Methods': 'OPTIONS,POST,GET,DELETE'
        },
    }


from __future__ import absolute_import

import logging
import numpy as np

from .import utils
from .import sampling

from sklearn.preprocessing import MultiLabelBinarizer, LabelBinarizer
from sklearn.model_selection import StratifiedShuffleSplit

logger = logging.getLogger(__name__)

class Dataset(object):

    def __init__(self, inputs, labels, test_indices=None, **kwargs):
        """Encapsulates all pieces of data to run an experiment. This is basically a bag of items that makes it
        easy to serialize and deserialize everything as a unit.

        Args:
            inputs: The raw model inputs. This can be set to None if you dont want
                to serialize this value when you save the dataset.
            labels: The raw output labels.
            test_indices: The optional test indices to use. Ideally, this should be generated one time and reused
                across experiments to make results comparable. `generate_test_indices` can be used generate first
                time indices.
            **kwargs: Additional key value items to store.
        """
        self.X = np.array(inputs)
        self.y = np.array(labels)
        for key, value in kwargs.items():
            setattr(self, key, value)

        self._test_indices = None
        self._train_indices = None
        self.test_indices = test_indices

        self.is_multi_label = isinstance(labels[0], (set, list, tuple))
        self.label_encoder = MultiLabelBinarizer() if self.is_multi_label else LabelBinarizer()
        self.y = self.label_encoder.fit_transform(self.y).flatten()

    def update_test_indices(self, test_size=0.1):
        """Updates `test_indices` property with indices of `test_size` proportion.

        Args:
            test_size: The test proportion in [0, 1] (Default value: 0.1)
        """
        if self.is_multi_label:
            self._train_indices, self._test_indices = sampling.multi_label_train_test_split(self.y, test_size)
        else:
            sss = StratifiedShuffleSplit(n_splits=1, test_size=test_size)
            self._train_indices, self._test_indices = next(sss.split(self.X, self.y))

    def save(self, file_path):
        """Serializes this dataset to a file.

        Args:
            file_path: The file path to use.
        """
        utils.dump(self, file_path)

    def train_val_split(self, split_ratio=0.1):
        """Generates train and validation sets from the training indices.

        Args:
            split_ratio: The split proportion in [0, 1] (Default value: 0.1)

        Returns:
            The stratified train and val subsets. Multi-label outputs are handled as well.
        """
        if self.is_multi_label:
            train_indices, val_indices = sampling.multi_label_train_test_split(self.y, split_ratio)
        else:
            sss = StratifiedShuffleSplit(n_splits=1, test_size=split_ratio)
            train_indices, val_indices = next(sss.split(self.X, self.y))
        return self.X[train_indices], self.X[val_indices], self.y[train_indices], self.y[val_indices]

    @staticmethod
    def load(file_path):
        """Loads the dataset from a file.

        Args:
            file_path: The file path to use.

        Returns:
            The `Dataset` instance.
        """
        return utils.load(file_path)

    @property
    def test_indices(self):
        return self._test_indices

    @test_indices.setter
    def test_indices(self, test_indices):
        if test_indices is None:
            self._train_indices = np.arange(0, len(self.y))
        else:
            self._test_indices = test_indices
            self._train_indices = np.setdiff1d(np.arange(0, len(self.y)), self.test_indices)

    @property
    def train_indices(self):
        return self._train_indices

    @property
    def labels(self):
        return self.label_encoder.classes_

    @property
    def num_classes(self):
        if len(self.y.shape) == 1:
            return 1
        else:
            return len(self.labels)


from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import NoSuchElementException
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.common.by import By
from selenium import webdriver
import time

'''
제가 실행하면 보이나요
python3 test.py
쳐보실래요??
저만 실행 되나요?? 안되는데 그러면 ㅠㅠ

맥으로 하는거라 리팩토링하면서 하고있어요

지금 날짜 찾는거까진 했거든요 이제 몇일인지 찾아야해요
ㅋㅋㅋㅋㅋㅋㅋㅋㅋㅋㅋ 함 봐볼까여
'''

'''
근데 그거 그냥 코드로 박으면 안되나여? 궁금하네여 예
'''

options = webdriver.ChromeOptions()
options.add_argument("window-size=800,600")

driver = webdriver.Chrome()
wait = WebDriverWait(driver, 10)
url = "https://ticket.interpark.com/Gate/TPLogin.asp"
driver.get(url)

def interpark_login():  # 인터파크 로그인
    driver.switch_to.frame(driver.find_element(By.TAG_NAME, "iframe"))

    driver.find_element(By.ID, "userId").send_keys("chlwldnjs0416")
    driver.find_element(By.ID, "userPwd").send_keys("#Chl4689056")
    driver.find_element(By.ID, "btn_login").click()

def booking_number_site():  # 예약번호 입력 후, 입장
    driver.get(
        "http://poticket.interpark.com/Book/BookSession.asp?GroupCode="
        # + showcode_entry.get()
        + "23002291"
    )

# def date_select():
#     # Select date
#     while True:
#         driver.switch_to.frame(driver.find_element(By.ID, "ifrmBookStep"))
#         # if int(calender_entry.get()) > 0:  # 날짜 설정
#         if int(1) > 0:  # 날짜 설정
#             # for i in range(int(calender_entry.get())): # 해당 월 아닐시 +1씩 증가하여 해당 월 찾기.
#             for i in range(int(1)):  # +1씩 증가하여 해당 월 찾기.:) 날짜 설정과 같은 넘버로 해야함.
#                 driver.find_element(
#                     By.XPATH, "/html/body/div/div[1]/div[1]/div/span[3]").click()
#                 try:
#                     '''
#                     회차 클릭 해야함.
#                     '''
#                     time.sleep(100)
#                     driver.find_element(
#                         # By.XPATH, '(//*[@id="CellPlayDate"])' + "[" + date_entry.get() + "]"  # 회차 설정
#                         # By.XPATH, '(//*[@id="CellPlayDate"])' + "[" + 21 + "]").click()  # 회차 설정
#                         By.XPATH, '(//*[@id="CellPlayDate"])' + "[" + 21 + "]").click()  # 회차 설정
#                     break
#                 except NoSuchElementException:
#                     print("NoSearch")
#                     #         # link_go()
#                     #         # go()

#                     # # Select round
#                     # # round_xpath = f"/html/body/div/div[3]/div[1]/div/span/ul/li[{round_entry.get()}]/a"
#                     # round_xpath = f"/html/body/div/div[3]/div[1]/div/span/ul/li['10']/a"
#                     # wait.until(EC.element_to_be_clickable((By.XPATH, round_xpath))).click()

#                     # # Click next button
#                     # driver.switch_to.default_content()
#                     # driver.find_element(By.ID, "LargeNextBtnImage").click()

def date_select():
    day_value = 23
    # 날짜
    while True:
        driver.switch_to.frame(driver.find_element(By.ID, "ifrmBookStep"))
        if int(1) == 0:
            pass
        elif int(1) >= 1:
            for i in range(1, int(1) + 1):
                driver.find_element(
                    By.XPATH, "/html/body/div/div[1]/div[1]/div/span[3]"
                ).click()
        try:
            driver.find_element(
                By.XPATH, '(//*[@id="CellPlayDate"])' +
                "[" + day_value + "]"
            ).click()
            break
        except NoSuchElementException:
            #     # link_go()
            #     # go()
            #     print("Element 못찾음.")
            time.sleep(1111)

    # 회차
    wait.until(
        EC.element_to_be_clickable(
            (
                By.XPATH,
                "/html/body/div/div[3]/div[1]/div/span/ul/li["
                + round_entry.get()
                + "]/a",
            )
        )
    ).click()
    driver.switch_to.default_content()
    driver.find_element(By.ID, "LargeNextBtnImage").click()

def find_random_seat():  # 좌석 무작위로 설정
    driver.switch_to.default_content()
    seat_frame = driver.find_element(By.NAME, "ifrmSeat")
    driver.switch_to.frame(seat_frame)

    # wait.until(EC.presence_of_element_located(
    # ))

interpark_login()  # 예약번호 입력 후, 입장
booking_number_site()  # 예약번호 입력 후, 입장
date_select()  # 상품 날짜 찾기.

# find_random_seat() # 좌석 무작위로 설정


import pygame
import config
import math
from unit import Unit
from unit_move import UnitMove

class Enemy(Unit):
    def __init__(self, x, y):
        super().__init__(x, y)
        self.start_position = [x, y]
        self.time_till_damage = 0
        self.look = []              # [center_x, center_y, radius]
        self.score = 0

    def move_to_position(self, move_to_position):
        position = self.get_position()

        # Find direction vector (dx, dy) between enemy and player.
        dx, dy = move_to_position[0] - position[0], move_to_position[1] - position[1]
        dist = math.hypot(dx, dy)

        if dx < 0:
            self.set_direction(UnitMove.LEFT)
        else:
            self.set_direction(UnitMove.RIGHT)

        if dist > 1:
            dx, dy = dx / dist, dy / dist  # Normalize.
            # Move along this normalized vector towards the player at current speed.
            position[0] += dx * config.BAT_VELOCITY
            position[1] += dy * config.BAT_VELOCITY
            self.set_position(position[0], position[1])

    def contains_look(self, player):
        player_corners = player.get_hitbox_corners()
        enemy_look = self.get_look()
        player_center = [player_corners[0][0] + (player_corners[1][0] - player_corners[0][0])/2,
                         player_corners[0][1] + (player_corners[3][1] - player_corners[0][1])/2]
        dx = enemy_look[0] - player_center[0]
        dy = enemy_look[1] - player_center[1]
        if dx * dx + dy * dy <= enemy_look[2] * enemy_look[2]:
            return True
        return False

    def contains(self, player):
        player_hitbox = player.get_hitbox()
        enemy_hitbox = self.get_hitbox()
        time_till_damage = self.get_time_till_damage()
        if time_till_damage == 0:
            # Check player up left corner is in enemy hitbox
            if enemy_hitbox[0] <= player_hitbox[0] <= enemy_hitbox[0] + enemy_hitbox[2]:
                if enemy_hitbox[1] <= player_hitbox[1] <= enemy_hitbox[1] + enemy_hitbox[3]:
                    self.set_time_till_damage(time_till_damage + 1)
                    return True

            # Check player up right corner is in enemy hitbox
            if enemy_hitbox[0] <= player_hitbox[0] <= enemy_hitbox[0] + enemy_hitbox[2]:
                if enemy_hitbox[1] <= player_hitbox[1] + player_hitbox[3] <= enemy_hitbox[1] + enemy_hitbox[3]:
                    self.set_time_till_damage(time_till_damage + 1)
                    return True

            # Check player down right corner is in enemy hitbox
            if enemy_hitbox[0] <= player_hitbox[0] + player_hitbox[2] <= enemy_hitbox[0] + enemy_hitbox[2]:
                if enemy_hitbox[1] <= player_hitbox[1] + player_hitbox[3] <= enemy_hitbox[1] + enemy_hitbox[3]:
                    self.set_time_till_damage(time_till_damage + 1)
                    return True

            # Check player down left corner is in enemy hitbox
            if enemy_hitbox[0] <= player_hitbox[0] + player_hitbox[2] <= enemy_hitbox[0] + enemy_hitbox[2]:
                if enemy_hitbox[1] <= player_hitbox[1] <= enemy_hitbox[1] + enemy_hitbox[3]:
                    self.set_time_till_damage(time_till_damage + 1)
                    return True
        else:
            time_till_damage += 1
            if time_till_damage < config.TIME_TILL_DAMAGE:
                self.set_time_till_damage(time_till_damage + 1)
            else:
                self.set_time_till_damage(0)

        return False

    def render_look(self, screen, camera):
        look = self.get_look()
        pygame.draw.circle(screen, (0, 0, 255), [look[0] - camera[0], look[1] - camera[1]], look[2], 1)

    def set_look(self, hitbox):
        self.look = [hitbox[0] + hitbox[2] / 2, hitbox[1] + hitbox[3] / 2, config.RADIUS]

    def get_look(self):
        return self.look

    def get_start_position(self):
        return self.start_position

    def set_time_till_damage(self, time_till_damage):
        self.time_till_damage = time_till_damage

    def get_time_till_damage(self):
        return self.time_till_damage

    def move_directrion(self, dx, dy):
        pass

    def get_score(self):
        return self.score


import scrapy

class CjSpider(scrapy.Spider):
    name = 'cj'
    # allowed_domains = ['caijing.com']
    start_urls = ['https://www.dyxhw.com/']

    def parse(self, response):
        typess = response.xpath('//div[@class="nav clearfix"]/a[@class="j_ch_nav _block_news_menu"]/@href').getall()
        for one_type in typess:
            # print(one_type)
            yield scrapy.Request(url=one_type, callback=self.parse_types)

    def parse_types(self, response):
        news_links = response.xpath('//ul[@class="list14 ml10"]/li/a/@href').getall()
        for news_link in news_links:
            # print(news_link)
            yield scrapy.Request(url=news_link, callback=self.parse_detial)

    def parse_detial(self, response):
        title = response.xpath('//h1[@class="title"]/text()').get()
        contents = response.xpath('//div[@class="clearfix"]/p/text()').getall()
        content = '\n'.join(x for x in contents)
        recurse = response.xpath('//div[@class="info fl"]//tr/td/text()').get().strip()
        pubtime = response.xpath('//div[@class="info fl"]//span[@class="pubTime"]/text()').get()

        item = dict()
        item['title'] = title
        item['content'] = content
        item['pubtime'] = pubtime
        yield item

        rela_article = response.xpath('//div[@class="pic-list clearfix"]//h3/a/@href').getall()
        if rela_article:
            for rela in rela_article:
                yield scrapy.Request(url=rela, callback=self.parse_detial)


def bin(l, h):
    # dap 이라는 변수에 최종 출력값 담기
    global dap
    # 종료 조건
    if l > h:
        return
    # 중간 값 설정
    mid = (l + h) // 2
    # 반복문 돌려서 문제 조건에 따라 절단기 높이 설정 후
    # 나무 높이에서 절단기 높이를 빼준 값들을 다 더해서 fin으로 값 받기
    fin = 0
    for a in trees:
        if a > mid:
            fin += a - mid
    # 가져가야 할 나무 높이만큼 가져왔으면
    # 절단기의 높이(mid)를 dap 이라는 변수에 담기
    if fin == M:
        dap = mid
        return
    # 가져가야 할 나무의 높이보다 덜 가져오면
    # 절단기의 높이를 낮추기 --> 여기서 반대로 생각해서 계속 헤맴
    elif fin < M:
        bin(l, mid - 1)
    # 위의 케이스와 반대로 적용
    elif fin > M:
        dap = mid
        bin(mid + 1, h)

N, M = map(int, input().split())
trees = list(map(int, input().split()))
# 이진 탐색을 할 것이기 때문에 정렬
trees.sort()
dap = 0
bin(0, trees[-1])
print(dap)


import sys
from PyQt6 import QtCore, QtGui, QtWidgets
from CurConUi import Ui_MainWindow
from currency_converter import CurrencyConverter

class CurrencyConv(QtWidgets.QMainWindow):
    def __init__(self):
        super(CurrencyConv, self).__init__()

        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.init_ui()

    def init_ui(self):
        self.ui.line_new_currency.setPlaceholderText("В какую валюту перевести")
        self.ui.line_old_currency.setPlaceholderText("Из какой валюты перевести")
        self.ui.line_old_amount.setPlaceholderText("У вас было")
        self.ui.button_convert.clicked.connect(self.convert)

    #        self.ui.button_convert.setObjectName()

    def convert(self):
        converter = CurrencyConverter()
        old_currency = self.ui.line_old_currency.text().upper()
        new_currency = self.ui.line_new_currency.text().upper()
        old_amount = self.ui.line_old_amount.text()
        if old_amount.isdigit() and old_currency and new_currency:
            new_amount = round(converter.convert(int(old_amount), f"{old_currency}", f"{new_currency}"), 2)
            self.ui.line_new_amount.setText(str(new_amount))
        else:
            self.ui.line_new_amount.setText("Ошибка ввода")

if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    application = CurrencyConv()
    application.show()

    sys.exit(app.exec())


#Intro
print('This program tells you how far an object will fall in a number of seconds.')
#Input
time = int(input('Enter the falling time in seconds: '))
#Defining Function
def fallingDistance(time):
    gravity = 9.8
    distance = 1 / 2 * gravity * time**2
    return round(distance, 1)
#Loop
while time > 0:
    print('The distance the object will fall in', time, 'seconds is:', fallingDistance(time), '\n')
    time = int(input('Enter the falling time in seconds: '))


import torch
import drjit as dr
import mitsuba as mi
import sys,os,json
import importlib
sys.path.append(".")
import cv2
import numpy as np
if torch.cuda.is_available():
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)
else:
    device = torch.device("cpu")
from utils.logger import Logger
from utils.matcher import Matcher
from mitsuba.scalar_rgb import Transform4f as T
from tqdm.std import tqdm
mi.set_variant('cuda_ad_rgb')

log_level = 1

Pooler = torch.nn.AvgPool2d(kernel_size=2)
@dr.wrap_ad(source='drjit', target='torch')
def down_res_loss(st, img, img_ref):
    img = img[None,...].permute(0,3,1,2)
    img_ref = img_ref[None,...].permute(0,3,1,2)
    while st>0:
        img = Pooler(img)
        img_ref = Pooler(img_ref)
        st = st-1
    if log_level>0:
        Logger.save_img("down_res.png",img.permute(0,2,3,1)[0])
    return torch.mean((img-img_ref)**2)

if __name__=="__main__":

    method = sys.argv[1]
    config = sys.argv[2]
    Logger.init(exp_name=config+"/"+method, show=False, debug=False, path="results/",add_time=False)

    tasks = importlib.import_module(f'exp.{config}') # import specific task
    resolution = tasks.resolution #resolution
    spp = tasks.spp # spp
    scene = tasks.scene # scene
    thres = tasks.thres # for hybrid scheme
    max_depth = tasks.max_depth
    match_res = tasks.match_res

    # get target image
    if hasattr(tasks,"gt_img")==True:
        gt_img = torch.from_numpy(cv2.cvtColor(cv2.imread(tasks.gt_img),cv2.COLOR_BGR2RGB)).to(device)/255.0
        img_ref = mi.TensorXf(gt_img.reshape(-1,3))
    else:
        if hasattr(tasks,"gt_scene")==True:
            img_ref = mi.render(tasks.gt_scene, seed=0, spp=8192, sensor=0)
        else:
            img_ref = mi.render(scene, seed=0, spp=8192, sensor=0)
        img_ref = img_ref[...,:3]
        img_np = np.array(mi.util.convert_to_bitmap(img_ref))
        gt_img = torch.from_numpy(img_np).to(device)/255.0

    if log_level>0:
        Logger.save_img("gt_img.png",gt_img)

    gt_img_low= torch.from_numpy(cv2.resize(np.array(mi.util.convert_to_bitmap(img_ref)),(match_res,match_res))).to(device)/255.0

    # pixel matcher using optimal transport(Sinkhorn)
    matcher = Matcher(match_res, device)

    # get optimized parameter and transformation
    opt, apply_transformation, output, params = tasks.optim_settings()
    apply_transformation(params, opt)

    for key in opt.keys():
        dr.enable_grad(opt[key])
    params = mi.traverse(scene)

    # get init image
    img_init = mi.render(scene, params, seed=0, spp=512, sensor=0)
    init_img = torch.from_numpy(np.array( mi.util.convert_to_bitmap(img_init[...,:3]))).to(device)/255.0

    if log_level>0:
        Logger.save_img("init_img.png",init_img)

    # deal with hybrid scheme
    if method.endswith("hybrid"):
        method = method[:-7]
        integrator2 = mi.load_dict({
            'type': "prb_reparam",
            'max_depth': max_depth
        })
    else:
        thres = 10000

    # define integrator
    integrator1 = mi.load_dict({
        'type': method,
        'max_depth': max_depth
    })

    # camera settings are slightly different between EPSM and PRB.
    if method.startswith("manifold"):
        sensor_id = 1
    else:
        sensor_id = 0

    loop = tqdm(range(tasks.it))

    for it in loop:

        apply_transformation(params, opt)
        if it<thres:
            img = mi.render(scene, params, seed=it, spp=spp, integrator=integrator1, sensor=sensor_id)
        else:
            if it==thres:
                for key in opt.keys():
                    opt.reset(key)
            img = mi.render(scene, params, seed=it, spp=spp, integrator=integrator2, sensor=0)

        imgs = np.array(mi.util.convert_to_bitmap(img[...,:3]))

        if log_level>0:
            Logger.save_img(f"optim.png",imgs/255.0,flip=False)
            Logger.add_image(f"optim",imgs/255.0,flip=False)
        if log_level>1:
            Logger.save_img_2(f"optim{it}.png",imgs/255.0,flip=False)

        if img.shape[-1]==5:
            render_img = torch.from_numpy(cv2.resize(imgs,(match_res,match_res))).to(device)/255.0
            grad_ = matcher.match_Sinkhorn(render_img[...,:3].reshape(-1,3), gt_img_low[...,:3].reshape(-1,3))
            grad_ = grad_.reshape(match_res,match_res,5)
            grad_ = grad_.repeat(resolution//match_res,resolution//match_res,1)
            grad = mi.TensorXf(grad_)
            dr.backward(img*grad)
        else:
            # whether using multi-resolution loss
            # loss = down_res_loss(6-((7*it)//tasks.it),img,img_ref[...,:3])
            loss = dr.sum(dr.sqr(img - img_ref[...,:3])) / len(img)
            dr.backward(loss)

        try:
            # remove nan in grad
            dic = {}
            for key in opt.keys():
                x = dr.grad(opt[key])
                x[dr.isnan(x)] = 0
                dr.set_grad(opt[key],x)
                dic[key] = float(opt[key].torch().item())#.item()
            if log_level>1:
                Logger.save_param(f"param{it}.npy",dic)
        except:
            pass

        opt.step()
        loop.set_description(f"Iteration {it:02d}: error={output(opt)}")

    Logger.exit()
    img_final = mi.render(scene, params, seed=0, spp=8192, sensor=0)
    img_final = torch.from_numpy(np.array( mi.util.convert_to_bitmap(img_final[...,:3]))).to(device)/255.0
    if log_level>0:
        Logger.save_img(f"{sys.argv[1]}.png",img_final)

    print("finish optim")


#!/usr/bin/env python3
#
# 10. Bayesian History Matching technique (advanced use)
#
import os
from pathlib import Path

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

from GPErks.constants import DEFAULT_TMP_OUTFILE_DIR
from GPErks.perks.history_matching import Wave
from GPErks.serialization.labels import read_labels_from_file
from GPErks.serialization.path import posix_path
from GPErks.utils.array import get_minmax
from GPErks.utils.plotting import interp_col, get_col
from GPErks.utils.sampling import Sampler
from gpytorch.kernels import MaternKernel, ScaleKernel
from gpytorch.likelihoods import GaussianLikelihood
from torchmetrics import MeanSquaredError, R2Score

from GPErks.gp.data.dataset import Dataset
from GPErks.gp.experiment import GPExperiment
from GPErks.gp.mean import LinearMean
from GPErks.log.logger import get_logger
from GPErks.train.emulator import GPEmulator
from GPErks.utils.random import set_seed

def main():
    # Set logger and enforce reproducibility
    log = get_logger()
    seed = 8
    set_seed(seed)

    # Load experimental values (mean and variance) you aim to match
    exp_data_file = posix_path(os.getcwd(), "data", "example_10", "expdata.json")
    expdata = {}
    with open(exp_data_file, "r") as f:
        expdata = json.load(f)
    exp_mean = [val["mean"] for val in expdata.values()]
    exp_var = [val["var"] for val in expdata.values()]

    # Load input parameters and output features' names
    dataset_dir = Path(posix_path(os.getcwd(), "datasets", "stefano", "8p", "sham"))
    xlabels = read_labels_from_file(dataset_dir / "xlabels.txt")
    ylabels = read_labels_from_file(dataset_dir / "ylabels.txt")
    feature_idx = {key: val for val, key in enumerate(ylabels)}
    active_features = list(expdata.keys())
    active_indices = [feature_idx[f] for f in active_features]

    # Train list of univariate emulators (one for each feature to match)
    X = np.loadtxt(dataset_dir / "X.txt", dtype=float)
    Y = np.loadtxt(dataset_dir / "Y.txt", dtype=float)

    emulators = []
    for idx, feature in zip(active_indices, active_features):
        y = Y[:, idx]

        dataset = Dataset(X, y, x_labels=xlabels, y_label=feature)
        likelihood = GaussianLikelihood()
        mean = LinearMean(degree=1, input_size=dataset.input_size, bias=True)
        covar = ScaleKernel(MaternKernel(ard_num_dims=dataset.input_size))
        metrics = [MeanSquaredError(), R2Score()]

        experiment = GPExperiment(
            dataset,
            likelihood,
            mean,
            covar,
            metrics=metrics,
            seed=seed
        )
        device = "cpu"

        emulator = GPEmulator(experiment, device)
        emulator.train_auto()
        emulators.append(emulator)

    minmax = get_minmax(X)

    waveno = 1  # number of iteration we are at (wave id if you want)
    cutoff = 3.0  # threshold value for the implausibility criterion
    maxno = 1  # explained below
    # The univariate GPE of each output feature will give for each point x_i a specific implausibility measure.
    # With the current implausibility criterion, for each x_i we take the maximum implausibility across all the output
    # features. With maxno=1, the maximum is calculated across all the output features (i.e., till the last worse
    # implausibility measure). If maxno=2 --> till the previous-to-last worse implausibility measure and so on.
    # With this criterion, the worse-performing emulator (the output feature which is the least well captured) will
    # dominate the entire analysis and thus determine if a point is non-implausible or implausible

    w = Wave(
        emulator=emulators,
        Itrain=minmax,
        cutoff=cutoff,
        maxno=maxno,
        mean=exp_mean,
        var=exp_var,
    )  # instantiate the wave object

    sampler = Sampler(design="lhs", dim=X.shape[1], seed=seed)
    n_samples = 100000
    X = sampler.sample(
        n_samples,
        l_bounds=list(minmax[:, 0]),
        u_bounds=list(minmax[:, 1]),
    )

    # Run one iteration of HM, which is: apply the implausibility criterion to detect regions of non-implausible
    # and of implausible points starting from the initial samples in X
    w.find_regions(X)
    w.print_stats()  # show statistics about the two obtained spaces
    w.plot_wave(xlabels=xlabels, display="impl")  # plot the current wave of history matching (impl. measure plot)
    w.plot_wave(xlabels=xlabels, display="var")  # we can also check the accuracy of the GPEs for the current wave
    # note: if filepath=<path_to_file> flag is provided, the plot will be saved to <path_to_file>

    # How to continue on the next wave in 5 steps
    #
    # (0) Save an exact copy of the wave. We always recommend saving each wave right on completion before manipulating
    # its internal structure as you might need it later for other purposes (see Appendix - A2)
    outfiles_dir = Path(DEFAULT_TMP_OUTFILE_DIR)
    outfiles_dir.mkdir(parents=True, exist_ok=True)

    w0 = w.copy()
    w0.print_stats()
    w0.save(outfiles_dir / f"wave_{waveno}.json")

    # (1) From the current non-implausible region, select points to be simulated and points to be used as tests
    # for the next wave
    n_tests = 10000  # number of TEST points we want for the next wave
    n_simuls = 128  # number of current NIMP points we want to simulate to augment training dataset for the next wave
    n_avail_nimps = len(w0.nimp_idx)  # we currently have available only this number of NIMP points
    if n_tests + n_simuls > n_avail_nimps:  # if they are not enough
        n_total_points = n_tests + n_simuls
        w.augment_nimp(n_total_points)  # use 'cloud technique' to generate new NIMP points starting from existing ones

    # Get the requested datasets
    X_simul, X_test = w.get_nimps(n_simuls)

    # We now have all the necessary data to run the next wave: a dataset to simulate to augment the training dataset
    # and build new emulators, and new TEST points to be evaluated with the new emulators. Saving the data to files.
    np.savetxt(outfiles_dir / f"X_simul_{waveno}.txt", X_simul, fmt="%.6f")
    np.savetxt(outfiles_dir / f"X_test_{waveno}.txt", X_test, fmt="%.6f")

    w.print_stats()  # quick check on TESTS, IMP, and NIMP sets' sizes after augmentation

    # (2) Simulate the selected points X_simul
    # (3) Add the simulated points and respective results to the training dataset used in the previous wave
    # (3) Train GPEs on the new, augmented training dataset
    # (4) Start a new wave of HM, where the initial parameter space to be split into non-implausible and
    # implausible regions is no more a Latin Hypercube design but is now the non-implausible region obtained
    # (and augmented) in the previous wave (i.e., X_test)

    # Appendix
    #
    # (A1) Visual check on the datasets generated for the next wave
    X_nimp = w.NIMP
    X_test = np.loadtxt(outfiles_dir / f"X_test_{waveno}.txt", dtype=float)
    X_simul = np.loadtxt(outfiles_dir / f"X_simul_{waveno}.txt", dtype=float)

    # We will inspect only 2 dimensions of the full 8D parameter space to keep it simple
    param = [4, 5]  # select 2 dimensions
    subset_idx = list(np.random.randint(0, X_test.shape[0], size=10*X_simul.shape[0]))  # select an example portion

    colors = interp_col(get_col("blue"), 4)  # getting some blue colour variants

    # Plotting current wave NIMP + next wave TEST + next wave SIMUL
    fig, axis = plt.subplots(1, 1)
    axis.scatter(X_nimp[:, param[0]], X_nimp[:, param[1]], fc=colors[1], ec=colors[1], label=f"X_nimp of wave {waveno}")
    axis.scatter(X_test[subset_idx, param[0]], X_test[subset_idx, param[1]], fc=colors[-1], ec=colors[-1], label=f"X_test for wave {waveno+1}")
    axis.scatter(X_simul[:, param[0]], X_simul[:, param[1]], fc='r', ec='r', label=f"X_simul for wave {waveno+1}")
    axis.set_xlabel(xlabels[param[0]], fontsize=12)
    axis.set_ylabel(xlabels[param[1]], fontsize=12)
    axis.legend()
    fig.tight_layout()
    plt.show()  # TEST + SIMUL points for NEXT wave are all within NIMP space CURRENT wave

    # (A2) Loading a wave object
    # You can load a wave object by providing the same data used to instantiate the wave (emulator, Itrain, cutoff,
    # maxno, mean, var). This is normally done when you need to re-run the wave differently. Alternatively, you can load
    # the wave object with no arguments. This is normally done when you need to examine the wave internal structure.
    # Let's try loading with no arguments.
    w = Wave()
    w.load(Path(outfiles_dir) / f"wave_{waveno}.json")
    w.print_stats()  # notice that TESTS, IMP, and NIMP sets' sizes are the same as pre-augmentation

    # You can get a list of all wave object attributes by printing:
    # print(w.__dict__.keys())

    # Noteworthy attributes are:
    # W.I: implausibility measure obtained for each point in the test set
    # W.PV: percentage emulator variance over experimental variance at each point (given as a fraction)
    # W.NIMP: non-implausible region
    # W.nimp_idx: indices of the initial test set which resulted to be non-implausible
    # W.IMP: implausible region
    # W.imp_idx: indices of the initial test set which resulted to be implausible
    # W.simul_idx: indices of W.NIMP that were selected to be simulated for the next wave
    # W.nsimul_idx: indices of W.NIMP which were not selected for simulations
    # (the respective points will appear in the test set of the next wave instead)

    # The original test set is not stored as an attribute to save space. However, this information can still be
    # retrieved from the stored attributes as:
    X_test = w.reconstruct_tests()
    print((np.equal(X_test, X)).all())  # the test set of first wave was the LHD we generated initially in this script

if __name__ == "__main__":
    main()


import sys
sys.path.append('../../preprocess')
from make_pca import load_landmarks
import numpy as np
import tensorflow as tf
from pfld import predict_landmarks as pfld_predict_landmarks
from pfld_custom import predict_landmarks as pfld_custom_predict_landmarks
from skimage.color import rgb2gray
import cv2
import dlib
from skimage.transform import resize
from prepare_data import IMAGE_SIZE, view_img, resize_lmks
from skimage.transform import resize
import matplotlib
from train_pfld import normalize_data
import os
from detector import get_face_detector

matplotlib.use("TkAgg")
# IMAGE_SIZE = 224

class Rect:
    def __init__(self, t, b, l, r):
        self.t = t
        self.b = b
        self.l = l
        self.r = r
    def top(self):
        return self.t
    def bottom(self):
        return self.b
    def right(self):
        return self.r
    def left(self):
        return self.l

def predict(data, model_path, predict_fn, image_size=IMAGE_SIZE, depth_multiplier=1.0, **kwargs):
    input_shape = [None, image_size, image_size, 3]
    inputs = tf.placeholder(tf.float32, shape=input_shape, name='input_images')
    preds, _, _ = predict_fn(inputs, image_size, is_training=False, depth_multiplier=depth_multiplier, **kwargs)
    print('predict tensor = ', preds)
    saver = tf.train.Saver()
    # g = tf.get_default_graph()
    # tf.contrib.quantize.create_eval_graph(input_graph=g)
    with tf.Session() as sess:
        saver.restore(sess, model_path)
        # sess.run(tf.global_variables_initializer())
        results = sess.run(preds, feed_dict={inputs: data})
        print('landmarks = ', results)
        # print('S1 > ')
        return results

def predict_tflite(data, model_path):
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    print('input_details ', input_details[0], ' data shape ', data.shape)
    interpreter.set_tensor(input_details[0]['index'], data)
    interpreter.invoke()
    landmarks = interpreter.get_tensor(output_details[0]['index'])
    print('landmarks = ', landmarks)
    return landmarks

def crop(img, box):
    return img[box.top(): box.bottom(), box.left(): box.right()]

def crop_landmarks(landmarks, box):
    return landmarks - np.array([box.left(), box.top()])

def predict_single(img_path, model_path,
                image_size=IMAGE_SIZE,
                depth_multiplier=1.0,
                predict_fn=pfld_predict_landmarks,
                zero_mean=True,
                box_detector='dlib',
                **kwargs):
    img_size = image_size
    gt_landmark = None
    if box_detector == 'gt':
        points, imgs_sizes, imgs = load_landmarks('/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/labels_ibug_300W_train.xml')
        fn = os.path.basename(img_path)
        gt_landmark = None
        for idx, img in enumerate(imgs):
            if img.endswith(fn):
                gt_landmark = points[idx]
                break
        if gt_landmark is not None:
            min_y, max_y = gt_landmark[:,1].min(), gt_landmark[:,1].max()
            min_x, max_x = gt_landmark[:,0].min(), gt_landmark[:,0].max()
            box = Rect(min_y, max_y, min_x, max_x)
            # _, gt_landmark = crop_and_resize(, gt_landmark, image_size)
    elif box_detector == 'tf':
        detector = get_face_detector()
        l, t, r, b = detector.detect(img_path)
        box = Rect(t, b, l, r)

    # get face bound
    else:
        img = dlib.load_rgb_image(img_path)
        detector = dlib.get_frontal_face_detector()
        box = detector(img, 1)[0]

    oridata = cv2.imread(img_path)
    # if image_size ==80:
    #     oridata = cv2.cvtColor(oridata,cv2.COLOR_BGR2RGB)
    data = crop(oridata, box)
    data = resize(data, (img_size, img_size), anti_aliasing=True, mode='reflect')
    # view_img(data, None)
    normalized_data = normalize_data(data)
    if model_path.endswith('.tflite'):
        # print('using tflite model ', model_path)
        # is_unint8 = model_path.find('uint8') >= 0
        # if is_unint8:
        #     print('int model')
        #     lmks = predict_tflite((np.reshape(data, (1, *data.shape)) * 255).astype(np.uint8), model_path)[0]
        # else:
        print('float model')
        lmks = predict_tflite(np.reshape(normalized_data, (1, *normalized_data.shape)).astype(np.float32), model_path)[0]
    else:
        lmks = predict(np.reshape(normalized_data, (1, *normalized_data.shape)), model_path, predict_fn,
                        image_size=image_size,
                        depth_multiplier=depth_multiplier,
                        **kwargs)[0]
    # print('landmark = ', lmks)
    if zero_mean:
        for i in range(0, 68):
            lmks[i*2] = (lmks[i*2]/2+0.5)* image_size# (lmks[i*2]/2+0.5)*image_size
            lmks[i*2+1] = (lmks[i*2+1]/2 + 0.5) * image_size# (lmks[i*2+1]/2 + 0.5)*image_size
    else:
        for i in range(0, 68):
            lmks[i*2] = (lmks[i*2])* image_size# (lmks[i*2]/2+0.5)*image_size
            lmks[i*2+1] = (lmks[i*2+1]) * image_size# (lmks[i*2+1]/2 + 0.5)*image_size
        # print('landmarks after denorm', lmks)
    lmks = lmks.reshape((68, 2))

    view_img(data, lmks)

if __name__ == '__main__':
    # 2960256451_1.jpg
    # '/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/30427236_1.jpg'
    use_tflite = False
    model = 'pfld-custom-80-025m-saux7-x3'
    # model = 'ailab'
    if model == 'pfld-64':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-64-05m/pfld-311400' if not use_tflite else '../../data/pfld-64-quant.tflite',
            depth_multiplier=0.5,
            image_size=64)
    elif model == 'pfld-112':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-112/pfld-1983600' if not use_tflite else '../../data/pfld-112-quant.tflite',
            # '../../data/pfld-64.tflite',
            image_size=112)
    elif model == 'pfld-80':
        predict_single('/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-80-025m/pfld-449100',
            # '../../data/pfld-64.tflite',
            zero_mean=False,
            depth_multiplier=0.25,
            image_size=80)
    elif model == 'pfld-custom-80':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom/pfld-183000',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=1,
            zero_mean=True,
            image_size=80)
    elif model == 'pfld-custom-80-025m':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m/pfld-314100',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80)
    elif model == 'pfld-custom-80-025m-aux7':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-aux7/pfld-376500',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80,
            aux_start_layer='layer_7')

    elif model == 'pfld-custom-80-025m-aux7-x3':
        predict_single( '/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-aux7-x3/pfld-220000',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80,
            fc_x_n=3,
            box_detector='tf',
            aux_start_layer='layer_7')
    elif model == 'pfld-custom-80-025m-saux7-x3':
        predict_single( '/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-saux7-x3/pfld-310500',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            simple_aux=True,
            zero_mean=True,
            image_size=80,
            fc_x_n=3,
            box_detector='dlib',
            aux_start_layer='layer_7')
    elif model == 'pfld-custom-80-025m-aux7-x4-m3':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg',# '/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-aux7-x4-m3/pfld-131500',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80,
            fc_x_n=4,
            mid_conv_n=3,
            box_detector='tf',
            aux_start_layer='layer_7')
    elif model == 'pfld-custom-80-025m-aux8':
        predict_single('/home/tamvm/Downloads/test_face_tamvm_2.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/checkpoints-pfld-custom-80-025m-aux8/pfld-445500',
            predict_fn=pfld_custom_predict_landmarks,
            # '../../data/pfld-64.tflite',
            depth_multiplier=0.25,
            zero_mean=True,
            image_size=80,
            aux_start_layer='layer_8')
    else:
        use_tflite = True
        predict_single('/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/testset/3035796193_1_mirror.jpg', #'/home/tamvm/Downloads/ibug_300W_large_face_landmark_dataset/helen/trainset/2960256451_1.jpg',
            '../../data/landmark_80pose.tflite',
            normalize_lmks=True,
            # '../../data/pfld-64.tflite',
            image_size=80)


# This is necessary to find the main code
import operator
import sys

from Bomberman.bomberman.entity import MonsterEntity
from Bomberman.bomberman.sensed_world import SensedWorld

sys.path.insert(0, '../bomberman')
# Import necessary stuff
from entity import CharacterEntity
from colorama import Fore, Back
from queue import PriorityQueue

import math
from enum import Enum

class TestCharacter(CharacterEntity):

    destination = (0, 0)
    expectiDepth = 3
    minimaxDepth = 4
    bound = 4

    def do(self, wrld):
        # Your code here
        loc = (self.x, self.y)
        wrldState = self.evaluateState(wrld)
        characterState = wrldState[0]
        # exit is first destination
        self.destination = wrld.exitcell

        # If the exit is right next to us, just pick it
        if wrld.exitcell in self.getNeighbors(loc, wrld, [obstacles.EXIT, obstacles.PLAYER]):
            move = self.calculateD(loc, wrld.exitcell)
            self.move(move[0], move[1])
            return

        # There is a monster close to us
        if characterState == state.UNSAFE:
            self.place_bomb()

            # running away from stupid
            if wrldState[1][0] == 'stupid':
                v, action = self.maxvalue(wrld, loc, 0, 'stupid')
                next_move = self.calculateD(loc, action)
                self.move(next_move[0], next_move[1])

            # Running away from aggressive
            if wrldState[1][0] == 'aggressive':
                v, action = self.miniMaxvalue(wrld, -math.inf, math.inf, loc, 0, 'aggressive')
                next_move = self.calculateD(loc, action)
                self.move(next_move[0], next_move[1])

            # Running away from selfpreserving
            if wrldState[1][0] == 'selfpreserving':
                v, action = self.miniMaxvalue(wrld, -math.inf, math.inf, loc, 0, 'selfpreserving')
                next_move = self.calculateD(loc, action)
                self.move(next_move[0], next_move[1])

        # What to do when there is a bomb near us
        if characterState == state.NEAR_BOMB:
            next_move = (0, 0)
            max = 0
            name = ''
            flag = True
            if wrldState[1]:
                name = wrldState[1][0]
                flag = False
            if self.bomb_check(loc, wrld):
                for cell in self.getNeighbors(loc, wrld, [obstacles.EXIT]):
                    if not self.bomb_check(cell, wrld):
                        # predict one step ahead
                        next_move = self.calculateD(loc, cell)
                        newWrld = SensedWorld.from_world(wrld)
                        character = next(iter(newWrld.characters.values()))[0]
                        new_move = self.calculateD((character.x, character.y), (cell[0], cell[1]))
                        character.move(new_move[0], new_move[1])
                        if name != '':
                            monster = self.getMonster(newWrld, name)
                            monster.move(0, 0)
                        newerWrld = newWrld.next()[0]
                        if flag:
                            test = self.exit_utility(newerWrld)
                        else:
                            test = self.utility(newerWrld, name)
                        if test > max:
                            max = test
                            next_move = new_move
                self.move(next_move[0], next_move[1])
            else:
                self.move(0, 0)

        # What to do if we cannot currently reach the exit
        if characterState == state.BLOCKED:
            walls = []
            route = []
            reachable = False
            # Map a direct course to the exit, ignoring walls
            came_from, cost_so_far = self.AStar(wrld, loc, wrld.exitcell, [obstacles.EXIT, obstacles.WALL])
            path = wrld.exitcell
            while path != loc:
                path = came_from[path]
                route.append(path)

            # Find all the walls you have to go through
            for stepnum, step in enumerate(route):
                self.set_cell_color(step[0], step[1], Fore.RED + Back.GREEN)
                if wrld.wall_at(step[0], step[1]):
                    walls.append(route[stepnum+1])
            # Choose the closest reachable wall to the exit
            closest_wall = (0,0)
            for wall in (walls):
                new_goal = wall
                came_from, cost_so_far = self.AStar(wrld, loc, new_goal, [obstacles.EXIT])
                for path in came_from:
                    if path == new_goal:
                        closest_wall = new_goal
                        reachable = True
                        break
                if reachable: break

            self.destination = closest_wall
            # Navigate to that location
            came_from, cost_so_far = self.AStar(wrld, loc, closest_wall, [obstacles.EXIT])
            path = closest_wall
            next_m = (0, 0)
            while path != loc:
                temp = path
                path = came_from[path]
                if path == loc:
                    next_m = temp
                    break
            next_move = self.calculateD(loc, next_m)

            # Place bomb at wall -- deal with diagonal!?!
            if loc == closest_wall:
                self.place_bomb()
            else:
                self.move(next_move[0], next_move[1])

        # What to do if there are no monsters near us and we can reach the exit
        if characterState == state.SAFE:
            # Just do A star
            came_from, cost_so_far = self.AStar(wrld, loc, self.destination, [obstacles.EXIT])
            path = self.destination
            next_m = (0, 0)
            while path != loc:
                temp = path
                path = came_from[path]
                if path == loc:
                    next_m = temp
                    break
            next_move = self.calculateD(loc, next_m)

            self.move(next_move[0], next_move[1])

    # Max Value function of expecitmax
    def maxvalue(self, wrld, curr, d, name):
        # Terminal state
        if self.evaluateState(wrld)[0] == state.SAFE or d == self.expectiDepth:
            return self.utility(wrld, name), curr
        if self.evaluateState(wrld)[0] == state.DEAD:
            return -10000, curr
        v = -math.inf
        action = (0, 0)
        for a in self.getNeighbors(curr, wrld, [obstacles.EXIT]):
            # simulate a new world where we make the move
            newWrld = SensedWorld.from_world(wrld)
            character = next(iter(newWrld.characters.values()))[0]
            new_move = self.calculateD((character.x, character.y), (a[0], a[1]))
            character.move(new_move[0], new_move[1])
            monster = self.getMonster(newWrld, name)
            monster.move(0, 0)
            newerWrld = newWrld.next()[0]
            val = self.expvalue(newerWrld, a, d + 1, name)
            if val > v:
                v = val
                action = a
        return v, action

    # Expected Value part of expectimax
    def expvalue(self, wrld, act, d, name):
        if self.evaluateState(wrld)[0] == state.SAFE or d == self.expectiDepth:
            return self.utility(wrld, name)
        v = 0
        mcurr = self.getMonster(wrld, name)
        possible_moves = self.getNeighbors((mcurr.x, mcurr.y), wrld, [obstacles.PLAYER])
        for a in possible_moves:
            p = 1.0/len(possible_moves)
            # Predict a step ahead using simulated world
            newWrld = SensedWorld.from_world(wrld)
            monster = self.getMonster(newWrld, name)
            new_move = self.calculateD((monster.x, monster.y), (a[0], a[1]))
            monster.move(new_move[0], new_move[1])
            try:
                character = next(iter(newWrld.characters.values()))[0]
            except(IndexError, StopIteration):
                return -10000
            character.move(0, 0)
            newerWrld = newWrld.next()[0]
            value = self.maxvalue(newerWrld, act, d+1, name)[0]
            v = v + p*value
        return v

    # Alpha Beta Minimax
    # Max value for Alpha-Beta Pruning
    def miniMaxvalue(self, wrld, alpha, beta, curr, d, name):
        # Terminal State is we are safe or depth reached
        if self.evaluateState(wrld)[0] == state.SAFE or d == self.minimaxDepth:
            return self.utility(wrld, name), curr
        if self.evaluateState(wrld)[0] == state.DEAD:
            return -10000, curr
        v = -math.inf
        action = (0, 0)
        for a in self.getNeighbors(curr, wrld, [obstacles.EXIT, obstacles.PLAYER]):
            # Simulate a new world where we made that action
            newWrld = SensedWorld.from_world(wrld)
            character = next(iter(newWrld.characters.values()))[0]
            new_move = self.calculateD((character.x, character.y), (a[0], a[1]))
            monster = self.getMonster(newWrld, name)
            character.move(new_move[0], new_move[1])
            monster.move(0, 0)
            newerWrld = newWrld.next()[0]
            val = self.minvalue(newerWrld, alpha, beta, a, d+1, name)
            if val > v:
                v = val
                action = a
            if v >= beta:
                return v, a
            alpha = max(alpha, v)
        return v, action

    # Min value for Minimax Alpha-Beta Pruning
    def minvalue(self, wrld, alpha, beta, act, d, name):
        # Terminal State is we are safe or depth reached
        if self.evaluateState(wrld)[0] == state.SAFE or d == self.minimaxDepth:
            return self.utility(wrld, name)
        v = math.inf
        mcurr = self.getMonster(wrld, name)
        possible_moves = self.getNeighbors((mcurr.x, mcurr.y), wrld, [obstacles.PLAYER, obstacles.EXIT, obstacles.MONSTER])
        for a in possible_moves:
            # Simulate a new world where we made that action
            newWrld = SensedWorld.from_world(wrld)
            monster = self.getMonster(newWrld, name)
            new_move = self.calculateD((monster.x, monster.y), (a[0], a[1]))
            monster.move(new_move[0], new_move[1])
            try:
                character = next(iter(newWrld.characters.values()))[0]
            except(IndexError, StopIteration):
                return -10000
            character.move(0, 0)
            newerWrld = newWrld.next()[0]
            val, act = self.miniMaxvalue(newerWrld, alpha, beta, act, d + 1, name)
            v = min(v, val)
            if v <= alpha:
                return v
            beta = min(beta, v)
        return v

    # Main utility function for terminal states
    def utility(self, wrld, name):
        # Utility for stupid monster
        if name == 'stupid':
            return 6*(1/(1 + self.exit_utility(wrld))) - 1*(1/((1 + self.monster_utility(wrld, name))**2))
        # Utility for non-stupid monster
        else:
            return 20 * (1 / (1 + self.exit_utility(wrld))) - 50 * (1 / ((1 + self.monster_utility(wrld, name)) ** 2)) + self.dpangle(wrld, name)

    # Calculate Vector between us, the monster, and the exit
    def dpangle(self, wrld, name):
        try:
            chara = next(iter(wrld.characters.values()))
            character = chara[0]
        except (IndexError, StopIteration):
            return -10
        # Vector for character to exit
        e = self.destination
        loc = (character.x, character.y)
        ce = tuple(map(operator.sub, e, loc))
        eu = self.calculateH(e, loc)

        if ce == (0, 0) or eu == 0:
            return 10000
        # Vector for character to monster
        monster = self.getMonster(wrld, name)
        mu = self.calculateH((monster.x, monster.y), loc)
        cm = tuple(map(operator.sub, (monster.x, monster.y), loc))
        if cm == (0, 0) or mu == 0:
            return -10000
        # Dot product
        dp = (ce[0] * cm[0]) + (ce[1] * cm[1])
        cosangle = dp / (eu * mu)

        try:
             angle = math.degrees(math.acos(cosangle))
        except(ValueError):
            return -10

        if self.exit_utility(wrld) <= 4:
            return 10
        # Return values based on if it is higher or lower than 90 degrees
        if angle >= 90:
            return eu
        else:
            return -mu

    # Gets the monster in the current world with a name
    def getMonster(self, wrld, name):
        for monster in list(wrld.monsters.values()):
            if monster[0].name == name:
                return monster[0]
        return MonsterEntity('dead', [0], 0, 0)

    # Utility function for the distance to the exit
    def exit_utility(self, wrld):
        try:
            chara = next(iter(wrld.characters.values()))
            character = chara[0]
        except (IndexError, StopIteration):
            return 10
        loc = (character.x, character.y)
        e = self.destination
        exit_came_from, exit_cost_so_far = self.AStar(wrld, loc, (e[0], e[1]), [obstacles.EXIT])
        counter = 0
        path = (e[0], e[1])
        while path != loc:
            try:
                path = exit_came_from[path]
            except (KeyError):
                return self.calculateH(loc, e)
            counter += 1
        if counter == -1:
            return counter
        return counter

    # Utility function for the distance to the monster
    def monster_utility(self, wrld, name):
        try:
            chara = next(iter(wrld.characters.values()))
            character = chara[0]
        except (IndexError, StopIteration):
            return -10
        m = self.getMonster(wrld, name)
        if m.name == 'dead':
            return 100
        loc = (character.x, character.y)
        mloc = (m.x, m.y)
        monster_came_from, monster_cost_so_far = self.AStar(wrld, loc, mloc, [obstacles.MONSTER, obstacles.PLAYER, obstacles.EXIT])
        counter = 0
        path = mloc
        while path != loc:
            try:
                path = monster_came_from[path]
            except (KeyError):
                return 100
            counter += 1
        return counter

    # A Star algorithm
    def AStar(self, wrld, start, goal, list_of_e):
        frontier = PriorityQueue()
        frontier.put((0, start))
        came_from = {}
        cost_so_far = {}
        came_from[start] = None
        cost_so_far[start] = 0

        while not frontier.empty():

            current = frontier.get()[1]
            if current == goal:
                break

            for next in self.getNeighbors(current, wrld, list_of_e):
                new_cost = cost_so_far[current] + self.calculateH(next, current)
                if next not in cost_so_far or new_cost < cost_so_far[next]:
                    cost_so_far[next] = new_cost
                    priority = new_cost + self.calculateH(goal, next)
                    frontier.put((priority, next))
                    came_from[next] = current
        return came_from, cost_so_far

    # Heuristic calculation - returns euclidean distance
    def calculateH(self, loc1, loc2):
        (x1, y1) = loc1
        (x2, y2) = loc2
        return math.sqrt(((loc1[0] - loc2[0]) ** 2) + ((loc1[1] - loc2[1]) ** 2))

    # Calculates the dx and dy between two locations
    def calculateD(self, loc1, loc2):
        (x1, y1) = loc1
        (x2, y2) = loc2
        return ((x2 - x1), (y2 - y1))

    # Returns the neighbors of a particular location according to the obstacles passed in - obstacles passed in ARE AVAILABLE to be considered neighbors
    def getNeighbors(self, loc, wrld, list_of_e):
        list_of_N = []
        for dx in [-1, 0, 1]:
            # Avoid out-of-bound indexing
            if (loc[0] + dx >= 0) and (loc[0] + dx < wrld.width()):
                # Loop through delta y
                for dy in [-1, 0, 1]:
                    # Make sure the monster is moving
                    if (dx != 0) or (dy != 0):
                        # Avoid out-of-bound indexing
                        if (loc[1] + dy >= 0) and (loc[1] + dy < wrld.height()):
                            # No need to check impossible moves
                            if obstacles.EXIT in list_of_e:
                                if wrld.exit_at(loc[0] + dx, loc[1] + dy):
                                    list_of_N.append((loc[0] + dx, loc[1] + dy))
                            if obstacles.MONSTER in list_of_e:
                                if wrld.monsters_at(loc[0] + dx, loc[1] + dy):
                                    list_of_N.append((loc[0] + dx, loc[1] + dy))
                            if obstacles.PLAYER in list_of_e:
                                if wrld.characters_at(loc[0] + dx, loc[1] + dy):
                                    list_of_N.append((loc[0] + dx, loc[1] + dy))
                            if obstacles.WALL in list_of_e:
                                if wrld.wall_at(loc[0] + dx, loc[1] + dy):
                                    list_of_N.append((loc[0] + dx, loc[1] + dy))
                            if wrld.empty_at(loc[0] + dx, loc[1] + dy):
                                list_of_N.append((loc[0] + dx, loc[1] + dy))
        return list_of_N

    # Checks if location is in range of a bomb
    def bomb_check(self, loc, wrld):
        bomb_range = wrld.expl_range
        for dx in range(-bomb_range, bomb_range):
            # Avoid out-of-bound indexing
            if (loc[0] + dx >= 0) and (loc[0] + dx < wrld.width()):
                if wrld.bomb_at((loc[0] + dx), loc[1]):
                    return True
        for dy in range(-bomb_range, bomb_range):
            # Avoid out-of-bound indexing
            if (loc[1] + dy >= 0) and (loc[1] + dy < wrld.height()):
                if wrld.bomb_at(loc[0], (loc[1] + dy)):
                    return True
        return False

    # Checks if location is in range of an explosion
    def expl_check(self, loc, wrld):
        bomb_range = wrld.expl_range
        for dx in range(-bomb_range, bomb_range):
            # Avoid out-of-bound indexing
            if (loc[0] + dx >= 0) and (loc[0] + dx < wrld.width()):
                if wrld.explosion_at((loc[0] + dx), loc[1]):
                    return True
        for dy in range(-bomb_range, bomb_range):
            # Avoid out-of-bound indexing
            if (loc[1] + dy >= 0) and (loc[1] + dy < wrld.height()):
                if wrld.explosion_at(loc[0], (loc[1] + dy)):
                    return True
        return False

    #Returns states and potentially a list of threats
    def evaluateState(self, wrld):
        monsters = []
        try:
            chara = next(iter(wrld.characters.values()))
            character = chara[0]
        except (IndexError, StopIteration):
            return state.DEAD, []
        try:
            monsters = list(wrld.monsters.values())
        except (StopIteration):
            pass

        loc = (character.x, character.y)

        counters = {}
        #Calculate each distance to the monster
        for monster in monsters:
            m = monster[0]
            monsterType = m.name
            mloc = (m.x, m.y)
            monster_came_from, monster_cost_so_far = self.AStar(wrld, loc, mloc, [obstacles.MONSTER, obstacles.PLAYER, obstacles.EXIT])
            counter = 0
            path = mloc
            while path != loc:
                try:
                    path = monster_came_from[path]
                except (KeyError):
                    counter = 100
                    break
                counter += 1
            counters[monsterType] = counter
        counts = [(k, v) for k, v in counters.items() if v <= 4]
        flag = False
        monsterTypes = []
        for count in counts:
            if count[1] <= self.bound:
                flag = True
                monsterTypes.append((count[0], count[1]))
        threats = []
        # Sort the monster list in order of closest
        monsterTypes.sort(key=lambda x: x[1])
        for monster in monsterTypes:
            threats.append(monster[0])
        if flag:
            return state.UNSAFE, threats

        if (wrld.bombs or wrld.explosions):
            return state.NEAR_BOMB, []

        # Does safe path exist?
        came_from, cost_so_far = self.AStar(wrld, loc, wrld.exitcell, [obstacles.EXIT])
        for path in came_from:
            if (path == wrld.exitcell):
                return state.SAFE, []
        return state.BLOCKED, []

class state(Enum):
    SAFE = 1
    UNSAFE = 2
    DEAD = 3
    NEAR_BOMB = 4
    BLOCKED = 5

class obstacles(Enum):
    EXIT = 1
    MONSTER = 2
    WALL = 3
    BOMB = 4
    EXPLOSION = 5
    PLAYER = 6


def factorial(n):
    """Return th factorial of N, a positive integer."""
    if n == 1:
        return 1
    return n * factorial(n-1)

def recursive_multiplication(m, n):
    if n == 1:
        return m
    return m + recursive_multiplication(m, n-1)

def is_prime(n):
    def helper(n, m):
        if m == 1:
            return True
        return n % m != 0 and helper(n, m-1)

    if n == 1:
        return False
    else:
        return helper(n, n - 1)

def hailstone(n):
    def helper(n, count):
        print(n)
        count += 1
        if n == 1:
            return count
        if n % 2 == 0:
            return helper(n//2, count)
        else:
            return helper(n*3+1, count)

    return helper(n, 0)

def merge(n1, n2):
    if n1 == 0:
        return n2
    if n2 == 0:
        return n1
    if n1 % 10 < n2 % 10:
        return n1 % 10 + 10 * merge(n1//10, n2)
    else:
        return n2 % 10 + 10 * merge(n1, n2//10)


from pynvml import *
import logging

from datasets import load_dataset
from datasets import ClassLabel
from transformers import LukeTokenizer, LukeModel, LukeForEntityPairClassification, TrainingArguments, Trainer
import torch
from tqdm import trange
# construir função que converta spans de relativos a frase para globais
import random
import os
import json

def print_gpu_utilization():
    nvmlInit()
    handle = nvmlDeviceGetHandleByIndex(0)
    info = nvmlDeviceGetMemoryInfo(handle)
    print(f"GPU memory occupied: {info.used//1024**2} MB.")

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: val[idx].clone().detach() for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

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

def convert_spans(item):
    sents = []
    sent_map = []

    entities = item["vertexSet"]
    entity_start, entity_end = [], []
    mention_types = []
    entity_spans = []
    for entity in entities:
        for mention in entity:
            if mention["sent_id"] != 0:
                current_id = mention["sent_id"]
                mention["pos"] = [sum(len(s) for s in item["sents"][:current_id])+mention["pos"][0],
                    sum(len(s) for s in item["sents"][:current_id])+mention["pos"][1]]
                mention["sent_id"] = 0

            pos = mention["pos"]
            mention_types.append(mention['type'])
            entity_spans.append(pos)
    item["vertexSet"] = entities
    return item, entity_spans

def load_examples_test(dataset):
    examples = []

    for i, item in enumerate(dataset["validation"]):
        concat_tokens = []
        counter = 0
        converted_item, entity_spans = convert_spans(item)
        tokens = item["sents"]
        for j in range(len(tokens)):
            concat_tokens += tokens[j]
        del j
        tokens = concat_tokens
        del concat_tokens

        # new
        text = ""
        cur = 0
        new_char_spans = [0]*len(entity_spans)
        entity_spans.sort(key=lambda y:y[0])
        for target_entity in entity_spans:
            tamanho_texto = len(text)
            text += " ".join(tokens[cur: target_entity[0]])
            if text:
                text += " "
            char_start = len(text)
            text += " ".join(tokens[target_entity[0]: target_entity[1]])
            char_end = len(text)
            new_char_spans[counter] = (char_start, char_end)
            text += " "
            cur = target_entity[1]
            counter+=1
        text += " ".join(tokens[cur:])
        text = text.rstrip()
        # get true labels
        labels_pairs = tuple(zip(item["labels"]["head"], item["labels"]["tail"], item["labels"]["relation_id"]))
        entity_spans = [tuple(l) for l in entity_spans]
        oldToNewPos =  dict(zip(entity_spans, new_char_spans))
        entities = item["vertexSet"]
        correlations = []
        for pair in labels_pairs:
            for head in entities[pair[0]]:
                if tuple(head["pos"]) in oldToNewPos:
                    head["pos"]=oldToNewPos[tuple(head["pos"])]
                for tail in entities[pair[1]]:
                    if tuple(tail["pos"]) in oldToNewPos:
                        tail["pos"] = oldToNewPos[tuple(tail["pos"])]
                    pack = tuple((head["pos"], tail["pos"], pair[2]))
                    correlations += (pack),
            item["vertexSet"] = entities
            examples.append(dict(
                text=text,
                entity_spans= [d[:][:-1] for d in correlations],
                labels = [d[:][-1] for d in correlations]
            ))
        return examples

def load_examples_competition(dataset):
    examples = []

    for i, item in enumerate(dataset["test"]):
        concat_tokens = []
        counter = 0
        converted_item, entity_spans = convert_spans(item)
        tokens = item["sents"]
        for j in range(len(tokens)):
            concat_tokens += tokens[j]
        del j
        tokens = concat_tokens
        del concat_tokens

        # new
        text = ""
        cur = 0
        new_char_spans = [0]*len(entity_spans)
        entity_spans.sort(key=lambda y:y[0])
        for target_entity in entity_spans:
            tamanho_texto = len(text)
            text += " ".join(tokens[cur: target_entity[0]])
            if text:
                text += " "
            char_start = len(text)
            text += " ".join(tokens[target_entity[0]: target_entity[1]])
            char_end = len(text)
            new_char_spans[counter] = (char_start, char_end)
            text += " "
            cur = target_entity[1]
            counter+=1
        text += " ".join(tokens[cur:])
        text = text.rstrip()
        aux_head = 0
        aux_tail = 0

        labels_pairs = []
        # get true labels
        for head_id in range(len(item["vertexSet"])):
            for tail_id in range(len(item["vertexSet"])):
                if (head_id!=tail_id):
                    labels_pair = tuple([head_id, tail_id , "Na"])
                    labels_pairs.append(labels_pair)

        entity_spans = [tuple(l) for l in entity_spans]
        oldToNewPos =  dict(zip(entity_spans, new_char_spans))
        entities = item["vertexSet"]
        correlations = []
        for pair in labels_pairs:
            head = random.choice(entities[pair[0]])
            tail = random.choice(entities[pair[1]])
            entity_head_id = pair[0]
            entity_tail_id = pair[1]
            rel = pair[2]

            if tuple(head["pos"]) in oldToNewPos:
                head["pos"]=oldToNewPos[tuple(head["pos"])]
            if tuple(tail["pos"]) in oldToNewPos:
                tail["pos"] = oldToNewPos[tuple(tail["pos"])]
            pack = tuple((head["pos"], tail["pos"], pair[2], tuple([entity_head_id, entity_tail_id]), item["title"]))

            item["vertexSet"] = entities
            examples.append(dict(
                text=text,
                entity_spans= pack[:2],
                labels = pack[2],
                idxs_entity_pair = pack[3],
                title = pack[4]
            ))
    return examples

torch.cuda.empty_cache()

dataset = load_dataset("docred")
max_value = 0
#for i, item in enumerate(dataset["train_annotated"]):
#    total_text_len = 0
#    tokens = item["sents"]
#    num_relations = len(item["labels"]["head"])

class ModifiedClassicLuke(LukeForEntityPairClassification):
    def __init__(self, config):
        super().__init__(config)
        self.classifier = torch.nn.Linear(in_features = 2048, out_features = 97, bias = True)

logging.info("Loading data and finetuned dataset for CLASSIC LUKE")
# FAZER LOAD DO MODEL FINETUNED DE 3 EPOCHS
model = ModifiedClassicLuke.from_pretrained("model_finetuned_classic")
tokenizer = LukeTokenizer.from_pretrained("studio-ousia/luke-large-finetuned-tacred")

test_examples = load_examples_competition(dataset)
maximum = 0
max_seq = 0

logging.info("Memory before choosing GPU")
#torch.cuda.empty_cache()

########################## Choose GPU ########################
# set the GPU device to use
cuda_device= 0 # mudar para 0 para dar o cuda
if cuda_device < 0:
    device = torch.device("cpu")
else:
    device = torch.device(f"cuda:{cuda_device}")
#model = model.to(device)
#model.eval()

# Convert to inputs
for batch_start_idx in trange(0, len(test_examples), len(test_examples)):
    batch_examples = test_examples[batch_start_idx:batch_start_idx+len(test_examples)]
    texts = [example["text"] for example in batch_examples]
    entity_spans = [example["entity_spans"] for example in batch_examples]
    #gold_labels = [example["labels"] for example in batch_examples]
    idxs_entity_pair = [example["idxs_entity_pair"] for example in batch_examples]
    titles = [example["title"] for example in batch_examples]

for i in range(len(entity_spans)):
    entity_spans[i] = list(entity_spans[i])

del batch_examples

logging.info("Removing too big examples!!")
num_rejected = 0
clean_texts = []
clean_ents = []
clean_idxs_entity_pairs = []
clean_titles = []
tokenizer2 = LukeTokenizer.from_pretrained("studio-ousia/luke-large")
for ix in range(len(texts)):
    input = tokenizer2(texts[ix])
    if len(input.data["input_ids"]) > 500:
        num_rejected+=1
        continue
    clean_texts.append(texts[i])
    clean_ents.append(entity_spans[ix])
    clean_idxs_entity_pairs.append(idxs_entity_pair)
    clean_titles.append(titles)
texts = clean_texts
entity_spans = clean_ents
idxs_entity_pair = clean_idxs_entity_pairs
titles = clean_titles

torch.cuda.empty_cache()

relations_code_list = ["P1376",
                "P607",
                "P136",
                "P137",
                "P131",
                "P527",
                "P1412",
                "P206",
                "P205",
                "P449",
                "P127",
                "P123",
                "P86",
                "P840",
                "P355",
                "P737",
                "P740",
                "P190",
                "P576",
                "P749",
                "P112",
                "P118",
                "P17",
                "P19",
                "P3373",
                "P6",
                "P276",
                "P1001",
                "P580",
                "P582",
                "P585",
                "P463",
                "P676",
                "P674",
                "P264",
                "P108",
                "P102",
                "P25",
                "P27",
                "P26",
                "P20",
                "P22",
                "Na",
                "P807",
                "P800",
                "P279",
                "P1336",
                "P577",
                "P570",
                "P571",
                "P178",
                "P179",
                "P272",
                "P170",
                "P171",
                "P172",
                "P175",
                "P176",
                "P39",
                "P30",
                "P31",
                "P36",
                "P37",
                "P35",
                "P400",
                "P403",
                "P361",
                "P364",
                "P569",
                "P710",
                "P1344",
                "P488",
                "P241",
                "P162",
                "P161",
                "P166",
                "P40",
                "P1441",
                "P156",
                "P155",
                "P150",
                "P551",
                "P706",
                "P159",
                "P495",
                "P58",
                "P194",
                "P54",
                "P57",
                "P50",
                "P1366",
            "P1365",
            "P937",
            "P140",
            "P69",
            "P1198",
            "P1056"]

c2l = ClassLabel(num_classes = 97, names = relations_code_list)
label_list_ids = [c2l.str2int(label) for label in relations_code_list]
#gold_labels_ids = [c2l.str2int(label) for label in gold_labels]
#aa = [c2l.int2str(label) for label in gold_labels_ids] # convert ints to CODE of label!! USE IN EVAL

#inputs = tokenizer(text=texts[0], entity_spans = entity_spans[0], padding = "max_length", max_length = 1024, task = "entity_pair_classification", return_tensors = "pt")
#torch.save(inputs, 'inputs_eval.pt')
#test_dataset = MyDataset(inputs, gold_labels_ids)

logging.info("Beginning of evaluation batching")
output_dir = "evalClassic_17Out"
if not os.path.exists(output_dir):
    os.makedirs(output_dir)
output_filename = os.path.join(output_dir, 'results.json')
output_file = open(output_filename, 'w')

batch_size = 10
rel2word = {
            "Na": "Na",
            "P6": "head of government",
            "P17": "country",
            "P19": "place of birth",
            "P20": "place of death",
            "P22": "father",
            "P25": "mother",
            "P26": "spouse",
            "P27": "country of citizenship",
            "P30": "continent",
            "P31": "instance of",
            "P35": "head of state",
            "P36": "capital",
            "P37": "official language",
            "P39": "position held",
            "P40": "child",
            "P50": "author",
            "P54": "member of sports team",
            "P57": "director",
            "P58": "screenwriter",
            "P69": "educated at",
            "P86": "composer",
            "P102": "member of political party",
            "P108": "employer",
            "P112": "founded by",
            "P118": "league",
            "P123": "publisher",
            "P127": "owned by",
            "P131": "located in the administrative territorial entity",
            "P136": "genre",
            "P137": "operator",
            "P140": "religion",
            "P150": "contains administrative territorial entity",
            "P155": "follows",
            "P156": "followed by",
            "P159": "headquarters location",
            "P161": "cast member",
            "P162": "producer",
            "P166": "award received",
            "P170": "creator",
            "P171": "parent taxon",
            "P172": "ethnic group",
            "P175": "performer",
            "P176": "manufacturer",
            "P178": "developer",
            "P179": "series",
            "P190": "sister city",
            "P194": "legislative body",
            "P205": "basin country",
            "P206": "located in or next to body of water",
            "P241": "military branch",
            "P264": "record label",
            "P272": "production company",
            "P276": "location",
            "P279": "subclass of",
            "P355": "subsidiary",
            "P361": "part of",
            "P364": "original language of work",
            "P400": "platform",
            "P403": "mouth of the watercourse",
            "P449": "original network",
            "P463": "member of",
            "P488": "chairperson",
            "P495": "country of origin",
            "P527": "has part",
            "P551": "residence",
            "P569": "date of birth",
            "P570": "date of death",
            "P571": "inception",
            "P576": "dissolved, abolished or demolished",
            "P577": "publication date",
            "P580": "start time",
            "P582": "end time",
            "P585": "point in time",
            "P607": "conflict",
            "P674": "characters",
            "P676": "lyrics by",
            "P706": "located on terrain feature",
            "P710": "participant",
            "P737": "influenced by",
            "P740": "location of formation",
            "P749": "parent organization",
            "P800": "notable work",
            "P807": "separated from",
            "P840": "narrative location",
            "P937": "work location",
            "P1001": "applies to jurisdiction",
            "P1056": "product or material produced",
            "P1198": "unemployment rate",
            "P1336": "territory claimed by",
            "P1344": "participant of",
            "P1365": "replaces",
            "P1366": "replaced by",
            "P1376": "capital of",
            "P1412": "languages spoken, written or signed",
            "P1441": "present in work",
            "P3373": "sibling"}

num_predicted = 0
num_gold = 0
num_correct = 0
this_pair = []
all_pairs = []
list_of_dicts = []
torch.cuda.empty_cache()

logging.info("Evaluation will start now!:")
model.eval()
model.to(device)
for batch_start_idx in trange(0, len(test_examples), batch_size):# len(test_examples) 100
    batch_examples = test_examples[batch_start_idx:batch_start_idx + batch_size]
    texts = [example["text"] for example in batch_examples]
    entity_spans = [example["entity_spans"] for example in batch_examples]
    idxs_entity_pair = [example["idxs_entity_pair"] for example in batch_examples]
    titles = [example["title"] for example in batch_examples]
    #gold_labels = [example["labels"] for example in batch_examples]

    #gold_labels_ids = [c2l.str2int(label) for label in gold_labels]

    for i in range(len(entity_spans)):
        entity_spans[i] = list(entity_spans[i])

    inputs = tokenizer(text=texts, entity_spans=entity_spans, truncation=True, padding = "max_length", max_length = 512, task = "entity_pair_classification", return_tensors = "pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    predicted_indices = outputs.logits.argmax(-1)
    predicted_labels = [c2l.int2str(pred) for pred in predicted_indices.tolist()]
    predicted_relation = [rel2word.get(rel) for rel in predicted_labels]

    for i in range(len(predicted_relation)):
        list_of_dicts.append(dict(
            title=titles[i],
            h_idx=idxs_entity_pair[i][0],
            t_idx = idxs_entity_pair[i][1],
            r = predicted_relation[i]
        ))
    torch.cuda.empty_cache()

json_object = json.dumps(list_of_dicts, indent = 4)
with open("results_classic.json", "w") as outfile:
    outfile.write(json_object)