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

import threading
import sys, os
import time, datetime , json
import shutil, hashlib
import inspect
import re
import requests
import boto3
from multiprocessing.pool import ThreadPool
from halo import Halo
from boto3.s3.transfer import TransferConfig
from botocore.handlers import disable_signing
from timeit import default_timer

### 상위 모듈 import 하기 위한 경로 추가(?)
parent_dir = os.path.dirname(os.path.abspath(os.path.dirname(__file__)))
sys.path.append(parent_dir)
sys.path.append(parent_dir+"/..")

from config.configure import Configure
from common import converter, output, base
from common.converter import region_info as region_s3_info     ### regeion config
from common.converter import region_cf_info                    ### regeion config

sys.excepthook = output.exception_handler

## SSL Warnning
base.disable_ssl_warnings()

filename = __file__
get_filename = filename.split('/')[-1]

def line_info(return_type=None):
    '''
       common/output.py 안에 추가 할 예정
       추가 할 경우  import 필요
       from common.output import line_info
       - line number 및 fucntion 위치를 출력하기 위함
    '''
    ### line number
    cf = inspect.currentframe()
    linenumber = cf.f_back.f_lineno

    ### Call to Function name
    func_name = cf.f_back.f_code.co_name

    ### file name
    frame = inspect.stack()[1]
    module = inspect.getmodule(frame[0])
    #filename = module.__file__
    #get_filename = filename.split('/')[-1]

    if return_type == 'filename':
        result_info = f'{get_filename}'
    elif return_type == 'function' or return_type == 'func_name' or return_type == 'f_name':
        result_info = f'{func_name}'
    elif return_type == 'lineinfo' or return_type == 'linenum' or return_type == 'lineno':
        result_info = f'{linenumber}'
    elif return_type == 'info_all' or return_type is None:
        result_info = f'{get_filename}({func_name}.{linenumber})'
    else:
        result_info = f'{get_filename}({func_name}.{linenumber})'

    return result_info

def log_print(msg, color=None, level=None):
    '''
      log 및 print를 하기 위한 함수
    '''
    color = color if color else f'green'
    level = level if level else f'INFO'

    now_time = (datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3])
    output.c_print(f'[{now_time}] [{level:5}] | {msg}', color)

class Restore:
    def __init__(self,
        region="kr",                              ### 사용안함
        db_path=None,                             ### .settings.icon2.GOLOOP_NODE_DIR
        network="MainNet",                        ### .settings.env.SERVICE
        send_url=None,                            ###  사용자 입력   /  없어도무방   .settings.env.RESTORE_SEND_URL
        bucket_name_prefix="icon2-backup",        ###  사용자 입력(docker 변수로 ) 없을 경우 defaut 사용  : false
        download_type="multi",                    ### 변경 필요 없음
        download_path='restore',                  ### Download directory | 사용자 입력(docker 변수로) 없을 경우 defaut 사용  : false
        download_url='download.solidwallet.io',   ### Downlaod type check  s3 or cloud front , Default : s3 download
        download_url_type='s3',                   ### Downlaod type check  s3 or cloud front , Default : s3 download , "s3" or "[cf | cloudfront]"
        download_tool='axel',                     ### Download Tool command option .  Default Axel command  [ axel | aria2(aria2c) ]
        download_force=False                      ### Download file이 동일할경우에 대한 액션 True일 경우 모두 삭제 후 download /
                                                  ### False일 경우 동일 파일 비교 후 다를 경우 삭제 |  사용자 입력(docker 변수로 ) 없을 경우 defaut 사용  : false
    ):
        self.region = region                            ### 사용안함
        self.db_path = db_path                          # 필수 Default => icon2 : /app/goloop
        self.network = network                          # 필수  Blockchain  node network name
        self.bucket_name_prefix = bucket_name_prefix    # 필수

        self.download_path = download_path              # Change variable name restore_path to download_path
        self.download_filename = ""
        self.download_url = re.sub('(http|https)://','',download_url)
        self.download_url_type = download_url_type
        self.download_type = download_type              ### 사용안함
        self.download_force = download_force
        self.download_tool = download_tool              ####  Default axel
        self.verbose = None

        ### Send url check
        self.send_url = send_url                        ### 없어도 무방
        if send_url:
            self.is_send = True
        else:
            self.is_send = False

        ##  Icon2 node DB Path check
        if self.db_path is None or not output.is_file(self.db_path):
            raise ValueError(f"-- {line_info()} | | db_path not found - '{self.db_path}'")

        self.get_nproc = ''.join(base.run_execute('nproc', capture_output=True)['stdout'])

        ### disk usage check percent
        self.diskcheck_per = 70

        self.run()

    ####  class main run function
    def run(self):
        ### running start time check
        run_start_time = default_timer()

        ## Downlaod(Restore) WorkDir Create
        self.restore_path = os.path.join(self.db_path, self.download_path)
        self.createDirectory(self.restore_path)

        ####  S3 또는 Cloudfront URL check
        self.download_url_type_check()

        ## fastest region check
        last_latency_rst = self.find_fastest_region()
        log_print(f'++ {line_info()} | | {last_latency_rst}', 'grey')

        self.s3_bkurl = last_latency_rst.replace('/route_check','')
        ### ex)  s3         : https://icon2-backup-kr.amazon.com/
        ### ex)  Cloudfront : https://download.soliwallet.io/kr/bk/
        log_print(f'++ {line_info()} | {self.s3_bkurl}', 'grey')

       	## get index.txt
        index_download_url = f'{self.s3_bkurl}/{self.network}/index.txt'
        log_print(f'++ {line_info()} | index file url => {index_download_url}','grey')

        ## dl , download   url_addr == self.dl_url
        self.dl_dict, self.url_addr = self.get_filelist(index_download_url)

        ### CID & Disk free size check & as_Download BackupFile Delete
        self.run_prejob(self.dl_dict)

        ## Backup file download
        self.get_bkfile(self.url_addr, self.dl_dict)

        self.db_file_decompress()

        ### Finished time check
        run_elapsed = default_timer() - run_start_time
        run_time_completed_at = "{:5.3f}s".format(run_elapsed)
        log_print(f"\n\n>>>> for_time_completed_at = {run_time_completed_at}\n", 'yellow')

    ####### >>>>>>>
    ### send result slack
    def send_slack(self, msg_text=None, msg_level="info"):
        if self.is_send and msg_text:
            output.send_slack(url=self.send_url, msg_text=msg_text, msg_level=msg_level)

    # Download(restore) Work Directory Create
    def createDirectory(self, dir_name):
        try:
            if not os.path.exists(dir_name):
                os.makedirs(dir_name)
                log_print(f"++ {line_info()} | create : {dir_name}", "green")
        except OSError:
            raise OSError(f"-- {line_info()} | Error: Creating directory fail. : {dir_name}")

    def download_url_type_check(self):
        if self.download_url_type == 's3' or self.download_url_type is None:
            self.dl_url = f'https://{self.bucket_name_prefix}_regioncode.amazonaws.com'
            self.region_info = region_s3_info
            self.match_keyword = f'^MainNet'
        elif self.download_url_type == 'cf' or self.download_url_type == 'cloudfront':
            self.dl_url = f'https://{self.download_url}/_regioncode'
            self.region_info = region_cf_info
            self.match_keyword = f'^kr/.*MainNet'

    ### s3 response time check
    def get_connect_time(self, url, name="NULL"):
        status_code = 999
        try:
            response = requests.get(f'{url}', timeout=5)
            response_text = response.text
            elapsed_time = response.elapsed.total_seconds()
            status_code = response.status_code
        except Exception as e:
            elapsed_time = None
            response_text = None
            log_print(f"-- {line_info()} | get_connect_time error : {url} -> {sys.exc_info()[0]} / {e}", "red")
        return {"url": url, "time": elapsed_time, "name": name, "text": response_text, "status": status_code}

    ### Check fast aws region
    def find_fastest_region(self):
        ## worker Thread Setting
        worker_thread_cnt = 8
        pool = ThreadPool(worker_thread_cnt)

        results = {}
        i = 0

        spinner = Halo(text=f"Finding fastest region", spinner='dots')
        spinner.start()

        check_file = 'route_check'

        for region_name, region_code in self.region_info.items():
            dl_addr = re.sub('_regioncode',region_code,self.dl_url)  ## region code 치환
            URL = f'{dl_addr}' + f'/{check_file}'
            #exec_func = "get_connect_time"
            exec_args = (f"{URL}", f"name={region_name}")
            results[i] = {}
            results[i]["data"] = pool.apply_async(self.get_connect_time, args = exec_args)
            i += 1
        pool.close()
        pool.join()

        last_latency = {}
        for i, p in results.items():
            data = p['data'].get()
            print(f"data => {data}") if self.verbose else False
            if time is not None:
                if len(last_latency) == 0:
                    last_latency = data
                if last_latency.get("time") and data.get("time"):
                    if last_latency.get("time", 99999) >= data.get("time"):
                        last_latency = data
            print(data) if self.verbose else False
        spinner.succeed(f'++ {line_info()} | [Done] Finding fastest region')

        #return last_latency
        return last_latency["url"]

    #### get download backup file list
    def get_filelist(self, file_url):
        print(file_url)
        try:
            ## get filelist.txt
            res = requests.get(file_url)
            filelist_text = []
            temp_list = res.text.strip().split('\n')     ### 마지막 라인이 공백일 경우 제거
            print(temp_list)

            for flist in temp_list:
                if re.match(self.match_keyword, flist):
                    filelist_text.append(flist)
            #file_text = filelist_text[-1].replace(f'{self.region_code}/bk/','')
            self.dl_url =  re.sub('/[\w]+/bk','',file_url).replace(f'/{self.network}/index.txt','')
            filelist_url = f'{self.dl_url}/{filelist_text[-1]}'
            log_print(f"++ {line_info()} | {filelist_url}","blue")
        except Exception as e:
            log_print(f"-- {line_info()} | {e}","red")

        try:
            #log_print(f"++ {line_info()} | {filelist_url}","green")
            dl_file_list = requests.get(filelist_url)
            self.f_dict = dl_file_list.json()
        except Exception as e:
            log_print(f"-- {line_info()} | {e}","red")

        return  self.f_dict, self.dl_url

    def run_prejob(self, dl_info):
        ### chain id check
        if 'CID' in dl_info:
            chain_id = dl_info['CID']
            del_dir_list = ['contract','db', 'wal', 'genesis.zip']
            for del_item in del_dir_list:
               dellist = os.path.join(self.db_path, chain_id, del_item)
               self.as_file_remove(dellist, file_opt=False)
            del dl_info['CID']

        ### free disk size check
        if 'Total_size' in dl_info:
            total, used, p_free = shutil.disk_usage(self.db_path)
            DL_Total_Size = dl_info['Total_size'].split(' ')[0]
            if float(DL_Total_Size) * 1.5 < float(p_free) :
                log_print(f"++ {line_info()} | You have enough disk space ", 'yellow')
                del dl_info['Total_size']
            else:
                raise ValueError(f"++ {line_info()} | Not enough disk - Download size : {DL_Total_Size} Byte , Disk Free Size : {p_free} Byte")
        else :
            dirusage = self.dir_free_size(self.db_path)
            ## dir size free size check
            if dirusage > self.diskcheck_per:
                self.send_slack(f"Not enough disk space : {dirusage:.2f}", "error")
                raise ValueError(f"-- {line_info()} | Not enough disk : {dirusage:.2f}")
            else:
                log_print(f"++ {line_info()} | You have enough disk space : {dirusage:.2f} % ", 'yellow')

        ### Download force delete ( Old Backup File Delete)
        if self.download_force:
            dl_path = os.path.join(self.db_path, self.download_path)
            del_list = os.listdir(dl_path)
            for delfile in del_list:
                if f'tar.zst' in delfile :
                    log_print(f'Delete Old Backup file  : {os.path.join(dl_path,delfile)}', 'red')
                    os.remove(os.path.join(dl_path,delfile))

    def get_bkfile(self, get_url, dl_info):
        ### backup file download
        for f_url, cksum_value in dl_info.items():
            download_url = f'{get_url}/{f_url}'
            ### Thread off
            #self.file_download(download_url,os.path.join(self.db_path,self.restore_path),cksum_value)

            ## Thread on
            t = threading.Thread(
                    target=self.file_download,
                    args = (
                        download_url,
                        os.path.join(self.db_path,self.restore_path),
                        cksum_value
                    )
                )
            t.start()

        log_print(f"++ {line_info()} | Download job wait ....", 'green')

        mainThread = threading.currentThread()
        for thread in threading.enumerate():
            log_print(f'{thread}', 'red')
            if thread is not mainThread:
                thread.join()
        log_print(f"++ {line_info()} | Download job finished", 'green')

    ##  파일 다운로드
    def file_download(self, download_url, download_path, hash_value):

        log_print(f'++ {line_info()} | download_url => {download_url} |\n\t\t\tdownload_path => {download_path}','yellow')
        try :
            local_dl_file = os.path.join(download_path, download_url.split("/")[-1])

            ## Old Download file Delete check
            if self.download_force:
                if os.path.isfile(local_dl_file):
                    log_print(f'>>> delete |    {local_dl_file}', 'red')
                    os.remove(local_dl_file)

            if os.path.isfile(local_dl_file) :
                diff_rst = self.download_diff(local_dl_file,hash_value)
                if diff_rst == 'nok' :
                    os.remove(local_dl_file)
                    log_print(f'-- {threading.get_ident()} | {download_url.split("/")[-1]} |{diff_rst} |delete : {local_dl_file} ', 'red')
                    diff_rst = None
            else:
                diff_rst = None

            if diff_rst is None:
                start_time = default_timer()

                if self.download_tool == 'axel' or self.download_tool is None:
                    #axel_option = f"-k -n {get_nproc} --verbose"    ###  axel 2.4 버전에서는 -k 옵션이 제외됨.
                    axel_option = f"-n {self.get_nproc} --verbose"
                    download_cmd = f'axel {axel_option} {download_url} -o "{download_path}"'
                elif self.download_tool == 'aria2' or  self.download_tool == 'aria2c':
                    aria2_option = f'-j {self.get_nproc}'
                    download_cmd = f'aria2c {aria2_option} {download_url} -d {download_path}'

                run_stat = base.run_execute(download_cmd, capture_output=True)

                elapsed = default_timer() - start_time
                time_completed_at = "{:5.3f}s".format(elapsed)

                diff_rst = self.download_diff(local_dl_file,hash_value)
                if diff_rst == 'ok':
                    log_print(f'\n\t{threading.get_ident()} | {download_url.split("/")[-1]} |{diff_rst} | time_completed_at = {time_completed_at}', 'green')
                else:
                    raise ValueError(f"++ {line_info()} | download file Checksum Check Fail - '{download_file}'")
            else:
                log_print(f'\n\t{threading.get_ident()} | {download_url.split("/")[-1]} |exist file (checksum is Same)', 'green')
        except Exception as e:
            log_print(f"-- {line_info()} | {e}","red")

    ### Old file remove
    def as_file_remove(self, delete_path, file_opt=False):
        ## /app/goloop/data/{CID} , ## /app/goloop/data/[restore]/*.tar*
        delete_dir = os.path.join(delete_path)
        log_print(f'-- {line_info()} | Delete_job_dir : {delete_dir}', 'yellow')

        try:
            if file_opt :
                ### Delete directory file
                filelist = os.listdir(f"{delete_dir}")
                for item in filelist:
                    if f".tar.zst" in item:
                        os.remove(os.path.join(delete_dir, item))
                        log_print(f"-- {line_info()}| delete : {delete_dir} {item}", "red")
            else:
               p_rst = f'deleted : {delete_dir}"'
               p_color = 'red'
               ## Delete directory
               if os.path.isdir(delete_dir):
                   shutil.rmtree(delete_dir)
                   lineinfo = f'{line_info()}'
               elif os.path.isfile(delete_dir):
                   os.remove(os.path.join(delete_dir))
                   lineinfo = f'{line_info()}'
               else:
                   p_rst = f'Not found {delete_dir}'
                   p_color = 'yellow'
                   lineinfo = f'{line_info()}'
               log_print(f"{lineinfo} | {p_rst}", p_color)

        except OSError:
            raise OSError(f"-- {line_info()} | Error: file_remove ({file_opt})")

    ## S3에 업로드된 파일과 다운로드 된 파일 의 MD5 체크섬 비교
    def download_diff(self, download_file, orig_hash):
        localfile_hash = self.getHash(download_file)
        if localfile_hash == orig_hash :
            hash_rst = "ok"
        else:
            hash_rst = "nok"
        return hash_rst

    ## MD5 HASH
    def getHash(self, path, blocksize=65536):
        afile = open(path, 'rb')
        hasher = hashlib.md5()
        buf = afile.read(blocksize)
        while len(buf) > 0:
            hasher.update(buf)
            buf = afile.read(blocksize)
        afile.close()
        return hasher.hexdigest()

    def db_file_decompress(self):
        log_print(f'++ {line_info()} | decompress start', 'yellow')
        comp_algorithm = 'pzstd'
        old_run_path = os.getcwd()

        if os.path.isdir(self.db_path):
            os.chdir(self.db_path)
            log_print(f'+++ {line_info()} | change job directory is \"{self.db_path}\"', 'yellow')
            sour_dir = os.path.join(self.db_path, self.restore_path)
            cmd = f"cat {sour_dir}/*tar.* | {comp_algorithm} -cd | tar -xf - -C {self.db_path}"

        log_print(f'++ {line_info()} | cmd = {cmd}', 'yellow')

        if base.run_execute(cmd).get("stderr"):
            raise OSError(f"-- {line_info()} | Failed Decompression - '{cmd}'")

        ### 이전  디렉토리로  변경
        os.chdir(old_run_path)

    ####### <<<<<<

class ProgressPercentage(object):
    def __init__(self, filename):
        self._filename = filename
        self._size = float(os.path.getsize(filename))
        self._seen_so_far = 0
        self._lock = threading.Lock()
        self._prevent_bytes = 0

    def __call__(self, bytes_amount):
        # To simplify we'll assume this is hooked up
        # to a single filename.
        with self._lock:
            self._seen_so_far += bytes_amount
            percentage = (self._seen_so_far / self._size) * 100
            # tx = bytes_amount - self._prevent_bytes
            sys.stdout.write(
                "\r \t %s  %s / %s  (%.2f%%) " % (
                    self._filename, converter.convert_bytes(
                        self._seen_so_far), converter.convert_bytes(self._size),
                    percentage))
            sys.stdout.flush

def main():
    from config.configure import Configure as CFG

    #test config  json download url  : https://d1hfk7wpm6ar6j.cloudfront.net/SejongNet/default_configure.json
    ###  test path & file   : /goloop/default_configure.json
    use_file=True
    cfg = CFG(use_file=use_file)
    config = cfg.config
    '''
    +  /goloop/configure.yaml  파일 내 참고 dict value
    settings
       > env
          > COMPOSE_ENV
              -  "SERVICE": "MainNet",
              -  "BASE_DIR": "/goloop",
              -  "DOWNLOAD_FORCE": "true",
              -  "RESTORE_PATH": "restore",
              -  "DOWNLOAD_TOOL" : "aria2c",
              -  "DOWNLOAD_URL_TYPE" : "file_list",
              -  "DOWNLOAD_URL" : "https://download.solidwallet.io"

    setting
       > icon2
           - "GOLOOP_NODE_DIR" : "/goloop/data"
    '''
    icon2_config = config['settings']['icon2']
    env_config = config['settings']['env']
    # compose_env_config = config['settings']['env']['COMPOSE_ENV']

    ### Goloop DB PATH
    if icon2_config.get('GOLOOP_NODE_DIR') :
        db_path = icon2_config['GOLOOP_NODE_DIR']
    else :
        default_db_path = 'data'
        # base_dir = compose_env_config['BASE_DIR']
        base_dir = env_config['BASE_DIR']
        db_path = os.path.join(base_dir, default_db_path)

    ### Restore Options
    ### network  =  MainNet | SejongNet ....
    # network = env_config['SERVICE']  if env_config.get('SERVICE') else compose_env_config['SERVICE']
    # restore_path = env_config['RESTORE_PATH']  if env_config.get('RESTORE_PATH') else compose_env_config['RESTORE_PATH']
    # dl_force = env_config['DOWNLOAD_FORCE']  if env_config.get('DOWNLOAD_FORCE') else compose_env_config['DOWNLOAD_FORCE']
    # download_tool = env_config['DOWNLOAD_TOOL']  if env_config.get('DOWNLOAD_TOOL') else compose_env_config['DOWNLOAD_TOOL']
    # download_url = env_config['DOWNLOAD_URL']  if env_config.get('DOWNLOAD_URL') else compose_env_config['DOWNLOAD_URL']
    # download_url_type = env_config['DOWNLOAD_URL_TYPE']  if env_config.get('DOWNLOAD_URL_TYPE') else compose_env_config['DOWNLOAD_URL_TYPE']
    network = env_config['SERVICE']
    restore_path = env_config['RESTORE_PATH']
    dl_force = env_config['DOWNLOAD_FORCE']
    download_tool = env_config['DOWNLOAD_TOOL']
    download_url = env_config['DOWNLOAD_URL']
    download_url_type = env_config['DOWNLOAD_URL_TYPE']

    print(f'db_path = {db_path}')
    print(f'network = {network}')
    print(f'restore_path = {restore_path}')

    print(f'dl_force = {dl_force}')
    print(f'download_tool = {download_tool}')
    print(f'download_url = {download_url}')
    print(f'download_url_type = {download_url_type}')

    Restore(
        db_path=db_path,
        network=network,
        download_path=restore_path,
        download_force=dl_force,
        download_url=download_url,
        download_tool=download_tool,
        download_url_type=download_url_type
    )

if __name__ == '__main__' :
    main()


# -*- coding: utf-8 -*-
import logging
import asyncio
from .utils import dotdict
from .streaming import Streaming
from .ohlcvbuilder import OHLCVBuilder
from .inventory import Inventory
from .exchange import Exchange, ExchangeError
from .board import Board
from collections import deque, defaultdict
from datetime import datetime, timedelta, timezone
from time import time

class Strategy:

    def __init__(self, yourlogic=None, interval=60):

        # トレーディングロジック設定
        self.yourlogic = yourlogic

        # 設定
        self.settings = dotdict()
        self.settings.apiKey = ''
        self.settings.secret = ''
        self.settings.symbol = 'BTC/JPY'

        # 動作タイミング
        self.settings.interval = interval
        self.settings.minimum_interval = 0

        # OHLCV設定
        self.settings.max_ohlcv_size = 1000
        self.settings.disable_rich_ohlcv = False

        # その他設定
        self.settings.enable_board = False
        self.settings.enable_board_api = False

        # ログ設定
        self.logger = logging.getLogger(__name__)

    async def start(self):
        self.logger.info('Start Trading')

        # ペア情報
        self.pair = self.settings.symbol.replace('/','_').lower()
        self.spec = Exchange.ProductSpecs[self.pair]

        # APIセットアップ
        self.api = Exchange(self.settings.apiKey, self.settings.secret)

        # ストリーム配信
        self.streaming = Streaming(Streaming.SocketioSource())
        self.executions_ep = await self.streaming.get_trades_endpoint(self.pair, 5000)

        # OHLCVビルダー設定
        self.ohlcvbuilder = OHLCVBuilder(
            maxlen=self.settings.max_ohlcv_size,
            disable_rich_ohlcv=self.settings.disable_rich_ohlcv)

        # 注文管理
        self.inventory = Inventory(self.spec)
        self.latest_trade_id = None

        # 板情報（配信）
        if self.settings.enable_board:
            self.board = Board(self.pair)
            ob = await self.api.orderbooks(self.pair)
            self.board.sync(ob)
            await self.board.attach(self.streaming)
        # 板情報（API）
        elif self.settings.enable_board_api:
            self.board = Board(self.pair)

        # ロジック実行
        await asyncio.wait([
            self.balance_polling(),
            self.cancel_nonactive_orders(),
            self.standard_logic(),
            self.inventory.start(),
            self.streaming.start()])

    def get_order(self, myid):
        return self.inventory.get_order(myid)

    async def order(self, myid, side, size, limit=None, cancel_after_seconds=None, limit_mask=0):
        # 注文がオープンならキャンセル
        o = self.inventory.get_order(myid)
        if o['status'] in Inventory.OPEN_STATUS:
            if abs(o['rate']-limit)>limit_mask or abs(o['amount']-size)>0:
                try:
                    self.logger.info('CANCEL {myid} {status} {order_type} {rate} {executed_amount}/{amount} {id}'.format(**o))
                    await self.api.cancel(o)
                except ExchangeError as e:
                    self.logger.warning(type(e).__name__ + ": {0}".format(e))
            else:
                # 価格・サイズが同じなら注文しない
                return

        # 新規注文
        try:
            res = await self.api.order(self.pair,side,size,limit)
            self.inventory.new_order(myid,res)
            o = self.inventory.get_order(myid)
            self.logger.info('NEW {myid} {status} {order_type} {rate} {executed_amount}/{amount} {id}'.format(**o))
            # 後でキャンセル
            if cancel_after_seconds is not None:
                asyncio.ensure_future(self._cancel_later(o,cancel_after_seconds))
        except ExchangeError as e:
            self.logger.warning(type(e).__name__ + ": {0}".format(e))

    async def _cancel_later(self, o, seconds):
        try:
            await asyncio.sleep(seconds)
            if o['status'] in Inventory.OPEN_STATUS:
                self.logger.info('CANCEL LATER {myid} {status} {order_type} {rate} {executed_amount}/{amount} {id}'.format(**o))
                await self.api.cancel(o)
        except ExchangeError as e:
            self.logger.warning(type(e).__name__ + ": {0}".format(e))
        except Exception as e:
            self.logger.exception(e)

    async def cancel(self, myid):
        # 注文がオープンならキャンセル
        o = self.inventory.get_order(myid)
        if o['status'] in Inventory.OPEN_STATUS:
            try:
                self.logger.info('CANCEL {myid} {status} {order_type} {rate} {executed_amount}/{amount} {id}'.format(**o))
                await self.api.cancel(o)
            except ExchangeError as e:
                self.logger.warning(type(e).__name__ + ": {0}".format(e))

    async def cancel_order_all(self):
        pass

    async def check_balance(self):
        try:
            ticker = await self.api.ticker(pair=self.pair)
            balance = await self.api.balance()
            btc = balance['btc']+balance['btc_reserved']
            jpy = balance['jpy']+balance['jpy_reserved']
            total = int(jpy + ticker['bid']*btc)
            balance['jpy_total'] = total
            self.balance = dotdict(balance)
            self.logger.info(f'btc {btc:.8f} jpy {jpy:.0f} total {total}')
        except ExchangeError as e:
            self.logger.warning(type(e).__name__ + ": {0}".format(e))

    async def check_trades(self):
        try:
            trades = await self.api.get_my_trades(end=self.latest_trade_id)
            if len(trades):
                self.latest_trade_id = max(t['id'] for t in trades)
            self.inventory.check_my_trades(trades)
        except ExchangeError as e:
            self.logger.warning(type(e).__name__ + ": {0}".format(e))

    async def balance_polling(self):
        while True:
            await asyncio.sleep(300)
            try:
                # 定期的に資産情報取得
                await self.check_balance()
            except Exception as e:
                self.logger.exception(e)

    async def cancel_nonactive_orders(self):
        while True:
            await asyncio.sleep(15)
            try:
                # 最近の注文情報取得
                recent_orders = self.inventory.get_nonactive_orders()
                if len(recent_orders):
                    # 未決済の注文取得
                    open_orders = await self.api.get_orders()
                    for o in open_orders:
                        self.logger.info('{id} {order_type} {rate} {pair} {pending_amount}'.format(**o))
                    open_orders = {o['id']:o for o in open_orders}
                    # アクティブ注文を全てキャンセル
                    cancel_needed = [o for o in recent_orders if o['id'] in open_orders]
                    if len(cancel_needed):
                        for o in cancel_needed:
                            try:
                                self.logger.info('FORCED CANCEL {myid} {status} {order_type} {rate} {executed_amount}/{amount} {id}'.format(**o))
                                await self.api.cancel(o)
                            except ExchangeError as e:
                                self.logger.warning(type(e).__name__ + ": {0}".format(e))

            except ExchangeError as e:
                self.logger.warning(type(e).__name__ + ": {0}".format(e))
            except Exception as e:
                self.logger.exception(e)

    async def standard_logic(self):
        try:
            await self.check_balance()
            await self.executions_ep.wait()
        except Exception as e:
            self.logger.exception(e)
        last_entry_time = time()
        while True:
            try:
                # 待ち
                if self.settings.interval:
                    await asyncio.sleep((-time() % self.settings.interval) or self.settings.interval)
                else:
                    await self.executions_ep.wait()

                # 最小インターバル
                if self.settings.minimum_interval:
                    t1 = last_entry_time // self.settings.minimum_interval
                    t2 = time() // self.settings.minimum_interval
                    can_entry = t2 > t1
                else:
                    can_entry = True

                # 注文情報更新
                if can_entry:
                    if self.settings.enable_board_api:
                        ob = await self.api.orderbooks(self.pair)
                        self.board.sync(ob)
                    await self.check_trades()

                # ポジション情報コピー
                self.long_size = self.inventory.position.long_size
                self.short_size = self.inventory.position.short_size
                self.position_size = self.inventory.position.position_size
                self.position_avg_price = self.inventory.position.position_avg_price

                # 約定履歴取得
                executions = await self.executions_ep.get_data()
                ohlcv = self.ohlcvbuilder.create_boundary_ohlcv(executions)

                # ロジックコール
                if can_entry:
                    last_entry_time = time()
                    if self.settings.enable_board:
                        board = self.board
                        bids, asks = board.sort()
                        bid, ask = bids[0]['price'], asks[0]['price']
                        spread = ask - bid
                        if spread<-50:
                            self.logger.warning(f'orderbooks needs to be sync. spr {spread} bid {bid} ask {ask}')
                            ob = await self.api.orderbooks(self.pair)
                            board.sync(ob)
                    elif self.settings.enable_board_api:
                        board = self.board
                    else:
                        board = None
                    await self.yourlogic(
                        executions=executions,
                        ohlcv=ohlcv,
                        board=board,
                        strategy=self)
            except Exception as e:
                self.logger.exception(e)


# %%
import json
import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from PIL import Image
import PIL
from glob import glob
import decord
from decord import VideoReader
import random
import logging

try:
    from lavis.datasets.datasets.base_dataset import BaseDataset
except ModuleNotFoundError:
    import sys
    sys.path.append(os.path.realpath(os.path.join(__file__, '../LAVIS')))
    from lavis.datasets.datasets.base_dataset import BaseDataset

from torch.utils.data.dataloader import default_collate
import pandas as pd
from collections import defaultdict
from typing import Literal, get_args, List

logging.getLogger().setLevel(logging.INFO)

decord.bridge.set_bridge('native')

### === ActionBenchDataset === ###
class ActionBenchDatasetBase(BaseDataset):
    def __init__(self,
        annotation,
        task = "video_text_matching",
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps=None,
        k=None,
        neg_sampling_same_clip=0
    ):
        """
            annotation: a list of annotation samples
        """
        self.task = task
        self.annotation = annotation
        print("annotation size:", len(self.annotation))

        self.vis_processor = vis_processor
        self.text_processor = text_processor

        self.frm_sampling_strategy = frm_sampling_strategy
        self.frame_height = frame_height
        self.frame_width = frame_width
        self.num_frm = num_frm
        self.fps = fps
        print("using fps:", self.fps)

        # if using negative sampling
        self.neg_sampling_same_clip = neg_sampling_same_clip
        if self.neg_sampling_same_clip > 0:
            print("Using negative sampling from same clip, number:", neg_sampling_same_clip)

    def _load_video_from_path_decord(self, video_path, fps, start_time=None, end_time = None):
        """
            returns frames as PIL images
        """
        frm_sampling_strategy = self.frm_sampling_strategy
        num_frm = self.num_frm
        task = self.task
        try:
            # load video
            if not self.frame_height or not self.frame_width:
                vr = VideoReader(video_path)
            else:
                vr = VideoReader(video_path, width=self.frame_width, height=self.frame_height)

            # get start and end frame idx
            vlen = len(vr)
            if start_time or end_time:
                assert fps > 0, 'must provide video fps if specifying start and end time.'
                start_idx = max(int(start_time * fps), 0)
                end_idx = min(int(end_time * fps), vlen-1)
                if start_idx == end_idx:
                    end_idx = min(start_idx + 1, vlen-1)
            else:
                start_idx, end_idx = 1, vlen-1

            # task

            if frm_sampling_strategy == 'uniform':
                frame_indices = np.linspace(
                    start_idx,
                    end_idx,
                    num = num_frm,
                    endpoint = True,
                    retstep = False,
                    dtype = int
                )
            else:
                raise NotImplementedError('Invalid sampling strategy {} '.format(frm_sampling_strategy))

            raw_sample_frms = vr.get_batch(frame_indices).asnumpy() # (num_frm, H, W, C)
            # raw_sample_frms = vr.get_batch(frame_indices) # tensor (num_frm, H, W, C)
            raw_sample_frms = [Image.fromarray(item, mode="RGB") for item in raw_sample_frms] # PIL Images

        except Exception as e:
            logging.info(e)
            return None
        # raw_sample_frms = raw_sample_frms.permute(0, 3, 1, 2) # torch tensor
        # raw_sample_frms = transforms.ToPILImage()(raw_sample_frms).convert('RGB')
        return raw_sample_frms

    def collater(self, samples, additional_fields: List[str] = [], text_fields: List[str] = []):
        if self.vis_processor is None:
            return samples
        if self.neg_sampling_same_clip == 0:
            return default_collate(samples)
        else:
            # additional work for batching the positive and negative samples
            ret = {
                "video_input":[],
                "video_input_reversed":[],
                "video_input_before":[],
                "video_input_after":[],
                "text_input":[],
                "action_antonym_text_input":[],
                "text_input_before":[],
                "text_input_after":[],
                **{f : [] for f in additional_fields}
            }
            for sample in samples:
                pos_instance = sample['pos_instance']
                neg_instances = sample['neg_instances']
                for att in ret.keys():
                    if att in pos_instance:
                        ret[att] += [pos_instance[att]] + [neg[att] for neg in neg_instances]

            empty_field_keys = []
            for att in ret.keys():
                if att not in text_fields:
                    if ret[att] != []:
                        ret[att] = torch.stack(ret[att])
                if ret[att] == []:
                    empty_field_keys.append(att)
            for remove_key in empty_field_keys:
                del ret[remove_key]
            return ret

    def __getitem__(self, index):
        if self.neg_sampling_same_clip == 0:
            # if not using any negative sampling strategy
            return self._get_single_item(index)
        else:
            # additional work if using the negative sampling strategy from the same clip
            ann_instance = self.annotation[index]
            indices_from_same_clip = list(self.clipuid_2_idx[ann_instance['clip_uid']])
            indices_from_same_clip.remove(index)
            assert len(indices_from_same_clip) == len(self.clipuid_2_idx[ann_instance['clip_uid']]) - 1
            if len(indices_from_same_clip) == 0:
                # the clip has only this one instance, then randomly sample another instance as negative
                sampled_neg_indices = random.sample(list(range(len(self.annotation))), self.neg_sampling_same_clip)
                print("Warning: no other instances from the same clip can be sampled, randomly sample negative instances:", sampled_neg_indices)
            else:
                while len(indices_from_same_clip) < self.neg_sampling_same_clip:
                    indices_from_same_clip += indices_from_same_clip
                sampled_neg_indices = random.sample(indices_from_same_clip, self.neg_sampling_same_clip)

            pos_instance = self._get_single_item(index) # Dict
            neg_instances = [self._get_single_item(idx) for idx in sampled_neg_indices] # List[Dict]
            return {
                "pos_instance":pos_instance,
                "neg_instances":neg_instances
            }

class ActionBenchDataset_Ego4D(ActionBenchDatasetBase):
    def __init__(self,
        vis_root,
        ann_path,
        state_change_filtering_json = None,
        state_change_filtering = False, # if true, use the filtered videos only during loading annotation
        task = "video_text_matching", # ["video_text_matching", "action_antonym", "object_shuffle", "reversed_video"]
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps=None,
        k=None,
        neg_sampling_same_clip=0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of EgoClip annotations
        params:
            vis_root (string): Root directory of video clips (e.g., datasets/Ego4D/video_clips/clips_downsampled_5fps_downsized_224x224)
            ann_path (string): path to the egoclip style ann jsonl, each line is one instance (e.g., ActionBench/ego4d/egoclip_subset_action_antonyms_train_val_test_split/val.jsonl)
            state_change_filtering_json: path to the json file that contains indication of whether a video is state-change salient (e.g., ActionBench/ego4d/egoclip_subset_action_antonyms_train_val_test_split/state_change_heavy_instance_filtering_val.json)
            task (string): what's the evaluation task: choose from ["video_text_matching", "action_antonym", "object_shuffle", "reversed_video"]
                - "video_text_matching": for general vl tasks; returns a video_input and the corresponding text
                - "action_antonym": returns a video_input and two text input: one original annotation text, one modified annotation text with all verbs repleced with their antonym
                - "object_shuffle": returns a video_input and two text input: one original annotation text, one modified annotation text with action's objects repleced with another random object
                - "reversed_video": returns two video input: an original video_input and a reversed video_input with all frames reversed; and a text input of the original text annotation
                Note: all video_input returned are a sequence of sampled frames with size (num_frm, C, H, W)
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        # Ego4d nouns and verbs taxonomy size: see ActionBench/ego4d/*.csv for the index mapping
        self.noun_dim = 582  # num of nouns of ego4d taxonomy dictionary, not used for now, could be used for fine-grained negative sampling
        self.verb_dim = 118  # num of verbs of ego4d taxonomy dictionary, not used for now, could be used for fine-grained negative sampling

        self.vis_root = vis_root
        if isinstance(ann_path, list):
            ann_path = ann_path[0]

        self.clipuid_2_idx = defaultdict(list)

        if state_change_filtering_json is not None:
            self.state_change_filtering_dict = json.load(open(state_change_filtering_json))
        else:
            self.state_change_filtering_dict = None

        ann_full = []
        with open(ann_path) as f:
            for line in f:
                ann_instance = json.loads(line)
                clip_uid = ann_instance['clip_uid']
                clip_relative_start = ann_instance['clip_relative_start']
                clip_relative_end = ann_instance['clip_relative_end']
                # filter out missing video clips
                if os.path.exists(os.path.join(vis_root, f'{clip_uid}.mp4')):
                    # check if filter
                    if self.state_change_filtering_dict is not None and state_change_filtering:
                        key_ = clip_uid + "__" + str(clip_relative_start) + "__" + str(clip_relative_end)
                        if self.state_change_filtering_dict[key_]["if_keep"]:
                            self.clipuid_2_idx[clip_uid].append(len(ann_full))
                            ann_full.append(ann_instance)
                    else:
                        self.clipuid_2_idx[clip_uid].append(len(ann_full))
                        ann_full.append(ann_instance)

                    if k is not None and len(ann_full) >= k:
                        break

        super().__init__(
            ann_full,
            task = task, # ["video_text_matching", "action_antonym", "object_shuffle", "reversed_video"]
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps=fps,
            k=k,
            neg_sampling_same_clip=neg_sampling_same_clip
        )

        print("annotation size:", len(self.annotation))

    def _get_noun_verb_vec(self, sample):
        noun_vec = torch.zeros(self.noun_dim)
        verb_vec = torch.zeros(self.verb_dim)
        if 'tag_noun' in sample and 'tag_verb' in sample:
            noun_idx = eval(sample['tag_noun'])
            verb_idx = eval(sample['tag_verb'])
            for i in noun_idx:
                noun_vec[i] = 1
            for i in verb_idx:
                verb_vec[i] = 1
        else:
            noun_vec = None
            verb_vec = None
        return noun_vec, verb_vec # if output zero tensor, meaning not used for deciding loss mask

    def _get_single_item(self, index):
        """otuput format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann_instance = self.annotation[index]
        narration_text = ann_instance['clip_text']

        clip_uid = ann_instance['clip_uid']
        video_path = os.path.join(self.vis_root, f"{clip_uid}.mp4")
        start_time = ann_instance['clip_relative_start']
        end_time = ann_instance['clip_relative_end']
        if self.fps is None:
            fps = ann_instance['clip_fps']
        else:
            fps = self.fps # overiding the original fps in ann

        # get state_change_saliency_flag
        key_ = clip_uid + "__" + str(start_time) + "__" + str(end_time)
        if self.state_change_filtering_dict is not None:
            state_change_saliency_flag = self.state_change_filtering_dict[key_]["if_keep"]
        else:
            state_change_saliency_flag = True
        state_change_saliency_flag = torch.tensor(state_change_saliency_flag)

        # get video frames
        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            fps,
            start_time=start_time,
            end_time=end_time
        )

        noun_vec, verb_vec = self._get_noun_verb_vec(ann_instance) # potentially useful for negative sampling (Egoclip annotation specific)

        if raw_sampled_frames is None:
            print(f"ERROR: cannot loading video clip: {ann_instance['clip_uid']}; try load another random instance...")
            return self.__getitem__(index+1)

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        if self.task == "action_antonym":
            narration_text_action_antonym = ann_instance['action_antonym_clip_text']
            if self.text_processor is not None:
                narration_text_action_antonym = self.text_processor(narration_text_action_antonym)
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "action_antonym_text_input": narration_text_action_antonym,
                "noun_vec": noun_vec,
                "verb_vec": verb_vec,
                "state_change_saliency_flag": state_change_saliency_flag
            }
        elif self.task == "object_shuffle":
            narration_text_object_shuffled = ann_instance['object_shuffled_clip_text']
            if self.text_processor is not None:
                narration_text_object_shuffled = self.text_processor(narration_text_object_shuffled)
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "object_shuffled_text_input": narration_text_object_shuffled,
                "noun_vec": noun_vec,
                "verb_vec": verb_vec,
                "state_change_saliency_flag": state_change_saliency_flag
            }
        elif self.task == "reversed_video":
            raw_sampled_frames_reverse = torch.flip(raw_sampled_frames, dims=[0])
            return {
                "video_input_reversed": raw_sampled_frames_reverse,
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "noun_vec": noun_vec,
                "verb_vec": verb_vec,
                "state_change_saliency_flag": state_change_saliency_flag
            }
        else: # "video_text_matching"
            # default output for general video-text matching tasks
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "noun_vec": noun_vec,
                "verb_vec": verb_vec,
                "state_change_saliency_flag": state_change_saliency_flag
            }

    ## overwrite collater here
    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=['noun_vec', 'verb_vec', "object_shuffled_text_input", "state_change_saliency_flag"],
            text_fields=["text_input","action_antonym_text_input","text_input_before","object_shuffled_text_input","text_input_after"]
        )

Split = Literal['train', 'val', 'test']
class ActionBenchDataset_SSv2(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: str,
        ann_root: str,
        split: Split,
        state_change_filtering = False, # if True, only keep the clips with state-change saliency
        use_templates_as_labels = False, # if True, use templates instead of the original annotations for text_input
        task = "video_text_matching", # ["video_text_matching", "action_antonym", "reversed_video"]
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm =  8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2
        params:
            vis_root (string): Root directory of video clips (e.g., datasets/SSv2/video_clips/clips_downsampled_5fps_downsized_224x224)
            ann_root (string): path to the annotations root directory (e.g., ActionBench/ssv2/shuffled_object_and_action_antonyms)
            use_templates_as_labels (bool): Whether to output the raw label template instead of the in-filled template. The test set only has raw templates as annotations.
            task (string): what's the evaluation task: choose from ["video_text_matching", "action_antonym", "reversed_video"]
                - "video_text_matching": for general vl tasks; returns a video_input and the corresponding text
                - "action_antonym": returns a video_input and two text input: one original annotation text, one modified annotation text with all verbs repleced with their antonym
                - "reversed_video": returns two video input: an original video_input and a reversed video_input with all frames reversed; and a text input of the original text annotation
                Note: all video_input returned are a sequence of sampled frames with size (num_frm, C, H, W)
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        if split not in get_args(Split):
            raise ValueError(f'split must be one of {get_args(Split)}')

        if split == 'test' and not use_templates_as_labels:
            raise ValueError('Test set only has raw templates as annotations')

        self.split = split

        self.vis_root = vis_root
        self.ann_root = ann_root
        self.use_templates_as_labels = use_templates_as_labels

        self.templates_to_ints = self._load_templates_to_ints()

        self.task = task # _load_annotations needs this to be set
        self.state_change_filtering = state_change_filtering # _load_annotations needs this to be set
        annotations = self._load_annotations(split)

        self.clipuid_2_idx = {d['clip_uid'] : [i] for i, d in enumerate(annotations)} # Each clip has only one annotation

        self.neg_sampling_same_clip = 0 # set neg_sampling_same_clip to 0 for SSv2 since there is only one annotation per video

        super().__init__(
            annotations,
            task = task, # ["video_text_matching", "action_antonym", "reversed_video"]
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = 0
        )

    def _load_templates_to_ints(self):
        '''
        Loads mapping from templates to integers
        '''
        with open(os.path.join(self.ann_root, 'labels.json')) as f:
            j = json.load(f)

        return  {k : int(v) for k, v in j.items()}

    def _load_annotations(self, split: Split):
        '''
        Builds list of annotations each containing a video id, label, template, and placeholders
        from the annotation files.

        Args:
            split (Split): The stage to load the annotations for.

        Returns: A list of annotations with the following structure:
            [
                {
                    'clip_uid': int,
                    'label' : Optional[str],
                    'template' : str,
                    'placeholders' : Optional[List[str]]
                }
            ]
        '''
        # load train_antonym if exists; for look up action antonym for the original training set
        # for training instances that doesn't have an antonym, we will randomly pick another instance as action antonym text
        id_to_train_antonym_anns = {}
        train_antonym_path = os.path.join(self.ann_root, f'train_antonym.json')
        if os.path.exists(train_antonym_path):
            train_antonym_anns = json.load(open(train_antonym_path))
            for item in train_antonym_anns:
                id_to_train_antonym_anns[item['id']] = item

        if split in ['train', 'val']:
            # load state_change_filtering_json if any
            state_change_filtering_json_path = os.path.join(self.ann_root, f"state_change_heavy_instance_filtering_{split}.json")
            if os.path.exists(state_change_filtering_json_path):
                self.state_change_filtering_dict = json.load(open(state_change_filtering_json_path))
            else:
                self.state_change_filtering_dict = None

            # load main ann
            split = 'validation' if split == 'val' else 'train' # File is called "validation.json"

            with open(os.path.join(self.ann_root, f'{split}.json')) as f:
                j = json.load(f)

            annots = []
            for d in j:
                annot = {
                    'clip_uid': int(d['id']),
                    'label': d['label'].capitalize(),
                    'template': d['template'].replace('[something]', 'something'),
                    'placeholders': d['placeholders']
                }

                if self.state_change_filtering_dict is not None and self.state_change_filtering:
                    if not self.state_change_filtering_dict[str(d['id'])]['if_keep']:
                        continue

                if self.task == 'action_antonym':
                    if split == 'validation':
                        # assert that validation instances contains antonym items
                        annot.update({
                            'label_action_antonym_clip_text': d['label_action_antonym_clip_text'],
                            'template_action_antonym_clip_text': d['template_action_antonym_clip_text']
                        })
                    else: # split == 'train'
                        id_ = int(d['id'])
                        if id_ in id_to_train_antonym_anns:
                            # if current id has an action antonym
                            train_antonym_ann = id_to_train_antonym_anns[id_]
                            annot.update({
                                'label_action_antonym_clip_text': train_antonym_ann['label_action_antonym_clip_text'],
                                'template_action_antonym_clip_text': train_antonym_ann['template_action_antonym_clip_text']
                            })
                        else:
                            # if not, randomly pick an item in train_antonym, and using the original text as antonym
                            random_id = random.choice(list(id_to_train_antonym_anns.keys()))
                            train_antonym_ann = id_to_train_antonym_anns[random_id]
                            annot.update({
                                'label_action_antonym_clip_text': train_antonym_ann['label'],
                                'template_action_antonym_clip_text': train_antonym_ann['template']
                            })

                if self.task == 'object_shuffle':
                    annot.update({
                        'label_object_shuffled_clip_text': d['label_object_shuffled_clip_text']
                    })

                annots.append(annot)
        else:
            # columns = ['id', 'template'] if self.task not in ['action_antonym','object_shuffle','reversed_video'] else None # Infer from csv if action_antonym
            # df = pd.read_csv(os.path.join(self.ann_root, 'test-answers.csv'), names=columns, sep=';')

            df = pd.read_csv(os.path.join(self.ann_root, 'test-answers.csv'), header=0, sep=';')

            annots = []
            for r in df.itertuples():
                annot = {
                    'clip_uid': int(r.id),
                    'label': None,
                    'template': r.template,
                    'placeholders': None
                }

                if self.task == 'action_antonym':
                    annot.update({
                        'label_action_antonym_clip_text': None,
                        'template_action_antonym_clip_text': r.template_action_antonym_clip_text
                    })

                annots.append(annot)

        return annots

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        narration_text = ann['template'] if self.use_templates_as_labels else ann['label']

        id = ann['clip_uid']

        # get state_change_saliency_flag
        if self.state_change_filtering_dict is not None:
            state_change_saliency_flag = self.state_change_filtering_dict[str(id)]['if_keep']
        else:
            state_change_saliency_flag = True
        state_change_saliency_flag = torch.tensor(state_change_saliency_flag)

        # get video frames
        video_path = os.path.join(self.vis_root, f"{id}.mp4")
        start_time = None # ann_instance['clip_relative_start']
        end_time = None # ann_instance['clip_relative_end']

        assert os.path.exists(video_path)
        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        if self.task == "action_antonym":
            narration_text_action_antonym = ann[f'{"template" if self.use_templates_as_labels else "label"}_action_antonym_clip_text']
            if self.text_processor is not None:
                narration_text_action_antonym = self.text_processor(narration_text_action_antonym)
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "action_antonym_text_input": narration_text_action_antonym,
                "state_change_saliency_flag": state_change_saliency_flag
            }

        elif self.task == "object_shuffle":
            assert not self.use_templates_as_labels
            narration_text_object_shuffled = ann['label_object_shuffled_clip_text']
            if self.text_processor is not None:
                narration_text_object_shuffled = self.text_processor(narration_text_object_shuffled)
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "object_shuffled_text_input": narration_text_object_shuffled,
                "state_change_saliency_flag": state_change_saliency_flag
            }

        elif self.task == "reversed_video":
            raw_sampled_frames_reverse = torch.flip(raw_sampled_frames, dims=[0])
            return {
                "video_input_reversed": raw_sampled_frames_reverse,
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "state_change_saliency_flag": state_change_saliency_flag
            }
        else:
            # default output for general video-text matching tasks
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "state_change_saliency_flag": state_change_saliency_flag
            }

    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=["object_shuffled_text_input", "state_change_saliency_flag"],
            text_fields=["text_input","action_antonym_text_input","text_input_before","object_shuffled_text_input","text_input_after"]
        )
        # return super().collater(
        #     samples,
        #     text_fields=["text_input","action_antonym_text_input","text_input_before","text_input_after"]
        # )

### === Downstream Datasets === ###
class DownstreamTask_Retrieval_SSv2(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: str,
        ann_root: str,
        split: Literal['train', 'val'],
        task = Literal['ssv2_label','ssv2_template'],
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2 for downstream retrieval tasks
        params:
            vis_root (string): Root directory of video clips (e.g., datasets/SSv2/video_clips/clips_downsampled_5fps_downsized_224x224)
            ann_root (string): path to the annotations root directory (e.g., datasets/SSv2/ssv2_label_ssv2_template)
            use_templates_as_labels (bool): Whether to output the raw label template instead of the in-filled template. The test set only has raw templates as annotations.
            task (string): ssv2_label or ssv2_template as described in https://arxiv.org/abs/2206.03428
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        self.split = split

        self.vis_root = vis_root
        self.ann_root = ann_root

        # self.templates_to_ints = self._load_templates_to_ints()

        self.task = task # _load_annotations needs this to be set

        (
            annotations,
            self.videos,
            self.texts,
            self.v2t_targets,
            self.t2v_targets
        ) = self._load_annotations(split)

        self.clipuid_2_idx = {d['clip_uid'] : [i] for i, d in enumerate(annotations)} # Each clip has only one annotation

        self.neg_sampling_same_clip = 0 # set neg_sampling_same_clip to 0 for SSv2 since there is only one annotation per video

        super().__init__(
            annotations,
            task = task,
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = 0
        )

    def _load_annotations(self, split):
        if split == "train" and self.task == "ssv2_label":
            ann_path = os.path.join(self.ann_root,"ssv2_ret_label_train.json")
        elif split == "val" and self.task == "ssv2_label":
            ann_path = os.path.join(self.ann_root,"ssv2_ret_label_val_small.json")
        elif split == "train" and self.task == "ssv2_template":
            ann_path = os.path.join(self.ann_root,"ssv2_ret_template_train.json")
        elif split == "val" and self.task == "ssv2_template":
            ann_path = os.path.join(self.ann_root,"ssv2_ret_template_val_small.json")

        j = json.load(open(ann_path))

        annots = []
        videos = []
        texts = []
        v2t_targets = []
        t2v_targets = []
        for item in j:
            video_names = item['video']
            if isinstance(video_names, str):
                video_names = [video_names]
            caption = item['caption'].replace('[something]', 'something')
            gt_video_idx_this_caption = []
            for video_name in video_names:
                video_id = video_name[:-5]
                if not self._check_if_video_exists(video_id):
                    logging.info("!!!! video {} does not exist, skip...".format(video_id))
                    continue
                # if video exists
                gt_video_idx_this_caption.append(len(videos)) # add current video idx
                videos.append(video_id)
                v2t_targets.append([len(texts)]) # one video only have one gt text
                annot = {
                    'clip_uid': int(video_id),
                    'caption': caption
                }
                annots.append(annot)
            texts.append(caption)
            t2v_targets.append(gt_video_idx_this_caption)

        assert len(videos) == len(v2t_targets)
        assert len(texts) == len(t2v_targets)
        return annots, videos, texts, v2t_targets, t2v_targets

    def _check_if_video_exists(self, video_id):
        return os.path.exists(os.path.join(self.vis_root, f"{video_id}.mp4"))

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        narration_text = ann['caption']

        id = ann['clip_uid']
        video_path = os.path.join(self.vis_root, f"{id}.mp4")
        start_time = None # ann_instance['clip_relative_start']
        end_time = None # ann_instance['clip_relative_end']

        assert os.path.exists(video_path)
        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        return {
            "video_input": raw_sampled_frames,
            "text_input": narration_text,
            "idx": int(index),
        }

    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=["idx"],
            text_fields=["text_input"]
        )

class DownstreamTask_MomentsInTime(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: str,
        ann_root: str,
        split: Literal['train', 'val'],
        task = Literal['video_action_retrieval','video_action_retrieval_2k'],
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2 for downstream retrieval tasks
        params:
            vis_root (string): Root directory of video clips (e.g., datasets/Moments_In_Time/videos)
            ann_root (string): path to the annotations root directory (e.g., datasets/Moments_In_Time/ann)
            task (string): video to action label
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        self.split = "training" if split == "train" else "validation"
        self.vis_root = os.path.join(vis_root, self.split)
        self.ann_root = ann_root

        self.task = task # only one task for now

        (
            annotations,
            self.videos,
            self.texts,
            self.v2t_targets,
            self.t2v_targets
        ) = self._load_annotations()

        self.clipuid_2_idx = {d['clip_uid'] : [i] for i, d in enumerate(annotations)} # Each clip has only one annotation

        self.neg_sampling_same_clip = 0 # set neg_sampling_same_clip to 0 for SSv2 since there is only one annotation per video

        super().__init__(
            annotations,
            task = task,
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = 0
        )

    def _load_annotations(self):
        if self.split == "training":
            ann_path = os.path.join(self.ann_root, "trainingSet.csv")
        else:
            if self.task == "video_action_retrieval_2k":
                ann_path = os.path.join(self.ann_root, "validationSet_2k.csv")
            else:
                ann_path = os.path.join(self.ann_root, "validationSet.csv")

        # get label texts & index mapping
        label_txt = os.path.join(self.ann_root, "moments_categories.txt")
        self.text_2_index = {}
        texts = []
        with open(label_txt, 'r') as f:
            for line in f:
                l, idx = line.split(',')
                l = l.strip().replace('+',' ')
                texts.append(l)
                self.text_2_index[l] = int(idx.strip())
        self.index_2_text = {value:key for key,value in self.text_2_index.items()}

        annotations = []
        videos = []
        v2t_targets = []
        t2v_targets = [[] for i in range(len(texts))]

        with open(ann_path, 'r') as f:
            for line in f:
                video_name, label, res_1, res_2 = line.split(',')
                vid_path = os.path.join(self.vis_root, video_name)
                if os.path.exists(vid_path):
                    text = label.strip().replace('+',' ')
                    annotations.append(
                        {
                            'clip_uid': video_name,
                            'caption': text
                        }
                    )
                    assert text in self.text_2_index
                    t_idx = self.text_2_index[text]
                    v2t_targets.append([t_idx])
                    t2v_targets[t_idx].append(len(videos))
                    videos.append(video_name)

        assert len(videos) == len(v2t_targets)
        assert len(texts) == len(t2v_targets)
        return annotations, videos, texts, v2t_targets, t2v_targets

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        narration_text = ann['caption']
        label_index = self.text_2_index[narration_text]

        video_path = os.path.join(self.vis_root, ann['clip_uid'])
        start_time = None # ann_instance['clip_relative_start']
        end_time = None # ann_instance['clip_relative_end']

        assert os.path.exists(video_path)
        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        return {
            "video_input": raw_sampled_frames,
            "text_input": narration_text,
            "label_index": label_index,
            "idx": int(index),
        }

    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=["idx","label_index"],
            text_fields=["text_input"]
        )

class DownstreamTask_Temporal(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: dict,
        ann_root: str,
        split: Literal['val'],
        task = Literal['v1.0','v1.0_2.4k'],
        subset = Literal['kinetics','ssv2'],
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2 for downstream retrieval tasks
        params:
            vis_root (dict): for example {
                'kinetics':'datasets/Temporal/video_clips/kinetics400/clips_downsampled_5fps_downsized_224x224',
                'ssv2': 'datasets/SSv2/video_clips/clips_downsampled_5fps_downsized_224x224'
            }
            ann_root (string): path to the annotations root directory (e.g. datasets/Temporal/ann)
            task (string): provided two versions of val set https://arxiv.org/abs/2301.02074
            subset: a list specifing the sub-datasets to use ['kinetic', 'ssv2'] or ['kinetic'] or ['ssv2']
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        self.split = split
        self.vis_root = vis_root[subset]
        self.ann_root = ann_root

        self.subset = subset # _load_annotations needs this to be set
        self.task = task # _load_annotations needs this to be set

        (
            annotations,
            self.videos,
            self.texts,
            self.v2t_targets,
            self.t2v_targets
        ) = self._load_annotations(split)

        self.clipuid_2_idx = {d['clip_uid'] : [i] for i, d in enumerate(annotations)} # Each clip has only one annotation

        self.neg_sampling_same_clip = 0 # set neg_sampling_same_clip to 0 for SSv2 since there is only one annotation per video

        super().__init__(
            annotations,
            task = task,
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = 0
        )

    def _load_annotations(self, split):
        if split == "val":
            if self.task == "v1.0":
                ann_path = os.path.join(self.ann_root, "val-v1.0.csv")
            elif self.task == "v1.0_2.4k":
                ann_path = os.path.join(self.ann_root, "val-v1.0-2.4k.csv")

            annots = []
            videos = []
            texts = []

            df = pd.read_csv(ann_path, header=0)

            text_2_video_idx = defaultdict(list)
            for index, row in df.iterrows():
                dataset_name = row['dataset'].lower()
                if dataset_name != self.subset:
                    continue
                if self.subset == "ssv2":
                    video_id = row['video_id'].strip()
                    start_time = None
                    end_time = None
                else:
                    video_id = row['video_id'].split("/")[1].strip()[:-14]
                    time_string = row['video_id'].split("/")[1].strip()[-13:]
                    start_time = float(time_string.split('_')[0])
                    end_time = float(time_string.split('_')[1])

                if not self._check_if_video_exists(video_id):
                    logging.info("!!!! video {} does not exist, skip...".format(video_id))
                    continue
                text = row['text'].replace('[something]', 'something')
                text_2_video_idx[text].append(len(videos))
                videos.append(video_id)
                annots.append(
                    {
                        'clip_uid': video_id,
                        'caption': text,
                        'start_time': start_time,
                        'end_time': end_time,
                    }
                )

            v2t_targets = [[] for i in range(len(videos))]
            t2v_targets = []
            for text,video_indices in text_2_video_idx.items():
                for v in video_indices:
                    v2t_targets[v] = [len(texts)]
                t2v_targets.append(video_indices)
                texts.append(text)

            assert len(videos) == len(v2t_targets)
            assert len(texts) == len(t2v_targets)
            return annots, videos, texts, v2t_targets, t2v_targets
        else:
            raise NotImplementedError("Temporal dataset for training not implemented")

    def _check_if_video_exists(self, video_id):
        return os.path.exists(os.path.join(self.vis_root, f"{video_id}.mp4"))

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        narration_text = ann['caption']
        start_time = ann['start_time']
        end_time = ann['end_time']

        id = ann['clip_uid']
        video_path = os.path.join(self.vis_root, f"{id}.mp4")
        assert os.path.exists(video_path)

        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        return {
            "video_input": raw_sampled_frames,
            "text_input": narration_text,
            "idx": int(index),
        }

    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=["idx"],
            text_fields=["text_input"]
        )

class DownstreamTask_QA_NextQA(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: dict,
        ann_root: str,
        split: Literal['train','val','test'],
        task = "mc", # multiple choice QA
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2 for downstream retrieval tasks
        params:
            vis_root (dict): path to videos (datasets/NextQA/video_clips/NExTVideo_downsampled_5fps_downsized_224x224)
            ann_root (string): path to the annotations root directory (datasets/NextQA/ann/nextqa)
            split: train, val, test
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        self.split = split
        self.vis_root = vis_root
        self.ann_root = ann_root

        self.map_vid_vidorID = json.load(open(os.path.join(self.ann_root, "map_vid_vidorID.json")))
        annotations, clipuid_2_idx = self._load_annotations(split)

        self.clipuid_2_idx = clipuid_2_idx # each clip_uid can have multiple annotations

        super().__init__(
            annotations,
            task = task,
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = neg_sampling_same_clip
        )

    def _load_annotations(self, split):
        csv = os.path.join(self.ann_root, f"{split}.csv")
        df = pd.read_csv(csv, header=0)

        annots = []
        clipuid_2_idx = defaultdict(list)
        for index, row in df.iterrows():
            vid = str(row['video'])
            if self._check_if_video_exists(vid):
                clipuid_2_idx[vid].append(len(annots))
                annots.append(
                    {
                        'clip_uid': vid,
                        'question': row['question'],
                        'answer': int(row['answer']),
                        'choices': [row[a] for a in ['a0','a1','a2','a3','a4']],
                        'action_antonym_choices': eval(row['action_antonym_choices']),
                        'qid': int(row['qid']),
                        'type': row['type'],
                        'start_time': None,
                        'end_time': None,
                    }
                )
        return annots, clipuid_2_idx

    def _check_if_video_exists(self, video_id):
        return os.path.exists(os.path.join(self.vis_root, self.map_vid_vidorID[video_id] + ".mp4"))

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        concat_qa_text = [ann['question']+"? "+ ans for ans in ann['choices']]
        action_antonym_concat_qa_text = [ann['question']+"? "+ ans for ans in ann['action_antonym_choices']]
        answer = ann['answer']
        start_time = ann['start_time']
        end_time = ann['end_time']

        id = ann['clip_uid']
        video_path = os.path.join(self.vis_root, self.map_vid_vidorID[id] + ".mp4")
        assert os.path.exists(video_path)

        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            concat_qa_text = [self.text_processor(t) for t in concat_qa_text]
        if self.text_processor is not None:
            action_antonym_concat_qa_text = [self.text_processor(t) for t in action_antonym_concat_qa_text]

        # task specific output
        return {
            "video_input": raw_sampled_frames,
            "text_input": concat_qa_text,
            "action_antonym_text_input": action_antonym_concat_qa_text,
            "answer": int(answer),
            "idx": int(index),
        }

    def collater(self, samples, text_fields = ['text_input', 'action_antonym_text_input']):
        if self.vis_processor is None:
            return samples
        ret = {
            "video_input":[],
            "text_input":[],
            "action_antonym_text_input":[],
            "answer": [],
            "idx": []
        }
        if self.neg_sampling_same_clip == 0:
            for sample in samples:
                for att in ret.keys():
                    ret[att].append(sample[att])
        else:
            # additional work for batching the positive and negative samples
            for sample in samples:
                pos_instance = sample['pos_instance']
                neg_instances = sample['neg_instances']
                for att in ret.keys():
                    if att in pos_instance:
                        ret[att] += [pos_instance[att]] + [neg[att] for neg in neg_instances]
        empty_field_keys = []
        for att in ret.keys():
            if att not in text_fields:
                if ret[att] != []:
                    if not isinstance(ret[att][0], torch.Tensor):
                        ret[att] = [torch.tensor(item) for item in ret[att]]
                    ret[att] = torch.stack(ret[att])
            if ret[att] == []:
                empty_field_keys.append(att)
        for remove_key in empty_field_keys:
            del ret[remove_key]
        return ret


from PyQt5 import QtGui
from PyQt5.QtWidgets import QMainWindow, QApplication, QLabel, QFileDialog
from PyQt5.uic import loadUi
from SearchEngine import textSearch, imageSearch, Freeman, DCT, os, Image, np
import glob, shutil

from tools import is_grayScale

class widget(QMainWindow):

    def __init__(self):
        QMainWindow.__init__(self)
        loadUi("assets/Search Engine.ui",self)
        self.setWindowTitle("Search Engine")
        self.startSearch.clicked.connect(self.deleteFiles)
        self.startSearch.clicked.connect(self.search)
        self.imgUpload.clicked.connect(self.deleteFiles)
        self.imgUpload.clicked.connect(self.UploadImage)

        self.path = os.getcwd()+'\\Download'
        self.searchPath = os.getcwd()+'\\Images'
        self.DBPath = os.getcwd()+'\\assets\\SearchEngine.db'
        self.pixel = None
        self.DCT = None
        self.Chains = None

    def UploadImage(self):
        try:
            imageFile = QFileDialog.getOpenFileName(None, "Open image", self.searchPath, "Image Files (*.png *.jpg *.bmp *.jpeg *.png *.jfif)")
            self.imagename = str(imageFile[0])
            self.pixel = np.array(Image.open(self.imagename).convert('RGB'))
            GrayScale = is_grayScale(self.pixel)
            if GrayScale:
                print('Method: Freeman')
                self.Chains = Freeman(self.pixel)
            else:
                print('Method: DCT')
                self.DCT = DCT(self.pixel)

            imageSearch(self.searchPath, self.path, self.DBPath, self.DCT, self.Chains)
            self.showImage()
            self.DCT = None
            self.Chains = None

        except: pass

    def showImage(self):
        os.chdir(self.path)
        row, column = 0, -1
        for file in glob.glob("*"):
            column += 1
            if column >= 4:
                row += 1
                column = 0
            pixmap = QtGui.QPixmap(self.path+"/"+file)
            pixmap = pixmap.scaled(270, 270)
            label = QLabel()
            label.setPixmap(pixmap)
            self.gridL.addWidget(label, row, column)

    def search(self):
        textSearch(self.path, self.DBPath, self.searchBar.text())
        self.showImage()

    def deleteFiles(self):
        for filename in os.listdir(self.path):
            file_path = os.path.join(self.path, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print('Failed to delete %s. Reason: %s' % (file_path, e))

        for i in reversed(range(self.gridL.count())):
            self.gridL.itemAt(i).widget().setParent(None)

app = QApplication([])
window = widget()
window.show()
app.exec_()


from decimal import Decimal

import pytest

from abusor.events.requirement_plugins import (
    ScoreIsAbove,
    ScoreIsBelow,
    SubjectContains,
)
from abusor.rules.plugins import RequirementPluginError

pytestmark = pytest.mark.django_db

@pytest.mark.parametrize(
    "value, required_value, expected",
    [
        ("0", "0", False),
        ("0", "5", False),
        ("5", "0", True),
        ("5", "5", False),
        ("10", "5", True),
    ],
)
def test_events_requirement_score_is_above(
    event_factory, value, required_value, expected
):
    event = event_factory(score=Decimal(value))
    req = ScoreIsAbove()
    outcome = req(event, Decimal(required_value))
    assert outcome == expected

def test_events_requirement_score_is_above_invalid_object():
    obj = object()
    req = ScoreIsAbove()
    with pytest.raises(RequirementPluginError) as excinfo:
        req(obj, Decimal("5"))
    assert "Object with type <class 'object'> has no attribute 'score'." in str(
        excinfo.value
    )

@pytest.mark.parametrize(
    "value, required_value, expected",
    [
        ("0", "0", False),
        ("0", "5", True),
        ("5", "0", False),
        ("5", "5", False),
        ("10", "5", False),
    ],
)
def test_events_requirement_score_is_below(
    event_factory, value, required_value, expected
):
    event = event_factory(score=Decimal(value))
    req = ScoreIsBelow()
    outcome = req(event, Decimal(required_value))
    assert outcome == expected

def test_events_requirement_score_is_below_invalid_object():
    obj = object()
    req = ScoreIsBelow()
    with pytest.raises(RequirementPluginError) as excinfo:
        req(obj, Decimal("5"))
    assert "Object with type <class 'object'> has no attribute 'score'." in str(
        excinfo.value
    )

@pytest.mark.parametrize(
    "value, required_value, expected",
    [("foo", "foo", True), ("bar", "foo", False), ("FOO", "foo", True)],
)
def test_events_requirement_subject_contains(
    event_factory, value, required_value, expected
):
    event = event_factory(subject=value)
    req = SubjectContains()
    outcome = req(event, required_value)
    assert outcome == expected

def test_events_requirement_subject_contains_invalid_object():
    obj = object()
    req = SubjectContains()
    with pytest.raises(RequirementPluginError) as excinfo:
        req(obj, "foo")
    assert "Object with type <class 'object'> has no attribute 'subject'." in str(
        excinfo.value
    )