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

# Let's import what we need
import os
import sys
import time
import shutil
import random
import urllib3
import argparse
import threading
from re import search, findall
from requests import get, post, codes
try:
    from urllib.parse import urlparse # for python3
except ImportError:
    from urlparse import urlparse # for python2

colors = True # Output should be colored
machine = sys.platform # Detecting the os of current system
if machine.startswith('os') or machine.startswith('win') or machine.startswith('darwin') or machine.startswith('ios'):
    colors = False # Colors shouldn't be displayed in mac & windows
if not colors:
    end = red = white = green = yellow = run = bad = good = info = que =  ''
else:
    end = '\033[1;m'
    red = '\033[91m'
    white = '\033[1;97m'
    green = '\033[1;32m'
    yellow = '\033[1;33m'
    run = '\033[1;97m[~]\033[1;m'
    bad = '\033[1;31m[-]\033[1;m'
    good = '\033[1;32m[+]\033[1;m'
    info = '\033[1;33m[!]\033[1;m'
    que =  '\033[1;34m[?]\033[1;m'

# Just a fancy ass banner
print ('''%s      ____  __          __
     / %s__%s \/ /_  ____  / /_____  ____
    / %s/_/%s / __ \/ %s__%s \/ __/ %s__%s \/ __ \\
   / ____/ / / / %s/_/%s / /_/ %s/_/%s / / / /
  /_/   /_/ /_/\____/\__/\____/_/ /_/ %s\n''' %
  (red, white, red, white, red, white, red, white, red, white, red, white, red, end))

urllib3.disable_warnings() # Disable SSL related warnings

# Processing command line arguments
parser = argparse.ArgumentParser()
parser.add_argument('-u', '--url', help='root url', dest='root')
parser.add_argument('-c', '--cookie', help='cookie', dest='cook')
parser.add_argument('-r', '--regex', help='regex pattern', dest='regex')
parser.add_argument('-s', '--seeds', help='additional seed urls', dest='seeds')
parser.add_argument('-l', '--level', help='levels to crawl', dest='level', type=int)
parser.add_argument('-t', '--threads', help='number of threads', dest='threads', type=int)
parser.add_argument('-n', '--ninja', help='ninja mode', dest='ninja', action='store_true')
parser.add_argument('-d', '--delay', help='delay between requests', dest='delay', type=int)
args = parser.parse_args()

if args.root: # if the user has supplied a url
    main_inp = args.root
    if main_inp.endswith('/'): # if the url ends with '/'
        main_inp = main_inp[:-1] # we will remove it as it can cause problems later in the code
else: # if the user hasn't supplied a url
    print ('\n' + parser.format_help().lower())
    quit()

delay = 0 # Delay between requests
cook = None # Cookie
ninja = False # Ninja mode toggle
crawl_level = 2 # Crawling level
thread_count = 2 # Number of threads

if args.cook:
    cook = args.cook
if args.ninja:
    ninja = True
if args.delay:
    delay = args.delay
if args.level:
    crawl_level = args.level
if args.threads:
    thread_count = args.threads

# Variables we are gonna use later to store stuff
files = set() # pdf, css, png etc.
intel = set() # emails, website accounts, aws buckets etc.
robots = set() # entries of robots.txt
custom = set() # string extracted by custom regex pattern
failed = set() # urls that photon failed to crawl
storage = set() # urls that belong to the target i.e. in-scope
scripts = set() # javascript files
external = set() # urls that don't belong to the target i.e. out-of-scope
fuzzable = set() # urls that have get params in them e.g. example.com/page.php?id=2
endpoints = set() # urls found from javascript files
processed = set() # urls that have been crawled

seeds = '' # custom urls to crawl
if args.seeds: # if the user has supplied custom seeds
    seeds = args.seeds
    for seed in seeds.split(','): # we will convert them into a list
        storage.add(seed) # and them to storage for crawling

# If the user hasn't supplied the root url with http(s), we will handle it
if main_inp.startswith('http'):
    main_url = main_inp
else:
    try:
        get('https://' + main_inp)
        main_url = 'https://' + main_inp
    except:
        main_url = 'http://' + main_inp

storage.add(main_url) # adding the root url to storage for crawling

name = urlparse(main_url).netloc # Extracts domain out of the url

####
# This function makes requests to webpage and returns response body
####

# list of user agents
user_agents = ['Mozilla/5.0 (X11; Linux i686; rv:60.0) Gecko/20100101 Firefox/60.0',
'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/60.0.3112.113 Safari/537.36'
'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/56.0.2924.87 Safari/537.36 OPR/43.0.2442.991']

def requester(url):
    processed.add(url) # mark the url as crawled
    time.sleep(delay) # pause/sleep the program for specified time
    try:
        def normal(url):
            headers = {
            'Host' : name, # ummm this is the hostname?
            'User-Agent' : random.choice(user_agents), # selecting a random user-agent
            'Accept' : 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
            'Accept-Language' : 'en-US,en;q=0.5',
            'Accept-Encoding' : 'gzip',
            'DNT' : '1',
            'Connection' : 'close'}
            # make request and return response
            return get(url, cookies=cook, headers=headers, verify=False)

        # pixlr.com API
        def pixlr(url):
            if url == main_url:
                url = main_url + '/' # because pixlr throws error if http://example.com is used
            # make request and return response
            return get('https://pixlr.com/proxy/?url=' + url, headers={'Accept-Encoding' : 'gzip'}, verify=False)

        # codebeautify.org API
        def code_beautify(url):
            headers = {
            'User-Agent' : 'Mozilla/5.0 (X11; Linux x86_64; rv:61.0) Gecko/20100101 Firefox/61.0',
            'Accept' : 'text/plain, */*; q=0.01',
            'Accept-Encoding' : 'gzip',
            'Content-Type' : 'application/x-www-form-urlencoded; charset=UTF-8',
            'Origin' : 'https://codebeautify.org',
            'Connection' : 'close'
            }
            # make request and return response
            return post('https://codebeautify.com/URLService', headers=headers, data='path=' + url, verify=False)

        # www.photopea.com API
        def photopea(url):
            # make request and return response
            return get('https://www.photopea.com/mirror.php?url=' + url, verify=False)

        if ninja: # if the ninja mode is enabled
            # select a random request function i.e. random API
            response = random.choice([photopea, normal, pixlr, code_beautify])(url)

            #check if link is not broken
            if response.status_code == codes.ok:
                return response.text  # return response body
            else:
                failed.add(url) # add it to the failed list

        else:
            response = normal(url)

            #check if link is not broken
            if response.status_code == codes.ok:
                return response.text  # return response body
            else:
                failed.add(url) # add it to the failed list

    except: # if photon fails to connect to the url
        failed.add(url) # add it to the failed list

####
# This function extracts links from robots.txt and sitemap.xml
####

def zap(url):
    try:
        response = requester(url + '/robots.txt') # makes request to robots.txt
        if '<body' not in response: # making sure robots.txt isn't some fancy 404 page
            matches = findall(r'Allow: (.*)|Disallow: (.*)', response) # If you know it, you know it
            for match in matches: # iterating over the matches, match is a tuple here
                match = match[0] + match[1] # one item in match will always be empty so will combine both items
                if '/*/' not in match: # if the url doesn't use a wildcard
                    url = main_url + match
                    storage.add(url) # add the url to storage list for crawling
                    robots.add(url) # add the url to robots list
                    if '=' in url: # self-explanatory
                        fuzzable.add(url) # add url to the fuzzable urls list
            if len(matches) > 0: # if there are more than 0 matches
                print ('%s URLs retrieved from robots.txt: %s' % (good, len(robots)))

        response = requester(url + '/sitemap.xml') # makes request to sitemap.xml

        if '<body' not in response: # making sure robots.txt isn't some fancy 404 page
            matches = findall(r'<loc>[^<]*</loc>', response) # regex for extracting urls
            if matches: # if there are any matches
                print ('%s URLs retrieved from sitemap.xml: %s' % (good, len(matches)))
                for match in matches:
                    storage.add(match.split('<loc>')[1][:-6]) #cleaning up the url & adding it to the storage list for crawling
    except:
        pass
####
# This functions checks whether a url should be crawled or not
####

def is_link(url):
    # file extension that don't need to be crawled and are files
    bad = ['xml', 'png', 'bmp', 'jpg', 'jpeg', 'pdf', 'css', 'ico', 'js', 'svg', 'json']
    conclusion = False # whether the the url should be crawled or not

    if url not in processed: # if the url hasn't been crawled already
        if '.' in url[-5:]: # url[-5:] = last 5 chars of url
            extension = url[-5:].split('.')[-1] # finding the extension of webpage
            if extension not in bad: # if the extension isn't "bad"
                conclusion = True # url can be crawled
            else:
                files.add(main_url + url) # its a file
        else: # it has no extension
            return True # url can be crawled
    return conclusion # return the conclusion :D

####
# This function extracts stuff from the response body
####

def extractor(url):
    response = requester(url) # make request to the url
    try:
        matches = findall(r'href=[\'"](.*?)["\']', response) # find all the linkx
        for link in matches: # iterate over the matches
            link = link.split('#')[0] # remove everything after a "#" to deal with in-page anchors
            if is_link(link): # checks if the urls should be crawled
                if link.startswith(main_url): # self-explanatory
                    storage.add(link)
                    if '=' in link:
                        fuzzable.add(link)
                if link.startswith('http') or link.startswith('//'):
                    if not link.startswith(main_url):
                        external.add(link)
                elif link.startswith('javascript') or link.startswith('{{'):
                    pass # we aren't fucking with javascript
                elif link.startswith('/') and not link.startswith('//'):
                    storage.add(main_url + link)
                    if '=' in link:
                        fuzzable.add(link)
                elif not link.startswith('http') and not link.startswith('//'):
                    storage.add(main_url + '/' + link)
                    if '=' in link:
                        fuzzable.add(link)

        matches = findall(r'src=[\'"](.*?\.js)["\']', response) # extract .js files
        for match in matches: # iterate over the matches
            if match.startswith(main_url):
                scripts.add(match)
            elif match.startswith('/') and not match.startswith('//'):
                scripts.add(main_url + match)
            elif not match.startswith('http') and not match.startswith('//'):
                scripts.add(main_url + '/' + match)
        # extract intel ;)
        matches = findall(r'''([\w\.-]+s[\w\.-]+\.amazonaws\.com)|(github\.com/[\w\.-/]+)|(facebook\.com/.*?)[\'" ]|
(youtube\.com/.*?)[\'" ]|(linkedin\.com/.*?)[\'" ]|(twitter\.com/.*?)[\'" ]|([\w\.-]+@[\w\.-]+\.[\.\w]+)''', response)
        for match in matches: # iterate over the matches
            for x in match: # iterate over the match because it's a tuple
                if x != '': # if the value isn't empty
                    intel.add(x) # add it to intel list

        def regxy(pattern):
            matches = findall(r'%s' % pattern, response)
            for match in matches:
                custom.add(match)
        if args.regex:
            regxy(args.regex)

    except: # if something bad happens
        failed.add(url)

####
# This function extracts endpoints from JavaScript Code
####

def jscanner(url):
    try:
        response = requester(url) # make request to the url
        matches = findall(r'[\'"](/.*?)[\'"]|[\'"](http.*?)[\'"]', response) # extract urls/endpoints
        for match in matches: # iterate over the matches, match is a tuple
            match = match[0] + match[1] # combining the items because one of them is always empty
            if not search(r'[}{><"\']', match) and not match == '/': # making sure it's not some js code
                endpoints.add(match) # add it to the endpoints list
    except:
        failed.add(url)

####
# This function starts multiple threads for a function
####

def threader(function, *urls):
    threads = [] # list of threads
    urls = urls[0] # because urls is a tuple
    for url in urls: # iterating over urls
        task = threading.Thread(target=function, args=(url,))
        threads.append(task)
    # start threads
    for thread in threads:
        thread.start()
    # wait for all threads to complete their work
    for thread in threads:
        thread.join()
    # delete threads
    del threads[:]

####
# This function processes the urls and sends them to "threader" function
####

def flash(function, links): # This shit is complicated, move on
    links = list(links) # convert links (set) to list
    begin = 0 # begining of slice
    end = thread_count # ending of slice
    # Okay I give up, if you know it, you know it
    for i in range((int(len(links)/thread_count)) + 1):
        splitted = links[begin:end]
        threader(function, splitted)
        begin += thread_count
        if (len(links) - end) >= thread_count:
            end += thread_count
        else:
            end += len(links) - begin
        progress = begin
        if progress > len(links):
            progress = len(links)
        sys.stdout.write('\r%s Progress: %i/%i' % (info, progress, len(links)))
        sys.stdout.flush()
    print ('')

then = time.time() # records the time at which crawling started

# Step 1. Extract urls from robots.txt & sitemap.xml
zap(main_url)

# Step 2. Crawl recursively to the limit specified in "crawl_level"
for level in range(crawl_level):
    links = storage - processed # links to crawl = all links - already crawled links
    if len(links) == 0: # if links to crawl are 0 i.e. all links have been crawled
        break
    if len(storage) <= len(processed): # if crawled links are somehow more than all links. Possible? ;/
        if len(storage) > 2 + len(seeds): # if you know it, you know it
            break
    print ('%s Level %i: %i URLs' % (run, level + 1, len(links)))
    try:
        flash(extractor, links)
    except KeyboardInterrupt:
        print ('')
        break

# Step 3. Scan the JavaScript files for enpoints
print ('%s Crawling %i JavaScript files' % (run, len(scripts)))
flash(jscanner, scripts)

now = time.time() # records the time at which crawling stopped
diff = (now  - then) # finds total time taken

def timer(diff):
    minutes, seconds = divmod(diff, 60) # Changes seconds into minutes and seconds
    time_per_request = diff / float(len(processed)) # Finds average time taken by requests
    return minutes, seconds, time_per_request
time_taken = timer(diff)
minutes = time_taken[0]
seconds = time_taken[1]
time_per_request = time_taken[2]
# Step 4. Save the results
if os.path.exists(name): # if the directory already exists
    shutil.rmtree(name, ignore_errors=True) # delete it, recursively
os.mkdir(name) # create a new directory

with open('%s/links.txt' % name, 'w+') as f:
    for x in storage:
        f.write(x + '\n')

with open('%s/files.txt' % name, 'w+') as f:
    for x in files:
        f.write(x + '\n')

with open('%s/intel.txt' % name, 'w+') as f:
    for x in intel:
        f.write(x + '\n')

with open('%s/robots.txt' % name, 'w+') as f:
    for x in storage:
        f.write(x + '\n')

with open('%s/failed.txt' % name, 'w+') as f:
    for x in failed:
        f.write(x + '\n')

with open('%s/custom.txt' % name, 'w+') as f:
    for x in custom:
        f.write(x + '\n')

with open('%s/scripts.txt' % name, 'w+') as f:
    for x in scripts:
        f.write(x + '\n')

with open('%s/fuzzable.txt' % name, 'w+') as f:
    for x in fuzzable:
        f.write(x + '\n')

with open('%s/external.txt' % name, 'w+') as f:
    for x in external:
        f.write(x + '\n')

with open('%s/endpoints.txt' % name, 'w+') as f:
    for x in endpoints:
        f.write(x + '\n')

# Printing out results
print ('''%s
%s URLs: %i
%s Intel: %i
%s Files: %i
%s Endpoints: %i
%s Fuzzable URLs: %i
%s Custom strings: %i
%s JavaScript Files: %i
%s External References: %i
%s''' % ((('%s-%s' % (red, end)) * 50), good, len(storage), good,
len(intel), good, len(files), good, len(endpoints), good, len(fuzzable), good,
len(custom), good, len(scripts), good, len(external),
(('%s-%s' % (red, end)) * 50)))

print ('%s Total time taken: %i:%i' % (info, minutes, seconds))
print ('%s Average request time: %s' % (info, str(time_per_request)[:4]))

if not colors: # if colors are disabled
    print ('%s Results saved in %s directory' % (good, name))
else:
    print ('%s Results saved in \033[;1m%s\033[0m directory' % (good, name))


#!/usr/bin/env python

import serial
import time

arduino = serial.Serial(
	port="/dev/cu.wchusbserialfa210",
	baudrate=9600,
	timeout=3,
	bytesize=serial.EIGHTBITS
)
time.sleep(3)
arduino.write('... --- ...\n')
time.sleep(1)
out=''
while arduino.inWaiting() > 0:
    out += arduino.read(40)

if out != '':
    print(">>" + out)

arduino.close()


"""
Collect a set of tweets based on keywords and save them to a specified file
"""

import sys
import tweepy
import json
import signal

from tweepy.streaming import StreamListener
from tweepy import OAuthHandler
from tweepy import Stream

consumer_key = ''
consumer_secret = ''
access_token = ''
access_token_secret = ''

# This is global so it can be access in the ctrl-c interrupt
tweet_writer = None

class TweetStreamWriter:
    def __init__(self, output_file):
        self.data = []
        self.output_file = output_file

    def add_tweet(self, tweet_data):
        self.data.append(tweet_data)

    def store_twitter_data(self):
        with open(self.output_file, 'w') as data_file:
            first_item = True
            data_file.write('[')

            for item in self.data:
                if first_item:
                    first_item = False
                else:
                    data_file.write(',')
                data_file.write(item)

            data_file.write(']')

        self.data = []

    def cleanup(self):
        # Write the extra data if there is any
        if self.data:
            self.store_twitter_data()

class TwitterStreamReader(StreamListener):
    def __init__(self, tweet_aggregator, max_tweets):
        self.tweet_aggregator = tweet_aggregator
        self.count = 0
        self.max_tweets = max_tweets

    def on_data(self, data):
        self.tweet_aggregator.add_tweet(data)
        self.count += 1
        return (self.count < self.max_tweets)

    def on_error(self, status):
        print(status)

def collect_tweets(output_file, max_tweets, keywords):
    def signal_handler(signal, frame):
        global tweet_writer
        tweet_writer.cleanup()
        sys.exit(0)

    # Register a signal handler for SIGINT
    signal.signal(signal.SIGINT, signal_handler)

    global tweet_writer
    tweet_writer = TweetStreamWriter(output_file)
    stream_reader = TwitterStreamReader(tweet_writer, max_tweets)

    auth = OAuthHandler(consumer_key, consumer_secret)
    auth.set_access_token(access_token, access_token_secret)
    stream = Stream(auth, stream_reader)
    stream.filter(track=keywords)
    tweet_writer.cleanup()

if __name__ == '__main__':
    if len(sys.argv) < 4:
        print('Usage: python stream.py output_file max_tweets keywords...')
        exit()
    collect_tweets(sys.argv[1], int(sys.argv[2]), sys.argv[3:])


import logging
import os

from osbenchmark.builder.downloaders.builders.binary_builder import BinaryBuilder
from osbenchmark.exceptions import ExecutorError, BuildError

class SourceBinaryBuilder(BinaryBuilder):
    def __init__(self, executor, path_manager, jdk_resolver, source_directory, build_jdk_version, log_directory):
        self.logger = logging.getLogger(__name__)
        self.executor = executor
        self.path_manager = path_manager
        self.jdk_resolver = jdk_resolver

        self.source_directory = source_directory
        self.build_jdk_version = build_jdk_version
        self.log_directory = log_directory

    def build(self, host, build_commands, override_source_directory=None):
        for build_command in build_commands:
            self._run_build_command(host, build_command, override_source_directory)

    def _run_build_command(self, host, build_command, override_source_directory):
        source_directory = self.source_directory if override_source_directory is None else override_source_directory

        self.path_manager.create_path(host, self.log_directory, create_locally=False)
        log_file = os.path.join(self.log_directory, "build.log")

        _, jdk_path = self.jdk_resolver.resolve_jdk_path(host, self.build_jdk_version)
        self.executor.execute(host, f"export JAVA_HOME={jdk_path}")

        self.logger.info("Running build command [%s]", build_command)
        try:
            self.executor.execute(host, f"{source_directory}/{build_command} > {log_file} 2>&1")
        except ExecutorError:
            raise BuildError(f"Executing {build_command} failed. The build log can be found at {log_file}")


from data import window, frameRow2, frameRow3, frameRow4

from PeriodiekSysteem import Nr, Symb, Naam, MolM, Dicht, SP, KP, zoek, frmPeriodiek
frmPeriodiek.pack(pady=10)

from MolaireMassa import frameMolM, entFormule, berekenMolM
frameMolM.grid(row=0, column=1, padx=5, sticky='n')

from MolaireMassa import frmCommas
frmCommas.grid(row=0, column=2, sticky='n',padx=5)

from WortelHerleiden import frmWortel, entGetal, entGetal2, wortel
frmWortel.grid(row=0, column=0, sticky='n', padx=5)

from ReactieVergelijking import frmReacTot, entReac, CalcReac
frmReacTot.grid(row=0, column=0, padx=5)

from TwentyFourGame import frm24Game, TwfGameEnt1, TwfGameEnt2, TwfGameEnt3, TwfGameEnt4, twfGame
frm24Game.grid(row=0, column=1, padx=5)

from GtaVoltHack import VLFrame, VLCinp1, VLCinp2, VLCinp3, VLFinp1, VLFinp2, VLFinp3, VLtargEnt, VLhack
VLFrame.grid(row=0, column=2, padx=5)

from FileDateChanger import FTFrame
FTFrame.grid(row=0, column=3, padx=5)

from nCrnPr import frmCP, entCPn, entCPr, CPber
frmCP.grid(row=0, column=0, padx=5)

def enter(key):
    if Nr.get() != '' or Symb.get() != '' or Naam.get() != '' or MolM.get() != '' or Dicht.get() != '' or SP.get() != '' or KP.get() != '':
        zoek()
    if entFormule.get() != '':
        berekenMolM(0, 0)
    if entGetal.get() != '' or entGetal2.get() != '':
        wortel()
    if entReac.get() != '':
        CalcReac(entReac.get())
    if TwfGameEnt1.get() != '' and TwfGameEnt2.get() != '' and TwfGameEnt3.get() != '' and TwfGameEnt4.get() != '':
        twfGame()
    if VLCinp1.get() != '' and VLCinp2.get() != '' and VLCinp3.get() != '' and VLFinp1.get() != '' and VLFinp2.get() != '' and VLFinp3.get() != '' and VLtargEnt.get() != '':
        VLhack()
    if entCPn.get() != '' and entCPr.get() != '':
        CPber()

frameRow2.pack(pady=5)
frameRow3.pack(pady=5)
frameRow4.pack(pady=5)

window.bind('<Return>', enter)

window.mainloop()


#!/usr/bin/env python

import os
import re
import sys
import nibabel
import pandas
import json

from glob import glob

# These arguments come from the command line. You should first
# run this interactively (eg with ipython in a terminal) and
# type/copy paste the commands one at a time.
# /home/vanessa/Documents/Work/pipelines/ryan/sub-01
data_folder = sys.argv[1]
output_folder = sys.argv[2]

# Here we are making sure that output folders and data folders exist
# and making them if necessary
if not os.path.exists(output_folder):
    os.mkdir(output_folder)
    print("Creating output folder %s" %(output_folder))

if not os.path.exists(data_folder):
    print("Cannot find %s. Check path and try again" %(data_folder))

# This is the subject id, and we are extracting from the folder name
sub_id = os.path.basename(data_folder) # You can also extract this from the data

# Here we are preparing output BIDS folders
bids_folder = "%s/%s" %(output_folder,sub_id)
if not os.path.exists(bids_folder):
    print("Making bids output folder %s" %(bids_folder))
    os.mkdir(bids_folder)

for folder in ['anat','pet']:
    modality_folder = "%s/%s" %(bids_folder,folder)
    if not os.path.exists(modality_folder):
        os.mkdir(modality_folder)

# S1 has the anatomicals
anat_files = glob('%s/S1/*.nii.gz' %data_folder)

# S2 has the pet
pet_file = glob('%s/S2/*.nii.gz' %data_folder)[0]

# Let's use this anatomical
anat_file = [x for x in anat_files if re.search('^co',os.path.basename(x))][0]
anat_bids = "%s/anat/%s_T1w.nii.gz" %(bids_folder,sub_id)
os.rename(anat_file,anat_bids)

# And rename the pet
pet_bids = "%s/pet/%s_pet.nii.gz" %(bids_folder,sub_id)
os.rename(pet_file,pet_bids)

# RYAN: here is where you start!
# Now we want to generate a json data structure for each of these files to be finished.
# Step 1: Download this example pet data in BIDS Format: https://drive.google.com/file/d/0B2JWN60ZLkgkOHdxRk5BUUlQelE/view?usp=sharing
# Step 2: Look at the json file in /pet_phno/sub-000005/pet/. This is the file that you will want to create for your pet data,
# and the equivalent for the T1 (see in the anat folder).
# Step 3: Read in your nifti output files generated above as follows:

anat = nibabel.load(anat_bids)
pet = nibabel.load(pet_bids)

# Take a look at pet.header.keys() that can be accessed in the header, and read about
# nibabel here: http://nipy.org/nibabel/
# Your goal is to generate that json file, name equivalently as the example, and save
# each one to the anat / pet folder respectively. Use the write_json function below
# to save your dictionary. So for example, you could do:

anat_dict = dict()
anat_dict["Info"] = dict()
anat_dict["Info"]["BodyPart"] = "brain"

# FILL IN THE REST HERE

# Save, each of anat and pet

# Final Step
# Generation of BIDS information

bids_description = { "BIDSVersion": "1.1.0-pet",
                     "Name": "FLUMAZENILAUTISM" }
bids_description_file = "%s/dataset_description.json" %bids_folder

# Here is a function to write a json file
def write_json(json_obj,filename,mode="w",print_pretty=True):
    '''write_json will (optionally,pretty print) a json object to file
    :param json_obj: the dict to print to json
    :param filename: the output file to write to
    :param pretty_print: if True, will use nicer formatting
    '''
    with open(filename,mode) as filey:
        if print_pretty == True:
            filey.writelines(json.dumps(json_obj, indent=4, separators=(',', ': ')))
        else:
            filey.writelines(json.dumps(json_obj))
    return filename

# Write the file
write_json(bids_description,bids_description_file)

# Write participants file
participants_file = "%s/participants.tsv" %(bids_folder)
if os.path.exists(participants_file):
    participants = pandas.read_csv(participants_file,index=True,sep="\t")
else:
    participants = pandas.DataFrame(columns=["participant_id","species","strain","population",
                                             "bodyWeight","bodyWeightUnits","age","ageUnits","gender"])

participants.loc[sub_id,"participant_id"] = sub_id
participants.to_csv(participants_file,sep="\t")


import sys
import numpy as np
from tqdm import tqdm

fil = sys.argv[1]
vocab_dir = sys.argv[2]

with open(vocab_dir + '/source_tokens.txt', 'r') as f:
    src_vocab_token2index = dict({x.strip('\n'): idx + 1 for idx, x in enumerate(f.readlines())})
    src_vocab_token2index['@@PADDING@@'] = 0
target_namespaces = ["answer", "program"]
dst_vocab_token2index = {}
for tn in target_namespaces:
    with open(vocab_dir + f'/{tn}.txt', 'r') as f:
        dst_vocab_token2index[tn] = dict({x.strip('\n'): idx + 1 for idx, x in enumerate(f.readlines())})
        dst_vocab_token2index[tn]['@@PADDING@@'] = 0

with open(fil, 'r') as f:
    lines = f.readlines()

QUES_MAX_LEN = 160 + 2
ANS_MAX_LEN = 30 + 3
PROG_MAX_LEN = 30 + 3
N = len(lines) // 2
Q = np.zeros((N, QUES_MAX_LEN + 1), dtype=np.ubyte)
A = np.zeros((N, ANS_MAX_LEN + 1), dtype=np.ubyte)
P = np.zeros((N, PROG_MAX_LEN + 1), dtype=np.ubyte)
get_src_idx = lambda x: src_vocab_token2index[x]
get_dst_idx = {k: dst_vocab_token2index[k].__getitem__ for k in dst_vocab_token2index}

for idx in tqdm(range(len(lines) // 2)):
    ques = lines[2 * idx].strip()
    ans, prog = lines[2 * idx + 1].strip().split('###')
    prog = prog.split()
    assert len(ques) <= QUES_MAX_LEN
    assert len(ans) <= ANS_MAX_LEN
    assert len(prog) <= PROG_MAX_LEN
    Q[idx][:len(ques)] = list(map(get_src_idx, ques))
    Q[idx][-1] = len(ques)
    A[idx][:len(ans) + 2] = [dst_vocab_token2index["answer"]["@start@"]] + list(map(get_dst_idx["answer"], ans)) + [
        dst_vocab_token2index["answer"]["@end@"]]
    A[idx][-1] = len(ans) + 2
    P[idx][:len(prog) + 2] = [dst_vocab_token2index["program"]["@start@"]] + list(map(get_dst_idx['program'], prog)) + [
        dst_vocab_token2index["program"]["@end@"]]
    P[idx][-1] = len(prog) + 2

data = {"Q": Q, "A": A, "P": P}

if Q.shape[0] > 10 ** 7:
    format = 'npz'
else:
    format = 'npy'

if fil.endswith('.txt'):
    save_path = fil[:-4] + '.' + format
else:
    save_path = f'{fil}.{format}'
if format == 'npy':
    np.save(save_path, data)
else:
    np.savez(save_path, Q=Q, A=A, P=P)


from datetime import datetime
from django.contrib import admin
from counsellor.forms import CounsellorAdminCreationForm, UserAdminCreationForm
from django.contrib import messages
from counsellor.models import *
from django.contrib.auth.admin import UserAdmin as BaseUserAdmin
from django.shortcuts import redirect

# Register your models here.
admin.site.register(Consent)
admin.site.register(CounsellingType)

class UserAdmin(BaseUserAdmin):
    # The forms to add and change user instances
    # form = UnPriviledgedUserForm
    add_form = UserAdminCreationForm
    # The fields to be used in displaying the User model.
    # These override the definitions on the base UserAdmin
    # that reference specific fields on auth.User
    list_display = [
        "email",
        "first_name",
        "last_name",
        "other_names",
        "is_active",
        "is_staff",
        "is_superuser",
        "user_type",
        "sex",
        "religion",
        "marital_status",
        "phone",
    ]

    list_display_links = ['email', 'first_name', 'last_name']

    search_fields = ['first_name', 'last_name', 'email', 'phone']

    ordering = ['-created_at']
    list_filter = []

    fieldsets = (
        (None, {'fields': ('email', 'password')}),
        ('Personal info', {'fields': (
            'first_name',
            'last_name',
            'other_names',
            'phone',
            'address',
            'sex',
            'religion',
            'marital_status',
        )}),
        ('Permissions', {'fields': (
            'is_active',
            'is_staff',
            'is_superuser',
        )}),
        (None, {'fields': (
            'user_type',
            'avatar',
        )}),
    )

    # add_fieldsets is not a standard ModelAdmin attribute. UserAdmin
    # overrides get_fieldsets to use this attribute when creating a user.
    add_fieldsets = (
        (None, {
            'classes': ('wide',),
            'fields': (
                'email',
                'first_name',
                'last_name',
                'password',
                'password2'
                )
            }
        ),
    )

class CounsellorAdmin(UserAdmin):
    add_form = CounsellorAdminCreationForm

    list_display = [
        "email",
        "first_name",
        "last_name",
        "other_names",
        "is_active",
        "is_staff",
        "is_superuser",
        "sex",
        "phone",
    ]

    list_display_links = ['email', 'first_name', 'last_name']

    search_fields = ['first_name', 'last_name', 'email', 'phone']
    filter_horizontal = ['fields','availability']
    ordering = ['-created_at']
    list_filter = []

    fieldsets = (
        (None, {'fields': ('email', 'password')}),
        ('Personal info', {'fields': (
            'title',
            'first_name',
            'last_name',
            'other_names',
            'phone',
            'address',
            'sex',
            'description',
        )}),
        ('Permissions', {'fields': (
            'is_active',
            'is_staff',
            'is_superuser',
        )}),
        (None, {'fields': (
            'available',
            'profile',
            'fields',
            'availability',
        )}),
    )

    # add_fieldsets is not a standard ModelAdmin attribute. UserAdmin
    # overrides get_fieldsets to use this attribute when creating a user.
    add_fieldsets = (
        (None, {
            'classes': ('wide',),
            'fields': (
                'email',
                'first_name',
                'last_name',
                'password',
                'password2'
                )
            }
        ),
    )

@admin.register(Booking)
class BookingAdmin(admin.ModelAdmin):
    list_display = [
        "client",
        "counsellor",
        "status",
        "created_at"
    ]

    list_display_links = ['client', 'counsellor', 'status']

    search_fields = ['client', 'counsellor', 'status']
    ordering = ['-created_at']
    list_filter = ["status"]

    def save_model(self, request, obj, form, change):
        start_time   = datetime.combine(obj.start_date, obj.start_time).replace(hour=0,minute=0, second=0, microsecond=0)
        current_time = datetime.now().replace(hour=0,minute=0, second=0, microsecond=0)

        if start_time < current_time:
            messages.error(request, "Starting date cannot be any day before current")
            return redirect('.')

        elif obj.start_time > obj.end_time:
            messages.error(request, "Starting time cannot be more than ending time.")
            return redirect('.')

        elif obj.start_time == obj.end_time:
            messages.error(request, "Starting time cannot be same as ending time.")
            return redirect('.')

        else:
            return super(BookingAdmin, self).save_model(request, obj, form, change)

@admin.register(AvailableDay)
class AvailableDayAdmin(admin.ModelAdmin):
    list_display = ['day']

    def has_change_permission(self, request, obj=None):
        return True

    def has_delete_permission(self, request, obj=None):
        return True

    def has_add_permission(self, request, obj=None):
        return True

admin.site.register(User, UserAdmin)
admin.site.register(Counsellor, CounsellorAdmin)


from scipy import io as sio
import numpy as np
import os
import dicom
import nibabel as nib
import datetime
from cafndl_fileio import *
from cafndl_utils import *
from cafndl_network import *
from cafndl_metrics import *
from keras.callbacks import ModelCheckpoint
from keras.optimizers import Adam

'''
convert dicom to nifti
$ mkdir DRF100_nifti
$ dicom2nifti DRF100 DRF100_nifti
'''

'''
dataset
'''
filename_checkpoint = '../ckpt/model_demo_0713.ckpt'
list_dataset_test =  [
				{
				 'inputs':['/data/enhaog/data_lowdose/GBM_Ex1496/DRF020_nifti/802_.nii.gz',
				 			'/data/enhaog/data_lowdose/GBM_Ex1496/DRF100_nifti/803_.nii.gz'],
				 'gt':'/data/enhaog/data_lowdose/GBM_Ex1496/DRF001_nifti/800_.nii.gz'
				},
				{
				 'inputs':['/data/enhaog/data_lowdose/GBM_Ex842/DRF020_nifti/802_.nii.gz',
				 			'/data/enhaog/data_lowdose/GBM_Ex842/DRF100_nifti/803_.nii.gz'],
				 'gt':'/data/enhaog/data_lowdose/GBM_Ex842/DRF001_nifti/800_.nii.gz'
				}
				]
filename_results = '../results/result_demo_0713'
num_dataset_test = len(list_dataset_test)
print('process {0} data description'.format(num_dataset_test))

'''
augmentation
'''
list_augments = []

'''
generate test data
'''
list_test_input = []
list_test_gt = []
for index_data in range(num_dataset_test):
	# directory
	list_data_test_input = []
	for path_test_input in list_dataset_test[index_data]['inputs']:
		# load data
		data_test_input = prepare_data_from_nifti(path_test_input, list_augments)
		list_data_test_input.append(data_test_input)
	data_test_input = np.concatenate(list_dataset_test_input, axis=-1)

	# get ground truth
	path_test_gt = list_dataset_test[index_data]['gt']
	data_test_gt = prepare_data_from_nifti(path_test_gt, [])

	# append
	list_test_input.append(data_test_input)
	list_test_gt.append(data_test_gt)

# generate test dataset
scale_data = 100.
data_test_input = scale_data * np.concatenate(list_test_input, axis = 0)
data_test_gt = scale_data * np.concatenate(list_test_gt, axis = 0)
data_test_residual = data_test_gt - data_test_input[:,:,:,0:1]
print('mean, min, max')
print(np.mean(data_test_input.flatten()),np.min(data_test_input.flatten()),np.max(data_test_input.flatten()))
print(np.mean(data_test_gt.flatten()),np.min(data_test_gt.flatten()),np.max(data_test_gt.flatten()))
print(np.mean(data_test_residual.flatten()),np.min(data_test_residual.flatten()),np.max(data_test_residual.flatten()))
print('generate test dataset with augmentation size {0},{1}'.format(
	data_test_input.shape, data_test_gt.shape))

'''
setup parameters
'''
keras_memory = 0.3
keras_backend = 'tf'
num_channel_input = data_test_input.shape[-1]
num_channel_output = data_test_gt.shape[-1]
img_rows = data_test_input.shape[1]
img_cols = data_test_gt.shape[1]
num_poolings = 3
num_conv_per_pooling = 3
lr_init = 0.001
with_batch_norm = True
ratio_validation = 0.1
batch_size = 4
always_retrain = 1
num_epoch = 100
print('setup parameters')

'''
define model
'''
setKerasMemory(keras_memory)
model = deepEncoderDecoder(num_channel_input = num_channel_input,
						num_channel_output = num_channel_output,
						img_rows = img_rows,
						img_cols = img_cols,
						lr_init = lr_init,
						num_poolings = num_poolings,
						num_conv_per_pooling = num_conv_per_pooling,
						with_bn = with_batch_norm, verbose=1)
print('train model:', filename_checkpoint)
print('parameter count:', model.count_params())

'''
load network
'''
model.load_weights(filename_checkpoint)
print('model load from' + filename_checkpoint)

'''
apply model
'''
t_start_pred = datetime.datetime.now()
data_test_output = model.predict(data_test_input, batch_size=batch_size)
# clamp
clamp_min = -0.5
clamp_max = 0.5
data_test_output = np.maximum(np.minimum(data_test_output, clamp_max), clamp_min)

# add
data_test_output += data_test_input[:,:,:,0:1]

t_end_pred = datetime.datetime.now()
print('predict on data size {0} using time {1}'.format(
	data_test_output.shape, t_end_pred - t_start_pred))

'''
export images
'''
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
num_sample = data_test_output.shape[0]
list_err_pred = []
list_err_input = []
aug_err = 10
for i in range(num_sample):
	# get image
	im_gt = np.squeeze(data_test_gt[i,:,:,0]).T
	im_input = np.squeeze(data_test_input[i,:,:,0]).T
	im_pred = np.squeeze(data_test_output[i,:,:,0]).T

	# get error
	err_pred = getErrorMetrics(im_pred, im_gt)
	err_input = getErrorMetrics(im_input, im_gt)
	list_err_pred.append(err_pred)
	list_err_input.append(err_input)

	# display
	im_toshow = [im_gt, im_input, (im_input-im_gt)*aug_err, im_pred, (im_pred-im_gt)*aug_err]
	im_toshow = np.abs(np.concatenate(im_toshow, axis=1))
	plt.figure(figsize=[20,8])
	plt.imshow(im_toshow, clim=[0,0.5], cmap='gray')
	im_title = 'sample #{0}, input PSNR {1:.4f}, SSIM {2:.4f}, predict PSNR {3:.4f}, SSIM {4:.4f}'.format(
				i, err_input['psnr'], err_input['ssim'], err_pred['psnr'], err_pred['ssim'])
	plt.title(im_title)
	print(im_title)
	path_figure = filename_results+'_{0}.png'.format(i)
	plt.savefig(path_figure)

'''
export results
'''
print('input average metrics:', {k:np.mean([x[k] for x in list_err_input]) for k in list_err_input[0].keys()})
print('prediction average metrics:', {k:np.mean([x[k] for x in list_err_pred]) for k in list_err_pred[0].keys()})# save history dictionary
import json
result_error = {'err_input':list_err_input,
				'err_pred':list_err_pred}
path_error = filename_results+'_error.json'
with open(path_error, 'w') as outfile:
    json.dump(result_error, outfile)
print('error exported to {0}'.format(path_error))


import datasets
import os
import torch
import numpy as np
import json
import random
import collections
import sentencepiece as spm
import time
import argparse

def get_time():
    return time.strftime("%Y%m%d %H:%M:%S", time.localtime())

OPENMODEL_PATH = 'your/model/path'
if OPENMODEL_PATH == 'your/model/path':
    raise('Model path is not valid')
CEVAL_DATA_PATH = 'your/ceval/data/path'
results_filename = 'results.txt'

def load_model(model_id):
    print('Loading', model_id)
    model_id = os.path.join(OPENMODEL_PATH, model_id)
    if 'gpt2' in model_id:
        from transformers import GPT2LMHeadModel, GPT2TokenizerFast
        model = GPT2LMHeadModel.from_pretrained(model_id, device_map='auto',)
        tokenizer = GPT2TokenizerFast.from_pretrained(model_id)
    elif 'glm' in model_id:
        from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        model = AutoModel.from_pretrained(model_id, trust_remote_code=True, device_map='auto', torch_dtype=torch.float16)
    elif 'bloomz' in model_id:
        from transformers import BloomForCausalLM
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = BloomForCausalLM.from_pretrained(model_id, device_map='auto', torch_dtype=torch.float16)
    elif 'tigerbot' in model_id:
        from transformers import AutoTokenizer, AutoModelForCausalLM
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto')
    elif 'Ziya-LLaMA' in model_id:
        from transformers import AutoTokenizer, LlamaForCausalLM
        model = LlamaForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
        tokenizer = AutoTokenizer.from_pretrained(model_id)
    elif 'falcon' in model_id:
        from transformers import AutoTokenizer, AutoModelForCausalLM
        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False,
                                                  trust_remote_code=True,
                                                  add_bos_token=False,
                                                  add_eos_token=False,
                                                  padding_dise='right')
        model = AutoModelForCausalLM.from_pretrained(model_id,
                                             trust_remote_code=True,
                                             device_map='auto',
                                             torch_dtype=torch.bfloat16)
    elif 'aquila' in model_id:
        from flagai.auto_model.auto_loader import AutoLoader
        from flagai.model.predictor.predictor import Predictor
        from flagai.data.tokenizer import Tokenizer
        import bminf
        model_id = model_id.split('/')[-1]
        print(model_id)
        loader = AutoLoader(
            "lm",
            model_dir=OPENMODEL_PATH,
            model_name=model_id,
            use_cache=True,
            device="cuda")
        model = loader.get_model()
        model.eval()
        model.half()
        tokenizer = loader.get_tokenizer()
    elif 'Open-Llama' in model_id:
        from transformers import AutoTokenizer, AutoModelForCausalLM
        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False, trust_remote_code=True, add_eos_token = False)
        model = AutoModelForCausalLM.from_pretrained(model_id,
                                             trust_remote_code=True,
                                             device_map='auto',
                                             torch_dtype=torch.bfloat16)
    else:
        from transformers import AutoTokenizer, AutoModelForCausalLM
        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(model_id,
                                             trust_remote_code=True,
                                             device_map='auto',
                                             torch_dtype=torch.bfloat16)
    candidates = ['A', 'B', 'C', 'D']
    for can in candidates:
        can_idx = tokenizer.encode(can)
        print(can, can_idx)

    return model, tokenizer

class BaseEval:
    def __init__(self, model_id, data_path, show_detail=False):
        self.data_path = data_path
        self.model_id = model_id
        self.task_names = self.get_task_names()
        self.show_detail = show_detail
        self.max_new_tokens = 3  # 最多生成的token数

    def get_task_names(self):
        pass

    def set_model_and_tokenizer(self, model, tokenizer):
        self.model = model
        self.tokenizer=tokenizer

    def build_an_example_prompt(self, data, with_answer=True, flag=0):
        pass

    def build_prompt(self, task_name, data, dataset, shot):
        pass

    def get_results_by_ABCD_prob(self, output, max_new_tokens):
        i = 0  # 用生成的第i个token进行预测
        if 'aquila' in self.model_id:
            scores = output.scores[i].to(torch.float32).detach()
        else:
            scores = output.scores[i][0].to(torch.float32)
        label_score = []
        candidates = ['A', 'B', 'C', 'D']
        for can in candidates:
            can_idx = self.tokenizer.encode(can)[-1]
            label_score.append(scores[can_idx].cpu().numpy())
        answer = candidates[np.argmax(label_score)]

        if self.show_detail:
            # 生成的具体内容
            if 'aquila' in self.model_id:
                outstr = output.text
            else:
                outstr = self.tokenizer.decode(output.sequences.cpu().numpy()[0][-max_new_tokens: ])
            print(f'Generate text (Start>>>){outstr}(<<<End)')
            print(f'Predict answer with ABCD probs: {answer}')
        return answer

    def predict_instance(self, prompt):
        if self.show_detail:
            print(100*'-')
            print(prompt)
        max_new_tokens = self.max_new_tokens
        if 'aquila' not in model_id:
            input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
            input_ids = input_ids.to('cuda')
        if 'glm' in self.model_id:
            # glm 需要特殊处理 https://github.com/mymusise/ChatGLM-Tuning/issues/150
            input_ids[0][-2] = 130001
            input_ids[0][-1] = 130004
        if 'aquila' in model_id:
            from aquila import aquila_generate
            output = aquila_generate(self.tokenizer, self.model, [prompt,], max_gen_len=1)
        else:
            output = self.model.generate(input_ids, max_new_tokens=max_new_tokens, return_dict_in_generate=True, output_scores=True, temperature=0.1, top_p=0.5, repetition_penalty=1.1)
        return self.get_results_by_ABCD_prob(output, max_new_tokens)

    def run_a_task(self, task_name, shot, split='test'):
        results = {}
        correct = []
        prompts = collections.OrderedDict()
        print(self.data_path, task_name)
        dataset = datasets.load_dataset(self.data_path, task_name)
        with torch.no_grad():
            for data in dataset[split]:
                prompt = self.build_prompt(task_name, data, dataset, shot)
                plabel = self.predict_instance(prompt)
                idx = data['id']
                results[str(idx)] = plabel
                correct.append(plabel == data['answer'] if 'answer' in data else 0)
                prompts[str(idx)] = prompt
        return {task_name: results}, {task_name: prompts}, {task_name: correct}

    def run(self, task_names, shot, split):
        results = dict()
        task_prompts = dict()
        task_acc = dict()
        filename_pfx = f"{self.model_id.split('/')[-1]}_{self.__class__.__name__}_{len(task_names)}_{split}_{shot}-shot"
        for name in task_names:
            print(100*'-')
            print(name)
            value, prompts, acc = self.run_a_task(name, shot, split)
            results.update(value)
            task_prompts.update(prompts)
            task_acc.update(acc)
            json.dump(results, open(f'{filename_pfx}_results.json', 'w'), indent=True, ensure_ascii=False)

        outputline = []
        outputline.append(f'{get_time()}\t{filename_pfx}')
        correct = []
        for k, v in task_acc.items():
            correct += v
        outputline.append(f'\tAverage example acc: {np.mean(correct):.4f}')
        acc = []
        for k, v in task_acc.items():
            acc.append(np.mean(v))
        outputline.append(f'\tAverage subject acc: {np.mean(acc):.4f}')
        [print(line) for line in outputline]
        self.save_results(outputline)

    def save_results(self, lines):
        with open(results_filename, 'a') as f:
            for line in lines:
                f.write(line)
                f.write('\n')

    def generate_sample(self, task_name='law', split='validatoin', shot=0):
        dataset = datasets.load_dataset(self.data_path, task_name)
        data = dataset[split][0]
        print('Data' + 100*'-')
        print(data)
        print(f'{shot}-shot prompt' + 100*'-')
        print('(Start>>>)'+self.build_prompt(task_name, data, dataset, shot)+'(<<<End)')

class CEval(BaseEval):
    task2desc = {
        "high_school_physics": "高中物理",
        "fire_engineer": "注册消防工程师",
        "computer_network": "计算机网络",
        "advanced_mathematics": "高等数学",
        "logic": "逻辑学",
        "middle_school_physics": "初中物理",
        "clinical_medicine": "临床医学",
        "probability_and_statistics": "概率统计",
        "ideological_and_moral_cultivation": "思想道德修养与法律基础",
        "operating_system": "操作系统",
        "middle_school_mathematics": "初中数学",
        "chinese_language_and_literature": "中国语言文学",
        "electrical_engineer": "注册电气工程师",
        "business_administration": "工商管理",
        "high_school_geography": "高中地理",
        "modern_chinese_history": "近代史纲要",
        "legal_professional": "法律职业资格",
        "middle_school_geography": "初中地理",
        "middle_school_chemistry": "初中化学",
        "high_school_biology": "高中生物",
        "high_school_chemistry": "高中化学",
        "physician": "医师资格",
        "high_school_chinese": "高中语文",
        "tax_accountant": "税务师",
        "high_school_history": "高中历史",
        "mao_zedong_thought": "毛泽东思想和中国特色社会主义理论概论",
        "high_school_mathematics": "高中数学",
        "professional_tour_guide": "导游资格",
        "veterinary_medicine": "兽医学",
        "environmental_impact_assessment_engineer": "环境影响评价工程师",
        "basic_medicine": "基础医学",
        "education_science": "教育学",
        "urban_and_rural_planner": "注册城乡规划师",
        "middle_school_biology": "初中生物",
        "plant_protection": "植物保护",
        "middle_school_history": "初中历史",
        "high_school_politics": "高中政治",
        "metrology_engineer": "注册计量师",
        "art_studies": "艺术学",
        "college_economics": "大学经济学",
        "college_chemistry": "大学化学",
        "law": "法学",
        "sports_science": "体育学",
        "civil_servant": "公务员",
        "college_programming": "大学编程",
        "middle_school_politics": "初中政治",
        "teacher_qualification": "教师资格",
        "computer_architecture": "计算机组成",
        "college_physics": "大学物理",
        "discrete_mathematics": "离散数学",
        "marxism": "马克思主义基本原理",
        "accountant": "注册会计师",
    }

    def get_task_names(self):
        task_names = list(sorted(self.task2desc.keys()))
        return task_names

    def build_an_example_prompt(self, data, with_answer=True):
        choice =  ['A. '+data["A"],'B. '+data["B"],'C. '+data["C"],'D. '+data["D"]]
        answer = data["answer"].strip().upper() if with_answer else ""
        question = data['question']
        cstr = '\n'.join(choice)
        return f"{question}\n{cstr}\n答案：{answer}"

    def build_prompt(self, task_name, data, dataset, shot):
        if shot == 0:
            prompt = f"以下是中国关于{self.task2desc[task_name]}考试的单项选择题，请选出其中的正确答案。\n\n{self.build_an_example_prompt(data, with_answer=False)}"
        else:
            shuffle_train = dataset['dev'].shuffle(seed=123)
            prompt = f"以下是中国关于{self.task2desc[task_name]}考试的单项选择题，请选出其中的正确答案。\n"
            for i in range(min(len(shuffle_train), shot)):
                prompt += '\n' + self.build_an_example_prompt(shuffle_train[i], with_answer=True)
            prompt += '\n' + self.build_an_example_prompt(data, with_answer=False)
        return prompt

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu_idx', type=str, required=True)
    parser.add_argument('--model_id', type=str, default=None)
    parser.add_argument('--task', type=str, required=True)
    parser.add_argument('--shot', type=int, default=5)
    parser.add_argument('--split', type=str, default='test')
    parser.add_argument('--show_detail', action='store_true')
    args = parser.parse_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_idx
    print(args.model_id, args.shot, args.split, args.task)
    model_id = args.model_id
    shot = args.shot
    split = args.split
    show_detail = args.show_detail

    model, tokenizer = None, None
    if model_id is not None:
        model, tokenizer = load_model(model_id)

    if args.task == 'ceval':
        data_path = CEVAL_DATA_PATH
        if data_path == 'your/ceval/data/path':
            raise('CEval data path is not valid')
        ceval = CEval(model_id, data_path, show_detail=show_detail)
        ceval.set_model_and_tokenizer(model, tokenizer)
        ceval.generate_sample(task_name='law', split=split, shot=3)
        ceval.run(task_names=ceval.task_names, shot=shot, split=split)
    else:
        raise(f'Task is not valid: {args.task}')


import os
import json
import gzip
import time
import tempfile
import warnings

import torch

from functools import partial

from markovdwp.source.datasets import utils

from markovdwp.runtime.trainer import get_trainer
from markovdwp.runtime.utils.common import linear
from markovdwp.runtime.utils.reconstruction_logger \
    import SliceReconstructionLogger

from markovdwp.utils.dicttools import flatten, unflatten
from markovdwp.utils.runtime import get_datasets, get_dataloaders
from markovdwp.utils.runtime import register, get_class, get_instance

from pytorch_lightning.loggers import WandbLogger
from markovdwp.utils.vendor.pytorch_lightning import GradInformation

class NoTargetWarning(UserWarning):
    pass

def get_vae(encoder, decoder, options):
    return {
        'encoder': get_instance(**encoder, **options),
        'decoder': get_instance(**decoder, **options)
    }

def apply_resampling(datasets, resampling):
    """Apply the specified resampling stratefy."""
    resampled = {}
    for name, spec in resampling.items():
        if isinstance(spec, dict):
            dataset, kind = datasets[spec['source']], spec['kind']
            if kind not in ('subset', 'undersample', 'stratified'):
                raise ValueError(f'Unrecognized resampling `{kind}`')

            if kind == 'stratified':
                dataset, _ = utils.stratified_split(
                    dataset, spec['size'], spec.get('random_state'))

            elif kind == 'undersample':
                dataset, _ = utils.undersampled_split(
                    dataset, spec['size'], spec.get('random_state'))

            elif kind == 'subset':
                dataset = utils.subset(
                    dataset, spec['size'], spec.get('random_state'))

        else:
            dataset = datasets[spec]

        resampled[name] = dataset
    return resampled

def configure_optimizers(pl_module):
    optim = torch.optim.Adam(pl_module.parameters(), lr=pl_module.lr)

    sched = torch.optim.lr_scheduler.LambdaLR(
        optim, partial(linear, t0=0, t1=pl_module.trainer.max_epochs))

    return [optim], [{'scheduler': sched, 'monitor': 'loss'}]

def train(gpus, config, logger=None):
    # select the specified runtime and attach the dynamic methods
    Runtime = register(
        '_VAERuntime', GradInformation, get_class(config['runtime']),
        configure_optimizers=configure_optimizers)

    vae = get_vae(**config['vae'])
    pl_module = Runtime(**vae, **config['options'])

    pl_trainer = get_trainer(gpus=gpus, logger=logger, **config['trainer'])

    datasets = get_datasets(config['dataset'])
    if 'resampling' in config:
        datasets.update(apply_resampling(datasets, config['resampling']))
    feeds = get_dataloaders(datasets, config['feeds'])

    # if 'reference' dataset exists, then we create a reconstruction logger
    if 'reference' in datasets:
        X, y, *rest = zip(*datasets['reference'])

        order = sorted(range(len(y)), key=y.__getitem__)
        ref_x = torch.stack([X[i] for i in order], dim=0)

        ref_z = pl_module.prior.sample([len(ref_x)])

        pl_trainer.callbacks.append(
            SliceReconstructionLogger(ref_x.unsqueeze(1), ref_z,
                                      n_draws_is=1000, scatter=False))

    # last-minute update of the config before training
    logger.experiment.config.update(
        flatten(config, delim='__'),
        allow_val_change=True)

    # do the training
    try:
        pl_trainer.fit(pl_module, train_dataloader=feeds.get('train'))

    except ValueError as e:
        # it would be better is pl raised special exceptions from
        # `detect_nan_tensors`, rather than the same ValueError...
        if 'nan or inf' in str(e) or 'nan and/or inf' in str(e):
            raise FloatingPointError(str(e)) from e

        raise

    pl_module.cpu().eval()
    return pl_module.encoder, pl_module.decoder, pl_module.prior

def pack(module):
    from inspect import signature

    parameters = {}
    # adopt sklearn's convention: save arguments in instance's `__dict__`
    for name, p in signature(type(module)).parameters.items():
        # ignore variable keyworded args
        if p.kind == p.VAR_KEYWORD:
            continue

        parameters[name] = getattr(module, p.name)

    return {
        'model': {'cls': str(type(module)), **parameters},
        'state': module.state_dict(),
    }

def main(manifest, target=None, gpus=[0], debug=False, tags=None):
    breakpoint() if debug else None

    assert tags is None or isinstance(tags, str)
    tags = tags.split(',') if tags is not None else []

    # the manifest must be complete and ready to go
    parameters = json.load(open(manifest, 'rt'))

    if target is None:
        # do not save anything
        warnings.warn('No target path specified, the model will not be saved',
                      NoTargetWarning)

    elif os.path.isdir(target):
        # consume the first tag
        tag, *tags = tags if tags else [None]

        # save under a random name if target is a directory
        fid, target = tempfile.mkstemp(dir=target, suffix='.gz',
                                       prefix=tag + '__' if tag else tag)
        os.close(fid)

    elif os.path.exists(target):
        raise ValueError(f'`{target}` already exists! Refusing to proceed.')

    # pl's Wandb logger uses reinit=true!
    # wandb.init(project='DWP Slice Replication Machine', reinit=False)
    logger = WandbLogger(project='Variatonal Autoencoders',
                         tags=tags)

    # sync with wandb's agent's arguments and rebuild the config
    # XXX wandb's config keys should not collide with meaningful manifest's keys
    config = unflatten({
        **flatten(parameters, delim='__'),  # defaults from the manifest
        **logger.experiment.config          # overrides from wandb
    }, delim='__')

    # train the model
    encoder, decoder, prior = train(gpus, config, logger)
    if target is not None:
        with gzip.open(target, 'wb', compresslevel=9) as fout:
            torch.save({
                '__dttm__': time.strftime('%Y%m%d %H%M%S'),
                'config': config,
                'encoder': pack(encoder),
                'decoder': pack(decoder),
                'prior': pack(prior) if isinstance(prior, torch.nn.Module) else {},
            }, fout)

    return encoder, decoder, prior

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
        description='Train a single VAE from .models on the specified dataset.',
        add_help=True)

    parser.add_argument(
        '--manifest', type=str, required=True,
        help='The manifest of the experiment.')

    parser.add_argument(
        '--target', type=str, required=False, default=None,
        help='The path where to store the trained model under a '
             'unique name, or the name of the snapshot wherein the '
             'model will be saved, or omitted, if model snapshot is '
             'not needed.')

    parser.add_argument(
        '--gpus', type=int, nargs='+', required=False, default=None,
        help='Device ids to occupy by the experiment.')

    parser.add_argument(
        '--tags', type=str, required=False,
        help='Optional comma-separated tags for the experiment being run. '
             'If `target` is a directory, then the first tag is used as a '
             'prefix to in model names.')

    parser.add_argument(
        '--debug', required=False, dest='debug', action='store_true',
        help='Enter trace mode.')

    parser.set_defaults(
        target=None, gpus=[0], tags=None, debug=False)

    args, _ = parser.parse_known_args()
    main(**vars(args))


from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt")
tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX" )

class txt2translate:
    def txt2translate(ref,text,language):
        article_en = text
        model_inputs = tokenizer(article_en, return_tensors="pt")
        generated_tokens = model.generate(**model_inputs, forced_bos_token_id=tokenizer.lang_code_to_id[language+"_IN"])
        translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
        # print(translation[0])
        translated = translation[0]
        return translated

# og = txt2translate()
# txt = "How are you ?"
# gt = og.txt2translate(txt)


import tweepy
import logging
from config import crearAPI
import time
import API.weatherAPI as weatherAPI

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger()

# since_id especifica resultados con una ID mayor que la id especificada (más reciente)

def verificarMenciones(api, palabrasClaves, idDesde):

    logger.info("Recuperando menciones...")

    idDesdeNuevo = idDesde

    # Cursor maneja la paginacion de resultados
    for tweet in tweepy.Cursor(api.mentions_timeline, since_id=idDesde).items():

        idDesdeNuevo = max(tweet.id, idDesdeNuevo)

        me = api.me()

        if any(palabraClave in tweet.text.lower() for palabraClave in palabrasClaves) and tweet.user.id != me.id:

            logger.info(f"Respondiendo al usuario {tweet.user.screen_name}")

            if not tweet.user.following:

                logger.info(f"Siguiendo al usuario {tweet.user.screen_name}")

                tweet.user.follow()

            if not tweet.favorited:

                logger.info("Dando me gusta al tweet...")

                tweet.favorite()

            usuario = tweet.user.screen_name

            ciudad = 'Mendoza, AR'

            clima = weatherAPI.getClima(ciudad)

            hora = time.strftime("%H:%M:%S")

            clima += f"\nHora: {hora}"

            # Se responde al tweet que nos mencionaron
            try:
                api.update_status(f"@{usuario} {clima}",
                                  in_reply_to_status_id=tweet.id)
            except Exception as e:
                logger.info(
                    "No se pudo responder, el tweet está duplicado.", exc_info=True)

    return idDesdeNuevo

def main():
    api = crearAPI()
    idDesde = 1

    while True:

        idDesde = verificarMenciones(
            api, ["getclima"], idDesde)

        logger.info("Esperando para volver a verificar en 1 minuto...")

        time.sleep(60)

if __name__ == "__main__":
    main()


import torch
import torch.nn as nn
from mmcv.cnn.bricks.registry import TRANSFORMER_LAYER_SEQUENCE
from mmcv.cnn.bricks.transformer import build_transformer_layer
from mmcv.runner.base_module import BaseModule
from projects.mmdet3d_plugin.models.utils.functional import (
    norm_points,
    pos2posemb2d,
    trajectory_coordinate_transform
)

@TRANSFORMER_LAYER_SEQUENCE.register_module()
class MotionTransformerDecoder(BaseModule):
    """Implements the decoder in DETR3D transformer.
    Args:
        return_intermediate (bool): Whether to return intermediate outputs.
        coder_norm_cfg (dict): Config of last normalization layer. Default：
            `LN`.
    """

    def __init__(self, pc_range=None, embed_dims=256, transformerlayers=None, num_layers=3, **kwargs):
        super(MotionTransformerDecoder, self).__init__()
        self.pc_range = pc_range
        self.embed_dims = embed_dims
        self.num_layers = num_layers
        self.intention_interaction_layers = IntentionInteraction()
        self.track_agent_interaction_layers = nn.ModuleList(
            [TrackAgentInteraction() for i in range(self.num_layers)])
        self.map_interaction_layers = nn.ModuleList(
            [MapInteraction() for i in range(self.num_layers)])
        self.bev_interaction_layers = nn.ModuleList(
            [build_transformer_layer(transformerlayers) for i in range(self.num_layers)])

        self.static_dynamic_fuser = nn.Sequential(
            nn.Linear(self.embed_dims*2, self.embed_dims*2),
            nn.ReLU(),
            nn.Linear(self.embed_dims*2, self.embed_dims),
        )
        self.dynamic_embed_fuser = nn.Sequential(
            nn.Linear(self.embed_dims*3, self.embed_dims*2),
            nn.ReLU(),
            nn.Linear(self.embed_dims*2, self.embed_dims),
        )
        self.in_query_fuser = nn.Sequential(
            nn.Linear(self.embed_dims*2, self.embed_dims*2),
            nn.ReLU(),
            nn.Linear(self.embed_dims*2, self.embed_dims),
        )
        self.out_query_fuser = nn.Sequential(
            nn.Linear(self.embed_dims*4, self.embed_dims*2),
            nn.ReLU(),
            nn.Linear(self.embed_dims*2, self.embed_dims),
        )

    def forward(self,
                track_query,
                lane_query,
                track_query_pos=None,
                lane_query_pos=None,
                track_bbox_results=None,
                bev_embed=None,
                reference_trajs=None,
                traj_reg_branches=None,
                agent_level_embedding=None,
                scene_level_ego_embedding=None,
                scene_level_offset_embedding=None,
                learnable_embed=None,
                agent_level_embedding_layer=None,
                scene_level_ego_embedding_layer=None,
                scene_level_offset_embedding_layer=None,
                **kwargs):
        """Forward function for `MotionTransformerDecoder`.
        Args:
            agent_query (B, A, D)
            map_query (B, M, D)
            map_query_pos (B, G, D)
            static_intention_embed (B, A, P, D)
            offset_query_embed (B, A, P, D)
            global_intention_embed (B, A, P, D)
            learnable_intention_embed (B, A, P, D)
            det_query_pos (B, A, D)
        Returns:
            None
        """
        intermediate = []
        intermediate_reference_trajs = []

        B, _, P, D = agent_level_embedding.shape
        track_query_bc = track_query.unsqueeze(2).expand(-1, -1, P, -1)  # (B, A, P, D)
        track_query_pos_bc = track_query_pos.unsqueeze(2).expand(-1, -1, P, -1)  # (B, A, P, D)

        # static intention embedding, which is imutable throughout all layers
        agent_level_embedding = self.intention_interaction_layers(agent_level_embedding)
        static_intention_embed = agent_level_embedding + scene_level_offset_embedding + learnable_embed
        reference_trajs_input = reference_trajs.unsqueeze(4).detach()

        query_embed = torch.zeros_like(static_intention_embed)
        for lid in range(self.num_layers):
            # fuse static and dynamic intention embedding
            # the dynamic intention embedding is the output of the previous layer, which is initialized with anchor embedding
            dynamic_query_embed = self.dynamic_embed_fuser(torch.cat(
                [agent_level_embedding, scene_level_offset_embedding, scene_level_ego_embedding], dim=-1))

            # fuse static and dynamic intention embedding
            query_embed_intention = self.static_dynamic_fuser(torch.cat(
                [static_intention_embed, dynamic_query_embed], dim=-1))  # (B, A, P, D)

            # fuse intention embedding with query embedding
            query_embed = self.in_query_fuser(torch.cat([query_embed, query_embed_intention], dim=-1))

            # interaction between agents
            track_query_embed = self.track_agent_interaction_layers[lid](
                query_embed, track_query, query_pos=track_query_pos_bc, key_pos=track_query_pos)

            # interaction between agents and map
            map_query_embed = self.map_interaction_layers[lid](
                query_embed, lane_query, query_pos=track_query_pos_bc, key_pos=lane_query_pos)

            # interaction between agents and bev, ie. interaction between agents and goals
            # implemented with deformable transformer
            bev_query_embed = self.bev_interaction_layers[lid](
                query_embed,
                value=bev_embed,
                query_pos=track_query_pos_bc,
                bbox_results=track_bbox_results,
                reference_trajs=reference_trajs_input,
                **kwargs)

            # fusing the embeddings from different interaction layers
            query_embed = [track_query_embed, map_query_embed, bev_query_embed, track_query_bc+track_query_pos_bc]
            query_embed = torch.cat(query_embed, dim=-1)
            query_embed = self.out_query_fuser(query_embed)

            if traj_reg_branches is not None:
                # update reference trajectory
                tmp = traj_reg_branches[lid](query_embed)
                bs, n_agent, n_modes, n_steps, _ = reference_trajs.shape
                tmp = tmp.view(bs, n_agent, n_modes, n_steps, -1)

                # we predict speed of trajectory and use cumsum trick to get the trajectory
                tmp[..., :2] = torch.cumsum(tmp[..., :2], dim=3)
                new_reference_trajs = torch.zeros_like(reference_trajs)
                new_reference_trajs = tmp[..., :2]
                reference_trajs = new_reference_trajs.detach()
                reference_trajs_input = reference_trajs.unsqueeze(4)  # BS NUM_AGENT NUM_MODE 12 NUM_LEVEL  2

                # update embedding, which is used in the next layer
                # only update the embedding of the last step, i.e. the goal
                ep_offset_embed = reference_trajs.detach()
                ep_ego_embed = trajectory_coordinate_transform(reference_trajs.unsqueeze(
                    2), track_bbox_results, with_translation_transform=True, with_rotation_transform=False).squeeze(2).detach()
                ep_agent_embed = trajectory_coordinate_transform(reference_trajs.unsqueeze(
                    2), track_bbox_results, with_translation_transform=False, with_rotation_transform=True).squeeze(2).detach()

                agent_level_embedding = agent_level_embedding_layer(pos2posemb2d(
                    norm_points(ep_agent_embed[..., -1, :], self.pc_range)))
                scene_level_ego_embedding = scene_level_ego_embedding_layer(pos2posemb2d(
                    norm_points(ep_ego_embed[..., -1, :], self.pc_range)))
                scene_level_offset_embedding = scene_level_offset_embedding_layer(pos2posemb2d(
                    norm_points(ep_offset_embed[..., -1, :], self.pc_range)))

                intermediate.append(query_embed)
                intermediate_reference_trajs.append(reference_trajs)

        return torch.stack(intermediate), torch.stack(intermediate_reference_trajs)

class TrackAgentInteraction(BaseModule):
    """
    Modeling the interaction between the agents
    """
    def __init__(self,
                 embed_dims=256,
                 num_heads=8,
                 dropout=0.1,
                 batch_first=True,
                 norm_cfg=None,
                 init_cfg=None):
        super().__init__(init_cfg)

        self.batch_first = batch_first
        self.interaction_transformer = nn.TransformerDecoderLayer(d_model=embed_dims,
                                                                  nhead=num_heads,
                                                                  dropout=dropout,
                                                                  dim_feedforward=embed_dims*2,
                                                                  batch_first=batch_first)

    def forward(self, query, key, query_pos=None, key_pos=None):
        '''
        query: context query (B, A, P, D)
        query_pos: mode pos embedding (B, A, P, D)
        key: (B, A, D)
        key_pos: (B, A, D)
        '''
        B, A, P, D = query.shape
        if query_pos is not None:
            query = query + query_pos
        if key_pos is not None:
            key = key + key_pos
        mem = key.expand(B*A, -1, -1)
        # N, A, P, D -> N*A, P, D
        query = torch.flatten(query, start_dim=0, end_dim=1)
        query = self.interaction_transformer(query, mem)
        query = query.view(B, A, P, D)
        return query

class MapInteraction(BaseModule):
    """
    Modeling the interaction between the agent and the map
    """
    def __init__(self,
                 embed_dims=256,
                 num_heads=8,
                 dropout=0.1,
                 batch_first=True,
                 norm_cfg=None,
                 init_cfg=None):
        super().__init__(init_cfg)

        self.batch_first = batch_first
        self.interaction_transformer = nn.TransformerDecoderLayer(d_model=embed_dims,
                                                                  nhead=num_heads,
                                                                  dropout=dropout,
                                                                  dim_feedforward=embed_dims*2,
                                                                  batch_first=batch_first)

    def forward(self, query, key, query_pos=None, key_pos=None):
        '''
        x: context query (B, A, P, D)
        query_pos: mode pos embedding (B, A, P, D)
        '''
        B, A, P, D = query.shape
        if query_pos is not None:
            query = query + query_pos
        if key_pos is not None:
            key = key + key_pos

        # N, A, P, D -> N*A, P, D
        query = torch.flatten(query, start_dim=0, end_dim=1)
        mem = key.expand(B*A, -1, -1)
        query = self.interaction_transformer(query, mem)
        query = query.view(B, A, P, D)
        return query

class IntentionInteraction(BaseModule):
    """
    Modeling the interaction between anchors
    """
    def __init__(self,
                 embed_dims=256,
                 num_heads=8,
                 dropout=0.1,
                 batch_first=True,
                 norm_cfg=None,
                 init_cfg=None):
        super().__init__(init_cfg)

        self.batch_first = batch_first
        self.interaction_transformer = nn.TransformerEncoderLayer(d_model=embed_dims,
                                                                  nhead=num_heads,
                                                                  dropout=dropout,
                                                                  dim_feedforward=embed_dims*2,
                                                                  batch_first=batch_first)

    def forward(self, query):
        B, A, P, D = query.shape
        # B, A, P, D -> B*A,P, D
        rebatch_x = torch.flatten(query, start_dim=0, end_dim=1)
        rebatch_x = self.interaction_transformer(rebatch_x)
        out = rebatch_x.view(B, A, P, D)
        return out


#%%
r"""
Hack into the MNIST Module
"""
import argparse
import sys

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as torchF

from mnist_module import MNISTClassifier
# from npxconv import outerconv_2d
from plot_utils import savefig
from record_utils import record_append, record_dump_to_file

_arg = argparse.ArgumentParser()
_arg.add_argument("--hack_layer",
                  type=int,
                  default=3,
                  help="hack layer in [1, 2, 3, 4]")
_arg.add_argument("--train_epho",
                  type=int,
                  default=256,
                  help="train epho for hacking network")
_arg.add_argument("--eng", type=float, default=0.4, help="alpha * Ep + X")
args = _arg.parse_args()

HACK_LAYER = args.hack_layer

# read data
# shape N, 28, 28
test_img = np.load("mnist/mnist-test-image.npy")
test_img = test_img.astype(np.float32) / 255.0
assert test_img.shape[1] == 28 and test_img.shape[2] == 28

# shape 6000, int8
test_lab = np.load("mnist/mnist-test-label.npy")

torch_test_img = torch.FloatTensor(test_img.reshape(-1, 1, 28, 28))

train_img = np.load("mnist/mnist-train-image.npy")
train_img = train_img.astype(np.float32) / 255.0
assert train_img.shape[1] == 28 and train_img.shape[2] == 28

#%%

state_dict = torch.load("result/mnist-classifier.torch")
classifier = MNISTClassifier().eval()

classifier.load_state_dict(state_dict)

# accuracy

classifier_preict = classifier.forward(torch_test_img)
classifier_label = torch.argmax(classifier_preict, dim=-1)
classifier_label = classifier_label.data.cpu().numpy()

acc = np.mean(classifier_label == test_lab)

record_append("classifier", "", accuracy=acc)

# compute order one proxy kernels
# the weights
# _W1 = classifier.conv1.weight  # (5, 1, 3, 3)
# _W2 = classifier.conv2.weight  # (10, 5, 3, 3)
# _W3 = classifier.conv3.weight  # (10, 10, 3, 3)
# # compute proxy kernel layer by layer

# # (1, 5, 3, 3) * (5, 10, 3, 3) -> (1, 10, 7, 7) s=4,p=3
# _W = torch.conv_transpose2d(_W1.transpose(0, 1),
#                             _W2.transpose(0, 1),
#                             stride=2,
#                             padding=0)
# # (1, 10, 7, 7) * (10, 10, 3, 3) -> (1, 10, 15, 15) s=8, p=3
# _W = torch.conv_transpose2d(_W, _W3.transpose(0, 1), stride=2, padding=0)
# _W = _W / torch.norm(_W, dim=(2, 3), keepdim=True)
# print(_W.shape)
# exit(0)

# emulate order1
class EMUo1(nn.Module):

    def __init__(self):
        super().__init__()
        if HACK_LAYER == 1:
            self.cv1 = nn.Conv2d(1, 5, 3, 2, 1)
        elif HACK_LAYER == 2:
            self.cv1 = nn.Conv2d(1, 10, 7, 4, 3)
        elif HACK_LAYER == 3:
            self.cv1 = nn.Conv2d(1, 10, 15, 8, 3)
        elif HACK_LAYER == 4:
            self.cv1 = nn.Conv2d(1, 10, 31, 8, 3)

    def forward(self, input):
        out = self.cv1(input)
        if HACK_LAYER == 4:
            out = out.reshape(out.shape[0], out.shape[1])
        return out

#%%

emu = EMUo1()
emuopt = torch.optim.Adam(emu.parameters(), lr=1e-2)

for i in range(args.train_epho):
    np.random.shuffle(train_img)
    for _i in range(1, train_img.shape[0], 64):
        emuopt.zero_grad()
        # _x = torch.rand(64, 1, 28, 28)
        # _x = torch.randn(64, 1 * 28 * 28)
        # _x /= torch.norm(_x, dim=(-1), keepdim=True)
        # _x = _x.reshape(64, 1, 28, 28)
        _x = torch.FloatTensor(train_img[_i:_i + 64].reshape(-1, 1, 28, 28))

        target = classifier.forward(_x, k=HACK_LAYER)
        predict = emu.forward(_x)
        loss = torchF.mse_loss(predict, target)
        loss.backward()
        emuopt.step()

    print(f"\repho {i:4d}/{args.train_epho} loss {loss :.8f}", end="")

print("")

#%%

#%% compute proxy kernel epilson
def compute_ker_perturbation(ker, x, scale):
    """
    ker: shape (n, n)
    x: shape (28, 28)
    """
    f_x = np.fft.fftn(x, axes=(0, 1))
    f_W = np.fft.fftn(ker, s=x.shape, axes=(0, 1))

    Ep = np.fft.ifftn(f_W - f_x, axes=(0, 1)).real
    Ep -= np.min(Ep)
    Ep /= np.max(Ep)

    return scale * Ep + x

# l2 norm of hacker network and classifier network
classifier_pred = classifier.forward(torch_test_img, k=HACK_LAYER)
hacker_pred = emu.forward(torch_test_img)

hacker_label = classifier.forward_rest(hacker_pred, k=HACK_LAYER)

hacker_label = torch.argmax(hacker_label, dim=-1).data.cpu().numpy()

record_append(
    "classifier_hacker",
    "",
    error=torchF.mse_loss(classifier_pred, hacker_pred),
    hacker_accuracy=np.mean(test_lab == hacker_label),
)
# tick_label = [str(_d) for _d in range(10)]
# y = np.linspace(0, 9, 10)
# barhig = 0.35
# y1 = y - barhig / 2
# y2 = y + barhig / 2

predWriter = open(f"result/mnist-hack-{HACK_LAYER}-predict.csv", "w")
# predWriter = sys.stdout  #open(f"result/mnist-hack-{HACK_LAYER}-predict-temp.csv", "w")

predWriter.write("label")
for i in range(10):
    predWriter.write(",%d" % i)
predWriter.write("\n")

# wemu = emu.cv1.weight.data.cpu().numpy()
# wfin = _W.transpose(0, 1).data.cpu().numpy()

# fig, axi = plt.subplots(2, 10, sharex=False, sharey=True)
# for i in range(10):
#     axi[0, i].imshow(wemu[i, 0])
#     axi[0, i].axis("off")
#     axi[1, i].imshow(wfin[i, 0])
#     axi[1, i].axis("off")

# plt.show()

wfin = emu.cv1.weight.data.cpu().numpy()
wfinlen = wfin.shape[0]

for tag in range(wfinlen + 2):

    while True:
        picked_idx = np.random.randint(0, test_img.shape[0])
        obj_x = test_img[picked_idx]
        real_label = test_lab[picked_idx]

        t_obj_x = torch.tensor(obj_x.reshape(1, 1, 28, 28),
                               dtype=torch.float32)
        obj_pred = classifier.forward(t_obj_x)
        print("predict", torch.argmax(obj_pred), real_label)
        # skip wrong prediction
        if torch.argmax(obj_pred[0]) != real_label:
            continue
        break

    if tag < wfinlen:
        weig = wfin[tag, 0].copy()
    elif tag == wfinlen:
        weig = np.random.rand(wfin.shape[2], wfin.shape[3])
    else:
        weig = np.random.randn(wfin.shape[2], wfin.shape[3])

    fig = plt.figure(figsize=(8, 8))
    plt.imshow(obj_x) #, cmap='gray')
    plt.axis('off')
    savefig(fig, f"mnist-hack-{HACK_LAYER}-img-{tag}")
    # plt.show()
    plt.close()

    fig = plt.figure(figsize=(8, 8))
    plt.imshow(weig) #, cmap='gray')
    plt.axis('off')
    savefig(fig, f"mnist-hack-{HACK_LAYER}-patch-{tag}")
    # plt.show()
    plt.close()

    # Epx = randX.copy()
    # Epx[0:weig.shape[0],0:weig.shape[1]] = weig

    Epx = compute_ker_perturbation(weig, obj_x, args.eng)

    fig = plt.figure(figsize=(8, 8))
    plt.imshow(Epx)  #, cmap='gray')
    plt.axis('off')
    savefig(fig, f"mnist-hack-{HACK_LAYER}-img-patch-{tag}")
    # plt.show()
    plt.close()

    # prepate data
    t_Epx = torch.tensor(Epx.reshape(1, 1, 28, 28), dtype=torch.float32)

    pke_pred = classifier.forward(t_Epx)

    # print(torch.topk(pke_pred, k=3)[1])
    obj_pred = obj_pred[0].data.cpu().numpy()
    pke_pred = pke_pred[0].data.cpu().numpy()

    record_append(
        f"hack-{tag}",
        "",
        label=real_label,
        hack_label=np.argmax(pke_pred),
        object_predict=obj_pred,
        patched_predict=pke_pred,
    )

    predWriter.write("preX %d" % tag)
    for x in obj_pred:
        predWriter.write(",%.8f" % x)
    predWriter.write("\npreE %d" % tag)
    for x in pke_pred:
        predWriter.write(",%.8f" % x)
    predWriter.write("\n")

print("all done")

predWriter.close()

record_dump_to_file(f"mnist-hack-{HACK_LAYER}-predict")
# record_dump_to_file(sys.stdout)


from typing import Iterator

from tqdm import tqdm

from datalabs.operations.aggregate import Aggregating, aggregating

class KGLinkPredictionAggregating(Aggregating):
    def __init__(self, *args, **kwargs):

        super(KGLinkPredictionAggregating, self).__init__(*args, **kwargs)
        self._data_type = "ag_news"

class kg_link_prediction_aggregating(aggregating):
    def __init__(self, *args, **kwargs):

        super(kg_link_prediction_aggregating, self).__init__(*args, **kwargs)
        # print(self.__dict__)

    def __call__(self, *param_arg):
        if callable(self.name):

            tf_class = KGLinkPredictionAggregating(
                name=self.name.__name__, func=self.name
            )
            return tf_class(*param_arg)
        else:
            f = param_arg[0]
            name = self.name or f.__name__
            tf_cls = KGLinkPredictionAggregating(
                name=name,
                func=f,
                resources=self.resources,
                contributor=self.contributor,
                task=self.task,
                description=self.description,
            )
            return tf_cls

"""
from datalabs import load_dataset
dataset = load_dataset('./ag_news')
from ag_news.featurize import *
get_number_of_tokens.__dict__
res = dataset['test'].apply(get_number_of_tokens)
print(next(res))
"""

@kg_link_prediction_aggregating(
    name="get_statistics",
    contributor="datalab",
    task="kg-link-prediction",
    description="aggregation function",
)
def get_statistics(samples: Iterator):
    dict_head = {}
    dict_link = {}
    dict_tail = {}

    for sample in tqdm(samples):

        if sample["tail"] not in dict_tail.keys():
            dict_tail[sample["tail"]] = 1
        else:
            dict_tail[sample["tail"]] += 1

        if sample["head"] not in dict_head.keys():
            dict_head[sample["head"]] = 1
        else:
            dict_head[sample["head"]] += 1

        if sample["link"] not in dict_link.keys():
            dict_link[sample["link"]] = 1
        else:
            dict_link[sample["link"]] += 1

    return {
        "head_fre": dict_head,
        "link_fre": dict_link,
        "tail_fre": dict_tail,
    }


from skimage.io import imread,imsave
import os, glob
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
day = '2012.10.17'
rl = '3R1'
filename = '/media/xuan/SP PHD U3/project/PROJECT/project/GT/slight/'+day+'/'+ rl +'/'
save_path = '/media/xuan/SP PHD U3/project/PROJECT/project/testgt/'
#filename1 = 'C:/Users/tommy/Desktop/AMDc/noextended/Non-AMD/'
file = []
file.append(filename)
#file.append(filename1)
X,y = [], []
for i in file:
    X.append(glob.glob(i + '/'+ '*.bmp'))
X = X[0]
for j in X:
    y.append((j.split('/')[-1]))
l = len(y)
c = 512
for i in range (0,l):
#    gt = imread(filename+'{}'.format(y[i]))
    ori = imread('/media/xuan/SP PHD U3/project/PROJECT/project/GT/slight/'+day+'/'+ rl +'/'+ y[i])
    imsave(save_path+'slight_'+str(c)+'.bmp', ori )
#    imsave(save_path+str(c) +'/masks_serious/'+str(c)+'.bmp', gt )
    c = c+1


import requests
def main():
    cep = input("Digite o Cep para consulta: ")
    if len(cep) != 8:
        print("quantidade de digitos inválido")
        main()

    request = requests.get('https://viacep.com.br/ws/{}/json/'.format(cep))
    endereco = request.json()

    if 'erro' not in endereco:
        print(" ==> Cep Encontrado <==")
        print("-"*20)
        print('CEP: {}'.format(endereco['cep']))
        print('Rua: {}'.format(endereco['logradouro']))
        print('Bairro: {}'.format(endereco['bairro']))
        print('Cidade: {}'.format(endereco['localidade']))
        print('Estado: {}'.format(endereco['uf']))
    else:
        print("{}: Cep invalido".format(cep))
    opcao = int(input("deseja realizar uma nova consulta?\n1. Sim\n2. Sair\n"))

    if opcao == 1:
        main()
    else:
        print('saindo...')
if __name__ == "__main__":
    main()


import random
from turtle import color
import pygame
width=600
height=500
window=pygame.display.set_mode((width,height))
pixel=50
player_x=width/2-pixel
player_y=450-pixel
player_xchange=0
ball_x=width/2-pixel
ball_y=height/2-pixel
ball_xchange=random.choice((1,-1))
vel=2
color='white'
ball_ychange=1
def collision():
    global ball_xchange
    global ball_ychange
    if player_y <(ball_y+pixel):
        if player_x > ball_x and player_x+pixel <ball_x+pixel or player_x+pixel>ball_x and player_x+pixel<ball_x+pixel:
            ball_xchange*=-1
            ball_ychange*=-1
run=True
while run:
    window.fill(color)
    # window.blit(text_1,)
    for event in pygame.event.get():
        if event.type==pygame.QUIT:
            run=False
    if ball_x>=width or ball_x<=0:
        ball_xchange*=-1
    if ball_y>=height or ball_y<=0:
        ball_ychange*=-1
    ball_x+=ball_xchange
    ball_y+=ball_ychange
    pygame.draw.circle(window,'red',(ball_x,ball_y),10)

    key=pygame.key.get_pressed()
    if key[pygame.K_LEFT] and player_x>0:
        player_x-=vel
    if key[pygame.K_RIGHT] and player_x<600-100:
        player_x+=vel
    pygame.draw.rect(window,'red',(player_x,player_y,100,30))
    collision()
    if ball_y==height:
        pygame.display.quit()
    # print(ball_x,ball_y)
    pygame.display.update()


import os
import tempfile
import shutil
import errno

def create_temp_folder():
    temp_folder = tempfile.mkdtemp(prefix='germanium_', suffix='install')

    def temp(subpath=''):
        return os.path.join(temp_folder, subpath)

    return temp

def get_germanium_folder(germanium_home):
    def ge_folder(subpath=''):
        return os.path.join(germanium_home, subpath)

    return ge_folder

def get_desktop_folder():
    home_folder = os.environ['USERPROFILE']

    def desktop_folder(subpath=''):
        return os.path.join(home_folder, 'Desktop', subpath)

    return desktop_folder

def mkdir_p(path):
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise

def rm_rf(*paths):
    for path in paths:
        if os.path.isdir(path) and not os.path.islink(path):
            shutil.rmtree(path)
        elif os.path.exists(path):
            os.remove(path)

def create_desktop_link(title, application, icon=None, workdir=None):
    desktop_folder = get_desktop_folder()

    if not icon:
        icon = application

    if not workdir:
        workdir = os.path.dirname(application)

    path = desktop_folder("%s.lnk" % title)

    try:
        from win32com.client import Dispatch

        shell = Dispatch('WScript.Shell')
        shortcut = shell.CreateShortCut(path)
        shortcut.Targetpath = application
        shortcut.WorkingDirectory = workdir
        shortcut.IconLocation = icon
        shortcut.save()
    except ImportError as e:
        pass


# https://www.acmicpc.net/problem/14502
import sys
from collections import deque
import copy

input = sys.stdin.readline

def dfs(depth, si, sj):
    global answer
    if depth == 3:
        arr = copy.deepcopy(board)
        q = copy.deepcopy(dq)
        d = [(1, 0), (-1, 0), (0, 1), (0, -1)]
        while q:
            x, y = q.popleft()
            for dx, dy in d:
                nx = x + dx
                ny = y + dy
                if -1 < nx < N and -1 < ny < M:
                    if arr[nx][ny] == 0:
                        q.append((nx, ny))
                        arr[nx][ny] = 2
        tmp = 0
        for a in arr:
            tmp += a.count(0)
        answer = max(answer, tmp)
        return
    for i in range(si, N):
        if i > si:
            sj = 0
        for j in range(sj, M):
            if board[i][j] == 0:
                board[i][j] = 1
                dfs(depth + 1, i, j)
                board[i][j] = 0

N, M = map(int, input().split())
board = [list(map(int, input().split())) for _ in range(N)]
dq = deque()
for i in range(N):
    for j in range(M):
        if board[i][j] == 2:
            dq.append((i, j))
answer = 0
dfs(0, 0, 0)
print(answer)