from flask import (
  Blueprint, request
)

from flaskr.db import get_db
import json

bp = Blueprint('products', __name__)

@bp.route('/products', methods=('GET', 'POST'))
def products():
  if request.method == 'GET':
    prod_id = request.args.get('prod_id')
    if not prod_id:
      return {'status': 'error', 'message': 'missing prod_id param'}

    db = get_db()
    product = db.execute(
      'SELECT * FROM products WHERE id = ?', (prod_id,)
    ).fetchone()

    if product is None:
      return {'status': 'error', 'message': 'product not found'}

    return {
      'status': 'success',
      'product': dict(zip(['id', 'name', 'price', 'quantity'], tuple(product)))
    }

  if request.method == 'POST':
    name = request.form['name']
    price = request.form['price']
    quantity = request.form['quantity']

    db = get_db()
    error = None

    if not name:
      error = 'name is required'
    elif not price:
      error = 'price is required'
    elif not quantity:
      error = 'price is required'

    if error is None:
      try:
        db.execute(
          'INSERT INTO products (name, price, quantity) VALUES (?, ?, ?)',
          (name, price, quantity)
        )
        db.commit()
      except:
        return {'status': 'error', 'message': 'an error happened and your product wasn\'t added'}
      else:
        return {'status': 'success', 'message': 'Your product as added successfully'}


import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

def summarize_text_with_hugging_face(
    text, model_name="facebook/bart-large-cnn", max_length=250, min_length=25
):
    """
    Summarize the given text using the specified model.

    Parameters:
    - text (str): The text to summarize.
    - model_name (str): The name of the pre-trained model to use.
    - max_length (int): The maximum length of the summary.
    - min_length (int): The minimum length of the summary.

    Returns:
    - str: The summary of the text.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    summarization_pipeline = pipeline(
        "summarization", model=model_name, tokenizer=tokenizer, device=0
    )  # device=0 if you have a GPU
    summary = summarization_pipeline(
        text, max_length=max_length, min_length=min_length, do_sample=False
    )
    return summary[0]["summary_text"]


import pygame as pg
from random import randrange

vec2 = pg.math.Vector2
snake1 = (227, 57, 25)
snake2 = (26, 147, 187)
snake3 = (122, 13, 89)
snake4 = (31, 3, 106)
green = (0, 255, 0)
blue = (0, 0, 255)
yellow = (255, 255, 0)
black = (0, 0, 0)
light_pink = (255, 182, 193)
pink = (255, 105, 180)
dark_pink = (255, 20, 147)
background = (245, 246, 246)

class Snake:
    def __init__(self, game, color, keys, snake_color, image):
        self.color = color
        self.l = keys[0]
        self.r = keys[1]
        self.u = keys[2]
        self.d = keys[3]
        self.game = game
        self.size = game.TILE_SIZE
        self.rect = pg.rect.Rect([0, 0,
                                  game.TILE_SIZE,
                                  game.TILE_SIZE])
        self.rect.center = self.get_random_position()
        self.direction = vec2(self.size, 0)
        self.step_delay = 100
        self.time = 0
        self.lenght = 1
        self.segments = []
        self.directions = {self.l: 1, self.r: 1, self.u: 1, self.d: 1}
        self.snake_color = snake_color
        self.img = import_image(image, 0.2, (0, 0, 0))
        self.head = self.img

    def end_of_snake(self):
        self.game.alive[self.snake_color] = 0
        self.game.table.append(self.snake_color)
        self.segments = []

    def check_collision(self):
        if self.snake_color != 'red':
            for segment in self.game.snake1.segments:
                if segment.center == self.rect.center:
                    print('a')
                    self.end_of_snake()
        if self.game.number_of_players > 1:
            if self.snake_color != 'light blue':
                for segment in self.game.snake2.segments:
                    if segment.center == self.rect.center:
                        print('b')
                        self.end_of_snake()
        if self.game.number_of_players > 2:
            if self.snake_color != 'violete':
                for segment in self.game.snake3.segments:
                    if segment.center == self.rect.center:
                        print('c')
                        self.end_of_snake()
        if self.game.number_of_players > 3:
            if self.snake_color != 'dark blue':
                for segment in self.game.snake4.segments:
                    if segment.center == self.rect.center:
                        print('d')
                        self.end_of_snake()

    def check_food(self):
        if self.rect.center == self.game.food.rect.center:
            self.game.food.rect.center = self.get_random_position()
            self.lenght += 4
        if self.rect.center == self.game.food1.rect.center:
            self.game.food1.rect.center = self.get_random_position()
            self.lenght += 3
        if self.rect.center == self.game.food2.rect.center:
            self.game.food2.rect.center = self.get_random_position()
            self.lenght += 2

    def check_border1(self):
        if self.rect.left < 0 or self.rect.right > self.game.WINDOW_SIZE:
            print('e')
            self.end_of_snake()
        if self.rect.top < 0 or self.rect.bottom > self.game.WINDOW_SIZE:
            print('f')
            self.end_of_snake()

    def check_border(self):
        if self.rect.left < 0:
            self.rect.right = self.game.WINDOW_SIZE
        if self.rect.right > self.game.WINDOW_SIZE:
            self.rect.left = 0
        if self.rect.top < 0:
            self.rect.bottom = self.game.WINDOW_SIZE
        if self.rect.bottom > self.game.WINDOW_SIZE:
            self.rect.top = 0

    def check_if_collision(self):
        if len(self.segments) > len(set([segment.center for segment in self.segments])):
            self.end_of_snake()

    def control(self, event):
        if event.type == pg.KEYDOWN:
            if event.key == self.r and self.directions[self.r]:
                self.direction = vec2(self.size, 0)
                self.directions = {self.l: 0, self.r: 0, self.u: 1, self.d: 1}
                self.head = pg.transform.rotate(self.img, 270)

            if event.key == self.u and self.directions[self.u]:
                self.direction = vec2(0, -self.size)
                self.directions = {self.l: 1, self.r: 1, self.u: 0, self.d: 0}
                self.head = self.img

            if event.key == self.l and self.directions[self.l]:
                self.direction = vec2(-self.size, 0)
                self.directions = {self.l: 0, self.r: 0, self.u: 1, self.d: 1}
                self.head = pg.transform.rotate(self.img, 90)

            if event.key == self.d and self.directions[self.d]:
                self.direction = vec2(0, self.size)
                self.directions = {self.l: 1, self.r: 1, self.u: 0, self.d: 0}
                self.head = pg.transform.rotate(self.img, 180)

    def delta_time(self):
        time_now = pg.time.get_ticks()
        if time_now - self.time > self.step_delay:
            self.time = time_now
            return True
        return False

    def get_random_position(self):
        return [randrange(self.size//2,
                          self.game.WINDOW_SIZE -self.size//2,
                          self.size)] * 2

    def move(self):
        if self.delta_time():
            self.rect.move_ip(self.direction)
            self.segments.append(self.rect.copy())
            self.segments = self.segments[-self.lenght:]

    def update(self):
        self.move()
        self.check_food()
        self.check_border()
        self.check_if_collision()
        self.check_collision()

    def draw(self):
        [pg.draw.rect(self.game.screen, self.color, segment) for segment in self.segments]
        self.game.screen.blit(self.head, (self.rect.left, self.rect.top))

class Food:
    def __init__(self, game):
        self.game = game
        self.size = game.TILE_SIZE
        self.image = import_image('images/fruit1.png', 0.08, (245, 246, 246))
        self.rect = pg.rect.Rect([0, 0,
                                  game.TILE_SIZE - 2,
                                  game.TILE_SIZE - 2])
        self.rect.center = self.get_random_position()

    def get_random_position(self):
        return [randrange(self.size//2,
                          self.game.WINDOW_SIZE -self.size//2,
                          self.size)] * 2

    def draw(self):
        #pg.draw.rect(self.game.screen, (255, 0, 0), self.rect)
        self.game.screen.blit(self.image, (self.rect.left, self.rect.top))

class Button:
    def __init__(self, game, x, y, image, scale):
        self.game = game
        width = image.get_width()
        height = image.get_height()
        self.image = pg.transform.scale(image, (int(width * scale), int(height * scale)))
        self.rect = self.image.get_rect()
        self.rect.topleft = (x, y)
        self.clicked = False

    def draw(self):
        action = False
        pos = pg.mouse.get_pos()
        if self.rect.collidepoint(pos):
            if pg.mouse.get_pressed()[0] == 1 and self.clicked == False:
                self.clicked = True
                action = True

        if pg.mouse.get_pressed()[0] == 0:
            self.clicked = False

        self.game.screen.blit(self.image, (self.rect.x, self.rect.y))
        return action

def draw_text(game, text, font, text_col, x, y):
    img = font.render(text, True, text_col)
    game.screen.blit(img, (x, y))

def import_image(image, scale, colorkey):
    im = pg.image.load(image)
    im.set_colorkey(colorkey)
    im = pg.transform.scale(im, (int(im.get_width() * scale), int(im.get_height() * scale)))
    return im


"""Configure for ASCAR DRL Evaluation

Copyright (c) 2016, 2017 The Regents of the University of California. All
rights reserved.

Created by Yan Li <yanli@tuneup.ai>, Kenneth Chang <kchang44@ucsc.edu>,
Oceane Bel <obel@ucsc.edu>. Storage Systems Research Center, Baskin School
of Engineering.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the Storage Systems Research Center, the
      University of California, nor the names of its contributors
      may be used to endorse or promote products derived from this
      software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
REGENTS OF THE UNIVERSITY OF CALIFORNIA BE LIABLE FOR ANY DIRECT,
INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
OF THE POSSIBILITY OF SUCH DAMAGE.
"""
from ascar import common
from ascar import LustreCommon
import glob
import logging
import os
import socket
from typing import List

__author__ = 'Yan Li'
__copyright__ = 'Copyright (c) 2016, 2017 The Regents of the University of California. All rights reserved.'

"""Maps hostname to MA ID
"""
# 主机名映射
nodeid_map = {
    '192.168.1.11': 1,
    '192.168.1.12': 2,
    '192.168.1.13': 3,
    '192.168.1.14': 4,
}

TICK_LEN = 1
servers = [
    '192.168.1.11',
    '192.168.1.12',
]
# client 观测参数数量
PI_PER_CLIENT_OBD = 11
# clint数量
OBD_PER_CLIENT_MA = len(servers)
clients = [
    '192.168.1.13',
    '192.168.1.14',
]
# todo 控制方法的含义
# 'mrif' max rpcs in flight osc到其ost的rpc的最大并发数
# control_method can be 'mrif', 'rules', or 'mrif_tau'
control_method = 'mrif_tau'
rule_archive_path = '/share/drl/rules'

if control_method == 'mrif':
    cpv_spec = [
        # name, initial value, min, max, step size
        ['mrif', 8, 1, 256, 4],
    ]
elif control_method == 'rules':
    cpv_spec = [
        # name, initial value, min, max, step size
        ['rtt_ratio_changepoint', 237, 37, 437, 10],
        ['ewma_changepoint_a', 41001, 20001, 60001, 1000],
        ['ewma_changepoint_b', 48427, 28427, 68427, 1000],
    ]
    rule_template_file = os.path.join(rule_archive_path, 'iorcp_alpha9999_472_3cpvs.csv')
    with open(rule_template_file, 'r') as f:
        rule_template = f.read()
elif control_method == 'mrif_tau':
    cpv_spec = [
        # name, initial value, min, max, step size
        ['mrif', 8, 1, 256, 4],
        ['tau', 32840, 0, 400000, 1500],
    ]
    rule_template_file = os.path.join(rule_archive_path, 'tau_only.csv')
    with open(rule_template_file, 'r') as f:
        tau_only_rule_template = f.read()
    _tau = 0
else:
    raise RuntimeError('Unknown control method: ' + control_method)

my_hostname = socket.gethostname()

# 收集到服务器的时间
def collect_ping_to_servers() -> List[float]:
    result = []
    for srv in servers:
        # PI 7
        result.append(common.get_ping_time(srv))
    return result

'''
收集了10个环境
max_rpcs_in_flight  最大并行rpc数
min_brw_rpc_gap
指数加权移动平均
ack_ewma
sent_ewma
rtt_ratio100
read_bandwidth
write_bandwidth
tau
cur_dirty_bytes
max_dirty_mb
'''

def collect_osc_pi_cpv(osc_path: str) -> List[float]:
    result = list()
    # PI 0
    result.append(LustreCommon.read_proc_file(os.path.join(osc_path, 'max_rpcs_in_flight')))
    # PI 1
    result.append(LustreCommon.read_proc_file(os.path.join(osc_path, 'min_brw_rpc_gap')))

    with open(os.path.join(osc_path, 'import'), 'r') as importfile:
        # import is a proc file and should be read as a whole, i.e., not using readline()
        import_data = importfile.read()

    # PI 2
    result.append(LustreCommon.extract_ack_ewma_from_import(import_data))
    # PI 3
    result.append(LustreCommon.extract_sent_ewma_from_import(import_data))
    # PI 4
    result.append(LustreCommon.extract_rtt_ratio100_from_import(import_data))
    # PI 5
    result.append(LustreCommon.extract_read_bandwidth_from_import(import_data))
    # PI 6
    result.append(LustreCommon.extract_write_bandwidth_from_import(import_data))
    # PI 7
    result.append(_tau)
    # PI 8
    result.append(LustreCommon.read_proc_file(os.path.join(osc_path, 'cur_dirty_bytes')))
    # PI 9
    result.append(LustreCommon.read_proc_file(os.path.join(osc_path, 'max_dirty_mb')))

    return result

# 从lustre收集所有数据
def lustre_collect_pi() -> List[float]:
    if my_hostname in clients:
        # collect ping time to all servers
        ping_times = collect_ping_to_servers()
        # collect PIs
        oscs = glob.glob('/proc/fs/lustre/osc/*/import')
        """ 收集信息大致如下
import:
    name: temp-OST0004-osc-ffff9068ea7c9000
    target: temp-OST0004_UUID
    state: IDLE
    connect_flags: [ write_grant, server_lock, version, request_portal, truncate_lock, max_byte_per_rpc, early_lock_cancel, adaptive_timeouts, lru_resize, alt_checksum_algorithm, fid_is_enabled, version_recovery, grant_shrink, full20, layout_lock, 64bithash, object_max_bytes, jobstats, einprogress, grant_param, lvb_type, short_io, lfsck, bulk_mbits, second_flags, lockaheadv2 ]
    connect_data:
       flags: 0xa0425af2e3440478
       instance: 1
       target_version: 2.12.7.0
       initial_grant: 8437760
       max_brw_size: 4194304
       grant_block_size: 4096
       grant_inode_size: 32
       grant_max_extent_size: 67108864
       grant_extent_tax: 24576
       cksum_types: 0xf7
       max_object_bytes: 17592186040320
    import_flags: [ replayable, pingable, connect_tried ]
    connection:
       failover_nids: [ 192.168.1.13@tcp ]
       current_connection: 192.168.1.13@tcp
       connection_attempts: 2
       generation: 2
       in-progress_invalidations: 0
       idle: 159490 sec
    rpcs:
       inflight: 0
       unregistering: 0
       timeouts: 0
       avg_waittime: 655647 usec
    service_estimates:
       services: 31 sec
       network: 1 sec
        """
        osc_paths = [os.path.dirname(p) for p in oscs]
        osc_paths.sort()
        osc_pis = list()
        for osc in osc_paths:
            osc_pis.extend(collect_osc_pi_cpv(osc))

        result = osc_pis + ping_times
        assert len(result) == len(servers) * PI_PER_CLIENT_OBD
        return result
    elif my_hostname in servers:
        # todo 服务器
        return list()
    else:
        raise RuntimeError('My hostname is neither client or server')

# lustre 根据不同的 control_method 进行控制
def lustre_controller(cpvs: List[float]):
    """Perform an action

    Applies the CPVs to the system

    :param cpvs: the new values of CPVs
    :return:
    """
    # cpvs[0] is the action id, remove it
    cpvs = cpvs[1:]
    if my_hostname in clients:
        if control_method == 'mrif':
            mrif = int(cpvs[0])
            # 必须位于最大最小值之间
            assert cpv_spec[0][2] <= mrif <= cpv_spec[0][3]
            LustreCommon.set_mrif(mrif, len(servers))
        elif control_method == 'rules':
            # convert cpvs list to a dict
            kv = dict([('cpv{0}'.format(i + 1), int(cpv)) for i, cpv in enumerate(cpvs)])
            rule = LustreCommon.gen_rule(rule_template, kv)
            LustreCommon.set_rule(rule, len(servers))
        elif control_method == 'mrif_tau':
            global _tau
            mrif = int(cpvs[0])
            assert cpv_spec[0][2] <= mrif <= cpv_spec[0][3]
            LustreCommon.set_mrif(mrif, len(servers))

            _tau = int(cpvs[1])
            rule = LustreCommon.gen_rule(tau_only_rule_template, {'tau': _tau})
            LustreCommon.set_rule(rule, len(servers))
        else:
            raise RuntimeError('Unknown control method: ' + control_method)
    elif my_hostname in servers:
        # todo 服务器优化
        pass
    else:
        raise RuntimeError('My hostname is neither client or server')

opt = {
    'loglevel': logging.INFO,
    'log_lazy_flush': True,
    # cProfile 性能分析工具
    # Pympler 内存诊断工具
    # level 1 enables cProfile, level 2 also enables Pympler
    'ma_debugging_level': 0,
    'dqldaemon_debugging_level': 0,
    # database location
    # todo: why not only use ram ?
    'dbfile': '/data/ascar/ascar_replay_db.sqlite',
    # todo 滴答声持续时间
    'tick_len': TICK_LEN,  # duration of a tick in second
    # todo 一次观察中有多少滴答声
    'ticks_per_observation': 10,  # how many ticks are in an observation
    # 节点映射表
    'nodeid_map': nodeid_map,
    # client 节点名
    'clients': clients,
    # server 节点名
    'servers': servers,
    # 可控制参数
    'cpvs': cpv_spec,
    # todo: 加减步长和不动？那也应该是3*len(cpv_spec)
    'num_actions': 2 * len(cpv_spec) + 1,
    'pi_per_client_obd': PI_PER_CLIENT_OBD,
    'obd_per_client_ma': OBD_PER_CLIENT_MA,
    'tick_data_size': PI_PER_CLIENT_OBD * OBD_PER_CLIENT_MA * len(clients),  # only four clients have PI so far

    # Collectors are functions that collect PIs. MA calls them in order and concatenate
    # their returns into a single list before passing them to IntfDaemon
    'collectors': [lustre_collect_pi],
    'controller': lustre_controller,
    # 用于接受，存储和聚合来自mas（monitor agents）的obs
    'intf_daemon_loc': '128.114.59.20:9123',
    # moniter 代理 日志文件
    'ascar.MonitorAgent.MonitorAgent_logfile': '/root/log/ma_log.txt',
    'ascar.IntfDaemon.IntfDaemon_logfile': '/data/ascar/intfdaemon_log.txt',
    'ascar.DQLDaemon.DQLDaemon_logfile': '/data/ascar/dqldaemon_log.txt',
    'pidfile_dir': '/tmp',

    'start_random_rate': 1,
    # How many actions should the exploration period include
    'exploration_period': 1000000,
    'random_action_probability': 0.05,

    'minibatch_size': 32,
    'enable_tuning': True,
}


# coding=utf-8
from unidecode import unidecode

from engine.models.record import *
from engine.modules.data_cleansing.data_cleansing_module import DataCleansingModule
from engine.modules.module import Module
import dateparser

class DateFormatting(DataCleansingModule):
    """
        Formatea como date los valores que logra parsear, los que no, quedán como estaban.

        Config:
        {
        languages:['en','es']
        }
    """

    def __init__(self, **kwargs):
        super(DateFormatting, self).__init__(**kwargs)

    @staticmethod
    def pretty_name():
        return "Date Formatting"

    def run(self, column):
        languages = self.config["languages"]
        for field in column.fields:
            if field.tipe.value == EnumType.string.value:
                parsed = dateparser.parse(field.value,languages=languages)
                if parsed is not None:
                    field.value = parsed
                    field.tipe = EnumType.date

        return column

    @staticmethod
    def config_json(**kwargs):
        return {
            "languages": {
                "label": "Languages",
                "type": "multipleselect",
                "options": [
                    "es",
                    "en",
                    "pt",
                    "fr",
                    "ru"
                ]
            }
        }


from django.conf.urls import url
from django.contrib.auth.decorators import login_required, permission_required

from . import views, users_pdf
from django.conf import settings
from django.conf.urls.static import static

urlpatterns = [
        url(r'^$', permission_required('userprofile.view_user', login_url='not_found')
                (views.users), name='users'),
        url(r'^add/$', permission_required('userprofile.add_user', login_url='not_found')
                (views.user_add), name='user-add'),
        url(r'^users_export_csv/$', views.users_export_csv, name='users_export_csv'),
        url(r'^user_process/$',  permission_required('userprofile.add_user')
        (views.user_process), name='user_process'),
        url(r'^detail/(?P<pk>[0-9]+)/$', permission_required('userprofile.change_user', login_url='not_found')(views.user_detail), name='user-detail'),
        url(r'user_trail/$', permission_required('userprofile.view_usertrail', login_url='not_found')
                (views.user_trails), name='user_trail'),
        url(r'^usertrail_paginate/', views.usertrail_paginate, name='usertrail_paginate'),
        url( r'^usertrail_search/$', views.usertrail_search, name = 'usertrail_search' ),
        url(r'^delete/(?P<pk>[0-9]+)/$', permission_required('userprofile.delete_user')
                (views.user_delete), name='user-delete'),

        url(r'^edit/(?P<pk>[0-9]+)/$', permission_required('userprofile.change_user', login_url='not_found')(views.user_edit), name='user-edit'),
        url(r'^user_update(?P<pk>[0-9]+)/$', permission_required('userprofile.change_user', login_url='not_found')
                (views.user_update), name='user-update'),
        url(r'^user_assign_permission/$', views.user_assign_permission, name='user_assign_permission'),
        url(r'^user_paginate/', views.user_paginate, name='user_paginate'),
        url( r'^users_search/$', views.user_search, name = 'user_search' ),
        url( r'^payload/pdf/$', users_pdf.pdf, name ='users_pdf'),

]

if settings.DEBUG:
    urlpatterns += static(settings.STATIC_URL, document_root=settings.STATIC_ROOT)
    urlpatterns += static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)


import math
import os
import sys
import itertools

import torch
import torch.nn.functional as F
import torch.optim as optim
from envs import create_atari_env
from model import ActorCritic
from torch.autograd import Variable
from torchvision import datasets, transforms

def ensure_shared_grads(model, shared_model):
    for param, shared_param in zip(model.parameters(), shared_model.parameters()):
        if shared_param.grad is not None:
            return
        shared_param._grad = param.grad

def is_dead(info):
    dead = False
    if is_dead.current_life > info['ale.lives']:
        dead = True
    is_dead.current_life = info['ale.lives']
    return dead

is_dead.current_life = 0

def train(rank, args, shared_model, optimizer=None):
    torch.manual_seed(args.seed + rank)

    env = create_atari_env(args.env_name)
    env.seed(args.seed + rank)

    model = ActorCritic(env.observation_space.shape[0], env.action_space, args.num_skips)

    if optimizer is None:
        optimizer = optim.Adam(shared_model.parameters(), lr=args.lr)

    model.train()

    state = env.reset()
    state = torch.from_numpy(state)
    done = True
    dead = False

    episode_length = 0
    for ep_counter in itertools.count(1):
        # Sync with the shared model
        model.load_state_dict(shared_model.state_dict())
        if done or dead:
            cx = Variable(torch.zeros(1, 256))
            hx = Variable(torch.zeros(1, 256))
        else:
            cx = Variable(cx.data)
            hx = Variable(hx.data)

        values = []
        log_probs = []
        rewards = []
        entropies = []

        for step in range(args.num_steps):
            value, logit, (hx, cx) = model(
                (Variable(state.unsqueeze(0)), (hx, cx)))
            prob = F.softmax(logit)
            log_prob = F.log_softmax(logit)
            entropy = -(log_prob * prob).sum(1)
            entropies.append(entropy)

            action = prob.multinomial().data
            log_prob = log_prob.gather(1, Variable(action))

            action_np = action.numpy()[0][0]
            if action_np < model.n_real_acts:
                state, reward, done, info = env.step(action_np)
                dead = is_dead(info)
                done = done or episode_length >= args.max_episode_length
                reward = max(min(reward, 1), -1)
                episode_length += 1
            else:
                reward = 0.
                for counter_skips in range(action_np - model.n_real_acts + 2):
                    state, rew, done, info = env.step(0)  # instead of random perform NOOP=0
                    dead = is_dead(info)
                    done = done or episode_length >= args.max_episode_length
                    rew = max(min(rew, 1), -1)

                    reward += rew
                    episode_length += 1
                    if done or dead:
                        break

                    if counter_skips != action_np - model.n_real_acts + 1: # maintain hx, cx for conseq frames
                        _, _, (hx, cx) = model((Variable(torch.from_numpy(state).unsqueeze(0)), (hx, cx)))

            if done:
                episode_length = 0
                state = env.reset()
                env.seed(args.seed + rank + (args.num_processes+1)*ep_counter)

            state = torch.from_numpy(state)
            values.append(value)
            log_probs.append(log_prob)
            rewards.append(reward)

            if done or dead:
                break

        R = torch.zeros(1, 1)
        if not done and not dead:
            value, _, _ = model((Variable(state.unsqueeze(0)), (hx, cx)))
            R = value.data

        values.append(Variable(R))
        policy_loss = 0
        value_loss = 0
        R = Variable(R)
        gae = torch.zeros(1, 1)
        for i in reversed(range(len(rewards))):
            R = args.gamma * R + rewards[i]
            advantage = R - values[i]
            value_loss = value_loss + 0.5 * advantage.pow(2)

            # Generalized Advantage Estimataion
            delta_t = rewards[i] + args.gamma * \
                values[i + 1].data - values[i].data
            gae = gae * args.gamma * args.tau + delta_t

            policy_loss = policy_loss - \
                log_probs[i] * Variable(gae) - 0.01 * entropies[i]

        optimizer.zero_grad()

        (policy_loss + 0.5 * value_loss).backward()
        torch.nn.utils.clip_grad_norm(model.parameters(), 40)

        ensure_shared_grads(model, shared_model)
        optimizer.step()


# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

import json
import logging
import math
import os
import platform
import random
import threading
import time
import unicodedata

import bpy
from bpy.app.handlers import persistent
from bpy.props import (  # TODO only keep the ones actually used when cleaning
    BoolProperty,
    StringProperty,
)
from bpy.types import Operator

from . import (
    addon_updater_ops,
    asset_bar_op,
    bkit_oauth,
    categories,
    colors,
    comments_utils,
    daemon_lib,
    global_vars,
    image_utils,
    paths,
    ratings_utils,
    reports,
    rerequests,
    resolutions,
    tasks_queue,
    ui,
    utils,
    version_checker,
)
from .daemon import tasks

bk_logger = logging.getLogger(__name__)

def check_errors(rdata):
  if rdata.get('statusCode') and int(rdata.get('statusCode')) > 299:
    utils.p(rdata)
    if rdata.get('detail') == 'Invalid token.':
      bkit_oauth.logout()
      return False, "Invalid token. You've been logged out. Use login panel to connect your profile."
    else:
      return False, rdata.get('detail')
  if rdata.get('statusCode') is None and rdata.get('results') is None:
    return False, 'Connection error'
  return True, ''

search_tasks = {}

def update_ad(ad):
  if not ad.get('assetBaseId'):
    try:
      ad['assetBaseId'] = ad['asset_base_id']  # this should stay ONLY for compatibility with older scenes
      ad['assetType'] = ad['asset_type']  # this should stay ONLY for compatibility with older scenes
      ad['verificationStatus'] = ad[
        'verification_status']  # this should stay ONLY for compatibility with older scenes
      ad['author'] = {}
      ad['author']['id'] = ad['author_id']  # this should stay ONLY for compatibility with older scenes
      ad['canDownload'] = ad['can_download']  # this should stay ONLY for compatibility with older scenes
    except Exception as e:
      bk_logger.error('BlenderKit failed to update older asset data')
  return ad

def update_assets_data():  # updates assets data on scene load.
  '''updates some properties that were changed on scenes with older assets.
  The properties were mainly changed from snake_case to CamelCase to fit the data that is coming from the server.
  '''
  data = bpy.data

  datablocks = [
    bpy.data.objects,
    bpy.data.materials,
    bpy.data.brushes,
  ]
  for dtype in datablocks:
    for block in dtype:
      if block.get('asset_data') != None:
        update_ad(block['asset_data'])

  dicts = [
    'assets used',
  ]
  for s in bpy.data.scenes:
    for bkdict in dicts:

      d = s.get(bkdict)
      if not d:
        continue;

      for asset_id in d.keys():
        update_ad(d[asset_id])
        # bpy.context.scene['assets used'][ad] = ad

@persistent
def undo_post_reload_previews(context):
  load_previews()

@persistent
def undo_pre_end_assetbar(context):
  ui_props = bpy.context.window_manager.blenderkitUI

  ui_props.turn_off = True
  ui_props.assetbar_on = False

@persistent
def scene_load(context):
  """Load categories , check timers registration, and update scene asset data.
  Should (probably) also update asset data from server (after user consent).
  """
  if not bpy.app.timers.is_registered(bkit_oauth.refresh_token_timer) and not bpy.app.background:
    bpy.app.timers.register(bkit_oauth.refresh_token_timer, persistent=True, first_interval=5)
    #bpy.app.timers.register(bkit_oauth.refresh_token_timer, persistent=True, first_interval=36000)
  # if utils.experimental_enabled() and not bpy.app.timers.is_registered(
  #         refresh_notifications_timer) and not bpy.app.background:
  #     bpy.app.timers.register(refresh_notifications_timer, persistent=True, first_interval=5)

  update_assets_data()

def fetch_server_data():
  '''Download profile, and refresh token if needed.'''
  if bpy.app.background:
    return
  user_preferences = bpy.context.preferences.addons['blenderkit'].preferences
  api_key = user_preferences.api_key
  # Only refresh new type of tokens(by length), and only one hour before the token timeouts.
  if api_key != '' and global_vars.DATA.get('bkit profile') == None:
    get_profile()
  # all_notifications_count = comments_utils.count_all_notifications()
  # comments_utils.get_notifications_thread(api_key, all_count = all_notifications_count)

first_time = True
first_search_parsing = True
last_clipboard = ''

def check_clipboard():
  '''
  Checks clipboard for an exact string containing asset ID.
  The string is generated on www.blenderkit.com as for example here:
  https://www.blenderkit.com/get-blenderkit/54ff5c85-2c73-49e9-ba80-aec18616a408/
  '''

  # clipboard monitoring to search assets from web
  if platform.system() != 'Linux':
    global last_clipboard
    if bpy.context.window_manager.clipboard != last_clipboard:
      last_clipboard = bpy.context.window_manager.clipboard
      instr = 'asset_base_id:'
      # first check if contains asset id, then asset type
      if last_clipboard[:len(instr)] == instr:
        atstr = 'asset_type:'
        ati = last_clipboard.find(atstr)
        # this only checks if the asset_type keyword is there but let's the keywords update function do the parsing.
        if ati > -1:
          search_props = utils.get_search_props()
          search_props.search_keywords = last_clipboard
          # don't run search after this - assigning to keywords runs the search_update function.
        # bpy.context.window_manager.clipboard = ''

def parse_result(r):
  '''
  needed to generate some extra data in the result(by now)
  Parameters
  ----------
  r - search result, also called asset_data
  '''
  scene = bpy.context.scene

  # TODO remove this fix when filesSize is fixed.
  # this is a temporary fix for too big numbers from the server.
  # try:
  #     r['filesSize'] = int(r['filesSize'] / 1024)
  # except:
  #     utils.p('asset with no files-size')
  asset_type = r['assetType']
  if len(r['files']) > 0:  # TODO remove this condition so all assets are parsed.
    get_author(r)

    r['available_resolutions'] = []
    allthumbs = []
    durl, tname, small_tname = '', '', ''

    if r['assetType'] == 'hdr':
      tname = paths.extract_filename_from_url(r['thumbnailLargeUrlNonsquared'])
    else:
      tname = paths.extract_filename_from_url(r['thumbnailMiddleUrl'])
    small_tname = paths.extract_filename_from_url(r['thumbnailSmallUrl'])
    allthumbs.append(tname)  # TODO just first thumb is used now.
    # if r['fileType'] == 'thumbnail':
    #     tname = paths.extract_filename_from_url(f['fileThumbnailLarge'])
    #     small_tname = paths.extract_filename_from_url(f['fileThumbnail'])
    #     allthumbs.append(tname)  # TODO just first thumb is used now.

    for f in r['files']:
      # if f['fileType'] == 'thumbnail':
      #     tname = paths.extract_filename_from_url(f['fileThumbnailLarge'])
      #     small_tname = paths.extract_filename_from_url(f['fileThumbnail'])
      #     allthumbs.append(tname)  # TODO just first thumb is used now.

      if f['fileType'] == 'blend':
        durl = f['downloadUrl'].split('?')[0]
        # fname = paths.extract_filename_from_url(f['filePath'])

      if f['fileType'].find('resolution') > -1:
        r['available_resolutions'].append(resolutions.resolutions[f['fileType']])

    # code for more thumbnails
    # tdict = {}
    # for i, t in enumerate(allthumbs):
    #     tdict['thumbnail_%i'] = t

    r['max_resolution'] = 0
    if r['available_resolutions']:  # should check only for non-empty sequences
      r['max_resolution'] = max(r['available_resolutions'])

    # tooltip = generate_tooltip(r)
    # for some reason, the id was still int on some occurances. investigate this.
    r['author']['id'] = str(r['author']['id'])

    # some helper props, but generally shouldn't be renaming/duplifiying original properties,
    # so blender's data is same as on server.
    asset_data = {'thumbnail': tname,
                  'thumbnail_small': small_tname,
                  # 'tooltip': tooltip,

                  }
    asset_data['downloaded'] = 0

    # parse extra params needed for blender here
    params = r['dictParameters']  # utils.params_to_dict(r['parameters'])

    if asset_type == 'model':
      if params.get('boundBoxMinX') != None:
        bbox = {
          'bbox_min': (
            float(params['boundBoxMinX']),
            float(params['boundBoxMinY']),
            float(params['boundBoxMinZ'])),
          'bbox_max': (
            float(params['boundBoxMaxX']),
            float(params['boundBoxMaxY']),
            float(params['boundBoxMaxZ']))
        }

      else:
        bbox = {
          'bbox_min': (-.5, -.5, 0),
          'bbox_max': (.5, .5, 1)
        }
      asset_data.update(bbox)
    if asset_type == 'material':
      asset_data['texture_size_meters'] = params.get('textureSizeMeters', 1.0)

    # asset_data.update(tdict)

    au = scene.get('assets used', {})
    if au == {}:
      scene['assets used'] = au
    if r['assetBaseId'] in au.keys():
      asset_data['downloaded'] = 100
      # transcribe all urls already fetched from the server
      r_previous = au[r['assetBaseId']]
      if r_previous.get('files'):
        for f in r_previous['files']:
          if f.get('url'):
            for f1 in r['files']:
              if f1['fileType'] == f['fileType']:
                f1['url'] = f['url']

    # attempt to switch to use original data gradually, since the parsing as itself should become obsolete.
    asset_data.update(r)
    return asset_data

def clear_searches():
  global search_tasks
  search_tasks.clear()

def cleanup_search_results():
  dicts = (
    'search results','bkit model search',
    'bkit scene search',
    'bkit hdr search',
    'bkit material search',
    'bkit texture search',
    'bkit brush search',
  )
  for sr in dicts:
    global_vars.DATA.pop(sr, None)
    global_vars.DATA.pop(f'{sr} orig', None)

def handle_search_task(task: tasks.Task) -> bool:
  '''parse search results, try to load all available previews.'''
  ##############original

  global search_tasks, first_search_parsing
  if len(search_tasks) == 0:
    # utils.p('end search timer')
    props = utils.get_search_props()
    props.is_searching = False
    return True

  # don't do anything while dragging - this could switch asset during drag, and make results list length different,
  # causing a lot of throuble.
  if bpy.context.window_manager.blenderkitUI.dragging:
    # utils.p('end search timer')
    return False
  # if original task was already removed (because user initiated another search), results are dropped- Returns True
  # because that's OK.
  orig_task = search_tasks.get(task.task_id)
  if orig_task is None:
    print('search task result not from active search', task.task_id, len(search_tasks), list(search_tasks.keys()))
    return True

  search_tasks.pop(task.task_id)

  # this fixes black thumbnails in asset bar, test if this bug still persist in blender and remove if it's fixed
  if bpy.app.version < (3, 3, 0):
    sys_prefs = bpy.context.preferences.system
    sys_prefs.gl_texture_limit = 'CLAMP_OFF'

  global first_search_parsing
  if first_search_parsing:
    comments_utils.check_notifications()
    first_search_parsing = False

  ###################

  asset_type = task.data['asset_type']
  props = utils.get_search_props()
  search_name = f'bkit {asset_type} search'

  if not task.data.get('get_next'):
    result_field = []
  else:
    result_field = []
    for r in global_vars.DATA[search_name]:
      result_field.append(r)

  ok, error = check_errors(task.result)
  if ok:
    ui_props = bpy.context.window_manager.blenderkitUI
    orig_len = len(result_field)

    for ri, r in enumerate(task.result['results']):
      asset_data = parse_result(r)
      if asset_data != None:
        result_field.append(asset_data)

    # Get ratings from BlenderKit server
    user_preferences = bpy.context.preferences.addons['blenderkit'].preferences
    api_key = user_preferences.api_key
    headers = utils.get_headers(api_key)
    if utils.profile_is_validator():
      for r in task.result['results']:
        if ratings_utils.get_rating_local(r['id']) is None:
          rating_thread = threading.Thread(target=ratings_utils.get_rating, args=([r['id'], headers]),
                                           daemon=True)
          rating_thread.start()

    global_vars.DATA[search_name] = result_field
    global_vars.DATA[search_name + ' orig'] = task.result

    if result_field and result_field[0]['assetType'] == ui_props.asset_type.lower():
      global_vars.DATA['search results'] = result_field
      global_vars.DATA['search results orig'] = task.result

    if len(result_field) < ui_props.scroll_offset or not (task.data.get('get_next')):
      # jump back
      ui_props.scroll_offset = 0
    props.report = f"Found {global_vars.DATA['search results orig']['count']} results."
    if len(global_vars.DATA['search results']) == 0:
      tasks_queue.add_task((reports.add_report, ('No matching results found.',)))
    else:
      tasks_queue.add_task((reports.add_report, (f"Found {global_vars.DATA['search results orig']['count']} results.",)))
    # show asset bar automatically, but only on first page - others are loaded also when asset bar is hidden.
    if not ui_props.assetbar_on and not task.data.get('get_next'):
      bpy.ops.view3d.run_assetbar_fix_context()

  else:
    props.report = error
    reports.add_report(error, 15, 'ERROR')

  if len(search_tasks) == 0:
    props.is_searching = False
  return True

def handle_preview_task(task: tasks.Task) -> bool:
  """Parse search results, try to load all available previews."""

  global_vars.DATA['images available'][task.data['image_path']] = True
  if asset_bar_op.asset_bar_operator is not None:
    if task.data['thumbnail_type'] =='small':
      asset_bar_op.asset_bar_operator.update_image(task.data['assetBaseId'])
    if task.data['thumbnail_type'] == 'full':
      asset_bar_op.asset_bar_operator.update_tooltip_image(task.data['assetBaseId'])
  return True

def load_preview(asset):
  # FIRST START SEARCH
  props = bpy.context.window_manager.blenderkitUI
  directory = paths.get_temp_dir('%s_search' % props.asset_type.lower())

  tpath = os.path.join(directory, asset['thumbnail_small'])
  tpath_exists = os.path.exists(tpath)
  if not asset['thumbnail_small'] or asset['thumbnail_small'] == '' or not tpath_exists:
    # tpath = paths.get_addon_thumbnail_path('thumbnail_notready.jpg')
    asset['thumb_small_loaded'] = False

  iname = f".{asset['thumbnail_small']}"
  # if os.path.exists(tpath):  # sometimes we are unlucky...
  img = bpy.data.images.get(iname)

  if img is None or len(img.pixels) == 0:
    if not tpath_exists:
      return False
    # wrap into try statement since sometimes
    try:
      img = bpy.data.images.load(tpath, check_existing=True)

      img.name = iname
      if len(img.pixels) > 0:
        return True
    except:
      pass
    return False
  elif img.filepath != tpath:
    if not tpath_exists:
      # unload loaded previews from previous results
      bpy.data.images.remove(img)
      return False
    # had to add this check for autopacking files...
    if bpy.data.use_autopack and img.packed_file is not None:
      img.unpack(method='USE_ORIGINAL')
    img.filepath = tpath
    try:
      img.reload()
    except:
      return False

  if asset['assetType'] == 'hdr':
    # to display hdr thumbnails correctly, we use non-color, otherwise looks shifted
    image_utils.set_colorspace(img, 'Non-Color')
  else:
    image_utils.set_colorspace(img, 'sRGB')
  asset['thumb_small_loaded'] = True
  return True

def load_previews():
  results = global_vars.DATA.get('search results')
  if results is not None:
    for i, result in enumerate(results):
      load_preview(result)

#  line splitting for longer texts...
def split_subs(text, threshold=40):
  if text == '':
    return []
  # temporarily disable this, to be able to do this in drawing code

  text = text.rstrip()
  text = text.replace('\r\n', '\n')

  lines = []

  while len(text) > threshold:
    # first handle if there's an \n line ending
    i_rn = text.find('\n')
    if 1 < i_rn < threshold:
      i = i_rn
      text = text.replace('\n', '', 1)
    else:
      i = text.rfind(' ', 0, threshold)
      i1 = text.rfind(',', 0, threshold)
      i2 = text.rfind('.', 0, threshold)
      i = max(i, i1, i2)
      if i <= 0:
        i = threshold
    lines.append(text[:i])
    text = text[i:]
  lines.append(text)
  return lines

def list_to_str(input):
  output = ''
  for i, text in enumerate(input):
    output += text
    if i < len(input) - 1:
      output += ', '
  return output

def writeblock(t, input, width=40):  # for longer texts
  dlines = split_subs(input, threshold=width)
  for i, l in enumerate(dlines):
    t += '%s\n' % l
  return t

def writeblockm(tooltip, mdata, key='', pretext=None, width=40):  # for longer texts
  if mdata.get(key) == None:
    return tooltip
  else:
    intext = mdata[key]
    if type(intext) == list:
      intext = list_to_str(intext)
    if type(intext) == float:
      intext = round(intext, 3)
    intext = str(intext)
    if intext.rstrip() == '':
      return tooltip
    if pretext == None:
      pretext = key
    if pretext != '':
      pretext = pretext + ': '
    text = pretext + intext
    dlines = split_subs(text, threshold=width)
    for i, l in enumerate(dlines):
      tooltip += '%s\n' % l

  return tooltip

def has(mdata, prop):
  if mdata.get(prop) is not None and mdata[prop] is not None and mdata[prop] is not False:
    return True
  else:
    return False

def generate_tooltip(mdata):
  col_w = 40
  if type(mdata['parameters']) == list:
    mparams = utils.params_to_dict(mdata['parameters'])
  else:
    mparams = mdata['parameters']
  t = ''
  t = writeblock(t, mdata['displayName'], width=int(col_w * .6))
  # t += '\n'

  # t = writeblockm(t, mdata, key='description', pretext='', width=col_w)
  return t

def generate_author_textblock(adata):
  t = ''

  if adata not in (None, ''):
    col_w = 2000
    if len(adata['firstName'] + adata['lastName']) > 0:
      t = '%s %s\n' % (adata['firstName'], adata['lastName'])
      t += '\n'
      if adata.get('aboutMe') is not None:
        t = writeblockm(t, adata, key='aboutMe', pretext='', width=col_w)
  return t

def write_gravatar(a_id, gravatar_path):
  '''
  Write down gravatar path, as a result of thread-based gravatar image download.
  This should happen on timer in queue.
  '''
  # print('write author', a_id, type(a_id))
  authors = global_vars.DATA['bkit authors']
  if authors.get(a_id) is not None:
    adata = authors.get(a_id)
    adata['gravatarImg'] = gravatar_path

def fetch_gravatar(adata=None):
  '''
  Gets avatars from blenderkit server
  Parameters
  ----------
  adata - author data from elastic search result

  '''

  # fetch new avatars if available already
  if adata.get('avatar128') is not None:
    avatar_path = paths.get_temp_dir(subdir='bkit_g/') + adata['id'] + '.jpg'
    if os.path.exists(avatar_path):
      tasks_queue.add_task((write_gravatar, (adata['id'], avatar_path)))
      return

    url = global_vars.SERVER + adata['avatar128']
    r = rerequests.get(url, stream=False, headers=utils.get_headers())
    # print(r.body)
    if r.status_code == 200:
      # print(url)
      # print(r.headers['content-disposition'])
      with open(avatar_path, 'wb') as f:
        f.write(r.content)
      tasks_queue.add_task((write_gravatar, (adata['id'], avatar_path)))
    elif r.status_code == '404':
      adata['avatar128'] = None
      utils.p('avatar for author not available.')
    return

  # older gravatar code
  if adata.get('gravatarHash') is not None:
    gravatar_path = paths.get_temp_dir(subdir='bkit_g/') + adata['gravatarHash'] + '.jpg'

    if os.path.exists(gravatar_path):
      tasks_queue.add_task((write_gravatar, (adata['id'], gravatar_path)))
      return;

    url = "https://www.gravatar.com/avatar/" + adata['gravatarHash'] + '?d=404'
    r = rerequests.get(url, stream=False, headers=utils.get_headers())
    if r.status_code == 200:
      with open(gravatar_path, 'wb') as f:
        f.write(r.content)
      tasks_queue.add_task((write_gravatar, (adata['id'], gravatar_path)))
    elif r.status_code == '404':
      adata['gravatarHash'] = None
      utils.p('gravatar for author not available.')

fetching_gravatars = {}

def get_author(r):
  ''' Writes author info (now from search results) and fetches gravatar if needed.
  this is now tweaked to be able to get authors from
  '''
  global fetching_gravatars

  a_id = str(r['author']['id'])
  preferences = bpy.context.preferences.addons['blenderkit'].preferences
  authors = global_vars.DATA.get('bkit authors', {})
  if authors == {}:
    global_vars.DATA['bkit authors'] = authors
  a = authors.get(a_id)
  if a is None:  # or a is '' or (a.get('gravatarHash') is not None and a.get('gravatarImg') is None):
    a = r['author']
    a['id'] = a_id
    a['tooltip'] = generate_author_textblock(a)

    authors[a_id] = a
    if fetching_gravatars.get(a['id']) is None:
      fetching_gravatars[a['id']] = True

    thread = threading.Thread(target=fetch_gravatar, args=(a.copy(),), daemon=True)
    thread.start()
  return a

def write_profile(adata):
  user = adata['user']
  # we have to convert to MiB here, numbers too big for python int type
  if user.get('sumAssetFilesSize') is not None:
    user['sumAssetFilesSize'] /= (1024 * 1024)
  if user.get('sumPrivateAssetFilesSize') is not None:
    user['sumPrivateAssetFilesSize'] /= (1024 * 1024)
  if user.get('remainingPrivateQuota') is not None:
    user['remainingPrivateQuota'] /= (1024 * 1024)

  if adata.get('canEditAllAssets') is True:
    user['exmenu'] = True
  else:
    user['exmenu'] = False

  global_vars.DATA['bkit profile'] = adata

def request_profile(api_key):
  a_url = f'{paths.BLENDERKIT_API}/me/'
  headers = utils.get_headers(api_key)
  r = rerequests.get(a_url, headers=headers)
  adata = r.json()
  if adata.get('user') is None:
    utils.p(adata)
    utils.p('getting profile failed')
    return None
  return adata

def fetch_profile(api_key):
  try:
    adata = request_profile(api_key)
    if adata is not None:
      tasks_queue.add_task((write_profile, (adata,)))
  except Exception as e:
    bk_logger.error(e)

def get_profile():
  preferences = bpy.context.preferences.addons['blenderkit'].preferences
  profile = global_vars.DATA.get('bkit profile')
  thread = threading.Thread(target=fetch_profile, args=(preferences.api_key,), daemon=True)
  thread.start()

  return profile

def query_to_url(query={}, params={}):
  # build a new request
  url = f'{paths.BLENDERKIT_API}/search/'

  # build request manually
  # TODO use real queries
  requeststring = '?query='
  #
  if query.get('query') not in ('', None):
    requeststring += query['query'].lower()
  for i, q in enumerate(query):
    if q != 'query':
      requeststring += '+'
      requeststring += q + ':' + str(query[q]).lower()

  # add dict_parameters to make results smaller
  # result ordering: _score - relevance, score - BlenderKit score
  order = []
  if params['free_first']:
    order = ['-is_free', ]
  if query.get('query') is None and query.get('category_subtree') == None:
    # assumes no keywords and no category, thus an empty search that is triggered on start.
    # orders by last core file upload
    if query.get('verification_status') == 'uploaded':
      # for validators, sort uploaded from oldest
      order.append('created')
    else:
      order.append('-last_upload')
  elif query.get('author_id') is not None and utils.profile_is_validator():

    order.append('-created')
  else:
    if query.get('category_subtree') is not None:
      order.append('-score,_score')
    else:
      order.append('_score')
  if requeststring.find('+order:') == -1:
    requeststring += '+order:' + ','.join(order)
  requeststring += '&dict_parameters=1'

  requeststring += '&page_size=' + str(params['page_size'])
  requeststring += '&addon_version=%s' % params['addon_version']
  if not (query.get('query') and query.get('query').find('asset_base_id')>-1):
    requeststring += '&blender_version=%s' % params['blender_version']
  if params.get('scene_uuid') is not None:
    requeststring += '&scene_uuid=%s' % params['scene_uuid']
  # print('params', params)
  urlquery = url + requeststring
  return urlquery

def parse_html_formated_error(text):
  report = text[text.find('<title>') + 7: text.find('</title>')]

  return report

def build_query_common(query, props):
  '''add shared parameters to query'''
  query_common = {}
  if props.search_keywords != '':
    # keywords = urllib.parse.urlencode(props.search_keywords)
    keywords = props.search_keywords.replace('&', '%26')
    query_common["query"] = keywords

  if props.search_verification_status != 'ALL' and utils.profile_is_validator():
    query_common['verification_status'] = props.search_verification_status.lower()

  if props.unrated_only and utils.profile_is_validator():
    query["quality_count"] = 0

  if props.search_file_size:
    query_common["files_size_gte"] = props.search_file_size_min * 1024 * 1024
    query_common["files_size_lte"] = props.search_file_size_max * 1024 * 1024

  if props.quality_limit > 0:
    query["quality_gte"] = props.quality_limit

  query.update(query_common)

def build_query_model():
  '''use all search input to request results from server'''

  props = bpy.context.window_manager.blenderkit_models
  query = {
    "asset_type": 'model',
    # "engine": props.search_engine,
    # "adult": props.search_adult,
  }
  if props.search_style != 'ANY':
    if props.search_style != 'OTHER':
      query["model_style"] = props.search_style
    else:
      query["model_style"] = props.search_style_other

  # the 'free_only' parametr gets moved to the search command and is used for ordering the assets as free first
  # if props.free_only:
  #     query["is_free"] = True

  if props.search_condition != 'UNSPECIFIED':
    query["condition"] = props.search_condition

  if props.search_design_year:
    query["designYear_gte"] = props.search_design_year_min
    query["designYear_lte"] = props.search_design_year_max
  if props.search_polycount:
    query["faceCount_gte"] = props.search_polycount_min
    query["faceCount_lte"] = props.search_polycount_max
  if props.search_texture_resolution:
    query["textureResolutionMax_gte"] = props.search_texture_resolution_min
    query["textureResolutionMax_lte"] = props.search_texture_resolution_max
  if props.search_animated:
    query["animated"] = True
  build_query_common(query, props)

  return query

def build_query_scene():
  '''use all search input to request results from server'''

  props = bpy.context.window_manager.blenderkit_scene
  query = {
    "asset_type": 'scene',
    # "engine": props.search_engine,
    # "adult": props.search_adult,
  }
  build_query_common(query, props)
  return query

def build_query_HDR():
  '''use all search input to request results from server'''

  props = bpy.context.window_manager.blenderkit_HDR
  query = {
    "asset_type": 'hdr',

    # "engine": props.search_engine,
    # "adult": props.search_adult,
  }
  if props.true_hdr:
    query["trueHDR"] = props.true_hdr
  build_query_common(query, props)
  return query

def build_query_material():
  props = bpy.context.window_manager.blenderkit_mat
  query = {
    "asset_type": 'material',

  }
  # if props.search_engine == 'NONE':
  #     query["engine"] = ''
  # if props.search_engine != 'OTHER':
  #     query["engine"] = props.search_engine
  # else:
  #     query["engine"] = props.search_engine_other
  if props.search_style != 'ANY':
    if props.search_style != 'OTHER':
      query["style"] = props.search_style
    else:
      query["style"] = props.search_style_other
  if props.search_procedural == 'TEXTURE_BASED':
    # todo this procedural hack should be replaced with the parameter
    query["textureResolutionMax_gte"] = 0
    # query["procedural"] = False
    if props.search_texture_resolution:
      query["textureResolutionMax_gte"] = props.search_texture_resolution_min
      query["textureResolutionMax_lte"] = props.search_texture_resolution_max

  elif props.search_procedural == "PROCEDURAL":
    # todo this procedural hack should be replaced with the parameter
    query["files_size_lte"] = 1024 * 1024
    # query["procedural"] = True

  build_query_common(query, props)

  return query

def build_query_texture():
  props = bpy.context.scene.blenderkit_tex
  query = {
    "asset_type": 'texture',

  }

  if props.search_style != 'ANY':
    if props.search_style != 'OTHER':
      query["search_style"] = props.search_style
    else:
      query["search_style"] = props.search_style_other

  build_query_common(query, props)

  return query

def build_query_brush():
  props = bpy.context.window_manager.blenderkit_brush

  brush_type = ''
  if bpy.context.sculpt_object is not None:
    brush_type = 'sculpt'

  elif bpy.context.image_paint_object:  # could be just else, but for future p
    brush_type = 'texture_paint'

  query = {
    "asset_type": 'brush',

    "mode": brush_type
  }

  build_query_common(query, props)

  return query

def add_search_process(query, params):
  global search_tasks

  if len(search_tasks) > 0:
    # just remove all running search tasks.
    # we can also kill them in daemon, but not so urgent now
    # TODO stop tasks in daemon?
    bk_logger.debug('Removing old search tasks')
    search_tasks = dict()

  tempdir = paths.get_temp_dir('%s_search' % query['asset_type'])
  headers = utils.get_headers(params['api_key'])

  if params.get('get_next'):
    urlquery = params['next']
  else:
    urlquery = query_to_url(query, params)

  data = {
    'PREFS': utils.get_prefs_dir(),
    'tempdir': tempdir,
    'urlquery': urlquery,
    'asset_type': query['asset_type'],
  }
  data.update(params)
  response = daemon_lib.search_asset(data)
  search_tasks[response['task_id']] = data

def get_search_simple(parameters, filepath=None, page_size=100, max_results=100000000, api_key=''):
  '''
  Searches and returns the

  Parameters
  ----------
  parameters - dict of blenderkit elastic parameters
  filepath - a file to save the results. If None, results are returned
  page_size - page size for retrieved results
  max_results - max results of the search
  api_key - BlenderKit api key

  Returns
  -------
  Returns search results as a list, and optionally saves to filepath

  '''
  headers = utils.get_headers(api_key)
  url = f'{paths.BLENDERKIT_API}/search/'
  requeststring = url + '?query='
  for p in parameters.keys():
    requeststring += f'+{p}:{parameters[p]}'

  requeststring += '&page_size=' + str(page_size)
  requeststring += '&dict_parameters=1'

  bk_logger.debug(requeststring)
  response = rerequests.get(requeststring, headers=headers)  # , params = rparameters)
  # print(response.json())
  search_results = response.json()

  results = []
  results.extend(search_results['results'])
  page_index = 2
  page_count = math.ceil(search_results['count'] / page_size)
  while search_results.get('next') and len(results) < max_results:
    bk_logger.info(f'getting page {page_index} , total pages {page_count}')
    response = rerequests.get(search_results['next'], headers=headers)  # , params = rparameters)
    search_results = response.json()
    # print(search_results)
    results.extend(search_results['results'])
    page_index += 1

  if not filepath:
    return results

  with open(filepath, 'w', encoding='utf-8') as s:
    json.dump(results, s, ensure_ascii=False, indent=4)
  bk_logger.info(f'retrieved {len(results)} assets from elastic search')
  return results

def get_single_asset(asset_base_id):
  preferences = bpy.context.preferences.addons['blenderkit'].preferences
  params = {
    'asset_base_id': asset_base_id
  }
  results = get_search_simple(params, api_key=preferences.api_key)
  if len(results) > 0:
    return results[0]
  return None

def search(category='', get_next=False, query = None, author_id=''):
  ''' initialize searching
  query : submit an already built query from search history
  '''

  if global_vars.DAEMON_ACCESSIBLE != True:
    reports.add_report('Cannot search, daemon is not accessible.', timeout = 2, type='ERROR')
    return

  # print(category,get_next,author_id)
  user_preferences = bpy.context.preferences.addons['blenderkit'].preferences

  scene = bpy.context.scene
  wm = bpy.context.window_manager
  ui_props = bpy.context.window_manager.blenderkitUI

  props = utils.get_search_props()
  # it's possible get_next was requested more than once.
  if props.is_searching and get_next == True:
    # print('return because of get next and searching is happening')
    return;

  if not query:
    if ui_props.asset_type == 'MODEL':
      if not hasattr(wm, 'blenderkit_models'):
        return;
      query = build_query_model()

    if ui_props.asset_type == 'SCENE':
      if not hasattr(wm, 'blenderkit_scene'):
        return;
      query = build_query_scene()

    if ui_props.asset_type == 'HDR':
      if not hasattr(wm, 'blenderkit_HDR'):
        return;
      query = build_query_HDR()

    if ui_props.asset_type == 'MATERIAL':
      if not hasattr(wm, 'blenderkit_mat'):
        return;

      query = build_query_material()

    if ui_props.asset_type == 'TEXTURE':
      if not hasattr(wm, 'blenderkit_tex'):
        return;
      # props = scene.blenderkit_tex
      # query = build_query_texture()

    if ui_props.asset_type == 'BRUSH':
      if not hasattr(wm, 'blenderkit_brush'):
        return;
      query = build_query_brush()

    # crop long searches
    if query.get('query'):
      if len(query['query']) > 50:
        query['query'] = strip_accents(query['query'])

      if len(query['query']) > 150:
        idx = query['query'].find(' ', 142)
        query['query'] = query['query'][:idx]

    if category != '':
      if utils.profile_is_validator() and user_preferences.categories_fix:
        query['category'] = category
      else:
        query['category_subtree'] = category

    if author_id != '':
      query['author_id'] = author_id

    elif props.own_only:
      # if user searches for [another] author, 'only my assets' is invalid. that's why in elif.
      profile = global_vars.DATA.get('bkit profile')
      if profile is not None:
        query['author_id'] = str(profile['user']['id'])

    if not get_next:
      #write to search history and check history length
      if len(global_vars.DATA['search history'])>0 and global_vars.DATA['search history'][-1] == query:
        # don't send same query again, when user clicks multiple times and waits e.t.c.
        return

      global_vars.DATA['search history'].append(query)
  # utils.p('searching')
  props.is_searching = True

  page_size = min(40, ui_props.wcount * user_preferences.max_assetbar_rows + 5)
  params = {
    'scene_uuid': bpy.context.scene.get('uuid', None),
    'addon_version': version_checker.get_addon_version(),
    'blender_version': version_checker.get_blender_version(),
    'api_key': user_preferences.api_key,
    'get_next': get_next,
    'free_first': props.free_only,
    'page_size': page_size,
  }

  orig_results = global_vars.DATA.get(f'bkit {ui_props.asset_type.lower()} search orig')
  if orig_results is not None and get_next:
    params['next'] = orig_results['next']
  add_search_process(query, params)

  props.report = 'BlenderKit searching....'

def clean_filters():
  '''cleanup filters in case search needs to be reset, typicaly when asset id is copy pasted'''
  sprops = utils.get_search_props()
  ui_props = bpy.context.window_manager.blenderkitUI
  sprops.property_unset('own_only')
  sprops.property_unset('search_texture_resolution')
  sprops.property_unset('search_file_size')
  sprops.property_unset('search_procedural')
  sprops.property_unset('free_only')
  sprops.property_unset('quality_limit')
  if ui_props.asset_type == 'MODEL':
    sprops.property_unset('search_style')
    sprops.property_unset('search_condition')
    sprops.property_unset('search_design_year')
    sprops.property_unset('search_polycount')
    sprops.property_unset('search_animated')
  if ui_props.asset_type == 'HDR':
    sprops.true_hdr = False

def update_filters():
  sprops = utils.get_search_props()
  ui_props = bpy.context.window_manager.blenderkitUI
  fcommon = sprops.own_only or \
            sprops.search_texture_resolution or \
            sprops.search_file_size or \
            sprops.search_procedural != 'BOTH' or \
            sprops.free_only or \
            sprops.quality_limit > 0

  if ui_props.asset_type == 'MODEL':
    sprops.use_filters = fcommon or \
                         sprops.search_style != 'ANY' or \
                         sprops.search_condition != 'UNSPECIFIED' or \
                         sprops.search_design_year or \
                         sprops.search_polycount or \
                         sprops.search_animated
  elif ui_props.asset_type == 'MATERIAL':
    sprops.use_filters = fcommon
  elif ui_props.asset_type == 'HDR':
    sprops.use_filters = sprops.true_hdr

def search_update(self, context):
  # if self.search_keywords != '':
  update_filters()
  ui_props = bpy.context.window_manager.blenderkitUI
  if ui_props.down_up != 'SEARCH':
    ui_props.down_up = 'SEARCH'

  # here we tweak the input if it comes form the clipboard. we need to get rid of asset type and set it in UI
  sprops = utils.get_search_props()
  instr = 'asset_base_id:'
  atstr = 'asset_type:'
  kwds = sprops.search_keywords
  id_index = kwds.find(instr)
  if id_index>-1:
    asset_type_index = kwds.find(atstr)
    # if the asset type already isn't there it means this update function
    # was triggered by it's last iteration and needs to cancel
    if asset_type_index > -1:
      asset_type_string = kwds[asset_type_index:].lower()
      # uncertain length of the remaining string -  find as better method to check the presence of asset type
      if asset_type_string.find('model') > -1:
        target_asset_type = 'MODEL'
      elif asset_type_string.find('material') > -1:
        target_asset_type = 'MATERIAL'
      elif asset_type_string.find('brush') > -1:
        target_asset_type = 'BRUSH'
      elif asset_type_string.find('scene') > -1:
        target_asset_type = 'SCENE'
      elif asset_type_string.find('hdr') > -1:
        target_asset_type = 'HDR'
      if ui_props.asset_type != target_asset_type:
        sprops.search_keywords = ''
        ui_props.asset_type = target_asset_type

      # now we trim the input copypaste by anything extra that is there,
      # this is also a way for this function to recognize that it already has parsed the clipboard
      # the search props can have changed and this needs to transfer the data to the other field
      # this complex behaviour is here for the case where the user needs to paste manually into blender,
      # Otherwise it could be processed directly in the clipboard check function.
      sprops = utils.get_search_props()
      clean_filters()
      sprops.search_keywords = kwds[:asset_type_index].rstrip()
      # return here since writing into search keywords triggers this update function once more.
      return

  if global_vars.DAEMON_ACCESSIBLE:
    reports.add_report(f"Searching for: '{kwds}'", 2)
  search()

# accented_string is of type 'unicode'
def strip_accents(s):
  return ''.join(c for c in unicodedata.normalize('NFD', s)
                 if unicodedata.category(c) != 'Mn')

class SearchOperator(Operator):
  """Tooltip"""
  bl_idname = "view3d.blenderkit_search"
  bl_label = "BlenderKit asset search"
  bl_description = "Search online for assets"
  bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

  esc: BoolProperty(name="Escape window",
                    description="Escape window right after start",
                    default=False,
                    options={'SKIP_SAVE'}
                    )

  own: BoolProperty(name="own assets only",
                    description="Find all own assets",
                    default=False,
                    options={'SKIP_SAVE'})

  category: StringProperty(
    name="category",
    description="search only subtree of this category",
    default="",
    options={'SKIP_SAVE'}
  )

  author_id: StringProperty(
    name="Author ID",
    description="Author ID - search only assets by this author",
    default="",
    options={'SKIP_SAVE'}
  )

  get_next: BoolProperty(name="next page",
                         description="get next page from previous search",
                         default=False,
                         options={'SKIP_SAVE'}
                         )

  keywords: StringProperty(
    name="Keywords",
    description="Keywords",
    default="",
    options={'SKIP_SAVE'}
  )

  # close_window: BoolProperty(name='Close window',
  #                            description='Try to close the window below mouse before download',
  #                            default=False)

  tooltip: bpy.props.StringProperty(default='Runs search and displays the asset bar at the same time')

  @classmethod
  def description(cls, context, properties):
    return properties.tooltip

  @classmethod
  def poll(cls, context):
    return True

  def execute(self, context):
    # TODO ; this should all get transferred to properties of the search operator, so sprops don't have to be fetched here at all.
    if self.esc:
      bpy.ops.view3d.close_popup_button('INVOKE_DEFAULT')
    sprops = utils.get_search_props()
    if self.author_id != '':
      sprops.search_keywords = ''
    if self.keywords != '':
      sprops.search_keywords = self.keywords

    search(category=self.category, get_next=self.get_next, author_id=self.author_id)
    # bpy.ops.view3d.blenderkit_asset_bar_widget()

    return {'FINISHED'}

  # def invoke(self, context, event):
  #     if self.close_window:
  #         context.window.cursor_warp(event.mouse_x, event.mouse_y - 100);
  #         context.area.tag_redraw()
  #
  #         context.window.cursor_warp(event.mouse_x, event.mouse_y);
  #     return self. execute(context)

class UrlOperator(Operator):
  """"""
  bl_idname = "wm.blenderkit_url"
  bl_label = ""
  bl_description = "Search online for assets"
  bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

  tooltip: bpy.props.StringProperty(default='Open a web page')
  url: bpy.props.StringProperty(default='Runs search and displays the asset bar at the same time')

  @classmethod
  def description(cls, context, properties):
    return properties.tooltip

  def execute(self, context):
    bpy.ops.wm.url_open(url=self.url)
    return {'FINISHED'}

class TooltipLabelOperator(Operator):
  """"""
  bl_idname = "wm.blenderkit_tooltip"
  bl_label = ""
  bl_description = "Empty operator to be able to create tooltips on labels in UI"
  bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

  tooltip: bpy.props.StringProperty(default='Open a web page')

  @classmethod
  def description(cls, context, properties):
    return properties.tooltip

  def execute(self, context):
    return {'FINISHED'}

classes = [
  SearchOperator,
  UrlOperator,
  TooltipLabelOperator
]

def register_search():
  bpy.app.handlers.load_post.append(scene_load)
  bpy.app.handlers.load_post.append(undo_post_reload_previews)
  bpy.app.handlers.undo_post.append(undo_post_reload_previews)
  bpy.app.handlers.undo_pre.append(undo_pre_end_assetbar)

  for c in classes:
    bpy.utils.register_class(c)

def unregister_search():
  bpy.app.handlers.load_post.remove(scene_load)

  for c in classes:
    bpy.utils.unregister_class(c)


import torch
if torch.__version__ >= '1.8':
    import torch_npu

import torch.nn.functional as F
# from SSD import _C as C
import numpy as np
import itertools

from math import sqrt

def calc_iou_tensor(box1, box2):
    """
        Calculation of IoU based on two boxes tensor,
        Reference to https://github.com/kuangliu/pytorch-ssd
        input:
            box1 (N, 4)
            box2 (M, 4)
        output:
            IoU (N, M)
    """
    N = box1.size(0)
    M = box2.size(0)

    be1 = box1.unsqueeze(1).expand(-1, M, -1)
    be2 = box2.unsqueeze(0).expand(N, -1, -1)

    # Left Top & Right Bottom
    lt = torch.max(be1[:, :, :2], be2[:, :, :2])
    rb = torch.min(be1[:, :, 2:], be2[:, :, 2:])

    delta = rb - lt
    delta[delta < 0] = 0
    intersect = delta[:, :, 0] * delta[:, :, 1]

    delta1 = be1[:, :, 2:] - be1[:, :, :2]
    area1 = delta1[:, :, 0] * delta1[:, :, 1]
    delta2 = be2[:, :, 2:] - be2[:, :, :2]
    area2 = delta2[:, :, 0] * delta2[:, :, 1]

    iou = intersect / (area1 + area2 - intersect)
    return iou

class DefaultBoxes(object):
    def __init__(self, fig_size, feat_size, steps, scales, aspect_ratios, \
                       scale_xy=0.1, scale_wh=0.2):

        self.feat_size = feat_size
        self.fig_size = fig_size

        self.scale_xy_ = scale_xy
        self.scale_wh_ = scale_wh

        # According to https://github.com/weiliu89/caffe
        # Calculation method slightly different from paper
        self.steps = steps
        self.scales = scales

        fk = fig_size/np.array(steps)
        self.aspect_ratios = aspect_ratios

        self.default_boxes = []
        # size of feature and number of feature
        for idx, sfeat in enumerate(self.feat_size):

            sk1 = scales[idx]/fig_size
            sk2 = scales[idx+1]/fig_size
            sk3 = sqrt(sk1*sk2)
            all_sizes = [(sk1, sk1), (sk3, sk3)]

            for alpha in aspect_ratios[idx]:
                w, h = sk1*sqrt(alpha), sk1/sqrt(alpha)
                all_sizes.append((w, h))
                all_sizes.append((h, w))
            for w, h in all_sizes:
                for i, j in itertools.product(range(sfeat), repeat=2):
                    cx, cy = (j+0.5)/fk[idx], (i+0.5)/fk[idx]
                    self.default_boxes.append((cx, cy, w, h))

        self.dboxes = torch.tensor(self.default_boxes, dtype=torch.float).cpu()
        self.dboxes.clamp_(min=0, max=1)
        # For IoU calculation
        self.dboxes_ltrb = self.dboxes.clone()
        self.dboxes_ltrb[:, 0] = self.dboxes[:, 0] - 0.5*self.dboxes[:, 2]
        self.dboxes_ltrb[:, 1] = self.dboxes[:, 1] - 0.5*self.dboxes[:, 3]
        self.dboxes_ltrb[:, 2] = self.dboxes[:, 0] + 0.5*self.dboxes[:, 2]
        self.dboxes_ltrb[:, 3] = self.dboxes[:, 1] + 0.5*self.dboxes[:, 3]

    @property
    def scale_xy(self):
        return self.scale_xy_

    @property
    def scale_wh(self):
        return self.scale_wh_

    def __call__(self, order="ltrb"):
        if order == "ltrb": return self.dboxes_ltrb
        if order == "xywh": return self.dboxes

# This class is from https://github.com/kuangliu/pytorch-ssd
class Encoder(object):
    """
        Inspired by https://github.com/kuangliu/pytorch-ssd
        Transform between (bboxes, lables) <-> SSD output

        dboxes: default boxes in size 8732 x 4,
            encoder: input ltrb format, output xywh format
            decoder: input xywh format, output ltrb format

        encode:
            input  : bboxes_in (Tensor nboxes x 4), labels_in (Tensor nboxes)
            output : bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732)
            criteria : IoU threshold of bboexes

        decode:
            input  : bboxes_in (Tensor 8732 x 4), scores_in (Tensor 8732 x nitems)
            output : bboxes_out (Tensor nboxes x 4), labels_out (Tensor nboxes)
            criteria : IoU threshold of bboexes
            max_output : maximum number of output bboxes
    """

    def __init__(self, dboxes):
        self.dboxes = dboxes(order="ltrb")
        self.dboxes_xywh = dboxes(order="xywh").unsqueeze(dim=0)
        self.nboxes = self.dboxes.size(0)
        #print("# Bounding boxes: {}".format(self.nboxes))
        self.scale_xy = dboxes.scale_xy
        self.scale_wh = dboxes.scale_wh

        self.dboxes = self.dboxes.cpu()
        self.dboxes_xywh = self.dboxes_xywh.cpu()
        self.dboxes_npu = self.dboxes.npu()
        self.dboxes_xywh_npu = self.dboxes_xywh.npu()
    def encode(self, bboxes_in, labels_in, criteria = 0.5):

        # try:
        # print("bboxes_in.shape, self.dboxes.shape", bboxes_in.shape, self.dboxes.shape)
        ious = calc_iou_tensor(bboxes_in, self.dboxes)
        best_dbox_ious, best_dbox_idx = ious.max(dim=0)
        best_bbox_ious, best_bbox_idx = ious.max(dim=1)

        # set best ious 2.0
        best_dbox_ious.index_fill_(0, best_bbox_idx, 2.0)

        idx = torch.arange(0, best_bbox_idx.size(0), dtype=torch.int64)
        best_dbox_idx[best_bbox_idx[idx]] = idx

        # filter IoU > 0.5
        masks = best_dbox_ious > criteria
        labels_out = torch.zeros(self.nboxes, dtype=torch.long).cpu()

        labels_out[masks] = labels_in[best_dbox_idx[masks]]
        bboxes_out = self.dboxes.clone()
        bboxes_out[masks, :] = bboxes_in[best_dbox_idx[masks], :]
        # Transform format to xywh format
        x, y, w, h = 0.5*(bboxes_out[:, 0] + bboxes_out[:, 2]), \
                    0.5*(bboxes_out[:, 1] + bboxes_out[:, 3]), \
                    -bboxes_out[:, 0] + bboxes_out[:, 2], \
                    -bboxes_out[:, 1] + bboxes_out[:, 3]
        bboxes_out[:, 0] = x
        bboxes_out[:, 1] = y
        bboxes_out[:, 2] = w
        bboxes_out[:, 3] = h
        # except:
        #     labels_out = torch.zeros(self.nboxes, dtype=torch.long)
        #     bboxes_out = torch.zeros(self.nboxes, 4)
        # print("bboxes_out.shape, labels_out.shape", bboxes_out.shape, labels_out.shape)
        return bboxes_out, labels_out

    def scale_back_batch(self, bboxes_in, scores_in):
        """
            Do scale and transform from xywh to ltrb
            suppose input Nx4xnum_bbox Nxlabel_numxnum_bbox
        """
        bboxes_in = bboxes_in.permute(0, 2, 1)
        scores_in = scores_in.permute(0, 2, 1)

        bboxes_in[:, :, :2] = self.scale_xy*bboxes_in[:, :, :2]
        bboxes_in[:, :, 2:] = self.scale_wh*bboxes_in[:, :, 2:]

        bboxes_in[:, :, :2] = bboxes_in[:, :, :2]*self.dboxes_xywh_npu[:, :, 2:] + self.dboxes_xywh_npu[:, :, :2]
        bboxes_in[:, :, 2:] = bboxes_in[:, :, 2:].exp()*self.dboxes_xywh_npu[:, :, 2:]

        # Transform format to ltrb
        l, t, r, b = bboxes_in[:, :, 0] - 0.5*bboxes_in[:, :, 2],\
                     bboxes_in[:, :, 1] - 0.5*bboxes_in[:, :, 3],\
                     bboxes_in[:, :, 0] + 0.5*bboxes_in[:, :, 2],\
                     bboxes_in[:, :, 1] + 0.5*bboxes_in[:, :, 3]

        bboxes_in[:, :, 0] = l
        bboxes_in[:, :, 1] = t
        bboxes_in[:, :, 2] = r
        bboxes_in[:, :, 3] = b

        return bboxes_in, F.softmax(scores_in, dim=-1)

    def decode_batch(self, bboxes_in, scores_in,  criteria = 0.45, max_output=200):
        bboxes, probs = self.scale_back_batch(bboxes_in, scores_in)
        N,A,C = probs.shape
        bboxes = bboxes.unsqueeze(-2).repeat([1,1,80,1])
        probs = probs[...,1:]
        nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_num = torch_npu.npu_batch_nms(bboxes.half(), probs.half(),
                                                                                0.05, criteria,
                                                                                max_output, max_output)

        return nmsed_boxes, nmsed_classes, nmsed_scores

    # perform non-maximum suppression###############################################################
    def decode_single(self, bboxes_in, scores_in, criteria, max_output, max_num=200):
    #     # Reference to https://github.com/amdegroot/ssd.pytorch
        scores_in = scores_in.transpose(1, 0)
        #
        #     # Sort every row (in hopefully a single kernel launch)
        #     # NOTE: Not masked out things yet
        #     # NOTE: descending sort is easier to reason about
        #     # NOTE: Indices are to _global_ bboxes, we're not going to mask them
        score_sorted, score_sorted_idx = scores_in.sort(dim=1, descending=True)

        #     # Now generate the mask on the sorted scores
        mask = score_sorted > 0.05

        #     # number of default boxes per class that have a score > 0.05
        splits = mask.sum(dim=1).tolist()

        # only keep scores & indices for default boxes that contribute to this class
        # NOTE: Not masking out bboxes, all indices are global
        score_sorted = score_sorted[mask].split(splits)
        score_sorted_idx = score_sorted_idx[mask].split(splits)

        # assemble prefix sum of splits
        offsets = torch.tensor([0] + list(itertools.accumulate(splits)), dtype=torch.int32, device=bboxes_in.device)
        bboxes_out, scores_out, labels_out = C.nms(1, # N
                                                  scores_in.shape[0],
                                                  offsets,
                                                  torch.cat(score_sorted),
                                                  torch.cat(score_sorted_idx),
                                                  bboxes_in.contiguous(), # VITAL otherwise we get bad results :(
                                                  criteria,
                                                  max_num)

        _, max_ids = scores_out.sort(dim=0)
        max_ids = max_ids[-max_output:]
        return bboxes_out[max_ids, :], labels_out[max_ids], scores_out[max_ids]

def npu_multiclass_nms(multi_bboxes,
                       multi_scores,
                       score_thr=0.05,
                       nms_thr=0.5,
                       max_num=200,
                       max_output = 200,
                       score_factors=None):
    """NMS for multi-class bboxes using npu api.

    Origin implement from mmdetection is
    https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/post_processing/bbox_nms.py#L7

    This interface is similar to the original interface, but not exactly the same.

    Args:
        multi_bboxes (Tensor): shape (n, #class, 4) or (n, 4)
        multi_scores (Tensor): shape (n, #class+1), where the last column
            contains scores of the background class, but this will be ignored.
            On npu, in order to keep the semantics unblocked, we will unify the dimensions
        score_thr (float): bbox threshold, bboxes with scores lower than it
            will not be considered.
        nms_thr (float): NMS IoU threshold. In the original implementation, a dictionary of {"iou_threshold": 0.45}
            was passed, which is simplified here.
        max_num (int): if there are more than max_num bboxes after NMS,
            only top max_num will be kept; if there are less than max_num bboxes after NMS,
            the output will zero pad to max_num. On the NPU, the memory needs to be requested in advance,
            so the current max_num cannot be set to -1 at present
        score_factors (Tensor): The factors multiplied to scores before applying NMS

    Returns:
        tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). Labels are 0-based.
    """

    num_classes = multi_scores.size(1)-1
    num_boxes = multi_scores.size(0)
    #print(num_classes,num_boxes)                     #####################
    if score_factors is not None:
        multi_scores = multi_scores[:, :-1] * score_factors[:, None]
    else:
        multi_scores = multi_scores[:, :-1]
    multi_bboxes = multi_bboxes.reshape(1, num_boxes, multi_bboxes.numel() // 4 // num_boxes, 4)
    multi_scores = multi_scores.reshape(1, num_boxes, num_classes)

    nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_num = torch_npu.npu_batch_nms(multi_bboxes.half(),
                                                                                multi_scores.half(),
                                                                                score_thr, nms_thr,
                                                                                max_num, max_num)
    print(nmsed_boxes.shape, nmsed_scores.shape, nmsed_classes.shape, nmsed_num.shape)
    nmsed_boxes = nmsed_boxes.reshape(nmsed_boxes.shape[1:])
    nmsed_scores = nmsed_scores.reshape(nmsed_scores.shape[1])
    nmsed_classes = nmsed_classes.reshape(nmsed_classes.shape[1])
    #print(nmsed_boxes.shape,nmsed_scores.shape,nmsed_classes.shape)
    #return torch.cat([nmsed_boxes, nmsed_scores[:, None]], -1), nmsed_classes
    _, max_ids = nmsed_scores.sort(dim=0)
    max_ids = max_ids[-max_output:]
    #print(type(nmsed_classes))
    ones = torch.ones(200).npu()
    nmsed_classes = nmsed_classes + ones
    #print(nmsed_boxes[max_ids, :], nmsed_classes[max_ids], nmsed_scores[max_ids])
    #return nmsed_boxes, nmsed_classes, nmsed_scores
    return nmsed_boxes[max_ids, :], nmsed_classes[max_ids], nmsed_scores[max_ids]

def npu_batched_multiclass_nms(
        multi_bboxes,
        multi_scores,
        score_thr=0.05,
        nms_thr=0.5,
        max_num=200,
        score_factors=None):
    """NMS for batched multi-class bboxes using npu api.

    Origin implement from mmdetection is
    https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/post_processing/bbox_nms.py#L7

    This interface is similar to the original interface, but not exactly the same.
    This interface implements the nms method under batch.

    Args:
        multi_bboxes (Tensor): shape (bs, n, #class, 4) or (bs, n, 4)
        multi_scores (Tensor): shape (bs, n, #class+1), where the last column
            contains scores of the background class, but this will be ignored.
            On npu, in order to keep the semantics unblocked, we will unify the dimensions
        score_thr (float): bbox threshold, bboxes with scores lower than it
            will not be considered.
        nms_thr (float): NMS IoU threshold. In the original implementation, a dictionary of {"iou_threshold": 0.45}
            was passed, which is simplified here.
        max_num (int): if there are more than max_num bboxes after NMS,
            only top max_num will be kept; if there are less than max_num bboxes after NMS,
            the output will zero pad to max_num. On the NPU, the memory needs to be requested in advance,
            so the current max_num cannot be set to -1 at present
        score_factors (Tensor): The factors multiplied to scores before applying NMS

    Returns:
        tuple: (bboxes, labels), tensors of shape (bs, k, 5) and (bs, k, 1). Labels are 0-based.
    """

    num_classes = multi_scores.size(2) - 1
    num_boxes = multi_scores.size(1)
    batch_size = multi_scores.size(0)
    if score_factors is not None:
        multi_scores = multi_scores[..., :-1] * score_factors[..., None]
    else:
        multi_scores = multi_scores[..., :-1]
    multi_bboxes = multi_bboxes.reshape(batch_size, num_boxes, multi_bboxes.numel() // 4 // num_boxes // batch_size, 4)
    multi_scores = multi_scores.reshape(batch_size, num_boxes, num_classes)

    nmsed_boxes, nmsed_scores, nmsed_classes, nmsed_num = torch_npu.npu_batch_nms(multi_bboxes.half(),
                                                                                multi_scores.half(),
                                                                                score_thr, nms_thr,
                                                                                max_num, max_num)

    return torch.cat([nmsed_boxes, nmsed_scores[..., None]], -1), nmsed_classes, nmsed_scores

def dboxes300_coco():
    figsize = 300
    feat_size = [38, 19, 10, 5, 3, 1]
    steps = [8, 16, 32, 64, 100, 300]
    # use the scales here: https://github.com/amdegroot/ssd.pytorch/blob/master/data/config.py
    scales = [21, 45, 99, 153, 207, 261, 315]
    aspect_ratios = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
    dboxes = DefaultBoxes(figsize, feat_size, steps, scales, aspect_ratios)
    return dboxes

def build_ssd300_coder():
    return Encoder(dboxes300_coco())


'''
Python 3.0
PyLucene 8.11.0
Create Required PyLucene Index
Retrieve Top-K Relevant Document for String Query
'''

# Importing Necessary Libraries
import os
import nltk
import math
from collections import defaultdict
import lucene
from java.io import File
from org.apache.lucene.document import Document, Field, FieldType
from org.apache.lucene.store import FSDirectory
from org.apache.lucene.analysis.standard import StandardAnalyzer
from org.apache.lucene.index import IndexWriter, IndexWriterConfig, IndexOptions, DirectoryReader, Term
from org.apache.lucene.search import IndexSearcher, TermQuery
from org.apache.lucene.util import BytesRefIterator

# Add the Relevant File/Folder Paths
stopwords_path = 'Enter The Full Path to Folder with File containing Stop Words'
stopwords_filename = 'File Name of the File with Stop Words (With Extention)'
collection_path = 'Enter The Full Path to Folder with Document Collection'
index_path = 'Enter The Full Path to Folder where PyLucene Index will be Stored'

# Tokenizer to keep only Alpha-Numeric Terms  [a-zA-Z0-9_]
tokenizer = nltk.RegexpTokenizer(r"\w+")

# Initializa JVM
lucene.initVM()

'''
Pre-Process Document Data, Query String
     - Remove All Punctuations
     - Remove Stopwords (Optional)
     - Convert to Lower Case
'''
def PreProcess(data):
    data = tokenizer.tokenize(data)
    data = [token for token in data if token not in stopwords] # Skip this Line if No Stopwords Are Imported
    data = ' '.join(data)
    data = data.lower()
    return data

# Transform Query to Vector
def QueryVectorGenerator(query):
        query = PreProcess(query)
        queryTerms = query.split()
        queryVector = defaultdict(lambda: 0)
        for term in queryTerms:
            queryVector[term] += 1
        return queryVector

# Transform Document to Vector in order of Query Vector and Terms
def DocumentVectorGenerator(docId, queryVector):
        documentVector = {}
        for term in queryVector.keys():
            documentVector[term] = TermFreequencyInDoc(term, docId)
        return documentVector

# Return a List of All Documents that contain the Term
def DocsWithTerm(term):
    termQuery = TermQuery(Term('DocData', term))
    hits = searcher.search(termQuery, lnCorpus).scoreDocs
    docIds = []
    i = 0
    while i < len(hits):
        docIds.append(hits[i].doc)
        i = i + 1
    return docIds

# Calculate TF Value
def TermFreequencyInDoc(termText, docId):
    termsList = indexReader.getTermVector(int(docId), 'DocData')
    terms = termsList.iterator()
    for term in BytesRefIterator.cast_(terms):
        dpEnum = terms.postings(None)
        dpEnum.nextDoc()
        tf = dpEnum.freq()
        if tf!=0:
            return tf
    return 0

# Calculate IDF Value
def idf(term):
    tf = indexReader.docFreq(Term('DocData', PreProcess(term)))
    return math.log((lnCorpus + 1) /(tf + 1), 10)

# Calculate TF-IDF Value
def tfidf(idf, wtd):
    return idf * wtd

# Transform Vector to TF-IDF for a given Document
def VectorToTFIDF(vector, docId):
    for term in vector.keys():
        # Calculate Term Weight
        tf = TermFreequencyInDoc(term, docId)
        if tf == 0:
            wt = 0
        else:
            wt = 1 + math.log(float(tf), 10)
        vector[term] = tfidf(idf(term), wt)
    return vector

# Transform Query Vector to TF-IDF
def QueryVectorToTFIDF(queryVector):
    for term, tf in queryVector.items():
        queryVector[term] = tfidf(idf(term), 1 + math.log(tf, 10))
    return queryVector

# Compute Cosine Similarity Score with L2 Normalization
def computeCosine(vector1, vector2):
    len_v1 = len(vector1)
    len_v2 = len(vector2)
    # Make Both Vectors of Same Length
    if len_v1 != len_v2:
        if len_v1 > len_v2:
            for i in range(len_v1 - len_v2):
                vector2.append(0)
        else:
            for i in range(len_v2 - len_v1):
                vector1.append(0)
    # Calculate Cross Product
    sq_v1 = 0
    sq_v2 = 0
    crossProduct = 0
    for i in range(len(vector1)):
        crossProduct = crossProduct + vector1[i] * vector2[i]
        sq_v1 += math.pow(vector1[i], 2)
        sq_v2 += math.pow(vector2[i], 2)
    # Calculate Cosine Similarity Score
    cs = crossProduct / ((math.sqrt(sq_v1) * math.sqrt(sq_v2)))
    return cs

# Create PyLucene Index
def Create_IndexFn(indexpath,collectionpath,stopwords_path,stopwords_filename):

    # Getting StopWords from a File (Optional)
    global stopwords
    '''
    For Windows Users, Use the following line of code to open file
        with open(path+'\'+filename) as f:
    '''
    with open(stopwords_path+stopwords_filename) as f:
        stopwords = f.read().split('\n')

    # Specify the Path where Index will be Stored
    global indexPath
    indexPath = File(indexpath).toPath()
    global indexDir
    indexDir = FSDirectory.open(indexPath)

    # Specify Analyser and Index Writer Configuration
    global writerConfig
    writerConfig = IndexWriterConfig(StandardAnalyzer())
    global writer
    writer = IndexWriter(indexDir, writerConfig)

    # Set Field Parameter for DocumentName,DocumentId
    field_filename = FieldType()
    field_filename.setStored(True)
    field_filename.setTokenized(False)
    field_filename.setIndexOptions(IndexOptions.DOCS)

    # Set Field Parameter for DocumentData
    field_filedata = FieldType()
    field_filedata.setStored(True)
    field_filedata.setTokenized(True)
    field_filedata.setStoreTermVectors(True)
    field_filedata.setStoreTermVectorPositions(True)
    field_filedata.setIndexOptions(IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS)

    # Change Working Directory to Collection Path
    os.chdir(collectionpath)

    # Import and Pre-Process the Document Collection
    for file in os.listdir():
        docId = 0
        if file.endswith(".txt"):
            file_path = f"./{file}"
            with open(file_path, "r") as f:
                doc_data = f.read()
            # Add DocumentId and Pre-Processed Document Data to PyLucene Index
            if len(doc_data)>0:
                doc_data = PreProcess(doc_data)
                doc = Document()
                doc.add(Field('DocName', str(file[:-4]), field_filename))
                doc.add(Field('DocId', docId, field_filename))
                doc.add(Field('DocData', doc_data, field_filedata))
                writer.addDocument(doc)
                docId += 1

    # Initialize Index Reader and Searcher
    global indexReader
    global searcher
    global lnCorpus
    indexReader = DirectoryReader.open(writer)
    searcher = IndexSearcher(indexReader)
    lnCorpus = indexReader.numDocs()

# Retrieve Top-K Relevant Documents for Specified Query
def TopKRetrieval(query, k):
    queryVector = QueryVectorGenerator(query)
    docIds = []
    docVectors = {}
    computedSimilarity = []
    topKDocInformation = {}

    # Get Documents with the Term
    for term in queryVector.keys():
        docIds += DocsWithTerm(term)

    # Create Vector for Documents
    for docId in docIds:
        docVectors[docId] = DocumentVectorGenerator(docId, queryVector)

    # Transform Document Vectors to TF-IDF Values
    for docId, vector in docVectors.items():
        docVectors[docId] = VectorToTFIDF(vector, docId)

    # Transform Query Vector to TF-IDF Values
    queryVector = QueryVectorToTFIDF(queryVector)

    # Compute Cosine Similarity Scores
    for docId, vector in docVectors.items():
        computedSimilarity.append({"DocId": docId, "CosineSimilarityScore": computeCosine(list(queryVector.values()), list(vector.values()))})
    computedSimilarity = sorted(computedSimilarity, key=lambda docIndex: docIndex["CosineSimilarityScore"],reverse=True)

    # Print Top-K RetrievedDocuments
    i = 0
    while i < k and k < len(computedSimilarity):
        DocName = searcher.doc(computedSimilarity[i]['DocId']).get('DocName')
        print(DocName, str(computedSimilarity[i]['DocId']))
        topKDocInformation[computedSimilarity[i]['DocId']] = docVectors[computedSimilarity[i]['DocId']]
        i += 1
    return topKDocInformation

Create_IndexFn(index_path,collection_path,stopwords_path,stopwords_filename)
query = 'Provide Query As a String Here'
k = 10 # Change this to Required Integer Value
Results = TopKRetrieval(query, k)