from flask import Flask, request, send_file, render_template, jsonify
from time import sleep, time
from queue import Queue
import pandas as pd
import threading
import pickle
from geobr import read_municipality, read_municipal_seat
from os import path
import uuid
import sys
import subprocess
from os.path import exists, realpath
import json
import base64

sys.setrecursionlimit(1024)

# Init global variables
queue = Queue()
max_threads = 2
threads = [None]*max_threads
outputs = [None]*max_threads
images  = [None]*max_threads
user_execution = {}
pd_df = None

start_time = time()

def read_csvs(pd_df: pd.DataFrame) -> pd.DataFrame:
    altura_idade = {
        -1: -1,
        0: 0,
        None: 0,
        'Muito baixa estatura para idade': 1,
        'Baixa estatura para idade': 2,
        'Estatura adequada para a idade': 3,
    }

    imc_idade = {
        -1: -1,
        0: 0,
        None: 0,
        'Magreza acentuada': 1,
        'Magreza': 2,
        'Eutrofia': 3,
        'Risco de sobrepeso': 4,
        'Sobrepeso': 5,
        'Obesidade': 6
    }

    raca_cor = {
        '01': 1, '02': 2, '03': 3, '04': 4, '05': 5, 'X': 0, '99': 99,
        0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 99: 99, None: 0, -1: -1
    }
    pd_df = pd.read_csv(realpath('./')+'/xaa', encoding='latin-1', sep=';')

    pd_df.drop(['CO_ACOMPANHAMENTO', 'CO_PESSOA_SISVAN', 'ST_PARTICIPA_ANDI',
                'NO_MUNICIPIO', 'DS_FASE_VIDA', 'DS_RACA_COR',
                'PESO X IDADE', 'PESO X ALTURA', 'CO_SISTEMA_ORIGEM_ACOMP',
                'SISTEMA_ORIGEM_ACOMP', 'NU_COMPETENCIA',
                'DS_ESCOLARIDADE', 'DS_POVO_COMUNIDADE'], axis=1, inplace=True)
    pd_df['NU_PESO'] = pd_df['NU_PESO'].apply(lambda x: float(x.replace(',', '.')) if isinstance(x, str) else -1)
    pd_df['NU_ALTURA'] = pd_df['NU_ALTURA'].apply(lambda x: float(x.replace(',', '.')) if isinstance(x, str) else -1)
    pd_df['DS_IMC'] = pd_df['DS_IMC'].apply(lambda x: float(x.replace(',', '.')) if isinstance(x, str) else -1)
    pd_df['DS_IMC_PRE_GESTACIONAL'] = pd_df['DS_IMC_PRE_GESTACIONAL'].apply(lambda x: int(float(x.replace(',', '.'))*100) if isinstance(x, str) else -1)
    pd_df['SG_SEXO'] = pd_df['SG_SEXO'].apply(lambda x: False if x=='F' else True)
    pd_df['SG_SEXO'] = pd_df['SG_SEXO'].astype('bool')
    pd_df['CO_CNES'] = pd_df['CO_CNES'].astype('category')
    pd_df['DT_ACOMPANHAMENTO'] =  pd.to_datetime(pd_df['DT_ACOMPANHAMENTO'], format='%d/%m/%Y')
    pd_df['NU_PESO'] = pd_df['NU_PESO'].astype('float16')
    pd_df['NU_ALTURA'] = pd_df['NU_ALTURA'].astype('float16')
    pd_df['DS_IMC'] = pd_df['DS_IMC'].astype('float16')
    pd_df['DS_IMC_PRE_GESTACIONAL'] = pd_df['DS_IMC_PRE_GESTACIONAL'].astype('int16')
    pd_df['CO_POVO_COMUNIDADE'] = pd_df['CO_POVO_COMUNIDADE'].astype('category')
    pd_df['CO_ESCOLARIDADE'] = pd_df['CO_ESCOLARIDADE'].astype('category')
    pd_df['CRI. ALTURA X IDADE'].fillna(-1, inplace=True)
    pd_df['CRI. ALTURA X IDADE'] = pd_df['CRI. ALTURA X IDADE'].apply(lambda x: altura_idade[x]).astype('int8')
    pd_df['ADO. ALTURA X IDADE'].fillna(-1, inplace=True)
    pd_df['ADO. ALTURA X IDADE'] = pd_df['ADO. ALTURA X IDADE'].apply(lambda x: altura_idade[x]).astype('int8')
    pd_df['CRI. IMC X IDADE'].fillna(-1, inplace=True)
    pd_df['CRI. IMC X IDADE'] = pd_df['CRI. IMC X IDADE'].apply(lambda x: imc_idade[x]).astype('int8')
    pd_df['ADO. IMC X IDADE'].fillna(-1, inplace=True)
    pd_df['ADO. IMC X IDADE'] = pd_df['ADO. IMC X IDADE'].apply(lambda x: imc_idade[x]).astype('int8')
    pd_df['CO_ESTADO_NUTRI_ADULTO'] = pd_df['CO_ESTADO_NUTRI_ADULTO'].astype('category')
    pd_df['CO_ESTADO_NUTRI_IDOSO'] = pd_df['CO_ESTADO_NUTRI_IDOSO'].astype('category')
    pd_df['CO_ESTADO_NUTRI_IMC_SEMGEST'] = pd_df['CO_ESTADO_NUTRI_IMC_SEMGEST'].astype('category')
    pd_df['NU_FASE_VIDA'] = pd_df['NU_FASE_VIDA'].astype('category')
    pd_df['CO_RACA_COR'].fillna(-1, inplace=True)
    pd_df['CO_RACA_COR'] = pd_df['CO_RACA_COR'].apply(lambda x: raca_cor[x]).astype('int8')
    pd_df['NU_IDADE_ANO'].fillna(-1, inplace=True)
    pd_df['NU_IDADE_ANO'] = pd_df['NU_IDADE_ANO'].apply(lambda x: -1 if x == None else x)
    pd_df['NU_IDADE_ANO'] = pd_df['NU_IDADE_ANO'].astype('int8')
    pd_df['SG_UF'] = pd_df['SG_UF'].astype('category')

    count = 0

    required_files = ['xab', 'xac', 'xad', 'xae', 'xaf', 'xag', 'xah', 'xai', 'xaj', 'xak', 'xal', 'xam', 'xan', 'xao']
    if(not all([exists(realpath('./')+'/'+f) for f in required_files])):
        print("ERROR: some files are missing, try running: python3 setup.py")
        exit()

    for splitted_file in required_files:
        temp_pd_df = pd.read_csv(realpath('.')+'/'+splitted_file,
                                sep=';', encoding='latin-1')
        temp_pd_df.drop(['CO_ACOMPANHAMENTO', 'CO_PESSOA_SISVAN', 'ST_PARTICIPA_ANDI',
                    'NO_MUNICIPIO', 'DS_FASE_VIDA', 'DS_RACA_COR',
                    'PESO X IDADE', 'PESO X ALTURA', 'CO_SISTEMA_ORIGEM_ACOMP',
                    'SISTEMA_ORIGEM_ACOMP', 'NU_COMPETENCIA',
                    'DS_ESCOLARIDADE', 'DS_POVO_COMUNIDADE'], axis=1, inplace=True)
        temp_pd_df['NU_PESO'] = temp_pd_df['NU_PESO'].apply(lambda x: float(x.replace(',', '.')) if isinstance(x, str) else -1)
        temp_pd_df['NU_ALTURA'] = temp_pd_df['NU_ALTURA'].apply(lambda x: float(x.replace(',', '.')) if isinstance(x, str) else -1)
        temp_pd_df['DS_IMC'] = temp_pd_df['DS_IMC'].apply(lambda x: float(x.replace(',', '.')) if isinstance(x, str) else -1)
        temp_pd_df['DS_IMC_PRE_GESTACIONAL'] = temp_pd_df['DS_IMC_PRE_GESTACIONAL'].apply(lambda x: int(float(x.replace(',', '.'))*100) if isinstance(x, str) else -1)
        temp_pd_df['SG_SEXO'] = temp_pd_df['SG_SEXO'].apply(lambda x: False if x=='F' else True)
        temp_pd_df['SG_SEXO'] = temp_pd_df['SG_SEXO'].astype('bool')
        temp_pd_df['CO_CNES'] = temp_pd_df['CO_CNES'].astype('category')
        temp_pd_df['DT_ACOMPANHAMENTO'] =  pd.to_datetime(pd_df['DT_ACOMPANHAMENTO'], format='%d/%m/%Y')
        temp_pd_df['NU_PESO'] = temp_pd_df['NU_PESO'].astype('float16')
        temp_pd_df['NU_ALTURA'] = temp_pd_df['NU_ALTURA'].astype('float16')
        temp_pd_df['DS_IMC'] = temp_pd_df['DS_IMC'].astype('float16')
        temp_pd_df['DS_IMC_PRE_GESTACIONAL'] = temp_pd_df['DS_IMC_PRE_GESTACIONAL'].astype('int16')
        temp_pd_df['CO_POVO_COMUNIDADE'] = temp_pd_df['CO_POVO_COMUNIDADE'].astype('category')
        temp_pd_df['CO_ESCOLARIDADE'] = temp_pd_df['CO_ESCOLARIDADE'].astype('category')
        temp_pd_df['CRI. ALTURA X IDADE'].fillna(-1, inplace=True)
        temp_pd_df['CRI. ALTURA X IDADE'] = temp_pd_df['CRI. ALTURA X IDADE'].apply(lambda x: altura_idade[x]).astype('int8')
        temp_pd_df['ADO. ALTURA X IDADE'].fillna(-1, inplace=True)
        temp_pd_df['ADO. ALTURA X IDADE'] = temp_pd_df['ADO. ALTURA X IDADE'].apply(lambda x: altura_idade[x]).astype('int8')
        temp_pd_df['CRI. IMC X IDADE'].fillna(-1, inplace=True)
        temp_pd_df['CRI. IMC X IDADE'] = temp_pd_df['CRI. IMC X IDADE'].apply(lambda x: imc_idade[x]).astype('int8')
        temp_pd_df['ADO. IMC X IDADE'].fillna(-1, inplace=True)
        temp_pd_df['ADO. IMC X IDADE'] = temp_pd_df['ADO. IMC X IDADE'].apply(lambda x: imc_idade[x]).astype('int8')
        temp_pd_df['CO_ESTADO_NUTRI_ADULTO'] = temp_pd_df['CO_ESTADO_NUTRI_ADULTO'].astype('category')
        temp_pd_df['CO_ESTADO_NUTRI_IDOSO'] = temp_pd_df['CO_ESTADO_NUTRI_IDOSO'].astype('category')
        temp_pd_df['CO_ESTADO_NUTRI_IMC_SEMGEST'] = temp_pd_df['CO_ESTADO_NUTRI_IMC_SEMGEST'].astype('category')
        temp_pd_df['NU_FASE_VIDA'] = temp_pd_df['NU_FASE_VIDA'].astype('category')
        temp_pd_df['CO_RACA_COR'].fillna(-1, inplace=True)
        temp_pd_df['CO_RACA_COR'] = temp_pd_df['CO_RACA_COR'].apply(lambda x: raca_cor[x]).astype('int8')
        temp_pd_df['NU_IDADE_ANO'].fillna(-1, inplace=True)
        temp_pd_df['NU_IDADE_ANO'] =  temp_pd_df['NU_IDADE_ANO'].apply(lambda x: -1 if x == None else x)
        temp_pd_df['NU_IDADE_ANO'] =  temp_pd_df['NU_IDADE_ANO'].astype('int8')
        temp_pd_df['SG_UF'] = temp_pd_df['SG_UF'].astype('category')
        pd_df = pd.concat([pd_df, temp_pd_df], copy=False, ignore_index=True)
        del temp_pd_df
        count += 1
        print(count,'/ 14')

    with open('./df.pickle', 'wb') as handle:
        pickle.dump(pd_df, handle, protocol=pickle.HIGHEST_PROTOCOL)

    return pd_df

if path.exists('./df.pickle'):
    with open('./df.pickle', 'rb') as handle:
        try:
            pd_df = pickle.load(handle)
        except ModuleNotFoundError:
            pd_df = pd.read_pickle(handle)
else:
    pd_df = read_csvs(pd_df)

def create_selectors() -> None:
    estado_cidade = {}
    for val in read_municipal_seat().values:
        if(not val[3] in estado_cidade):
            estado_cidade.update({val[3]: []})
        else:
            if '.' in str(val[0]):
                val[0] = str(val[0])[:str(val[0]).index('.')]
            estado_cidade[val[3]].append({val[1]: str(val[0])})
    with open('estado_cidade_codigo.json', 'w') as f:
        json.dump(estado_cidade, f)

if not path.exists('./estado_cidade_codigo.json'):
    create_selectors()

end_time = time()
time_taken = end_time - start_time
print("Time taken to start cached pandas: {:.2f} seconds".format(time_taken))

def thread_filter(filters: dict, id_thread: int, queue: Queue) -> dict:
    # Using global variables
    global pd_df
    global outputs
    global images

    pd_copy = None
    pd_copy = pd_df.copy(deep=True)
    outputs[id_thread-1] = pd_copy

    print('filters:', filters, 'columns', pd_copy.columns)

    # print('len', len(pd_copy))

    # If Brasil is selected, do not apply state filter
    if(filters['state'] != 'BR' and filters['state'] != ''):
        pd_copy = pd_copy[pd_copy['SG_UF'] == filters['state']]
        print('1')
    # print('len', len(pd_copy))
    if(filters['city'] != ''):
        pd_copy = pd_copy[pd_copy['CO_MUNICIPIO_IBGE'] == int(filters['city'][:-1])]
        print('2')
    # print('len', len(pd_copy))

    # Weight range filter
    if(filters['minWeight']):
        pd_copy = pd_copy[pd_copy['NU_PESO'] > filters['minWeight']]
        print('3')
    if(filters['maxWeight']):
        pd_copy = pd_copy[pd_copy['NU_PESO'] < filters['maxWeight']]
        print('4')

    # print(pd_copy['NU_PESO'].values.tolist())
    # pd_copy['NU_PESO'] = pd_copy['NU_PESO'].dropna()
    pd_copy['NU_PESO'] = pd_copy['NU_PESO'].astype('float32')

    # Height range filter
    if(filters['minHeight']):
        pd_copy = pd_copy[pd_copy['NU_ALTURA'] > filters['minHeight']]
        print('5')
    if(filters['maxHeight']):
        pd_copy = pd_copy[pd_copy['NU_ALTURA'] < filters['maxHeight']]
        print('6')

    # pd_copy['NU_ALTURA'] = pd_copy['NU_ALTURA'].dropna()

    # Date range filter
    if(filters['minDate']):
        pd_copy = pd_copy[pd_copy['DT_ACOMPANHAMENTO'] > filters['minDate']]
        print('7')
    if(filters['maxDate']):
        pd_copy = pd_copy[pd_copy['DT_ACOMPANHAMENTO'] < filters['maxDate']]
        print('8')

    # print('len', len(pd_copy))

    # Age range filter
    if(filters['minAge']):
        print('9')
        pd_copy = pd_copy[pd_copy['NU_IDADE_ANO'] > filters['minAge']]
    if(filters['maxAge']):
        print('10')
        pd_copy = pd_copy[pd_copy['NU_IDADE_ANO'] < filters['maxAge']]

    pd_copy['DS_IMC'] = pd_copy['DS_IMC'].astype('int32')

    # IMC range filter
    if(filters['minIMC']):
        print('11')
        pd_copy = pd_copy[pd_copy['DS_IMC'] > filters['minIMC']]
    if(filters['maxIMC']):
        print('12')
        pd_copy = pd_copy[pd_copy['DS_IMC'] < filters['maxIMC']]

    # CNES filter
    if(filters['CNES']):
        print('13')
        pd_copy = pd_copy[pd_copy['CO_CNES'] == filters['CNES']]

    # School level filter
    if(filters['education']):
        print('14')
        pd_copy = pd_copy[pd_copy['CO_ESCOLARIDADE'] == filters['education']]

    # Genre filter
    if(filters['genre']):
        genre = filters['genre']
        if genre != 'Ambos':
            print('15')
            if genre == 'Masculino':
                pd_copy = pd_copy[pd_copy['SG_SEXO'] == True]
            else:
                pd_copy = pd_copy[pd_copy['SG_SEXO'] == False]

    # Color/race
    if(filters['etnicity']):
        print('16')
        pd_copy = pd_copy[pd_copy['CO_RACA_COR'] == filters['genre']]

    if(filters['community']):
        print('17')
        pd_copy = pd_copy[pd_copy['CO_POVO_COMUNIDADE'] == filters['community']]

    # if(filters['group']):
    #     group = filters['group']
    #     if group != 'Todos':
    #         if group == 'Criança':
    #             pd_copy = pd_copy[pd_copy['SG_SEXO'] == 1]
    #         elif group == 'Adulto':
    #             pd_copy = pd_copy[pd_copy['SG_SEXO'] == 0]
    #         else:
    #             # Idoso

    cidade = None
    try:
        cidade = read_municipality(code_muni=int(filters['city']), year=2010)
    except:
       images[id_thread-1] = 'undefined muni'
       if(cidade != None):
            print('dei except')
            return {'msg': 'error: undefined muni'}

    # print('len', len(pd_copy))

    image_uuid = uuid.uuid4().__str__()
    images[id_thread-1] = image_uuid
    queue.put(cidade)

    # print('new file at', realpath('.')+'/matplotimages/'+str(images[id_thread-1]))
    if(filters['state'] == 'BR'):
        with open(realpath('.')+'/matplotimages/'+str(images[id_thread-1]), 'wb') as imagefile:
            pickle.dump(pd_copy, imagefile)
        subprocess.run(["python3", f"{realpath('.')}/national_map.py", str(image_uuid)])
        return {'msg': 'success'}

    with open(realpath('.')+'/matplotimages/'+str(images[id_thread-1]), 'wb') as imagefile:
        pickle.dump(cidade, imagefile)

    # Outside with because it will run faster than the file is closed
    subprocess.run(["python3", f"{realpath('.')}/image_builder.py", 'cidade', str(image_uuid)])

    with open(realpath('.')+'/citiesjsons/'+str(images[id_thread-1]) + '.json', 'w') as jsonfile:
        # json.dump(pd_copy.describe().to_json(), jsonfile)
        # print(pd_copy.describe())
        # print(pd_copy['NU_PESO'].describe())
        jsonfile.write(pd_copy.describe().to_json())

    with open(realpath('.')+'/citiesjsons/'+str(images[id_thread-1]) + 'CNES.json', 'w') as jsonfile:
        jsonfile.write(pd_copy['CO_CNES'].value_counts().to_json())

    return {'msg': 'success'}

app = Flask('Lineu')
@app.route('/data', methods=['POST'])
def get_slash():
    while(all(threads)):
        print('Waiting while all threads are busy:', threads)
        sleep(0.2)

    id_thread = -1

    # Find first not None in threads array
    for i, thr in enumerate(threads):
        if thr == None:
            id_thread = i
            print('New thread in:', i)
            break

    if id_thread == -1:
        return 'System is unavailable'

    # Verify if request.json is a valid format (only asks a state or city)
    t = threading.Thread(target=thread_filter, args=(request.json, id_thread, queue))
    threads.append(t)

    t.start()
    t.join()
    # output = outputs[id_thread-1]

    if(images[id_thread-1] == 'undefined muni'):
        return 'Undefined Municipality'

    # Free threads and outputs
    outputs[id_thread-1] = None
    threads[id_thread-1] = None

    # print(images)
    while(not exists(realpath('.')+'/matplotimages/'+images[id_thread-1]+'.png')):
        print('to dormindo')
        sleep(1)

    # Probably there's a better way to transfer a file to buffer without using PIL
    # img = Image.open('./'+images[id_thread-1]+'.png')
    # buffer = io.BytesIO()
    # img.save(buffer, 'png')
    # buffer.seek(0)
    # data = buffer.read()
    # data = base64.b64encode(data).decode()
    data = None
    with open(realpath('.')+'/matplotimages/'+ images[id_thread-1] + '.png', 'rb') as f:
        data = f.read()

    data = base64.b64encode(data).decode()

    json_data = None
    if(exists(realpath('.')+'/citiesjsons/'+ images[id_thread-1] + '.json')):
        with open(realpath('.')+'/citiesjsons/'+ images[id_thread-1] + '.json', 'r') as j:
            json_data = j.read()

    cnes_data = None
    if(exists(realpath('.')+'/citiesjsons/'+ images[id_thread-1] + 'CNES.json')):
        with open(realpath('.')+'/citiesjsons/'+ images[id_thread-1] + 'CNES.json', 'r') as j:
            cnes_data = j.read()

    response = jsonify({'msg': 'success', 'format': 'png', 'data': json_data, 'image': data, 'cnes': cnes_data})
    response.headers.add('Access-Control-Allow-Origin', '*')
    return response
    # return send_file(f'{images[id_thread-1]}.png', mimetype='image/png')
    return f"<img src='data:application/html;base64,'/>"
    return render_template('oi.html')

@app.route('/municipalities', methods=['GET'])
def get_municipalities():
    response = send_file('estado_cidade_codigo.json', mimetype='application/json')
    response.headers.add('Access-Control-Allow-Origin', '*')
    return response

@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')

@app.route('/frontend/scripts.js', methods=['GET'])
def scripts():
    return send_file('./frontend/scripts.js', mimetype='application/json')

@app.route('/frontend/static/Lineu.png', methods=['GET'])
def lineu_logo():
    return send_file('./frontend/static/Lineu.png', mimetype='image/png')

# @app.route('/frontend/static/<path:path>')
# def serve_static(path):
#     return send_from_directory('static', path)

# Adicionar HOST e PORT vindos do .env
app.run(port=5050)


from pprint import pprint
import json

json_files = open("task5_data.json", "r")
files =  json.load(json_files)
json_files.close()

def analyse_movies_genre():
    list1 = []
    for movie in files:
        for genre in movie["genre"]:
            list1.append(genre)
            temp = set(list1)
            movie_genre = {}
            for genre in list1:
                movie_genre[genre] = list1.count(genre)

    files_11 = open("task11_data.json", "w")
    json.dump(movie_genre,files_11, indent = 4)
    files_11.close()

    return movie_genre
task11 = analyse_movies_genre()
# pprint(task11)


# CONQUEST:GO
# https://github.com/Kandru/conquest-go
# author: Karl-Martin Minkner
# website: https://gameshare.community

import pymysql
import uuid
from queue import Queue
from threading import Lock
from listeners.tick import GameThread
from conquest.debug import msg
class mysql:
	def __init__(self):
		self.socket = None
		self.cursor = None
		self.connect()
		self.queue = Queue()
		self.queue_lock = Lock()
		self.queue_queries = {}
		t = GameThread(target=self.queueworker, args=())
		t.start()

	def connect(self):
		try:
			if self.socket is None:
				self.socket = pymysql.connect(
					host='localhost',
					user='csgo_1',
					password='',
					db='csgo_1',
					charset='utf8mb4',
					autocommit=True,
					cursorclass=pymysql.cursors.DictCursor)
				self.cursor = self.socket.cursor()
				msg('INFO', 'successfully connected to mysql server')
			else:
				msg('INFO', 'reconnecting to mysql server')
				self.disconnect()
				self.connect()
		except:
			msg('ERROR', 'could not connect to mysql server')

	def query(self, query, params=[], fetch_all=True, wait=True):
		uid = uuid.uuid4()
		self.queue.put({
			'uid': uid,
			'query': query,
			'params': params,
			'fetch_all': fetch_all
		})
		if wait is True:
			while True:
				self.queue_lock.acquire()
				if uid in self.queue_queries:
					tmp = self.queue_queries[uid]
					del self.queue_queries[uid]
					self.queue_lock.release()
					return tmp
				self.queue_lock.release()

	def do_query(self, uid, query, params=[], fetch_all=True):
		try:
			self.cursor.execute(query, params)
			if fetch_all is True:
				return self.cursor.fetchall()
			else:
				return self.cursor.fetchone()
		except:
			msg('ERROR', 'could not perform query - mysql connection lost?')
			self.connect()
			self.query(query, params, fetch_all)

	def queueworker(self):
		while True:
			item = self.queue.get()
			if item is None:
				break
			self.queue_lock.acquire()
			self.queue_queries[item['uid']] = self.do_query(item['uid'], item['query'], item['params'], fetch_all=item['fetch_all'])
			self.queue_lock.release()

	def disconnect(self):
		self.socket.close()
		self.socket = None
		msg('INFO', 'disconnected from mysql server')


import setuptools

# https://packaging.python.org/tutorials/packaging-projects/

with open("README.md", "r") as fh:
    long_description = fh.read()

setuptools.setup(
    name="mathfn", # Replace with your own username
    version="0.0.1",
    url='git@github.com:NandaCj/private-packages.git',
    author="Nanda",
    author_email="nanda@gmail.com",
    description="A small example package",
    long_description=long_description,
    long_description_content_type="text/markdown",
    packages=setuptools.find_packages(),
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],
    python_requires='>=3.6',
)


from django.urls import path
from .views import GetIdealUser, UserDetailAPI,RegisterUserAPIView

urlpatterns = [
  path('register',RegisterUserAPIView.as_view()),
  path("get-details",UserDetailAPI.as_view()),        # need Authentication
  path("ideal/user/",GetIdealUser.as_view(), name="get-ideal-user"),
]


"""Regression tests for bug fixes, needing the full stack running."""
import logging
import time
from pathlib import Path

import pytest
import requests

logger = logging.getLogger(__name__)

TEST_VIDEO = str(
    (
        Path(__file__).parent
        / "test_data"
        / "test-abbor-small[2021-01-02_00-00-00]-000.mp4"
    ).resolve(),
)

@pytest.mark.usefixtures("start_core", "detection_api", "tracing_api")
def test_regression_add_two_jobs_with_objects() -> None:
    """Test regression of objects from previous job shows up.

    Ref. github issue #81
    """
    response = requests.post(
        "http://127.0.0.1:8000/projects/",
        json={
            "name": "Some name",
            "number": "NINA-123",
            "description": "Some kind of project description",
        },
    )
    assert response.status_code == 201

    response_job = requests.post(
        "http://127.0.0.1:8000/projects/1/jobs/",
        json={
            "name": "Job 1",
            "description": "Job description",
            "location": "test",
            "videos": [TEST_VIDEO],
        },
    )

    assert response_job.status_code == 201
    response = requests.post(
        "http://127.0.0.1:8000/projects/1/jobs/1/start",
    )

    url = "http://127.0.0.1:8000/projects/1/jobs/1"
    while requests.get(url).json()["_status"] != "Done":
        logger.info("Waiting for job to finish.")
        time.sleep(5)

    response_job = requests.post(
        "http://127.0.0.1:8000/projects/1/jobs/",
        json={
            "name": "Job 2",
            "description": "Job description",
            "location": "test",
            "videos": [TEST_VIDEO],
        },
    )

    assert response_job.status_code == 201
    response = requests.post(
        "http://127.0.0.1:8000/projects/1/jobs/2/start",
    )

    url = "http://127.0.0.1:8000/projects/1/jobs/2"

    while requests.get(url).json()["_status"] != "Done":
        logger.info("Waiting for job to finish.")
        time.sleep(5)

    job1 = requests.get(
        "http://127.0.0.1:8000/projects/1/jobs/1",
    ).json()
    job2 = requests.get(
        "http://127.0.0.1:8000/projects/1/jobs/2",
    ).json()

    assert job1["stats"]["total_objects"] == 1
    assert job2["stats"]["total_objects"] == 1  # bug results this to be 2


# -- Import packages ---------------------------------------------------------------------------------------------------
from collections import OrderedDict
import os
import json
import numpy as np

class Diagnostics:
    """
    Diagnostics for an episode.
    """

    def __init__(self, runs):
        self.d = OrderedDict()
        self.runs = runs
        self.reset()

    def reset(self):
        """
        Reset the diagnostics dictionary.
        """
        self.d = OrderedDict()

        for r in range(self.runs):
            self.d[f'run_{r}'] = {'n_steps': [],
                                  'net_reward_to_org': [],
                                  'actionFromStart': [],
                                  'optimalPathLength': [],
                                  'optimalAction': [],
                                  'short_open_door_path': [],
                                  'long_path': [],
                                  'short_closed_door_path': [],
                                  'noGoal': [],
                                  'leaders': [],
                                  'coordinationCostsAccumulated': [],
                                  'opportunityCostsAccumulated': [],
                                  'policyInDomain_longpath': [],
                                  'policyInDomain_shortpath': [],
                                  'n_samplesModelDomain1': [],
                                  'n_samplesModelDomain2': [],
                                  'n_samplesModelDomain3': [],
                                  'n_samplesModelDomain4': [],
                                  'n_visits_per_state': np.zeros((13, 13)),
                                  'neutral_states_share': [],
                                  'shortpath_states_share': [],
                                  'longpath_states_share': [],
                                  'neutral_states_count': [],
                                  'shortpath_states_count': [],
                                  'longpath_states_count': []
                                  }

    def to_disk(self, par_name_, par_value_, start_date, start_time, runs):
        """
        Write diagnostics to disk.
        :param par_name_: parameter name
        :param par_value_: parameter value
        :param start_date: start date of the simulation run
        :param start_time: start time of the simulation run
        """
        assert type(self.d) == OrderedDict, 'ERROR: diagnostics variable must be a dictionary.'

        os.makedirs(f'outputs/traces/{start_date}/{start_time}', exist_ok=True)
        path = f'outputs/traces/{start_date}/{start_time}'
        filename = f'/dyna-q_{par_name_}_{par_value_}_{start_time}.json'

        # JSONify ndarray
        for r in range(runs):
            self.d[f'run_{r}']['n_visits_per_state'] = self.d[f'run_{r}']['n_visits_per_state'].tolist()

        with open(path + filename, "w") as f:
            json.dump(self.d, f, indent=4)

        print(f'\nSuccessfully saved to disk: {filename}\n')


"""
tests.test_component_device_tracker
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Tests the device tracker compoments.
"""
# pylint: disable=protected-access,too-many-public-methods
import unittest
from unittest.mock import patch
from datetime import datetime, timedelta
import os

from homeassistant.config import load_yaml_config_file
from homeassistant.loader import get_component
import homeassistant.util.dt as dt_util
from homeassistant.const import (
    ATTR_ENTITY_ID, ATTR_ENTITY_PICTURE, ATTR_FRIENDLY_NAME, ATTR_HIDDEN,
    STATE_HOME, STATE_NOT_HOME, CONF_PLATFORM, DEVICE_DEFAULT_NAME)
import homeassistant.components.device_tracker as device_tracker

from tests.common import (
    get_test_home_assistant, fire_time_changed, fire_service_discovered)

class TestComponentsDeviceTracker(unittest.TestCase):
    """ Tests homeassistant.components.device_tracker module. """

    def setUp(self):  # pylint: disable=invalid-name
        """ Init needed objects. """
        self.hass = get_test_home_assistant()
        self.yaml_devices = self.hass.config.path(device_tracker.YAML_DEVICES)

    def tearDown(self):  # pylint: disable=invalid-name
        """ Stop down stuff we started. """
        try:
            os.remove(self.yaml_devices)
        except FileNotFoundError:
            pass

        self.hass.stop()

    def test_is_on(self):
        """ Test is_on method. """
        entity_id = device_tracker.ENTITY_ID_FORMAT.format('test')

        self.hass.states.set(entity_id, STATE_HOME)

        self.assertTrue(device_tracker.is_on(self.hass, entity_id))

        self.hass.states.set(entity_id, STATE_NOT_HOME)

        self.assertFalse(device_tracker.is_on(self.hass, entity_id))

    def test_migrating_config(self):
        csv_devices = self.hass.config.path(device_tracker.CSV_DEVICES)

        self.assertFalse(os.path.isfile(csv_devices))
        self.assertFalse(os.path.isfile(self.yaml_devices))

        person1 = {
            'mac': 'AB:CD:EF:GH:IJ:KL',
            'name': 'Paulus',
            'track': True,
            'picture': 'http://placehold.it/200x200',
        }
        person2 = {
            'mac': 'MN:OP:QR:ST:UV:WX:YZ',
            'name': '',
            'track': False,
            'picture': None,
        }

        try:
            with open(csv_devices, 'w') as fil:
                fil.write('device,name,track,picture\n')
                for pers in (person1, person2):
                    fil.write('{},{},{},{}\n'.format(
                        pers['mac'], pers['name'],
                        '1' if pers['track'] else '0', pers['picture'] or ''))

            self.assertTrue(device_tracker.setup(self.hass, {}))
            self.assertFalse(os.path.isfile(csv_devices))
            self.assertTrue(os.path.isfile(self.yaml_devices))

            yaml_config = load_yaml_config_file(self.yaml_devices)

            self.assertEqual(2, len(yaml_config))

            for pers, yaml_pers in zip(
                (person1, person2), sorted(yaml_config.values(),
                                           key=lambda pers: pers['mac'])):
                for key, value in pers.items():
                    if key == 'name' and value == '':
                        value = DEVICE_DEFAULT_NAME
                    self.assertEqual(value, yaml_pers.get(key))

        finally:
            try:
                os.remove(csv_devices)
            except FileNotFoundError:
                pass

    def test_reading_yaml_config(self):
        dev_id = 'test'
        device = device_tracker.Device(
            self.hass, timedelta(seconds=180), 0, True, dev_id,
            'AB:CD:EF:GH:IJ', 'Test name', 'http://test.picture', True)
        device_tracker.update_config(self.yaml_devices, dev_id, device)
        self.assertTrue(device_tracker.setup(self.hass, {}))
        config = device_tracker.load_config(self.yaml_devices, self.hass,
                                            device.consider_home, 0)[0]
        self.assertEqual(device.dev_id, config.dev_id)
        self.assertEqual(device.track, config.track)
        self.assertEqual(device.mac, config.mac)
        self.assertEqual(device.config_picture, config.config_picture)
        self.assertEqual(device.away_hide, config.away_hide)
        self.assertEqual(device.consider_home, config.consider_home)

    def test_setup_without_yaml_file(self):
        self.assertTrue(device_tracker.setup(self.hass, {}))

    def test_adding_unknown_device_to_config(self):
        scanner = get_component('device_tracker.test').SCANNER
        scanner.reset()
        scanner.come_home('DEV1')
        self.assertTrue(device_tracker.setup(self.hass, {
            device_tracker.DOMAIN: {CONF_PLATFORM: 'test'}}))
        config = device_tracker.load_config(self.yaml_devices, self.hass,
                                            timedelta(seconds=0), 0)[0]
        self.assertEqual('dev1', config.dev_id)
        self.assertEqual(True, config.track)

    def test_discovery(self):
        scanner = get_component('device_tracker.test').SCANNER

        with patch.dict(device_tracker.DISCOVERY_PLATFORMS, {'test': 'test'}):
            with patch.object(scanner, 'scan_devices') as mock_scan:
                self.assertTrue(device_tracker.setup(self.hass, {
                    device_tracker.DOMAIN: {CONF_PLATFORM: 'test'}}))
                fire_service_discovered(self.hass, 'test', {})
                self.assertTrue(mock_scan.called)

    def test_update_stale(self):
        scanner = get_component('device_tracker.test').SCANNER
        scanner.reset()
        scanner.come_home('DEV1')

        register_time = datetime(2015, 9, 15, 23, tzinfo=dt_util.UTC)
        scan_time = datetime(2015, 9, 15, 23, 1, tzinfo=dt_util.UTC)

        with patch('homeassistant.components.device_tracker.dt_util.utcnow',
                   return_value=register_time):
            self.assertTrue(device_tracker.setup(self.hass, {
                'device_tracker': {
                    'platform': 'test',
                    'consider_home': 59,
                }}))

        self.assertEqual(STATE_HOME,
                         self.hass.states.get('device_tracker.dev1').state)

        scanner.leave_home('DEV1')

        with patch('homeassistant.components.device_tracker.dt_util.utcnow',
                   return_value=scan_time):
            fire_time_changed(self.hass, scan_time)
            self.hass.pool.block_till_done()

        self.assertEqual(STATE_NOT_HOME,
                         self.hass.states.get('device_tracker.dev1').state)

    def test_entity_attributes(self):
        dev_id = 'test_entity'
        entity_id = device_tracker.ENTITY_ID_FORMAT.format(dev_id)
        friendly_name = 'Paulus'
        picture = 'http://placehold.it/200x200'

        device = device_tracker.Device(
            self.hass, timedelta(seconds=180), 0, True, dev_id, None,
            friendly_name, picture, away_hide=True)
        device_tracker.update_config(self.yaml_devices, dev_id, device)

        self.assertTrue(device_tracker.setup(self.hass, {}))

        attrs = self.hass.states.get(entity_id).attributes

        self.assertEqual(friendly_name, attrs.get(ATTR_FRIENDLY_NAME))
        self.assertEqual(picture, attrs.get(ATTR_ENTITY_PICTURE))

    def test_device_hidden(self):
        dev_id = 'test_entity'
        entity_id = device_tracker.ENTITY_ID_FORMAT.format(dev_id)
        device = device_tracker.Device(
            self.hass, timedelta(seconds=180), 0, True, dev_id, None,
            away_hide=True)
        device_tracker.update_config(self.yaml_devices, dev_id, device)

        scanner = get_component('device_tracker.test').SCANNER
        scanner.reset()

        self.assertTrue(device_tracker.setup(self.hass, {
            device_tracker.DOMAIN: {CONF_PLATFORM: 'test'}}))

        self.assertTrue(self.hass.states.get(entity_id)
                            .attributes.get(ATTR_HIDDEN))

    def test_group_all_devices(self):
        dev_id = 'test_entity'
        entity_id = device_tracker.ENTITY_ID_FORMAT.format(dev_id)
        device = device_tracker.Device(
            self.hass, timedelta(seconds=180), 0, True, dev_id, None,
            away_hide=True)
        device_tracker.update_config(self.yaml_devices, dev_id, device)

        scanner = get_component('device_tracker.test').SCANNER
        scanner.reset()

        self.assertTrue(device_tracker.setup(self.hass, {
            device_tracker.DOMAIN: {CONF_PLATFORM: 'test'}}))

        state = self.hass.states.get(device_tracker.ENTITY_ID_ALL_DEVICES)
        self.assertIsNotNone(state)
        self.assertEqual(STATE_NOT_HOME, state.state)
        self.assertSequenceEqual((entity_id,),
                                 state.attributes.get(ATTR_ENTITY_ID))

    @patch('homeassistant.components.device_tracker.DeviceTracker.see')
    def test_see_service(self, mock_see):
        self.assertTrue(device_tracker.setup(self.hass, {}))
        mac = 'AB:CD:EF:GH'
        dev_id = 'some_device'
        host_name = 'example.com'
        location_name = 'Work'
        gps = [.3, .8]

        device_tracker.see(self.hass, mac, dev_id, host_name, location_name,
                           gps)

        self.hass.pool.block_till_done()

        mock_see.assert_called_once_with(
            mac=mac, dev_id=dev_id, host_name=host_name,
            location_name=location_name, gps=gps)


import subprocess
import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

def get_gsea_enrichment(input_file, library, out_folder):
    """Runs GSEA based on input set of genes with associated logFC values
    using local jar file and gene set library (.gmt file)

    Parameters
    ----------
    input_file : str
        path to 2-column tab-seperated file  of gene (or phospshosite)
        and associated log2 fold-change values. Extension should be .rnk
    library : str
        path to gene set library
    out_folder : str
        path to output folder where all output files will be saved

    Returns
    -------
    df : pandas dataframe
       Summary of results from enrichment analysis that includes
       the term enriched, enrichment value, p-value and FDR.
    """
    path = out_folder
    os.makedirs(path)
    if '.gmt' not in library:
        lib = '%s.gmt' % library
    else:
        lib = library
    msda_path = os.path.dirname(os.path.abspath(__file__))
    arg_list = ['java', '-cp', os.path.join(msda_path, 'jars/gsea2-2.2.3.jar'),
                '-Xmx1024m',  'xtools.gsea.GseaPreranked',
                '-param_file', os.path.join(msda_path, 'gsea_parameters.txt'),
                '-rnk', input_file,
                '-gmx', lib,
                '-out', path]
    # if out_folder:
    #     arg_list += ['-out', path]
    # if out_file:
    #     arg_list += ['-rpt_label', out_file]
    subprocess.call(arg_list)
    terminal_folder = os.listdir(path)[0]
    timestamp = terminal_folder.split('.')[2]
    terminal_path = "%s/%s" % (path, terminal_folder)

    file1 = '%s/gsea_report_for_na_neg_%s.xls' % (terminal_path, timestamp)
    file2 = '%s/gsea_report_for_na_pos_%s.xls' % (terminal_path, timestamp)
    dfl = []
    if os.path.isfile(file1):
        dfl.append(pd.read_table(file1))
    if os.path.isfile(file2):
        dfl.append(pd.read_table(file2))
    # try:
    #     df_pos = pd.read_table(file1)
    #     df_neg = pd.read_table(file2)
    # except IOError:
    #     pass
    try:
        df = pd.concat(dfl)
    except ValueError:
        df = pd.DataFrame()
    return df

def make_rnkfile(df, sample1, sample2, identifier):
    """ Generated 2-column dataframe of identifier (gene or phosphosite)
    and associated log2 fold-change values

    Parameters
    ----------
    df : pandas dataframe
        proteomics or phosphoproteomics dataframe of dataset.
    sample1 : str
        column name of test sample.
    sample2 : str
        column name of control sample.
    identifier : str
       name of column that contains identifiers.

    Results
    -------
    df2 : pandas dataframe
       2-column dataframe of identifier (gene or phosphosite)
       and associated log2 fold-change values
    """
    df2 = df.copy()
    df2[sample1] = df2[sample1].div(df2[sample2])
    df2[sample1] = df2[sample1].apply(np.log2)
    df2 = df2.sort_values([sample1])
    df2 = df2[[identifier, sample1]]
    df2[identifier] = [g.upper() for g in df2[identifier].tolist()]
    df2 = df2.replace([-np.inf, np.inf], np.nan)
    df2 = df2.dropna()
    return df2

def plot_nes(df, filter=True, top=10, outfile=None,
             pval_label='NOM p-val', fdr_label='FDR q-val',
             regulon='NAME',
             fontsize=12, show_pval=False):
    """Plots gene set terms and associated enrichment scores
       from enrichment analyses.

    Parameters
    ----------
    df : pandas dataframe
       output dataframe from running enrichment analysis.
       Should contain term name, enrichment score, associated p-value and FDR.
    filter : Optional[bool]
       If True, terms for which the p-value is less than 0.05
       and FDR is less than 0.2 are filtered out. Default is True.
    top : Optional[int]
       Number of terms from downregulated and upregulated terms
       to be included in the plot. Default is 10.
    outfile : Optional[str]
       name of file to save the plot. Default is None
    pval_label : Optional[str]
       Name of column that contains p-values. Default is 'NOM p-val'
    fdr_label : Optional[str]
       Name of column that contains FDR values. Default is 'FDR q-val'
    regulon : Optional[str]
       Name of column that contains name of terms. Default is 'NAME'.
    fontsize : Optional[int]
       Size of fonts in plot. Default is 10.
    show_pval : bool
       p-values shown on plot if True. Default is False.
    """
    sns.set_style('dark')
    font = {'family': 'sans-serif',
            'color':  'black',
            'weight': 'normal',
            'size': fontsize,
            }
    sns.despine()
    if filter is True:
        df = df[(df[pval_label] <= 0.05) & (df[fdr_label] < 0.2)]
    df1 = df[df.NES > 0].iloc[:top, :]
    df2 = df[df.NES < 0].iloc[:top, :].sort_values(by='NES', ascending=False)
    ylen = np.max([len(df1), len(df2)])
    f, axes = plt.subplots(nrows=1, ncols=2)
    if not df1.empty:
        df1.NES.plot(ax=axes[1], kind='barh', width=0.3,
                     color='Green', alpha=0.5)
        for pos, name in zip(range(ylen), df1[regulon].tolist()):
            axes[1].text(0, pos-0.25, name, horizontalalignment='left',
                         fontdict=font)
    if not df2.empty:
        df2.NES.plot(ax=axes[0], kind='barh', width=0.3,
                     color='Red', alpha=0.5)
        for pos, name in zip(range(ylen), df2[regulon].tolist()):
            axes[0].text(0, pos+0.25, name, horizontalalignment='right',
                         fontdict=font)
    axes[0].set_yticks([])
    axes[1].set_yticks([])
    plt.subplots_adjust(wspace=0.075)
    axes[0].set_ylim([-0.5, ylen-0.5])
    axes[1].set_ylim([-0.5, ylen-0.5])

    if show_pval:
        pvals1 = ["%.2f" % p if p > 0 else "%.1f" % p
                  for p in df1[pval_label].tolist()]
        pvals2 = ["%.2f" % p if p > 0 else "%.1f" % p
                  for p in df2[pval_label].tolist()]
        for x_pos, y_pos, pval in zip(df1.NES.tolist(), range(ylen),
                                      pvals1):
            axes[1].text(x_pos+0.1, y_pos, pval)

        for x_pos, y_pos, pval in zip(df2.NES.tolist(), range(ylen)[::-1],
                                      pvals2):
            axes[1].text(x_pos-0.5, y_pos, pval)
    plt.gca().invert_yaxis()
    if not outfile:
        plt.show()
        plt.clf()
    else:
        plt.savefig(outfile)
        plt.clf()


from keras.applications.inception_v3 import InceptionV3
from keras.preprocessing import image
from keras.models import Model
from keras.layers import Dense, GlobalAveragePooling2D
from keras import backend as K
from keras.optimizers import SGD, RMSprop, Adam

train_data_dir = "/ebs/user05/data/train"
validation_data_dir = "/ebs/user05/data/test"
img_width, img_height = 299, 299
nb_epoch = 20
nb_train_samples = 20385
nb_validation_samples = 8733

# Create the base pre-trained model
base_model = InceptionV3(weights='imagenet', include_top=False)

# Add a global spatial average pooling layer
x = base_model.output
x = GlobalAveragePooling2D()(x)
# Add a fully-connected layer
x = Dense(256, activation='relu', name='fc_1')(x)
predictions = Dense(1, activation='sigmoid')(x)

model = Model(input=base_model.input, output=predictions)

# Freeze convolutional layers
for layer in base_model.layers:
    layer.trainable = False

model.compile(optimizer=RMSprop(lr = .00001), loss = 'binary_crossentropy', metrics=['accuracy'])

train_datagen = image.ImageDataGenerator(rescale = 1./255,
                                  shear_range =.2,
                                  zoom_range = .2,
                                  horizontal_flip = True)

# Rescale image color scale
test_datagen = image.ImageDataGenerator(rescale=1./255)

generator_train = train_datagen.flow_from_directory(
        train_data_dir,
        target_size=(img_width, img_height),
        batch_size=32,
        class_mode='binary')

generator_test = test_datagen.flow_from_directory(
        validation_data_dir,
        target_size=(img_width, img_height),
        batch_size=32,
        class_mode='binary')

model.fit_generator(generator_train,
            samples_per_epoch = nb_train_samples,
            nb_epoch = nb_epoch,
            validation_data = generator_test,
            nb_val_samples = nb_validation_samples)

model_json = model.to_json()
with open("incep_filter.json", 'w') as json_file:
    json_file.write(model_json)
model.save_weights("incep_filter.h5")


import functools

from peewee import fn
from playhouse.shortcuts import model_to_dict

from nni.nas.benchmark.utils import load_benchmark
from .schema import NdsTrialStats, NdsTrialConfig, proxy

def query_nds_trial_stats(model_family, proposer, generator, model_spec, cell_spec, dataset,
                          num_epochs=None, reduction=None, include_intermediates=False):
    """
    Query trial stats of NDS given conditions.

    Parameters
    ----------
    model_family : str or None
        If str, can be one of the model families available in :class:`nni.nas.benchmark.nds.NdsTrialConfig`.
        Otherwise a wildcard.
    proposer : str or None
        If str, can be one of the proposers available in :class:`nni.nas.benchmark.nds.NdsTrialConfig`. Otherwise a wildcard.
    generator : str or None
        If str, can be one of the generators available in :class:`nni.nas.benchmark.nds.NdsTrialConfig`. Otherwise a wildcard.
    model_spec : dict or None
        If specified, can be one of the model spec available in :class:`nni.nas.benchmark.nds.NdsTrialConfig`.
        Otherwise a wildcard.
    cell_spec : dict or None
        If specified, can be one of the cell spec available in :class:`nni.nas.benchmark.nds.NdsTrialConfig`.
        Otherwise a wildcard.
    dataset : str or None
        If str, can be one of the datasets available in :class:`nni.nas.benchmark.nds.NdsTrialConfig`. Otherwise a wildcard.
    num_epochs : float or None
        If int, matching results will be returned. Otherwise a wildcard.
    reduction : str or None
        If 'none' or None, all trial stats will be returned directly.
        If 'mean', fields in trial stats will be averaged given the same trial config.
    include_intermediates : boolean
        If true, intermediate results will be returned.

    Returns
    -------
    generator of dict
        A generator of :class:`nni.nas.benchmark.nds.NdsTrialStats` objects,
        where each of them has been converted into a dict.
    """

    if proxy.obj is None:
        proxy.initialize(load_benchmark('nds'))

    fields = []
    if reduction == 'none':
        reduction = None
    if reduction == 'mean':
        for field_name in NdsTrialStats._meta.sorted_field_names:
            if field_name not in ['id', 'config', 'seed']:
                fields.append(fn.AVG(getattr(NdsTrialStats, field_name)).alias(field_name))
    elif reduction is None:
        fields.append(NdsTrialStats)
    else:
        raise ValueError('Unsupported reduction: \'%s\'' % reduction)
    query = NdsTrialStats.select(*fields, NdsTrialConfig).join(NdsTrialConfig)
    conditions = []
    for field_name in ['model_family', 'proposer', 'generator', 'model_spec', 'cell_spec',
                       'dataset', 'num_epochs']:
        if locals()[field_name] is not None:
            conditions.append(getattr(NdsTrialConfig, field_name) == locals()[field_name])
    if conditions:
        query = query.where(functools.reduce(lambda a, b: a & b, conditions))
    if reduction is not None:
        query = query.group_by(NdsTrialStats.config)
    for trial in query:
        if include_intermediates:
            data = model_to_dict(trial)
            # exclude 'trial' from intermediates as it is already available in data
            data['intermediates'] = [
                {k: v for k, v in model_to_dict(t).items() if k != 'trial'} for t in trial.intermediates
            ]
            yield data
        else:
            yield model_to_dict(trial)


from collections import defaultdict
import sys

import numpy as np
import gffutils
import pandas as pd

def miso_exon_to_gencode_exon(exon):
    """Convert a single miso exon to gencode gffutils database exon id

    >>> # Skipped exon (SE) or Mutually exclusive exon (MXE) ID
    >>> miso_exon_to_gencode_exon('chr2:9624561:9624679:+')
    'exon:chr2:9624561-9624679:+'
    >>> # Alt 5' splice site (pick first of alternative)
    >>> miso_exon_to_gencode_exon('chr15:42565276:42565087|42565161:-')
    'exon:chr15:42565276-42565087:-'
    >>> # Alt 3' splice site (pick first of alternative)
    >>> miso_exon_to_gencode_exon('chr2:130914199|130914248:130914158:-')
    'exon:chr2:130914199-130914158:-'
    >>> # Retained intron: start, stop separated by '-' instead of ':'
    >>> miso_exon_to_gencode_exon('chr1:906259-906386:+')
    'exon:chr1:906259-906386:+'
    """
    return 'exon:{}:{}-{}:{}'.format(*miso_exon_to_coords(exon))

def miso_id_to_exon_ids(miso_id):
    """Convert a MISO-style alternative event ID

    Split on the pipe ("|") to account for Alt 5'/3' splice site events

    # A skipped exon (SE) ID
    >>> miso_id_to_exon_ids('chr2:9624561:9624679:+@chr2:9627585:9627676:+@chr2:9628276:9628591:+')
    ['exon:chr2:9624561-9624679:+', 'exon:chr2:9627585-9627676:+', 'exon:chr2:9628276-9628591:+']
    >>> # A mutually exclusive (MXE) ID
    >>> miso_id_to_exon_ids('chr16:89288500:89288591:+@chr16:89289565:89289691:+@chr16:89291127:89291210:+@chr16:89291963:89292039:+')
    ['exon:chr16:89288500-89288591:+', 'exon:chr16:89289565-89289691:+', 'exon:chr16:89291127-89291210:+', 'exon:chr16:89291963-89292039:+']
    >>> # An Alt 5' splice site (A5SS) ID
    >>> miso_id_to_exon_ids("chr15:42565276:42565087|42565161:-@chr15:42564261:42564321:-")
    ['exon:chr15:42565276-42565087:-', 'exon:chr15:42564261-42564321:-']
    >>> # An Alt 3' splice site (A3SS) ID
    >>> miso_id_to_exon_ids('chr2:130914824:130914969:-@chr2:130914199|130914248:130914158:-')
    ['exon:chr2:130914824-130914969:-', 'exon:chr2:130914199-130914158:-']
    >>> # A retained intron (RI) ID
    >>> miso_id_to_exon_ids('chr1:906066-906138:+@chr1:906259-906386:+')
    ['exon:chr1:906066-906138:+', 'exon:chr1:906259-906386:+']
    """
    return map(miso_exon_to_gencode_exon, miso_id.split('@'))

def miso_exon_to_coords(exon):
    """Convert a miso exon to gffutils coordinates

    >>> miso_exon_to_coords('chr2:130914824:130914969:-')
    ('chr2', '130914824', '130914969', '-')
    >>> # Alt 5' SS - pick the first of the alternative ends
    >>> miso_exon_to_coords('chr15:42565276:42565087|42565161:-')
    ('chr15', '42565276', '42565087', '-')
    >>> # Alt 3' SS - pick the first of the alternative starts
    >>> miso_exon_to_coords('chr2:130914199|130914248:130914158:-')
    ('chr2', '130914199', '130914158', '-')
    >>> # Retained intron
    >>> miso_exon_to_coords('chr1:906066-906138:+')
    ('chr1', '906066', '906138', '+')
    """
    strand = exon[-1]
    coords = map(lambda x: x.split('|')[0],
                 exon.split(':'))
    if '-' in coords[1]:
        start, stop = coords[1].split('-')
        coords = coords[0], start, stop, strand
    return coords[0], coords[1], coords[2], strand

def convert_miso_ids_to_everything(miso_ids, db,
                                   event_type,
                                   out_dir):
    """Given a list of miso IDs and a gffutils database, pull out the
    ensembl/gencode/gene name/gene type/transcript names, and write files
    into the out directory. Does not return a value.

    Parameters
    ----------
    miso_ids : list of str
        Miso ids to convert
    db : gffutils.FeatureDB
        gffutils feature database created from a gtf file
    event_type : str
        The type of splicing event. This is used for naming only
    out_dir : str
        Where to write the files to.
    """
    out_dir = out_dir.rstrip('/')
    event_type = event_type.lower()

    miso_to_ensembl = {}
    miso_to_gencode = {}
    miso_to_gene_name = {}
    miso_to_gene_type = {}
    miso_to_ensembl_transcript = {}
    miso_to_gencode_transcript = {}

    ensembl_to_miso = defaultdict(list)
    gencode_to_miso = defaultdict(list)
    gene_name_to_miso = defaultdict(list)

    n_miso_ids = len(miso_ids)
    sys.stdout.write('Converting {} {} miso ids using {} gffutils database '
                     'into {}.'.format(n_miso_ids, event_type, str(db),
                                       out_dir))

    for i, miso_id in enumerate(miso_ids):
        if i % 100 == 0:
            sys.stdout.write('On {}/{} {} miso ids'.format(i, n_miso_ids,
                                                           event_type))

        exons = miso_id_to_exon_ids(miso_id)

        gencode = set([])
        ensembl = set([])
        gene_name = set([])
        gene_type = set([])
        gencode_transcript = set([])
        ensembl_transcript = set([])
        for e in exons:
            try:
                exon = db[e]
                gencode.update(exon.attributes['gene_id'])
                ensembl.update(
                    map(lambda x: x.split('.')[0],
                        exon.attributes['gene_id']))
                gene_name.update(exon.attributes['gene_name'])
                gene_type.update(exon.attributes['gene_type'])
                gencode_transcript.update(exon.attributes['transcript_id'])
                ensembl_transcript.update(
                    map(lambda x: x.split('.')[0], exon.attributes[
                        'transcript_id']))
            except gffutils.FeatureNotFoundError:
                try:
                    #  not an exon, look for any overlapping transcripts here
                    prefix, chrom, startstop, strand = e.split(':')
                    start, stop = startstop.split('-')
                    transcripts = list(db.features_of_type('transcript',
                                                           strand=strand,
                                                           limit=(
                                                           chrom, int(start),
                                                           int(stop))))
                    # if there are overlapping transcripts...
                    if len(transcripts) == 0:
                        continue
                    else:
                        for transcript in transcripts:
                            gencode.update(transcript.attributes['gene_id'])
                            ensembl.update(
                                map(lambda x: x.split('.')[0],
                                    transcript.attributes['gene_id']))
                            gene_name.update(transcript.attributes['gene_name'])
                            gene_type.update(transcript.attributes['gene_type'])
                            gencode_transcript.update(
                                transcript.attributes['transcript_id'])
                            ensembl_transcript.update(
                                map(lambda x: x.split('.')[0],
                                    transcript.attributes['transcript_id']))
                except:
                    continue
        if len(gencode) > 0:

            for ens in ensembl:
                ensembl_to_miso[ens].append(miso_id)
            for g in gene_name:
                gene_name_to_miso[g].append(miso_id)
            for g in gencode:
                gencode_to_miso[g].append(miso_id)

            ensembl = ','.join(ensembl)

            gencode = ','.join(gencode)
            gene_name = ','.join(gene_name)
            gene_type = ','.join(gene_type)

            gencode_transcript = ','.join(gencode_transcript)
            ensembl_transcript = ','.join(ensembl_transcript)

            miso_to_gencode[miso_id] = gencode
            miso_to_ensembl[miso_id] = ensembl
            miso_to_gene_name[miso_id] = gene_name
            miso_to_gene_type[miso_id] = gene_type

            miso_to_gencode_transcript[miso_id] = gencode_transcript
            miso_to_ensembl_transcript[miso_id] = ensembl_transcript

        else:
            miso_to_gencode[miso_id] = np.nan
            miso_to_ensembl[miso_id] = np.nan
            miso_to_gene_name[miso_id] = np.nan
            miso_to_gene_type[miso_id] = np.nan

    miso_tos = {'gencode_gene': miso_to_gencode,
                'ensembl_gene': miso_to_ensembl,
                'gene_name': miso_to_gene_name,
                'gene_type': miso_to_gene_type,
                'gencode_transcript': miso_to_gencode_transcript,
                'ensembl_transcript': miso_to_ensembl_transcript}

    to_misos = {'ensembl_gene': ensembl_to_miso,
                'gene_name': gene_name_to_miso,
                'gencode_gene': gencode_to_miso}

    for name, d in miso_tos.iteritems():
        df = pd.DataFrame.from_dict(d, orient='index')
        df.index.name = 'event_name'
        df.columns = [name]
        tsv = '{}/miso_{}_to_{}.tsv'.format(out_dir, event_type, name)
        df.to_csv(tsv, sep='\t')
        sys.stdout.write('Wrote {}\n'.format(tsv))

    for name, d in to_misos.iteritems():
        tsv = '{}/{}_to_miso_{}.tsv'.format(out_dir, name, event_type)
        with open(tsv, 'w') as f:
            for k, v in d.iteritems():
                f.write('{}\t{}\n'.format(k, '\t'.join(v)))
        sys.stdout.write('Wrote {}\n'.format(tsv))


import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from fairseq import utils
from fairseq.models import (
    FairseqEncoder,
)
from fairseq.modules import (
    LayerNorm,
    PositionalEmbedding,
    SinusoidalPositionalEmbedding,
)
from nonauto.model_utils import positional_encodings_like
from nonauto.modules.blocks import (
    SelfAttentionBlock,
    FeedForwardBlock
)

class NATEncoderLayer(nn.Module):
    """Encoder layer block.

    In the original paper each operation (multi-head attention or FFN) is
    postprocessed with: `dropout -> add residual -> layernorm`. In the
    tensor2tensor code they suggest that learning is more robust when
    preprocessing each layer with layernorm and postprocessing with:
    `dropout -> add residual`. We default to the approach in the paper, but the
    tensor2tensor approach can be enabled by setting
    *args.encoder_normalize_before* to ``True``.

    Args:
        args (argparse.Namespace): parsed command-line arguments
    """

    def __init__(self, args):
        super().__init__()
        embed_dim = args.encoder_embed_dim
        normalize_before = args.encoder_normalize_before
        head_nums = args.encoder_attention_heads
        dropout = args.dropout

        self.self_attn_block = SelfAttentionBlock(
            input_dim=embed_dim,
            num_heads=head_nums,
            dropout=args.attention_dropout,
            normalize_before=normalize_before,
        )

        activation_dropout = getattr(args, 'activation_dropout', 0)
        if activation_dropout == 0:
            # for backwards compatibility with models that use args.relu_dropout
            activation_dropout = getattr(args, 'relu_dropout', 0)

        self.ffn_block = FeedForwardBlock(
            input_dim=embed_dim,
            hidden_dim=args.encoder_ffn_embed_dim,
            dropout=dropout,
            activation_fn=getattr(args, 'activation_fn', 'relu'),
            activation_dropout=activation_dropout,
            normalize_before=normalize_before,
            residual_type=args.residual_type
        )

    def forward(self, x, encoder_padding_mask):
        """
        Args:
            x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
            encoder_padding_mask (ByteTensor): binary ByteTensor of shape
                `(batch, src_len)` where padding elements are indicated by ``1``.

        Returns:
            encoded output of shape `(batch, src_len, embed_dim)`
        """
        x = self.self_attn_block(
            x=x,
            padding_mask=encoder_padding_mask
        )
        x = self.ffn_block(x)
        return x

    def upgrade_state_dict_named(self, state_dict, name):
        """Upgrade a (possibly old) state dict for new versions of fairseq."""
        return state_dict

class NATEncoder(FairseqEncoder):
    """
    Non-autoregressive Transformer encoder consisting of *args.encoder_layers* layers. Each layer
    is a :class:`NATEncoderLayer`.

    Args:
        args (argparse.Namespace): parsed command-line arguments
        dictionary (~fairseq.data.Dictionary): encoding dictionary
        embed_tokens (torch.nn.Embedding): input embedding
    """

    def __init__(self, args, dictionary, embed_tokens):
        super().__init__(dictionary)
        self.share_embed = args.share_all_embeddings
        embed_dim = args.encoder_embed_dim
        self.padding_idx = dictionary.pad()
        self.max_source_positions = args.max_source_positions

        self.embed_tokens = embed_tokens
        self.embed_scale = math.sqrt(embed_dim)
        self.embed_positions = PositionalEmbedding(
            args.max_source_positions, embed_dim, self.padding_idx,
            learned=args.encoder_learned_pos,
        ) if not args.no_token_positional_embeddings else None

        self.layers = nn.ModuleList([])
        self.layers.extend([
            NATEncoderLayer(args)
            for _ in range(args.encoder_layers)
        ])
        self.dropout = nn.Dropout(args.dropout)
        self.register_buffer('version', torch.Tensor([2]))
        self.normalize = args.encoder_normalize_before
        if self.normalize:
            self.layer_norm = LayerNorm(embed_dim)

    def forward(self, src_tokens, src_lengths, **unused):
        """
        Args:
            src_tokens (LongTensor): tokens in the source language of shape
                `(batch, src_len)`
            src_lengths (torch.LongTensor): lengths of each source sentence of
                shape `(batch)`

        Returns:
            dict:
                - **encoder_out** (Tensor): the last encoder layer's output of
                  shape `(src_len, batch, embed_dim)`
                - **encoder_padding_mask** (ByteTensor): the positions of
                  padding elements of shape `(batch, src_len)`
        """
        # embed tokens and positions
        if self.share_embed:
            x = F.embedding(src_tokens, self.embed_tokens.weight * self.embed_scale)
        else:
            x = self.embed_tokens(src_tokens) * self.embed_scale
        # x = self.embed_scale * self.embed_tokens(src_tokens)
        # if self.embed_positions is not None:
        #     x += self.embed_positions(src_tokens)
        x += positional_encodings_like(x)
        encoder_history = [x]
        # x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.dropout(x)
        # B x T x C -> T x B x C

        x = x.transpose(0, 1)

        # compute padding mask
        encoder_padding_mask = src_tokens.eq(self.padding_idx)
        if not encoder_padding_mask.any():
            encoder_padding_mask = None

        # encoder layers
        for layer in self.layers:
            x = layer(x, encoder_padding_mask)
            encoder_history.append(x.transpose(0, 1))

        if self.normalize:
            x = self.layer_norm(x)

        return {
            'encoder_history': encoder_history,  # List<B X T X C>
            'encoder_out': x,  # T x B x C
            'encoder_padding_mask': encoder_padding_mask,  # B x T
        }

    def reorder_encoder_out(self, encoder_out, new_order):
        """
        Reorder encoder output according to *new_order*.

        Args:
            encoder_out: output from the ``forward()`` method
            new_order (LongTensor): desired order

        Returns:
            *encoder_out* rearranged according to *new_order*
        """
        if encoder_out['encoder_out'] is not None:
            encoder_out['encoder_out'] = \
                encoder_out['encoder_out'].index_select(1, new_order)
        if encoder_out['encoder_padding_mask'] is not None:
            encoder_out['encoder_padding_mask'] = \
                encoder_out['encoder_padding_mask'].index_select(0, new_order)
        return encoder_out

    def max_positions(self):
        """Maximum input length supported by the encoder."""
        if self.embed_positions is None:
            return self.max_source_positions
        return min(self.max_source_positions, self.embed_positions.max_positions())

    def upgrade_state_dict_named(self, state_dict, name):
        """Upgrade a (possibly old) state dict for new versions of fairseq."""
        if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
            weights_key = '{}.embed_positions.weights'.format(name)
            if weights_key in state_dict:
                del state_dict[weights_key]
            state_dict['{}.embed_positions._float_tensor'.format(name)] = torch.FloatTensor(1)
        for i in range(len(self.layers)):
            # update layer norms
            self.layers[i].upgrade_state_dict_named(state_dict, f"{name}.layers.{i}")

        version_key = '{}.version'.format(name)
        if utils.item(state_dict.get(version_key, torch.Tensor([1]))[0]) < 2:
            # earlier checkpoints did not normalize after the stack of layers
            self.layer_norm = None
            self.normalize = False
            state_dict[version_key] = torch.Tensor([1])
        return state_dict


from typing import Optional, Tuple, Union

import torch
from torch import nn as nn
from torch.autograd import Function

from ..utils import ext_loader
from .ball_query import ball_query
from .knn import knn

ext_module = ext_loader.load_ext(
    '_ext', ['group_points_forward', 'group_points_backward'])

class QueryAndGroup(nn.Module):
    """Groups points with a ball query of radius.

    Args:
        max_radius (float): The maximum radius of the balls.
            If None is given, we will use kNN sampling instead of ball query.
        sample_num (int): Maximum number of features to gather in the ball.
        min_radius (float, optional): The minimum radius of the balls.
            Default: 0.
        use_xyz (bool, optional): Whether to use xyz.
            Default: True.
        return_grouped_xyz (bool, optional): Whether to return grouped xyz.
            Default: False.
        normalize_xyz (bool, optional): Whether to normalize xyz.
            Default: False.
        uniform_sample (bool, optional): Whether to sample uniformly.
            Default: False
        return_unique_cnt (bool, optional): Whether to return the count of
            unique samples. Default: False.
        return_grouped_idx (bool, optional): Whether to return grouped idx.
            Default: False.
    """

    def __init__(self,
                 max_radius: float,
                 sample_num: int,
                 min_radius: float = 0.,
                 use_xyz: bool = True,
                 return_grouped_xyz: bool = False,
                 normalize_xyz: bool = False,
                 uniform_sample: bool = False,
                 return_unique_cnt: bool = False,
                 return_grouped_idx: bool = False):
        super().__init__()
        self.max_radius = max_radius
        self.min_radius = min_radius
        self.sample_num = sample_num
        self.use_xyz = use_xyz
        self.return_grouped_xyz = return_grouped_xyz
        self.normalize_xyz = normalize_xyz
        self.uniform_sample = uniform_sample
        self.return_unique_cnt = return_unique_cnt
        self.return_grouped_idx = return_grouped_idx
        if self.return_unique_cnt:
            assert self.uniform_sample, \
                'uniform_sample should be True when ' \
                'returning the count of unique samples'
        if self.max_radius is None:
            assert not self.normalize_xyz, \
                'can not normalize grouped xyz when max_radius is None'

    def forward(
        self,
        points_xyz: torch.Tensor,
        center_xyz: torch.Tensor,
        features: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, Tuple]:
        """
        Args:
            points_xyz (torch.Tensor): (B, N, 3) xyz coordinates of the
                points.
            center_xyz (torch.Tensor): (B, npoint, 3) coordinates of the
                centriods.
            features (torch.Tensor): (B, C, N) The features of grouped
                points.

        Returns:
            Tuple | torch.Tensor: (B, 3 + C, npoint, sample_num) Grouped
            concatenated coordinates and features of points.
        """
        # if self.max_radius is None, we will perform kNN instead of ball query
        # idx is of shape [B, npoint, sample_num]
        if self.max_radius is None:
            idx = knn(self.sample_num, points_xyz, center_xyz, False)
            idx = idx.transpose(1, 2).contiguous()
        else:
            idx = ball_query(self.min_radius, self.max_radius, self.sample_num,
                             points_xyz, center_xyz)

        if self.uniform_sample:
            unique_cnt = torch.zeros((idx.shape[0], idx.shape[1]))
            for i_batch in range(idx.shape[0]):
                for i_region in range(idx.shape[1]):
                    unique_ind = torch.unique(idx[i_batch, i_region, :])
                    num_unique = unique_ind.shape[0]
                    unique_cnt[i_batch, i_region] = num_unique
                    sample_ind = torch.randint(
                        0,
                        num_unique, (self.sample_num - num_unique, ),
                        dtype=torch.long)
                    all_ind = torch.cat((unique_ind, unique_ind[sample_ind]))
                    idx[i_batch, i_region, :] = all_ind

        xyz_trans = points_xyz.transpose(1, 2).contiguous()
        # (B, 3, npoint, sample_num)
        grouped_xyz = grouping_operation(xyz_trans, idx)
        grouped_xyz_diff = grouped_xyz - \
            center_xyz.transpose(1, 2).unsqueeze(-1)  # relative offsets
        if self.normalize_xyz:
            grouped_xyz_diff /= self.max_radius

        if features is not None:
            grouped_features = grouping_operation(features, idx)
            if self.use_xyz:
                # (B, C + 3, npoint, sample_num)
                new_features = torch.cat([grouped_xyz_diff, grouped_features],
                                         dim=1)
            else:
                new_features = grouped_features
        else:
            assert (self.use_xyz
                    ), 'Cannot have not features and not use xyz as a feature!'
            new_features = grouped_xyz_diff

        ret = [new_features]
        if self.return_grouped_xyz:
            ret.append(grouped_xyz)
        if self.return_unique_cnt:
            ret.append(unique_cnt)
        if self.return_grouped_idx:
            ret.append(idx)
        if len(ret) == 1:
            return ret[0]
        else:
            return tuple(ret)

class GroupAll(nn.Module):
    """Group xyz with feature.

    Args:
        use_xyz (bool): Whether to use xyz.
    """

    def __init__(self, use_xyz: bool = True):
        super().__init__()
        self.use_xyz = use_xyz

    def forward(self,
                xyz: torch.Tensor,
                new_xyz: torch.Tensor,
                features: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            xyz (Tensor): (B, N, 3) xyz coordinates of the features.
            new_xyz (Tensor): new xyz coordinates of the features.
            features (Tensor): (B, C, N) features to group.

        Returns:
            Tensor: (B, C + 3, 1, N) Grouped feature.
        """
        grouped_xyz = xyz.transpose(1, 2).unsqueeze(2)
        if features is not None:
            grouped_features = features.unsqueeze(2)
            if self.use_xyz:
                # (B, 3 + C, 1, N)
                new_features = torch.cat([grouped_xyz, grouped_features],
                                         dim=1)
            else:
                new_features = grouped_features
        else:
            new_features = grouped_xyz

        return new_features

class GroupingOperation(Function):
    """Group feature with given index."""

    @staticmethod
    def forward(ctx, features: torch.Tensor,
                indices: torch.Tensor) -> torch.Tensor:
        """
        Args:
            features (Tensor): (B, C, N) tensor of features to group.
            indices (Tensor): (B, npoint, nsample) the indices of
                features to group with.

        Returns:
            Tensor: (B, C, npoint, nsample) Grouped features.
        """
        features = features.contiguous()
        indices = indices.contiguous()

        B, nfeatures, nsample = indices.size()
        _, C, N = features.size()
        output = torch.cuda.FloatTensor(B, C, nfeatures, nsample)

        ext_module.group_points_forward(
            features,
            indices,
            output,
            b=B,
            c=C,
            n=N,
            npoints=nfeatures,
            nsample=nsample)

        ctx.for_backwards = (indices, N)
        return output

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor, None]:
        """
        Args:
            grad_out (Tensor): (B, C, npoint, nsample) tensor of the gradients
                of the output from forward.

        Returns:
            Tensor: (B, C, N) gradient of the features.
        """
        idx, N = ctx.for_backwards

        B, C, npoint, nsample = grad_out.size()
        grad_features = torch.cuda.FloatTensor(B, C, N).zero_()

        grad_out_data = grad_out.data.contiguous()
        ext_module.group_points_backward(
            grad_out_data,
            idx,
            grad_features.data,
            b=B,
            c=C,
            n=N,
            npoints=npoint,
            nsample=nsample)
        return grad_features, None

grouping_operation = GroupingOperation.apply


import random
from typing import List

class Solution:

    def __init__(self, w: List[int]):
        self.total = sum(w)
        self.cumulative = [0]
        for val in w:
            self.cumulative.append(self.cumulative[-1] + val)

    def pickIndex(self) -> int:
        val = random.randint(0, self.total-1)

        # TLE
        # for i in range(len(self.cumulative)-1):
        #     if self.cumulative[i] <= val < self.cumulative[i+1]:
        #         return i

        left, right = 0, len(self.cumulative)-1
        while left < right:
            mid = left + (right - left) // 2
            if self.cumulative[mid] <= val < self.cumulative[mid + 1]:
                return mid
            elif self.cumulative[mid] < val:
                left = mid + 1
            else:
                right = mid

        # return left

# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()

# ["Solution", "pickIndex", "pickIndex", "pickIndex", "pickIndex", "pickIndex"]
# [[[1, 3]], [], [], [], [], []]

def test(init_w: List[int], times: int):
    print(init_w, times)
    obj = Solution(init_w)
    print(obj.cumulative)
    for _ in range(times):
        print(obj.pickIndex())

if __name__ == "__main__":
    test([1], 1)
    test([1, 3], 5)
    test([3, 14, 1, 7], 10000)


import hashlib
import random
import sqlite3

import emojis

import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(name)s | %(levelname)s | %(" "message)s",
    handlers=[
            logging.FileHandler("oi.log"),
            logging.StreamHandler()
        ]
)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)

DB_NAME = "oi.db"

class OIDatabase:
    def __init__(self):
        """
        If you want to make a database change:
        - Increment LATEST_VERSION
        - Make your changes and wrap them around an if db_version < <previous version> function
        """
        self.connection = sqlite3.connect(DB_NAME)
        self.cursor = self.connection.cursor()
        self.cursor.execute(
            "CREATE TABLE IF NOT EXISTS schema_version (version INT NOT NULL)"
        )
        db_version = self.cursor.execute(
            "SELECT version FROM schema_version"
        ).fetchall()
        if len(db_version) == 0:
            db_version = 0
        else:
            db_version = db_version[0][0]
        if db_version < 1:
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS oi (id TEXT NOT NULL, messagecount INT NOT NULL)"
            )
        if db_version < 2:
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS cumcry (id TEXT NOT NULL, emoji TEXT NOT NULL, cumcount INT NOT NULL, crycount INT NOT NULL)"
            )
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS confessions (confession TEXT NOT NULL)"
            )
        if db_version < 3:
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS wyr (question TEXT NOT NULL)"
            )
        if db_version < 4:
            self.cursor.execute("ALTER TABLE confessions ADD COLUMN hash TEXT")
            self.cursor.execute("ALTER TABLE wyr ADD COLUMN hash TEXT")
        if db_version < 5:
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS later_deletion (member_id TEXT NOT NULL, remove_time TEXT NOT NULL)"
            )
        if db_version < 6:
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS cumtime (id TEXT NOT NULL, datetime TEXT NOT NULL)"
            )
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS crytime (id TEXT NOT NULL, datetime TEXT NOT NULL)"
            )
        if db_version < 7:
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS unlater (id TEXT NOT NULL, cmdcount INT NOT NULL)"
            )
        if db_version < 8:
            self.cursor.execute("ALTER TABLE later_deletion ADD COLUMN jobid TEXT")
        if db_version < 9:
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS later_min (id TEXT NOT NULL, minutes INT NOT NULL)"
            )
            self.cursor.execute(
                "CREATE TABLE IF NOT EXISTS last_latered (id TEXT NOT NULL, datetime TEXT NOT NULL)"
            )
        LATEST_VERSION = 9
        if db_version == 0:
            self.cursor.execute(
                "INSERT INTO schema_version VALUES (?)", (LATEST_VERSION,)
            )
        else:
            self.cursor.execute(
                "UPDATE schema_version SET version = ?", (LATEST_VERSION,)
            )
        self.connection.commit()

    def update_message_count(self, id):
        """
        Increment #hi msg count by 1.
        :param id: user id
        :return: None
        """
        self.cursor.execute("SELECT * FROM oi WHERE id = ?", (id,))
        if len(self.cursor.fetchall()) == 0:
            self.cursor.execute("INSERT INTO oi VALUES (?, 1)", (id,))
        else:
            self.cursor.execute(
                "UPDATE oi SET messagecount = messagecount + 1 WHERE id = ?", (id,)
            )
        self.connection.commit()

    def get_hi_leaderboard(self):
        """
        Get ID, message count for users in #hi.
        :return:
        """
        self.cursor.execute("SELECT * FROM oi ORDER BY messagecount DESC")
        return self.cursor.fetchall()

    def get_cum_leaderboard(self):
        """
        Gets all id, emoji, cumcount, crycount from cumcry table, order by cumcount.
        :return:
        """
        self.cursor.execute(
            "SELECT id, emoji, cumcount, crycount FROM cumcry ORDER BY cumcount DESC"
        )
        return self.cursor.fetchall()

    def get_cry_leaderboard(self):
        """
        Gets all id, emoji, cumcount, crycount from cumcry table, order by crycount.
        :return:
        """
        self.cursor.execute(
            "SELECT id, emoji, cumcount, crycount FROM cumcry ORDER BY crycount DESC"
        )
        return self.cursor.fetchall()

    def get_aggregated_cumcry_leaderboard(self):
        """
        Gets all id, emoji, cumcount + crycount from cumcry table, order by cumcount + crycount.
        :return:
        """
        self.cursor.execute(
            "SELECT id, emoji, cumcount + crycount FROM cumcry ORDER BY cumcount + crycount DESC, cumcount DESC"
        )
        return self.cursor.fetchall()

    def increment_cumcry_count(self, id, action):
        """
        Increment cum or cry count for user by 1. If user doesn't exist, creates entry for them with random emoji.
        :param id: user id
        :param action: either "cum" or "cry"
        :return:
        """
        self.cursor.execute("SELECT * FROM cumcry WHERE id = ?", (id,))
        if len(self.cursor.fetchall()) == 0:
            # choose random emoji from 'Animals & Nature' category
            cat = [*emojis.db.get_emojis_by_category("Animals & Nature")]
            e = random.choice(cat)
            # check if anyone else has same emoji
            self.cursor.execute("SELECT * FROM cumcry WHERE emoji = ?", (e.aliases[0],))
            # repeat until we find one that no one else has
            while len(self.cursor.fetchall()) > 0:
                e = random.choice(cat)
                self.cursor.execute(
                    "SELECT * FROM cumcry WHERE emoji = ?", (e.aliases[0],)
                )
            # create entry for user with their unique emoji
            self.cursor.execute(
                "INSERT INTO cumcry VALUES (?, ?, 0, 0)", (id, e.aliases[0])
            )
        # update counts
        if action == "cum":
            self.cursor.execute(
                "UPDATE cumcry SET cumcount = cumcount + 1 WHERE id = ?", (id,)
            )
        elif action == "cry":
            self.cursor.execute(
                "UPDATE cumcry SET crycount = crycount + 1 WHERE id = ?", (id,)
            )
        self.connection.commit()

    def clear_cumcry_counts(self):
        """
        Clear all cum and cry counts. Irreversible operation, do not execute.
        :return:
        """
        self.cursor.execute("DELETE FROM cumcry")
        self.connection.commit()

    def store_confession(self, confession):
        """
        Stores confession and hashed confession in database. Returns string to be sent to user.
        :param confession: confession message in the format "confession <actual confession>"
        :return: message to send to user with hash
        """
        if confession.strip() == "":
            return "you can't confess nothing, dummy"
        confession = confession[len("confess ") :]
        gen_hash = hashlib.sha1(confession.encode("utf-8")).hexdigest()
        h = (
            "Here is the hash for your confession: `"
            + gen_hash
            + '`. If you want to delete your confession, send the following message: "unconfess <hash>" and if deleted '
            "the oi bot will react with the trash can emoji."
        )
        self.cursor.execute(
            "INSERT INTO confessions VALUES (?, ?)", (confession, gen_hash)
        )
        self.connection.commit()
        return h

    def store_wyr(self, wyr):
        """
        Stores wyr and hashed wyr in database. Returns hash to be sent to user.
        :param wyr: wyr message in the format "wyr <actual wyr>"
        :return: message to send to user with hash
        """
        wyr = wyr[len("wyr ") :]
        gen_hash = hashlib.sha1(wyr.encode("utf-8")).hexdigest()
        h = (
            "Here is the hash for your wyr: `"
            + gen_hash
            + '`. If you want to delete your wyr, send the following message: "unwyr <hash>" and if deleted '
            "the oi bot will react with the trash can emoji."
        )
        self.cursor.execute("INSERT INTO wyr VALUES (?, ?)", (wyr, gen_hash))
        self.connection.commit()
        return h

    def get_random_confession(self):
        """
        Gets random confession.
        :return: rowid, confession
        """
        self.cursor.execute(
            "SELECT rowid, * FROM confessions ORDER BY RANDOM() LIMIT 1"
        )
        table = self.cursor.fetchall()
        rowid = -1
        confession = None
        log.info(f"Fetched {len(table)} of {self.get_num_confessions()} confessions")
        if len(table) > 0:
            rowid = table[0][0]
            confession = table[0][1]
            log.info(f"Confession {rowid}: \"{confession}\"")
        return rowid, confession

    def get_num_confessions(self):
        """
        :return: number of entries in confessions table
        """
        self.cursor.execute("SELECT COUNT(1) FROM confessions")
        num = self.cursor.fetchall()
        return num

    def get_random_wyr(self):
        """
        Gets random wyr.
        :return: rowid, wyr
        """
        self.cursor.execute("SELECT rowid, * FROM wyr ORDER BY RANDOM() LIMIT 1")
        table = self.cursor.fetchall()
        rowid = -1
        wyr = ""
        if len(table) > 0:
            rowid = table[0][0]
            wyr = table[0][1]
        return rowid, wyr

    def get_num_wyr(self):
        """
        :return: number of entries in confessions table
        """
        self.cursor.execute("SELECT COUNT(1) FROM wyr")
        num = self.cursor.fetchall()
        return num

    def delete_confession(self, rowid):
        """
        Deletes confession by rowid. Use if rowid is knownn.
        :param rowid:
        :return:
        """
        self.cursor.execute("DELETE FROM confessions WHERE rowid = ?", (rowid,))
        self.connection.commit()

    def delete_confession_by_hash(self, h):
        """
        Deletes confession by hash.
        :param h: hash
        :return: number of rows deleted
        """
        self.cursor.execute("DELETE FROM confessions WHERE hash = ?", (h,))
        num_rows = self.cursor.rowcount
        self.connection.commit()
        return num_rows

    def delete_wyr(self, rowid):
        """
        Deletes wyr by rowid. Use if rowid is known.
        :param rowid:
        :return:
        """
        self.cursor.execute("DELETE FROM wyr WHERE rowid = ?", (rowid,))
        self.connection.commit()

    def delete_wyr_by_hash(self, h):
        """
        Deletes wyr by hash.
        :param h: hash
        :return: number of rows deleted
        """
        self.cursor.execute("DELETE FROM wyr WHERE hash = ?", (h,))
        num_rows = self.cursor.rowcount
        self.connection.commit()
        return num_rows

    def create_later_delete_job(self, member_id, remove_time, jobid):
        """
        Creates job to remove later role
        :param member_id: ID of target member
        :param remove_time: string of datetime object representing remove tim
        :param jobid: ID of apscheduler job
        :return:
        """
        self.cursor.execute(
            "INSERT INTO later_deletion (member_id, remove_time, jobid) SELECT ?, ?, ? WHERE NOT EXISTS (SELECT 1 FROM later_deletion WHERE member_id = ? AND remove_time = ?);",
            (
                member_id,
                remove_time,
                jobid,
                member_id,
                remove_time,
            ),
        )
        self.connection.commit()

    def delete_later_delete_job(self, member_id, remove_time):
        """
        Deletes job to remove later role from row. Used if job is completed without restart.
        :param member_id:
        :param remove_time:
        :return:
        """
        self.cursor.execute(
            "DELETE FROM later_deletion WHERE member_id = ? AND remove_time = ?",
            (
                member_id,
                remove_time,
            ),
        )
        self.connection.commit()

    def get_all_later_deletion_jobs(self):
        """
        Gets all later deletion job entries in the db.
        :return: entries in the form of a list of tuples (member_id, remove_time)
        """
        self.cursor.execute("SELECT * FROM later_deletion")
        entries = self.cursor.fetchall()
        return entries

    def get_later_deletion_jobs_by_id(self, member_id):
        """
        Gets all later deletion job entries in the db.
        :return: entries in the form of a list of tuples (member_id, remove_time)
        """
        self.cursor.execute(
            "SELECT * FROM later_deletion where member_id = ?", (member_id,)
        )
        entries = self.cursor.fetchall()
        return entries

    def delete_later_jobs_by_id(self, user_id):
        """
        Deletes all later deletion jobs by user id.
        :param user_id: target user id
        :return:
        """
        self.cursor.execute(
            "DELETE FROM later_deletion WHERE member_id = ?", (user_id,)
        )
        self.connection.commit()

    def get_all_cumcry_entries(self):
        """
        Gets all cumcry entries in the db.
        :return: cumcry entries in the form of a list of tuples (id, emoji, cumcount, crycount)
        """
        self.cursor.execute("SELECT * FROM cumcry")
        entries = self.cursor.fetchall()
        return entries

    def update_cumcry_entry(self, id, num_cum, num_cry):
        """
        Updates cumcry entry with new cum and cry counts.
        :param id:
        :param num_cum:
        :param num_cry:
        :return:
        """
        self.cursor.execute(
            "UPDATE cumcry SET cumcount = ?, crycount = ? WHERE id = ?",
            (num_cum, num_cry, id),
        )
        self.connection.commit()

    def insert_cum_date_entry(self, member_id, date):
        """
        Inserts a new cum date entry into the db.
        :param member_id: ID of member
        :param date: datetime object or string
        :return:
        """
        self.cursor.execute("INSERT INTO cumtime VALUES (?, ?)", (member_id, date))
        self.connection.commit()

    def insert_cry_date_entry(self, member_id, date):
        """
        Inserts a new cry date entry into the db.
        :param member_id: ID of member
        :param date: datetime object or string
        :return:
        """
        self.cursor.execute("INSERT INTO crytime VALUES (?, ?)", (member_id, date))
        self.connection.commit()

    def count_cum_date_entries_by_id(self, member_id):
        """
        Counts the number of cum date entries for a given member.
        :param member_id:
        :return:
        """
        self.cursor.execute("SELECT COUNT(*) FROM cumtime WHERE id = ?", (member_id,))
        return self.cursor.fetchone()[0]

    def count_cry_date_entries_by_id(self, member_id):
        """
        Counts the number of cry date entries for a given member.
        :param member_id:
        :return:
        """
        self.cursor.execute("SELECT COUNT(*) FROM crytime WHERE id = ?", (member_id,))
        return self.cursor.fetchone()[0]

    def get_cumcry_id_emoji_pairs(self):
        """
        Gets a dictionary mapping ids to emojis.
        :return: dictionary of id->emoji pairs. Emojis are given by names, eg. "potted_plant".
        """
        self.cursor.execute("SELECT id, emoji FROM cumcry")
        return dict(self.cursor.fetchall())

    def get_cum_date_entries_by_id(self, id):
        """
        Gets all cum date entries for a given member.
        :param id: discord member id
        :return:
        """
        self.cursor.execute("SELECT datetime FROM cumtime WHERE id = ?", (id,))
        return self.cursor.fetchall()

    def get_cry_date_entries_by_id(self, id):
        """
        Gets all cry date entries for a given member.
        :param id: discord member id
        :return:
        """
        self.cursor.execute("SELECT datetime FROM crytime WHERE id = ?", (id,))
        return self.cursor.fetchall()

    def update_unlater_count(self, id):
        """
        Increment /unlater count by 1.
        :param id: user id
        :return: None
        """
        self.cursor.execute("SELECT * FROM unlater WHERE id = ?", (id,))
        if len(self.cursor.fetchall()) == 0:
            self.cursor.execute("INSERT INTO unlater VALUES (?, 1)", (id,))
        else:
            self.cursor.execute(
                "UPDATE unlater SET cmdcount = cmdcount + 1 WHERE id = ?", (id,)
            )
        self.connection.commit()

    def get_unlater_leaderboard(self):
        """
        Get ID, call count for /unlater.
        :return:
        """
        self.cursor.execute("SELECT * FROM unlater ORDER BY cmdcount DESC")
        return self.cursor.fetchall()

    def insert_later_time(self, id, curr_time):
        """
        Insert into unlater table.
        :param id: user id
        :return: None
        """
        self.cursor.execute(
            "INSERT INTO last_latered VALUES (?, ?)",
            (
                id,
                curr_time,
            ),
        )
        self.connection.commit()

    def get_last_later_time(self, id):
        """
        Get last time user used /later.
        :param id: user id
        :return: Datetime string of last latered time
        """
        self.cursor.execute("SELECT datetime FROM last_latered WHERE id = ?", (id,))
        ret = self.cursor.fetchone()
        if ret is None:
            return None
        else:
            return ret[0]

    def remove_later_time(self, id):
        """
        Remove last latered time from table.
        :param id: user id
        :return: None
        """
        self.cursor.execute("DELETE FROM last_latered WHERE id = ?", (id,))
        self.connection.commit()

    def update_later_min_count(self, id, minutes):
        """
        Increment /unlater count by 1.
        :param id: user id
        :return: None
        """
        self.cursor.execute("SELECT * FROM later_min WHERE id = ?", (id,))
        if len(self.cursor.fetchall()) == 0:
            self.cursor.execute(
                "INSERT INTO later_min VALUES (?, ?)",
                (
                    id,
                    minutes,
                ),
            )
        else:
            self.cursor.execute(
                "UPDATE later_min SET minutes = minutes + ? WHERE id = ?",
                (
                    minutes,
                    id,
                ),
            )
        self.connection.commit()

    def get_later_min_leaderboard(self):
        """
        Get ID, call count for /unlater.
        :return:
        """
        self.cursor.execute("SELECT * FROM later_min ORDER BY minutes DESC")
        return self.cursor.fetchall()


import torch
import torch.nn as nn
import numpy as np
import sys
sys.path.append("../../")

from general.mutils import get_param_val, create_transformer_mask, create_channel_mask
from layers.flows.flow_model import FlowModel
from layers.flows.activation_normalization import ActNormFlow
from layers.flows.permutation_layers import InvertibleConv
from layers.flows.coupling_layer import CouplingLayer
from layers.flows.mixture_cdf_layer import MixtureCDFCoupling
from layers.categorical_encoding.mutils import create_encoding

class FlowSetModeling(FlowModel):

	def __init__(self, model_params, dataset_class):
		super().__init__(layers=None, name="Set Modeling Flow")
		self.model_params = model_params
		self.dataset_class = dataset_class
		self.set_size = self.model_params["set_size"]
		self.vocab_size = self.dataset_class.get_vocab_size(self.set_size)

		self._create_layers()
		self.print_overview()

	def _create_layers(self):

		self.latent_dim = self.model_params["categ_encoding"]["num_dimensions"]
		model_func = lambda c_out : CouplingTransformerNet(c_in=self.latent_dim,
														   c_out=c_out,
														   num_layers=self.model_params["coupling_hidden_layers"],
														   hidden_size=self.model_params["coupling_hidden_size"])
		self.model_params["categ_encoding"]["flow_config"]["model_func"] = model_func
		self.model_params["categ_encoding"]["flow_config"]["block_type"] = "Transformer"
		self.encoding_layer = create_encoding(self.model_params["categ_encoding"],
											  dataset_class=self.dataset_class,
											  vocab_size=self.vocab_size)

		num_flows = self.model_params["coupling_num_flows"]
		if self.latent_dim > 1:
			coupling_mask = CouplingLayer.create_channel_mask(self.latent_dim,
															  ratio=self.model_params["coupling_mask_ratio"])
			coupling_mask_func = lambda flow_index : coupling_mask
		else:
			coupling_mask = CouplingLayer.create_chess_mask()
			coupling_mask_func = lambda flow_index : coupling_mask if flow_index%2==0 else 1-coupling_mask

		layers = []
		for flow_index in range(num_flows):
			layers += [
				ActNormFlow(self.latent_dim),
				InvertibleConv(self.latent_dim),
				MixtureCDFCoupling(c_in=self.latent_dim,
								   mask=coupling_mask_func(flow_index),
								   model_func=model_func,
								   block_type="Transformer",
								   num_mixtures=self.model_params["coupling_num_mixtures"])
			]

		self.flow_layers = nn.ModuleList([self.encoding_layer] + layers)

	def forward(self, z, ldj=None, reverse=False, length=None, **kwargs):
		if length is not None:
			kwargs["src_key_padding_mask"] = create_transformer_mask(length)
			kwargs["channel_padding_mask"] = create_channel_mask(length)
		return super().forward(z, ldj=ldj, reverse=reverse, length=length, **kwargs)

	def get_inner_activations(self, z, length=None, return_names=False, **kwargs):
		if length is not None:
			kwargs["length"] = length
			kwargs["src_key_padding_mask"] = create_transformer_mask(length)
			kwargs["channel_padding_mask"] = create_channel_mask(length)

		out_per_layer = []
		layer_names = []
		for layer_index, layer in enumerate(self.flow_layers):
			z = self._run_layer(layer, z, reverse=False, **kwargs)[0]
			out_per_layer.append(z.detach())
			layer_names.append(layer.__class__.__name__)

		if not return_names:
			return out_per_layer
		else:
			return out_per_layer, layer_names

	def initialize_data_dependent(self, batch_list):
		# Batch list needs to consist of tuples: (z, kwargs)
		print("Initializing data dependent...")
		with torch.no_grad():
			for batch, kwargs in batch_list:
				kwargs["src_key_padding_mask"] = create_transformer_mask(kwargs["length"])
				kwargs["channel_padding_mask"] = create_channel_mask(kwargs["length"])
			for layer_index, layer in enumerate(self.flow_layers):
				batch_list = FlowModel.run_data_init_layer(batch_list, layer)

class CouplingTransformerNet(nn.Module):

	def __init__(self, c_in, c_out, num_layers, hidden_size):
		super().__init__()
		self.input_layer = nn.Sequential(
				nn.Linear(c_in, hidden_size),
				nn.GELU(),
				nn.Linear(hidden_size, hidden_size)
			)
		self.transformer_layers = nn.ModuleList([
				nn.TransformerEncoderLayer(hidden_size,
										   nhead=4,
										   dim_feedforward=2*hidden_size,
										   dropout=0.0,
										   activation='gelu') for _ in range(num_layers)
			])
		self.output_layer = nn.Sequential(
				nn.LayerNorm(hidden_size),
				nn.Linear(hidden_size, hidden_size),
				nn.GELU(),
				nn.Linear(hidden_size, c_out)
			)

	def forward(self, x, src_key_padding_mask, **kwargs):
		x = x.transpose(0, 1) # Transformer layer expects [Sequence length, Batch size, Hidden size]
		x = self.input_layer(x)
		for transformer in self.transformer_layers:
			x = transformer(x, src_key_padding_mask=src_key_padding_mask)
		x = self.output_layer(x)
		x = x.transpose(0, 1)
		return x