import pygame
from lib.AnimatedSprite import AnimatedSprite
from lib.AutoMovingAnimatedSprite import AutoMovingAnimatedSprite
from lib import invaderutils
from lib.GameObjectKeeper import GameObjectKeeper

class Hero(AnimatedSprite):
    MIS_STEPS = 0  # Number of waits before we trigger another missile
    mis_steps = 0
    MOVE_STEPS = 0 # Number of waits before we allow another move
    move_steps = 0

    def __init__(self, ix, iy, filenames, steps=1, move_steps=1, mis_steps=50):
        """What is steps, and msteps?? """

        super().__init__(ix, iy, filenames, steps)
        Hero.MIS_STEPS = mis_steps
        Hero.MOVE_STEPS = move_steps

        self.effect = pygame.mixer.Sound('sounds/thump.wav')
        self.effect.set_volume(0.03)

        self.count = len(filenames)-1

    def process_events(self, evt):
        if (self.mis_steps >= Hero.MIS_STEPS):
            if evt.type == pygame.KEYDOWN and evt.key == pygame.K_SPACE:
                self.fire_missile()

    def process_keys(self, pressed):
        """constantly being called"""

        self.mis_steps += 1
        self.move_steps += 1

        if (self.mis_steps >= Hero.MIS_STEPS):
            if pressed[pygame.K_UP]:
                self.fire_missile()

        if (self.move_steps >= Hero.MOVE_STEPS):
            self.move_steps = 0
            if pressed[pygame.K_LEFT]:
                self.x -= self.STEPS
            if pressed[pygame.K_RIGHT]:
                self.x += self.STEPS

        self.myclipx()

    def myclipx(self):
        """ Keeps the ship from hanging out at the edge where aliens can't
        shoot it. """
        if (self.x < 40):
            self.x = 0+40

        if (self.x+self.get_width() + 40 > invaderutils.get_game_width()):
            self.x = invaderutils.get_game_width() - self.get_width()-40

    def fire_missile(self):
        self.mis_steps = 0

        newx = int(self.x + self.get_width()/2)
        newy = self.y

        missile = AutoMovingAnimatedSprite(
            newx, newy,
            invaderutils.missile_png(), movex=0, movey=-4,
            msteps=1)
        GameObjectKeeper.setupfriendly(missile, 45)

    def draw(self, screen):
        image = self.images[self.count]
        screen.blit(image, (self.x, self.y))

        # self.steps += 1
        # if (self.steps >= self.STEPS):
        #     self.count = self.count + 1
        #     self.count = self.count % len(self.images)
        #     self.steps = 0

    def terminate(self):

        if self.count > 0:
            self.count -= 1
        else:
            super().terminate()


import sqlite3

from config import CONF
from tools.log import *

# 实例化日志对象
logger = setLogger('databases')

def db(sql=None):
    """
    数据库操作
    :param sql:
    :return:
    """
    # 挂载链接
    connect = connect_database()
    logger.info(connect)
    # 传入查询时
    if sql:
        # 执行SQL
        result = do_query(connect, sql)
        connect.close()
        return result
    # 关闭链接
    connect.close()

def connect_database():
    """
    挂载数据库
    :return: sqlite3.Connection
    """
    # 挂载链接
    connect = sqlite3.connect(CONF.DB_PATH)

    # 强制启动外键约束
    if CONF.USE_FOREIGN:
        connect.execute("PRAGMA foreign_keys=1;")

    if CONF.CACHE_SIZE:
        sql = "PRAGMA cache_size=" + str(CONF.CACHE_SIZE) + ";"
        connect.execute(sql)

    return connect

def do_query(connect, sql):
    """
    执行SQL
    :type connect: sqlite3.Connection
    :type sql: string
    """
    sql = sql.strip()

    if sql.lstrip().upper().startswith("SELECT"):
        # 创建游标
        cursor = connect.cursor()
        # 查询
        result = select(cursor, sql)
        # 关闭游标
        cursor.close()
        return result
    elif sql.lstrip().upper().startswith("INSERT"):
        result = insert(connect, sql)
        return result
    elif sql.lstrip().upper().startswith("UPDATE"):
        result = update(connect, sql)
        return result
    elif sql.lstrip().upper().startswith("DELETE"):
        result = delete(connect, sql)
        return result
    else:
        return 0

def select(cursor, sql):
    """
    SELECT语句
    :param sql:
    :param cursor:
    :return:
    """
    # 执行SQL
    cursor.execute(sql)
    # 获取结果
    result = cursor.fetchall()

    return result

def insert(connect, sql):
    """
    INSERT语句
    :param sql:
    :param connect:
    :type connect: sqlite3.Connection
    :return:
    """
    # 执行SQL
    try:
        connect.execute(sql)
        connect.commit()
        result = 1
    except sqlite3.DatabaseError:
        result = 0
        logger.error(sqlite3.DatabaseError)
        print(sqlite3.DatabaseError)
    except sqlite3.Error:
        result = 0
        logger.error(sqlite3.Error)
        print(sqlite3.Error)

    return result

def update(connect, sql):
    """
    UPDATE语句
    :param connect:
    :param sql:
    :type connect: sqlite3.Connection
    :return:
    """
    # 执行SQL
    try:
        connect.execute(sql)
        connect.commit()
        result = 1
    except sqlite3.DatabaseError:
        result = 0
        logger.error(sqlite3.DatabaseError)
        print(sqlite3.DatabaseError)
    except sqlite3.Error:
        result = 0
        logger.error(sqlite3.Error)
        print(sqlite3.Error)

    return result

def delete(connect, sql):
    """
    DELETE语句
    :param connect:
    :param sql:
    :type connect: sqlite3.Connection
    :return:
    """
    # 执行SQL
    try:
        connect.execute(sql)
        connect.commit()
        result = 1
    except sqlite3.DatabaseError:
        result = 0
        logger.error(sqlite3.DatabaseError)
        print(sqlite3.DatabaseError)
    except sqlite3.Error:
        result = 0
        logger.error(sqlite3.Error)
        print(sqlite3.Error)

    return result

# if __name__ == '__main__':
#     # pass
#     sql = """
#     UPDATE bug SET tester_id=1, developer_id=1, phase_id=1, bug_type=1, category=8, title='作品展示页翻页后点击进入作品详情页后，点击退回上个页面保持页面筛选状态，即维持在进入此作品详情页的作品类表页数',
#     model='前端-作品展示区', create_time=1611147000, close_time=0, is_finished='False', is_closed='False', is_online='False' WHERE kb_id=71892;
#     """
#     db(sql)


# -*- coding: utf-8 -*-
# __author__ = 'zhouxuan'

from biocluster.api.database.base import Base, report_check
import re
import datetime
from bson import ObjectId, SON
from bson.objectid import ObjectId
from types import StringTypes
# from biocluster.config import Config

class MetaSourcetracker(Base):
    """
    微生物来源组成比例分析
    """
    def __init__(self, bind_object):
        super(MetaSourcetracker, self).__init__(bind_object) #
        self._project_type = 'meta'
        # self._db_name = Config().MONGODB
        self.task_id = ""

    @report_check
    def add_sg_sourcetracker_detail(self, id=None, file_path=None, stdev_file_path=None, name_1=None, name_2=None):
        insert_data = list()
        d_dict = dict()
        def addtwodimdict(thedict, key_a, key_b, val):
            if key_a in thedict:
                thedict[key_a].update({key_b: val})
            else:
                thedict.update({key_a: {key_b: val}})
        with open(stdev_file_path, "rb") as s:
            head_s = s.next().strip('\r\n')
            head_s = re.split('\t', head_s)
            source_name = head_s[1:]
            for line in s:
                line = line.rstrip("\r\n")
                line = re.split('\t', line)
                stdev_number = line[1:]
                for i in range(0, len(source_name)):
                    addtwodimdict(d_dict, line[0], source_name[i], float(stdev_number[i]))
        print(d_dict)
        with open(file_path, 'rb') as r:
            head = r.next().strip('\r\n')  # windows换行符
            head = re.split('\t', head)
            new_head = head[1:]
            for line in r:
                line = line.rstrip("\r\n")
                line = re.split('\t', line)
                group_num = line[1:]
                detail = dict()
                detail['sourcetracker_id'] = ObjectId(id)
                detail['file_name_1'] = name_1
                detail['file_name_2'] = name_2
                detail['sample_name'] = line[0]
                for i in range(0, len(group_num)):
                    # detail[new_head[i]] = float(group_num[i])
                    detail[new_head[i] + "_mean"] = float(group_num[i])
                    detail[new_head[i] + "_stdev"] = float(d_dict[line[0]][new_head[i]])
                    detail[new_head[i] + "_h"] = float(group_num[i]) + (d_dict[line[0]][new_head[i]])
                    detail[new_head[i] + "_l"] = float(group_num[i]) - (d_dict[line[0]][new_head[i]])
                    # detail[new_head[i] + "_stdev_plot"] = float(d_dict[line[0]][new_head[i]]*d_dict[line[0]][new_head[i]])
                    # detail[new_head[i] + "_h"] = float(group_num[i]) + (d_dict[line[0]][new_head[i]]*d_dict[line[0]][new_head[i]])
                    # detail[new_head[i] + "_l"] = float(group_num[i]) - (d_dict[line[0]][new_head[i]]*d_dict[line[0]][new_head[i]])
                insert_data.append(detail)
        try:
            collection = self.db['sg_sourcetracker_detail']
            collection.insert_many(insert_data)
        except Exception as e:
            self.bind_object.logger.error("导入sg_sourcetracker_detail表格信息出错:{}".format(e))
            self.bind_object.set_error("导入sg_sourcetracker_detail表格信息出错", code="51003801")
        else:
            self.bind_object.logger.info("导入sg_sourcetracker_detail表格成功")


# -*- coding: utf-8 -*-
"""
Created on Fri Oct 11 14:40:05 2019

@author: Fede
"""
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import numpy as np
import matplotlib.pyplot as plt
import os

def plot_results(vector_mean,vector_mean1,train_episode,negative_r1,negative_r2,pos_r,pos_r2,null_r,loss_mean,title,LOSSES,success_episodes):
    bound=[]
    for i in range(int(train_episode/100)):
        bound.append(1)
    print (bound)
    episodes=[]
    for i in range(int(train_episode/10)):
        episodes.append(i*100)
    base=[]
    base1=[]
    base2=[]
    base3=[]
    base4=[]

    for i in range(int(len(negative_r1))):
        base.append(negative_r1[i]*100/(negative_r1[i]+negative_r2[i]+pos_r[i]+pos_r2[i]+null_r[i]))
        base1.append(negative_r2[i]*100/(negative_r1[i]+negative_r2[i]+pos_r[i]+pos_r2[i]+null_r[i]))
        base2.append(pos_r[i]*100/(negative_r1[i]+negative_r2[i]+pos_r[i]+pos_r2[i]+null_r[i]))
        base3.append(pos_r2[i]*100/(negative_r1[i]+negative_r2[i]+pos_r[i]+pos_r2[i]+null_r[i]))
        base4.append(null_r[i]*100/(negative_r1[i]+negative_r2[i]+pos_r[i]+pos_r2[i]+null_r[i]))

    fig = make_subplots(rows=2,cols=2,subplot_titles=("Rewards distributions","Times both drones reach their targets [%]","Mean reward per episode per step","Mean loss"))

    fig.add_trace(go.Bar(x=episodes, y=base,
                        marker_color='yellow',
                        name='Same position'),row=1,col=1)
    fig.add_trace(go.Bar(x=episodes, y=base2,
                        marker_color='darkgreen',
                        name='Optimal position'
                        ),row=1,col=1)
    fig.add_trace(go.Bar(x=episodes, y=base3,
                        marker_color='lightgreen',
                        name='Only 1 drone correct'
                        ),row=1,col=1)
    fig.add_trace(go.Bar(x=episodes, y=base1,
                        marker_color='crimson',
                        name='Outside map position'
                        ),row=1,col=1)
    fig.add_trace(go.Bar(x=episodes, y=base4,
                        marker_color='grey',
                        name='Null rewards'
                        ),row=1,col=1)

    min_v=[]
    max_v=[]
    mean_v=[]
    min_v1=[]
    max_v1=[]
    mean_v1=[]
    min_loss=[]
    max_loss=[]
    mean_loss=[]

    for i in range(len(vector_mean[0])):
        temp=[]
        temp1=[]
        for j in range(len(vector_mean)):
            temp.append(vector_mean[j][i])
            temp1.append(vector_mean1[j][i])
        min_v.append(min(temp))
        max_v.append(max(temp))
        mean_v.append(np.mean(temp))
        min_v1.append(np.mean(temp1)-min(temp1))
        max_v1.append(max(temp1)-np.mean(temp1))
        mean_v1.append(np.mean(temp1))
    for i in range(len(loss_mean[0])):
        temp_loss=[]
        for j in range(len(loss_mean)):
            temp_loss.append(loss_mean[j][i])
        min_loss.append(min(temp_loss))
        max_loss.append(max(temp_loss))
        mean_loss.append(np.mean(temp_loss))

    fig.add_trace(go.Bar(x=episodes, y=mean_v1,
                        marker_color='orange',error_y=dict(
                                            symmetric=False,
                                            array=max_v1,
                                            arrayminus=min_v1),
                        name='Perfect episodes [%]'
                        ),row=1,col=2)

    fig.add_trace(go.Scatter(y=np.asarray(min_v),line=dict(width=0),mode='lines',name='Min reward'),row=2,col=1)
    fig.add_trace(go.Scatter(y=np.asarray(mean_v),line=dict(color='rgb(31, 119, 180)'),
                             fillcolor='rgba(68, 68, 68, 0.3)',fill='tonexty',name='Mean reward'),row=2,col=1)
    fig.add_trace(go.Scatter(y=np.asarray(max_v),line=dict(width=0),fillcolor='rgba(68, 68, 68, 0.3)',
                             fill='tonexty',name='Upper reward'),row=2,col=1)
    fig.add_trace(go.Scatter(y=bound,line=go.scatter.Line(color="crimson"),name='Bound reward'),
                  row=2,col=1)

    fig.add_trace(go.Scatter(y=np.asarray(min_loss),line=dict(width=0),mode='lines',name='Min reward'),row=2,col=2)
    fig.add_trace(go.Scatter(y=np.asarray(mean_loss),line=dict(color='rgb(68, 0, 0)'),
                             fillcolor='rgba(68, 0, 0, 0.3)',fill='tonexty',name='Mean reward'),row=2,col=2)
    fig.add_trace(go.Scatter(y=np.asarray(max_loss),line=dict(width=0),fillcolor='rgba(68, 0, 0, 0.3)',
                             fill='tonexty',name='Upper reward'),row=2,col=2)

    fig.update_xaxes(tickfont=dict(size=30))
    fig.update_yaxes(tickfont=dict(size=30))
    fig.update_layout(barmode='relative',
                          title=go.layout.Title(
                          text="Results"),
                          title_font_size=30,
                          legend=go.layout.Legend(
                          font=dict(
                          family="sans-serif",
                          size=30,
                          color="black"
                          )),
                          width=1920,height=1080)

    fig.write_image('General_Results'+str(title)+'.pdf')

    fig = go.Figure(data=[
            go.Bar(x=episodes, y=base1,
                        marker_color='crimson',
                        name='Outside map position'),
            go.Bar(x=episodes, y=base,
                        marker_color='yellow',
                        name='Same position'),
            go.Bar(x=episodes, y=base4,
                        marker_color='grey',
                        name='Null rewards'
                        ),
            go.Bar(x=episodes, y=base3,
                        marker_color='lightgreen',
                        name='Only 1 drone correct'),
            go.Bar(x=episodes, y=base2,
                        marker_color='darkgreen',
                        name='Optimal position')
                   ])

    fig.update_layout(barmode='stack',
                          title=go.layout.Title(
                          text="Greedy rewards"),
                          yaxis=dict(title='Reward distributions (%)'),
                          xaxis=dict(title='Trained episodes',tick0 = 0,dtick = 50))
    fig.write_image('Greedy_rewards_'+str(title)+'.pdf')

    fig = go.Figure(data=[
            go.Bar(x=episodes, y=mean_v1,
                        marker_color='orange',error_y=dict(
                                            symmetric=False,
                                            array=max_v1,
                                            arrayminus=min_v1),
                        name='Perfect episodes [%]'
                        )])

    fig.update_layout(barmode='stack',
                          title=go.layout.Title(
                          text="Correct episodes"),
                          yaxis=dict(title='Correct episodes (%)'),
                          xaxis=dict(title='Trained episodes',tick0 = 0,dtick = 50))
    fig.write_image('Correct_episodes_'+str(title)+'.pdf')

    fig = go.Figure(data=[
            go.Scatter(y=np.asarray(min_loss),line=dict(width=0.5),mode='lines',name='Min reward'),
            go.Scatter(y=np.asarray(mean_loss),line=dict(color='rgb(68, 0, 0)'),
                             fillcolor='rgba(68, 0, 0, 0.3)',fill='tonexty',name='Mean reward'),
            go.Scatter(y=np.asarray(max_loss),line=dict(width=0.5),fillcolor='rgba(68, 0, 0, 0.3)',
                             fill='tonexty',name='Upper reward')])

    fig.update_layout(title=go.layout.Title(
                          text="Loss"),
                          yaxis=dict(title='MSE'),
                          xaxis=dict(title='Times NN trained',tick0 = 0,dtick = 50))
    fig.write_image('Loss_'+str(title)+'.pdf')

    fig = go.Figure(data=[
            go.Scatter(y=np.asarray(min_v),line=dict(width=0.5),mode='lines',name='Min reward'),
            go.Scatter(y=np.asarray(mean_v),line=dict(color='rgb(31, 119, 180)'),
                             fillcolor='rgba(68, 68, 68, 0.3)',fill='tonexty',name='Mean reward'),
            go.Scatter(y=np.asarray(max_v),line=dict(width=0.5),fillcolor='rgba(68, 68, 68, 0.3)',
                             fill='tonexty',name='Upper reward'),
            go.Scatter(y=bound,line=go.scatter.Line(color="crimson"),name='Bound reward'),
                  ])

    fig.update_layout(barmode='stack',
                          title=go.layout.Title(
                          text="Mean reward per episodes per step"),
                          yaxis=dict(title='Reward'),
                          xaxis=dict(title='Trained episodes',tick0 = 0,dtick = 50))
    fig.write_image('Mean_reward_'+str(title)+'.pdf')

def plot_memory_replay(SMR,GSMR,TMR,GTMR,Map,title):
    fig = make_subplots(rows=2,cols=2,
                        subplot_titles=("Pre-training","Available gaussians in the map in the initial MR","Available gaussians in the map in the end in the MR","Map heterogeneity"),
                        specs=[[{"type": "domain"}, {"type": "domain"}],
                               [{"type": "domain"},{"type": "heatmap"}]])
    labels=["Outside map","Same position","0 reward", "1 Correct","All correct"]
    labels1=["0 available gaussians","1 available gaussians","2 available gaussians"]
    total=[0,0,0]
    for i in range(len(GTMR)):
        total[0]+=GTMR[i][0]
        total[1]+=GTMR[i][1]
        total[2]+=GTMR[i][2]
    fig.add_trace(go.Pie(values=GSMR[0],labels=labels),row=1,col=1)
    fig.add_trace(go.Pie(values=GTMR[0],labels=labels1),row=1,col=2)
    fig.add_trace(go.Pie(values=total,labels=labels1),row=2,col=1)
    maximum=max(max(Map[0]),max(Map[1]),max(Map[2]),max(Map[3]),max(Map[4]),max(Map[5]),max(Map[6]))
    fig.add_trace(go.Heatmap(z=[Map[0],Map[1],Map[2],Map[3],Map[4],Map[5],Map[6]],zmax=0,zmid=maximum/2,zmin=maximum,reversescale=False,showscale=False,colorscale='jet'),row=2,col=2)
    fig.update_xaxes(tickfont=dict(size=30))
    fig.update_yaxes(tickfont=dict(size=30))
    fig.update_layout(barmode='stack',
                          title=go.layout.Title(
                          text="Dataset"),
                          title_font_size=30,
                          legend=go.layout.Legend(
                          font=dict(
                          family="sans-serif",
                          size=30,
                          color="black"
                          )),
                          width=1920,height=1080)

    fig.write_image('Dataset'+str(title)+'.pdf')

def plot_qvalues(VAR,MIN,MAX,MEAN,title):
    fig = make_subplots(rows=2,cols=2,
                        subplot_titles=("Mean qvalues variance",
                                        "Mean qvalues minimum",
                                        "Mean qvalues maximum",
                                        "Mean qvalues mean"),
                        specs=[[{"type": "xy"}, {"type": "xy"}],
                               [{"type": "xy"},{"type": "xy"}]])
    VAR1=[]
    MEAN1=[]
    MINIMUM=[]
    MAXIMUM=[]
    for i in range(len(VAR)):
        VAR1.append(VAR[i][0])
        MEAN1.append(MEAN[i][0])
        MINIMUM.append(MIN[i][0])
        MAXIMUM.append(MAX[i][0])
    fig.add_trace(go.Scatter(y=VAR1,line=go.scatter.Line(color="red"),name='Variance'),row=1,col=1)
    fig.add_trace(go.Scatter(y=MINIMUM,line=go.scatter.Line(color="orange"),name='Minimum'),row=1,col=2)
    fig.add_trace(go.Scatter(y=MAXIMUM,line=go.scatter.Line(color="green"),name='Maximum'),row=2,col=1)
    fig.add_trace(go.Scatter(y=MEAN1,line=go.scatter.Line(color="darkgreen"),name='Mean'),row=2,col=2)
    fig.update_xaxes(tickfont=dict(size=30))
    fig.update_yaxes(tickfont=dict(size=30))
    fig.update_layout(title=go.layout.Title(
                          text="Dataset"),
                          title_font_size=30,
                          legend=go.layout.Legend(
                          font=dict(
                          family="sans-serif",
                          size=30,
                          color="black"
                          )),
                          width=1920,height=1080)

    fig.write_image('Qvalues'+str(title)+'.pdf')

def plot_IAHOS(y,ogp,ogp2,tgp,tgp2,model):
    """fig = make_subplots(rows=2, cols=2,subplot_titles=("Mean train. accuracy first round",
                                                   "Mean valid accuracy first round",
                                                  "Mean train accuracy last round",
                                                  "Mean valid accuracy last round"))"""

    fig = go.Figure()
    x = np.linspace(0,len(tgp[0])-1,len(tgp[0]))
    Colorscale = [[0, '#FF0000'],[0.5, '#F1C40F'], [1, '#00FF00']]
    """fig.add_trace(go.Heatmap(y=[y[i] for i in range(len(y))],
                       x=[0,1],
                       z=ogp2, colorscale = Colorscale),row=1,col=1)
    fig.add_trace(go.Heatmap(y=[y[i] for i in range(len(y))],
                       x=[0,1],
                       z=ogp,colorscale=Colorscale),row=1,col=2)
    fig.add_trace(go.Heatmap(y=[y[i] for i in range(len(y))],
                       x=x,
                       z=tgp2, colorscale = Colorscale),row=2,col=1)"""
    fig.add_trace(go.Heatmap(y=[y[i] for i in range(len(y))],
                       x=x,
                       z=tgp,colorscale=Colorscale))
    fig.update_layout(height=600, width=800,title=dict(text='IAHOS results'))
    if not os.path.exists("images"):
        os.mkdir("images")
    fig.write_image('images/IAHOS_'+str(model)+'.pdf')

def plot_pretrain(performances,y,params,model):
    fig = make_subplots(rows=2, cols=2,subplot_titles=("All performances",
                                                   "Performances w.r.t. Gamma",
                                                  "Performances w.r.t. malus same position",
                                                  ""))
    indeces=[]
    for i in range(len(performances)):
        indeces.append(performances[i][-1])
    index1=np.argmax(indeces)
    index2=np.argmin(indeces)

    bound=[]
    for i in range(len(performances[0])):
        bound.append(1)
    for i in range(len(performances)):
        if i!=index1 and i!=index2:
            fig.add_trace(go.Scatter(y=performances[i],line=go.scatter.Line(color="grey"),showlegend=False),row=1,col=1)
    fig.add_trace(go.Scatter(y=performances[index1],line=go.scatter.Line(color="green",width=4),name='Best performance'),row=1,col=1)
    fig.add_trace(go.Scatter(y=performances[index2],line=go.scatter.Line(color="red",width=4),name='Worst performance'),row=1,col=1)
    fig.add_trace(go.Scatter(y=bound,line=go.scatter.Line(color="green",dash='dashdot',width=4),name='Bound'),row=1,col=1)
    fig.add_trace(go.Bar(x=params[0], y=y[0]),row=1,col=2)
    fig.add_trace(go.Bar(x=params[1], y=y[1]),row=2,col=1)
    #fig.add_trace(go.Bar(x=params[2], y=y[2]),row=2,col=2)

    fig.update_layout(title=go.layout.Title(text="Results"),title_font_size=30,width=1920,height=1080)

    fig.write_image('Results_pretraining_'+str(model)+'.pdf')

def plot_training(VARIATIONS):
    VAR1=[]
    VAR2=[]
    VAR3=[]
    VAR4=[]
    for i in range(len(VARIATIONS)):
        VAR1.append(VARIATIONS[i][0]*100)
        VAR2.append(VARIATIONS[i][1]*100)
        VAR3.append(VARIATIONS[i][2]*100)
        VAR4.append(VARIATIONS[i][3]*100)
    fig = make_subplots(rows=2, cols=2,subplot_titles=("Outside map","Same position","Null rewards","+1 reward"))
    fig.add_trace(go.Scatter(y=VAR1,marker_color='crimson',name='Outside map position'),row=1,col=1)

    fig.add_trace(go.Scatter(y=VAR2,marker_color='orange', name='Same position'),row=1,col=2)
    fig.add_trace(go.Scatter(y=VAR3,marker_color='grey',name='Null rewards'),row=2,col=1)
    fig.add_trace(go.Scatter(y=VAR4,marker_color='lightgreen',name='+1 reward'),row=2,col=2)

    fig.update_layout(title=go.layout.Title(
                          text="Argmax q-values variations during training"),
                          yaxis=dict(title='Variation (%)'),
                          xaxis=dict(title='Trained episodes',tick0 = 0,dtick = 50)
                          ,width=1920,height=1080)
    fig.write_image('Argmax q-values variations during training.pdf')


"""
Pollard rho based hash mining algorithm.
"""
from random import randint
from typing import Optional
import time
from blockchain.block import Block
from crypto.pollard_rho import func_g, func_h
from mining.pollard_rho_solution import PRSolution

class SimplePRMiner:
    """A pollard rho miner uses textbook mapping functions to search private key."""
    def __init__(self, block: Block, block_time=0):
        self.block = block
        self.block_time = block_time

    def header_hash(self, nonce: int) -> int:
        self.block.nonce = nonce
        # if the block is restored from database or json, the full Tx info may missing
        # WARNING: this part of code is insecure and it is not based on original design but only for
        # simplification of code.
        header_hash = self.block.current_block_hash if self.block.current_block_hash else self.block.hash_header()
        hash_val = int.from_bytes(header_hash, byteorder='little', signed=True) % self.block.public_key.p
        return hash_val

    def func_f(self, hash_i, y_i):
        """x_(i+1) = func_f(hash(x_i))"""
        base = self.block.public_key.g
        h = self.block.public_key.h
        p = self.block.public_key.p
        if hash_i % 3 == 2:
            return (h * y_i) % p
        elif hash_i % 3 == 0:
            return pow(y_i, 2, p)
        elif hash_i % 3 == 1:
            return base * y_i % p
        else:
            raise ValueError("Input value has error!")

    def func_g(self, a, n, hash_i):
        """Wrapper of classic mapping function for a."""
        p = self.block.public_key.p
        return func_g(a, n, p, hash_i)

    def func_h(self, b, n, hash_i):
        """Wrapper of classic mapping function for b."""
        p = self.block.public_key.p
        return func_h(b, n, p, hash_i)

    @staticmethod
    def _calculate_y(a, b, g, h, p):
        return (pow(g, a, p) * pow(h, b, p)) % p

    def mining(self) -> tuple[int, Optional[PRSolution]]:
        """
        Refer to section 3.6.3 of Handbook of Applied Cryptography
        Computes `x` = a mod n for the DLP base**x % p == y
        in the Group G = {0, 1, 2, ..., n}
        given that order `n` is a prime number.

        Returns:
            nonce and paired private key
        """
        pubkey = self.block.public_key
        n = (pubkey.p - 1) // 2

        a_i = randint(0, n)
        b_i = randint(0, n)
        a_2i = a_i
        b_2i = b_i

        y_i = self._calculate_y(a_i, b_i, pubkey.g, pubkey.h, pubkey.p)
        y_2i = y_i

        i = 1
        init_time = time.time()
        while i <= n and (time.time() - init_time) < self.block_time:
            # Single Step calculations
            hash_1i = self.header_hash(y_i)
            a_i = self.func_g(a_i, n, hash_1i)
            b_i = self.func_h(b_i, n, hash_1i)
            y_i = self.func_f(hash_1i, y_i)
            # print(f"y: {y_i}, a: {a_i}, b: {b_i}")
            # assert(y_i == self._calculate_y(a_i, b_i, pubkey.g, pubkey.h, pubkey.p))

            # Double Step calculations
            hash_2mi = self.header_hash(y_2i)
            a_2i = self.func_g(a_2i, n, hash_2mi)
            b_2i = self.func_h(b_2i, n, hash_2mi)
            y_2mi = self.func_f(hash_2mi, y_2i)
            hash_2i = self.header_hash(y_2mi)
            a_2i = self.func_g(a_2i, n, hash_2i)
            b_2i = self.func_h(b_2i, n, hash_2i)
            y_2i = self.func_f(hash_2i, y_2mi)
            # assert (y_2i == self._calculate_y(a_2i, b_2i, pubkey.g, pubkey.h, pubkey.p))

            if y_i == y_2i:
                # print("left side = ", (pow(pubkey.g, a_i, pubkey.p) * pow(pubkey.h, b_i, pubkey.p)) % pubkey.p)
                # print("right side = ", (pow(pubkey.g, a_2i, pubkey.p) * pow(pubkey.h, b_2i, pubkey.p)) % pubkey.p)
                solution = PRSolution(a_i, a_2i, b_i, b_2i, n)
                solution.pubkey = pubkey
                return y_2mi, solution
            else:
                i += 1
                continue
        # failed to find the solution and nonce, return none
        print("Cannot seal a new block within block time! Try again")
        return 0, None

class PRMiner(SimplePRMiner):
    """The pollard rho miner with a mapping function which is hard to compute reversely."""
    def func_f(self, hash_i, y_i):
        base = self.block.public_key.g
        h = self.block.public_key.h
        p = self.block.public_key.p
        if hash_i % 3 == 2:
            return (pow(h, hash_i, p) * y_i) % p
        elif hash_i % 3 == 0:
            return pow(y_i, hash_i, p)
        elif hash_i % 3 == 1:
            return (pow(base, hash_i, p) * y_i) % p
        else:
            raise ValueError("Input value has error!")

    def func_g(self, a, n, hash_i):
        p = self.block.public_key.p
        if hash_i % 3 == 2:
            return a
        elif hash_i % 3 == 0:
            return a * hash_i % (p-1)
        elif hash_i % 3 == 1:
            return (a + hash_i) % (p-1)
        else:
            raise ValueError("Input value has error!")

    def func_h(self, b, n, hash_i):
        p = self.block.public_key.p
        if hash_i % 3 == 2:
            return (b + hash_i) % (p-1)
        elif hash_i % 3 == 0:
            return b * hash_i % (p-1)
        elif hash_i % 3 == 1:
            return b
        else:
            raise ValueError("Input value has error!")


# all model classes

import joblib
import numpy as np
import os
import random
import time
import datetime

from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (precision_score, recall_score, f1_score, roc_curve)
from sklearn.model_selection import GridSearchCV
from sklearn.svm import SVC
from xgboost import XGBClassifier

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
from transformers import BertForSequenceClassification, AdamW, get_linear_schedule_with_warmup

MODEL_DIR = 'models/'
MODEL_DIR_MAP = {
    "lr": 'logistic_regression',
    "svm": 'svm',
    "rf": 'random_forest',
    "xgb": 'xgboost',
    "bert": 'bert',
    "cnn": '1d_cnn'
}

class Classifier():

    def __init__(self, model_name, random_state):
        super().__init__()
        self.model_name = model_name
        self.model = None
        self.model_params = None
        self.random_state = random_state
        self.batch_size = None

    def fit(self, X_train, Y_train, X_test=None, word_embeddings=None, tune=False, save=True):
        if self.model_name not in ['bert', 'cnn']:
            model_file = os.path.join(MODEL_DIR + MODEL_DIR_MAP[self.model_name], 'tfidf_0.005/clf.joblib')
        else:
            model_file = os.path.join(MODEL_DIR + MODEL_DIR_MAP[self.model_name], 'saved_model')

        if self.model_name == 'lr':
            if tune:
                lr = LogisticRegression(random_state=self.random_state)
                parameters = {
                    "penalty": ['l2'],
                    "C": [0.5, 0.7, 0.8, 1.0],
                    }

                clf = GridSearchCV(lr, parameters, scoring='f1_macro', cv=5, n_jobs=-1, verbose=2)
                clf.fit(X_train, Y_train)
                if save:
                    joblib.dump(clf, model_file)

            else:
                clf = joblib.load(model_file)

            self.model = clf.best_estimator_
            self.model_params = clf.best_params_

        elif self.model_name == 'rf':
            if tune:
                rf = RandomForestClassifier(random_state=self.random_state)
                parameters = {
                    "n_estimators": [100, 200, 500],
                    "max_depth": [5, 7, 11, 17, 19, 23],
                    "max_features": ['auto', 'sqrt', 'log2'],
                    "criterion": ['gini', 'entropy']
                    }

                clf = GridSearchCV(rf, parameters, scoring='f1_macro', cv=5, n_jobs=-1, verbose=2)
                clf.fit(X_train, Y_train)
                if save:
                    joblib.dump(clf, model_file)

            else:
                clf = joblib.load(model_file)

            self.model = clf.best_estimator_
            self.model_params = clf.best_params_

        elif self.model_name == 'svm':
            if tune:
                svc = SVC(random_state=self.random_state)
                parameters = {
                    "C": [0.5, 0.8, 1.0],
                    "kernel": ['poly', 'rbf', 'sigmoid'],
                    "degree": [3, 5, 7, 9],
                    "gamma": ['scale', 'auto']
                    }

                clf = GridSearchCV(svc, parameters, scoring='f1_macro', cv=5, n_jobs=-1, verbose=2)
                clf.fit(X_train, Y_train)
                if save:
                    joblib.dump(clf, model_file)

            else:
                clf = joblib.load(model_file)

            self.model = clf.best_estimator_
            self.model_params = clf.best_params_

        elif self.model_name == 'xgb':
            if tune:
                xgb_classifier = XGBClassifier(random_state=self.random_state, n_jobs=-1, use_label_encoder=False)
                parameters = {
                    "n_estimators": [100, 200, 500],
                    "max_depth": [5, 9, 13, 23, 47],
                    "subsample": [0.8, 1.0],
                    "colsample_bytree": [0.5, 0.8, 1.0],
                }

                clf = GridSearchCV(xgb_classifier, parameters, scoring='f1_macro', cv=5, n_jobs=-1, verbose=2)
                clf.fit(X_train, Y_train)
                if save:
                    joblib.dump(clf, model_file)

            else:
                clf = joblib.load(model_file)

            self.model = clf.best_estimator_
            self.model_params = clf.best_params_

        elif self.model_name == 'cnn':
            model_file = os.path.join(model_file, 'model.pt')
            self.model = Classifier_1D_CNN(
                    embedding=word_embeddings,
                    dropout_size=0.2,
                    n_grams=[2, 3, 4],
                    out_channels=[50, 50, 50],
                    output_size_dense_layer=10,
                    classes=3,
                    num_epochs=30, lr=0.001, batch_size=32)
            if tune:
                self.model.fit(X_train, Y_train)

                if save:
                    torch.save(self.model.state_dict(), model_file)
            else:
                self.model.load_state_dict(torch.load(model_file))

        elif self.model_name == 'bert':
            self.model = CustomBERT()
            self.batch_size = 32
            if tune:
                self.model.finetune(
                    train_dataset=X_train, test_dataset=X_test, num_epochs=4, batch_size=self.batch_size, random_state=self.random_state)

                if save:
                    self.model.save(model_file)
            else:
                self.model.load(model_file)

        else:
            raise ValueError(f'Unsupported model type: {self.model_name}')

    def predict(self, X_test):

        if self.model_name != 'bert':
            y_pred = self.model.predict(X_test)
        else:
            y_pred = self.model.predict(X_test, batch_size=self.batch_size)

        return y_pred

    def interpret_results(self):
        pass

    def classification_results(self, true, pred, pred_prob=None, result_label='test'):

        result_dict = {"acc": (true==pred).mean(),
                    "precision": precision_score(true, pred, average='macro'),
                    "recall": recall_score(true, pred, average='macro'),
                    "f1": f1_score(true, pred, average='macro')}

        if pred_prob is not None:
            result_dict.update({"roc_auc": get_roc_auc(true, pred_prob)})

        return pd.DataFrame(result_dict, index=[result_label]).T

class Classifier_1D_CNN(nn.Module):

    def __init__(self, embedding, n_grams, out_channels, output_size_dense_layer, classes, batch_size, lr, num_epochs, dropout_size):

        super(Classifier_1D_CNN, self).__init__()

        self.lr = lr
        self.batch_size = batch_size
        self.num_epochs = num_epochs
        self.vocab_size, self.embed_dim = embedding.shape
        self.embedding = nn.Embedding.from_pretrained(embedding, freeze=True)
        self.out_channels = out_channels
        self.n_grams = n_grams
        self.conv1d_list = nn.ModuleList([
                                    nn.Conv1d(
                                        in_channels=self.embed_dim,
                                        out_channels=out_channels[i],
                                        kernel_size=n_grams[i])
                                    for i in range(len(self.n_grams))
                                    ]
                                    )

        self.dropout = nn.Dropout(dropout_size)
        self.dense1 = nn.Linear(np.sum(out_channels), output_size_dense_layer)
        self.output_dense = nn.Linear(output_size_dense_layer, classes)

    def forward(self, x):

        x = self.embedding(x).float().permute(0,2,1)
        x_conv_list = [F.relu(layer(x)) for layer in self.conv1d_list]

        x_pool_list = [F.max_pool1d(x_conv, kernel_size=x_conv.shape[2]) for x_conv in x_conv_list]

        x = torch.cat([x_pool.squeeze(dim=2) for x_pool in x_pool_list], dim=1)

        x = self.dense1(x)
        x = self.dropout(x)
        output = self.output_dense(x)
        output = F.softmax(output)

        return output

    def fit(self, train_data, train_labels):

        dataset = TensorDataset(train_data, torch.Tensor(train_labels))
        dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)

        loss_fn = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)

        self.train()

        for epoch in range(self.num_epochs):
            for id_batch, (x_batch, y_batch) in enumerate(dataloader):
                optimizer.zero_grad()
                y_pred = self(x_batch)
                loss = loss_fn(y_pred, y_batch)
                loss.backward()
                optimizer.step()

            print(f"Epoch {epoch}: training loss: {loss.item()}")

    def predict(self, test_data, prob=False):

        self.eval()

        pred = self(test_data).detach()

        if not prob:
            pred = torch.max(pred, dim=1)[-1]

        return pred.numpy()

class Classifier_NN(nn.Module):

    def __init__(self, input_size, output_size, hidden_layer_dims, dropouts):
        super().__init__()

        self.input_size = input_size
        self.output_size = output_size

        self.all_layer_dims = [self.input_size] + hidden_layer_dims + [self.output_size]
        self.linear_layers = nn.ModuleList([nn.Linear(self.all_layer_dims[i], self.all_layer_dims[i+1]) for i in range(len(self.all_layer_dims)-1)])

        self.dropouts = dropouts
        self.classes_ = None

    def forward(self, x):
        for layer in range(len(self.linear_layers)-1):
            x = F.relu(F.dropout(self.linear_layers[layer](x), p=self.dropouts[layer]))

        x = self.linear_layers[-1](x)
        x = F.softmax(x)

        return x

    def fit(self, train_data, train_labels, batch_size, lr, num_epochs):

        dataset = TensorDataset(torch.Tensor(train_data), torch.Tensor(train_labels))
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

        loss_fn = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(self.parameters(), lr=lr)

        self.train()

        for epoch in range(num_epochs):
            for id_batch, (x_batch, y_batch) in enumerate(dataloader):
                optimizer.zero_grad()
                y_pred = self(torch.Tensor(x_batch))
                loss = loss_fn(y_pred, y_batch)
                loss.backward()
                optimizer.step()

        print(f"Epoch {epoch}: traing loss: {loss.item()}")

    def predict(self, test_data, prob=False):
        self.eval()

        pred = self(torch.Tensor(test_data)).detach()

        if not prob:
            pred = torch.max(pred, dim=1)[-1]

        return pred.numpy()

class CustomBERT():

    def __init__(self, model_type="bert-base-uncased", to_cuda=True):

        self.model = BertForSequenceClassification.from_pretrained(
            model_type, num_labels=3, output_attentions=False, output_hidden_states=False)
        self.device = torch.device("cpu")
        if to_cuda:
            self.device = torch.device("cuda")
            self.model = self.model.cuda()

    def finetune(self, train_dataset, test_dataset, num_epochs, batch_size, random_state):

        train_loader = DataLoader(
            train_dataset, sampler=RandomSampler(train_dataset), batch_size=batch_size)

        optimizer = AdamW(self.model.parameters(), lr=2e-5, eps=1e-8)
        epochs = num_epochs
        total_steps = len(train_loader) * epochs
        scheduler = get_linear_schedule_with_warmup(
            optimizer, num_warmup_steps=0, num_training_steps=total_steps)

        random.seed(random_state)
        np.random.seed(random_state)
        torch.manual_seed(random_state)
        torch.cuda.manual_seed_all(random_state)

        training_stats = []
        total_t0 = time.time()
        for epoch_i in range(0, epochs):
            print("")
            print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))
            print('Training...')

            t0 = time.time()
            total_train_loss = 0

            self.model.train()

            for step, batch in enumerate(train_loader):

                if step % 40 == 0 and not step == 0:
                    elapsed = format_time(time.time() - t0)
                    print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(train_loader), elapsed))

                b_input_ids = batch[0].to(self.device)
                b_input_mask = batch[1].to(self.device)
                b_labels = batch[2].to(self.device)

                self.model.zero_grad()

                result = self.model(
                    b_input_ids, token_type_ids=None, attention_mask=b_input_mask,
                    labels=b_labels, return_dict=True)

                loss = result.loss
                logits = result.logits
                total_train_loss += loss.item()

                loss.backward()

                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)

                optimizer.step()
                scheduler.step()

            avg_train_loss = total_train_loss / len(train_loader)
            training_time = format_time(time.time() - t0)
            print("")
            print("  Average training loss: {0:.2f}".format(avg_train_loss))
            print("  Training epcoh took: {:}".format(training_time))

            print("")
            print("Running on Test Data...")
            t0 = time.time()

            test_loader = DataLoader(
                test_dataset, sampler=SequentialSampler(test_dataset), batch_size=batch_size)

            self.model.eval()

            total_eval_accuracy = 0
            total_eval_loss = 0
            nb_eval_steps = 0
            for batch in test_loader:
                b_input_ids = batch[0].to(self.device)
                b_input_mask = batch[1].to(self.device)
                b_labels = batch[2].to(self.device)

                with torch.no_grad():
                    result = self.model(
                        b_input_ids, token_type_ids=None, attention_mask=b_input_mask,
                        labels=b_labels, return_dict=True)

                loss = result.loss
                logits = result.logits
                total_eval_loss += loss.item()

                logits = logits.detach().cpu().numpy()
                label_ids = b_labels.to('cpu').numpy()
                total_eval_accuracy += self.flat_accuracy(logits, label_ids)

            avg_test_accuracy = total_eval_accuracy / len(test_loader)
            print("  Accuracy: {0:.2f}".format(avg_test_accuracy))

            avg_test_loss = total_eval_loss / len(test_loader)
            test_time = format_time(time.time() - t0)

            print("  Test Loss: {0:.2f}".format(avg_test_loss))
            print("  Run on Test Data took: {:}".format(test_time))

            training_stats.append(
                {
                    'epoch': epoch_i + 1,
                    'Training Loss': avg_train_loss,
                    'Test Loss': avg_test_loss,
                    'Test Accuracy': avg_test_accuracy,
                    'Training Time': training_time,
                    'Test Time': test_time
                    }
                    )

        print("")
        print("Training complete!")
        print("Total training took {:} (h:mm:ss)".format(format_time(time.time()-total_t0)))

    def predict(self, test_dataset, batch_size):

        test_loader = DataLoader(
            test_dataset, sampler=SequentialSampler(test_dataset), batch_size=batch_size)

        self.model.eval()

        predictions = []
        for batch in test_loader:
            b_input_ids = batch[0].to(self.device)
            b_input_mask = batch[1].to(self.device)
            b_labels = batch[2].to(self.device)

            with torch.no_grad():
                result = self.model(
                    b_input_ids, token_type_ids=None, attention_mask=b_input_mask, return_dict=True)

            logits = result.logits
            logits = logits.detach().cpu().numpy()
            predictions.append(logits)

        predictions = np.concatenate(predictions, axis=0)
        predictions = np.argmax(predictions, axis=1).flatten()

        return predictions

    def save(self, file_path):

        self.model.save_pretrained(file_path)

    def load(self, file_path):

        self.model = BertForSequenceClassification.from_pretrained(file_path)

        if self.device == torch.device("cuda"):
            self.model = self.model.cuda()

    def flat_accuracy(self, preds, labels):

        pred_flat = np.argmax(preds, axis=1).flatten()
        labels_flat = labels.flatten()

        return np.sum(pred_flat == labels_flat) / len(labels_flat)

def format_time(elapsed):

    elapsed_rounded = int(round((elapsed)))

    return str(datetime.timedelta(seconds=elapsed_rounded))


import sys
from os.path import dirname, realpath
import argparse

sys.path.insert(1, dirname(dirname(realpath(__file__))))
from src.utils.SQLUtils import CalibrationDB
from src.evaluation.Calibrator import Calibrator

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("cal1", help="cal path 1")
    parser.add_argument("cal2", help="cal path 2")
    parser.add_argument("calname", help="cal path 2")
    args = parser.parse_args()
    cal1 = CalibrationDB(args.cal1, args.calname)
    cal2 = CalibrationDB(args.cal2, args.calname)
    cal1 = Calibrator(cal1)
    cal2 = Calibrator(cal2)
    max_diff = 0
    max_det = 0
    for i in range(14):
        for j in range(11):
            for k in range(2):
                diff = (cal1.gains[i,j,k] - cal2.gains[i,j,k]) / cal1.gains[i,j,k]
                if(abs(diff) > max_diff):
                    max_det = 2*(14*j + i) + k
                    max_diff = abs(diff)
                print("det {0} gain diff {1}".format((2*(14*j + i) + k), diff))
    print("max diff is det {0} diff {1}".format(max_det, max_diff))
    #print((cal1.gains - cal2.gains) / cal1.gains)

if __name__=="__main__":
    main()


import codecs
import torch
import numpy as np
from nltk.corpus import wordnet as wn
import pickle as pkl

def get_trimmed_w2v_vectors(filename):
    """
    Args:
        filename: path to the npz file
    Returns:
        matrix of embeddings (np array)
    """
    with np.load(filename) as data:
        return data['embeddings']

def load_vocab(filename):
    """
    Args:
        filename: file with a word per line
    Returns:
        d: dict[word] = index
    """
    d = dict()
    with open(filename) as f:
        for idx, word in enumerate(f):
            word = word.strip()
            d[word] = idx + 1  # preserve idx 0 for pad_tok
    return d

def load_vocab_utf8(filename):
    """
    Args:
        filename: file with a word per line
    Returns:
        d: dict[word] = index
    """
    d = dict()
    with codecs.open(filename, encoding='utf-8') as f:
        for idx, word in enumerate(f):
            word = word.strip()
            d[word] = idx + 1  # preserve idx 0 for pad_tok
    return d

def load_pkl(path):
    with open(path, 'rb') as file:
        return pkl.load(file)

def dump_pkl(obj, path):
    with open(path, 'wb') as file:
        pkl.dump(obj, file)

def build_label_for_df(df):
    '''
    Build label for dataframe
    '''
    for i in range(len(df)):
        label = torch.tensor([df.iloc[i]['label']]).type(torch.FloatTensor)
        if i == 0:
            labels = label
        else:
            labels = torch.vstack((labels, label))

    return labels

def get_lookup(path):
    '''
    Get lookup table from file
    '''
    with open(path, 'r') as file:
        f = file.read().split('\n')
    return {f[i]: i + 1 for i in range(len(f))}

def lookup(word, dct):
    '''
    Get index of word in lookup table
    '''
    try:
        idx = dct[word]
    except:
        idx = 0
    return idx

def get_idx(sent, vocab_lookup, tag_lookup, direction_lookup, edge_lookup):
    '''
    Get index of features of tokens in sentence
    '''
    if sent == None:
        return None
    to_return = list()
    i = 0
    for dp in sent:
        word1_idx = lookup(dp[0][0], vocab_lookup)
        word2_idx = lookup(dp[2][0], vocab_lookup)
        tag1_idx = lookup(dp[0][1], tag_lookup)
        tag2_idx = lookup(dp[2][1], tag_lookup)
        direction_idx = lookup(dp[1][0], direction_lookup)
        edge_idx = lookup(dp[1][1], edge_lookup)
        pos1 = dp[0][2]
        pos2 = dp[2][2]
        v = torch.tensor([word1_idx, tag1_idx, direction_idx, edge_idx, word2_idx, tag2_idx])
        v = torch.hstack((v[:2], pos1, v[2:6], pos2))
        if i == 0:
            to_return = v.view(1, -1)
        else:
            to_return = torch.vstack((to_return, v))
        i += 1
    return to_return

def get_idx_dataset(data,
                    vocab_lookup,
                    tag_lookup,
                    direction_lookup,
                    edge_lookup):
    tmp = list()
    for i in data:
        tmp.append(get_idx(i, vocab_lookup, tag_lookup, direction_lookup, edge_lookup))
    return tmp

def prepare_dict_simul_paths(lookup, X, y):
    '''
    Get training tensors from lookup table and X for simultaneous paths approach
    Args:
        lookup: lookup table
        X: list of tensors
    Returns:
        list of tensors
    '''
    to_return_dict = dict()

    for k, v in lookup.items():
        tmp = []
        for i in v:
            if X[i] != None:
                tmp.append(X[i])
            label = y[i]
        if tmp != []:
            to_return_dict[k] = (tmp, label)

    return to_return_dict
def get_tensors_simul_paths(lookup: dict,
                            select_option: str,
                            top_k: int):
    '''
    Get top k paths and convert to tensors for simultaneous paths approach
    Args:
        lookup: lookup table
        select_option: select method, could be "random", "shortest", "longest"
        top_k: number of paths to select
    Returns:
        X: dataset of tensors
        y: labels
    '''
    from tqdm import tqdm

    X = list()
    y = list()
    for k, v in tqdm(lookup.items()):
        X_tmp = v[0]
        y_tmp = v[1]
        if select_option == "random":
            if top_k < len(X_tmp):
                X_tmp = X_tmp[:top_k]
        X.append(X_tmp)
        y.append(y_tmp)

    return X, y


# bot.py
import os
import random
from time import sleep
import asyncio
import time
import youtube_dl

import discord
from discord.ext import commands
from discord.utils import get

TOKEN = "[input token here]"
justUsed = ""
connected = False
playGachi = True
optoutusers = []

bot = commands.Bot(command_prefix='pog, ')
def openThing(fileName):
    global justUsed
    notGood = True
    with open('.\\media\\' + fileName) as Links:
        words = Links.read().split(",")
        while notGood:
            response = random.choice(words)
            if response != justUsed:
                notGood = False
        print(response)
    justUsed = response
    return response

async def mp3Player(fileName, message, length, directory):
    try:
	    global vc
	    global connected
	    global playGachi
	    user = message.author
	    print(user)
	    voice_channel = user.voice.channel
	    print(voice_channel)
	    channel = None
	    if voice_channel != None:
	        channel = voice_channel.name
	        if not connected:
	            vc = await voice_channel.connect()
	            connected = True
	        if directory == 0:
	            vc.play(discord.FFmpegPCMAudio(".\\mp3\\"+fileName))
	        elif directory == 1:
	            vc.play(discord.FFmpegPCMAudio("M:\\Gachi\\"+fileName))
	        elif directory == 2:
	        	vc.play(discord.FFmpegPCMAudio(".\\minecraft\\"+fileName))
	        elif directory == 3:
	        	vc.play(discord.FFmpegPCMAudio(".\\mp3\\yt\\"+fileName))

	        if (directory != 0 or directory != 3) and length == 0:
	            while vc.is_playing():
	                await asyncio.sleep(1)
	            vc.stop()
	        elif directory != 0 or directory != 3:
	            await asyncio.sleep(length)
	            vc.stop()
	        elif length == 0:
	            while vc.is_playing():
	                await asyncio.sleep(1)
	            await vc.disconnect()
	            connected = False
	        else:
	            await asyncio.sleep(length)
	            await vc.disconnect()
	            connected = False
    except:
        try:
            vc.stop()
            await vc.disconnect()
            connected = False
            playGachi = False
            pass
        except:
        	await vc.disconnect()
        	connected = False
        	pass
        print("whoop exception at mp3Player")
    await vc.disconnect()
    connected = False

def checkifplaying():
	return vc.is_playing()

with open('optout.txt') as optoutList:
	optoutusers = optoutList.read().split(",")

@bot.event
async def on_ready():
    print('Bot is now working!')
    await bot.change_presence(activity=discord.Game('with your dad ;)'))

@bot.event
async def on_message(message):
    with open('optout.txt') as optoutList:
       optoutusers = optoutList.read().split(",")
    global messagem
    if message.author == bot.user or ('.com' in message.content and not 'play' in message.content.lower()):
        return
    if message.author == 327606483741179914 and random.randint(0, 2) == 1:
        print("thompson denied")
        return
    if str(message.author) in optoutusers and not "pog, " in message.content.lower():
    	print(str(message.author), "is opted out")
    	return

    messagem = message

    #==============================poggers

    if 'poggers' in message.content.lower():

        response = openThing('poggers.tenor')

        unpoggers = ["https://tenor.com/view/poggers-deku-bakugo-my-hero-academia-my-hero-gif-18930451",
                     "https://tenor.com/view/unpoggers-pog-gif-19224204",
                     "https://tenor.com/view/anime-poggers-anime-poggers-anime-gif-18290524"]

        await message.channel.send(response)

        if response in unpoggers:
            await message.channel.send("whoa!")
            await message.channel.send("unpoggers...")

        else:
            await mp3Player("poggers.mp3", message, 1.2, 0)
            chance = random.randint(1, 50)
            print(chance)
            if chance == 24:
                await message.channel.send("poggers", tts=True)

        #==============================amogus

    elif 'amogus' in message.content.lower():

        a = random.choice(["amoguss", "aamogus", "sus", "amogus"])

        await mp3Player(a + ".mp3", message, 0, 0)

        await message.channel.send("ඞ")

    await bot.process_commands(message)

@bot.command(help='tells (un)funny joke')
async def funny(ctx):
    await mp3Player(random.choice(["eel.mp3", "congrats.mp3", "ahhhh.mp3"]), messagem, 0, 0)

@bot.command(help="play command, plays mp3's from playlist, opt. seconds")
async def playmp3(ctx, fileName: str, seconds: float = 0):
	fileName.replace(".mp3", "")
	fileName = fileName + ".mp3"
	print(fileName)
	await mp3Player(fileName, messagem, seconds, 0)

@bot.command(help="youtube play command, only accepts url")
async def play(ctx, url: str):
	if "list" in url:
		url = url[:url.index("list")-1]
	songDownloaded = os.path.isfile("ytsong.mp3")
	try:
		if songDownloaded:
			os.remove("ytsong.mp3")
	except PermissionError:
		await ctx.send("oop gotta wait for this one to end. or like use stop command idk")
	ytdl_opts = {
	'format': 'bestaudio/best',
	'outtmpl': '.\\mp3\\yt\\ytsong.%(ext)s',
	'postprocessors': [{
	'key': 'FFmpegExtractAudio',
	'preferredcodec': 'mp3',
	'preferredquality': '192'
		}]
	}
	print("downloading", url)
	with youtube_dl.YoutubeDL(ytdl_opts) as ydl:
		ydl.download([url])
	await mp3Player('ytsong.mp3', messagem, 0, 3)

@bot.command(help='audio files you can play')
async def playlist(ctx):
    files = os.listdir(".\\mp3\\")
    await messagem.channel.send(files)

@bot.command(help='stop playing')
async def stop(ctx):
    global playGachi
    global playMinecraft
    global connected
    playGachi = False
    playMinecraft = False
    vc.stop()
    await vc.disconnect()
    connected = False

@bot.command(help='triggers list')
async def triggers(ctx):
    with open("triggers.txt") as triggers:
        response = triggers.read().split(",")
    await messagem.channel.send("this is the simple version of this bot, so most of these wont work")
    await messagem.channel.send(response)

''' Check bot-og.py for why this is commented out
@bot.command(help='gachi')
async def gachi(ctx):
    global playGachi
    global connected
    global vc
    m = messagem
    playGachi = True
    while playGachi:
        toPlay = random.choice(os.listdir("M:\\Gachi"))
        await m.channel.send(toPlay)
        await mp3Player(toPlay, messagem, 0, 1)
'''

@bot.command(help='minecraft')
async def minecraft(ctx):
	global playMinecraft
	global connected
	global vc
	m = messagem
	playMinecraft = True
	listofsongs = os.listdir(".\\minecraft")
	if listofsongs == []:
		playMinecraft = False
	while playMinecraft:
		toPlay = random.choice(listofsongs)
		listofsongs.remove(toPlay)
		await m.channel.send(toPlay)
		await mp3Player(toPlay, messagem, 0, 2)

@bot.command(help='opt out using this')
async def optout(ctx):
	with open('optout.txt') as optoutList:
		optoutusers = optoutList.read().split(",")
	print(optoutusers)
	if not str(messagem.author) in optoutusers:
		with open('optout.txt', 'a') as optoutList:
			optoutList.write(str(messagem.author) + ',')
	await messagem.channel.send('opt back in using "pog, optin"')

@bot.command(help='opt back in')
async def optin(ctx):
	with open('optout.txt') as optoutList:
		optoutusers = optoutList.read().split(",")
	print(optoutusers)
	new = ''
	if str(messagem.author) in optoutusers:
		optoutusers.remove(str(messagem.author))
	for i in optoutusers:
		if i != '':
			new += i + ','
	with open('optout.txt', 'w') as optoutList:
		optoutList.write(new)
	await messagem.channel.send("welcome back!")

@bot.command(help='random screenshot')
async def randomsc(ctx, n: int = 1):
	letters = "qwertyuiopasdfghjklzxcvbnm"
	for i in range(n):
		await ctx.send("https://prnt.sc/" + random.choice(letters) + random.choice(letters) + str(random.randint(1000, 9999)))

bot.run(TOKEN)


import logging
import urllib
import yaml
from simplejson import JSONDecodeError
import simplejson as json
from twisted.internet import defer

from synapse.api.constants import Membership
from synapse.appservice import ApplicationService, AppServiceTransaction
from synapse.storage.roommember import RoomsForUser
from synapse.types import UserID
from ._base import SQLBaseStore

logger = logging.getLogger(__name__)

class ApplicationServiceStore(SQLBaseStore):

    def __init__(self, hs):
        super(ApplicationServiceStore, self).__init__(hs)
        self.hostname = hs.hostname
        self.services_cache = []
        self._populate_appservice_cache(
            hs.config.app_service_config_files
        )

    def get_app_services(self):
        return defer.succeed(self.services_cache)

    def get_app_service_by_user_id(self, user_id):
        """Retrieve an application service from their user ID.

        All application services have associated with them a particular user ID.
        There is no distinguishing feature on the user ID which indicates it
        represents an application service. This function allows you to map from
        a user ID to an application service.

        Args:
            user_id(str): The user ID to see if it is an application service.
        Returns:
            synapse.appservice.ApplicationService or None.
        """
        for service in self.services_cache:
            if service.sender == user_id:
                return defer.succeed(service)
        return defer.succeed(None)

    def get_app_service_by_token(self, token):
        """Get the application service with the given appservice token.

        Args:
            token (str): The application service token.
        Returns:
            synapse.appservice.ApplicationService or None.
        """
        for service in self.services_cache:
            if service.token == token:
                return defer.succeed(service)
        return defer.succeed(None)

    def get_app_service_rooms(self, service):
        """Get a list of RoomsForUser for this application service.

        Application services may be "interested" in lots of rooms depending on
        the room ID, the room aliases, or the members in the room. This function
        takes all of these into account and returns a list of RoomsForUser which
        represent the entire list of room IDs that this application service
        wants to know about.

        Args:
            service: The application service to get a room list for.
        Returns:
            A list of RoomsForUser.
        """
        return self.runInteraction(
            "get_app_service_rooms",
            self._get_app_service_rooms_txn,
            service,
        )

    def _get_app_service_rooms_txn(self, txn, service):
        # get all rooms matching the room ID regex.
        room_entries = self._simple_select_list_txn(
            txn=txn, table="rooms", keyvalues=None, retcols=["room_id"]
        )
        matching_room_list = set([
            r["room_id"] for r in room_entries if
            service.is_interested_in_room(r["room_id"])
        ])

        # resolve room IDs for matching room alias regex.
        room_alias_mappings = self._simple_select_list_txn(
            txn=txn, table="room_aliases", keyvalues=None,
            retcols=["room_id", "room_alias"]
        )
        matching_room_list |= set([
            r["room_id"] for r in room_alias_mappings if
            service.is_interested_in_alias(r["room_alias"])
        ])

        # get all rooms for every user for this AS. This is scoped to users on
        # this HS only.
        user_list = self._simple_select_list_txn(
            txn=txn, table="users", keyvalues=None, retcols=["name"]
        )
        user_list = [
            u["name"] for u in user_list if
            service.is_interested_in_user(u["name"])
        ]
        rooms_for_user_matching_user_id = set()  # RoomsForUser list
        for user_id in user_list:
            # FIXME: This assumes this store is linked with RoomMemberStore :(
            rooms_for_user = self._get_rooms_for_user_where_membership_is_txn(
                txn=txn,
                user_id=user_id,
                membership_list=[Membership.JOIN]
            )
            rooms_for_user_matching_user_id |= set(rooms_for_user)

        # make RoomsForUser tuples for room ids and aliases which are not in the
        # main rooms_for_user_list - e.g. they are rooms which do not have AS
        # registered users in it.
        known_room_ids = [r.room_id for r in rooms_for_user_matching_user_id]
        missing_rooms_for_user = [
            RoomsForUser(r, service.sender, "join") for r in
            matching_room_list if r not in known_room_ids
        ]
        rooms_for_user_matching_user_id |= set(missing_rooms_for_user)

        return rooms_for_user_matching_user_id

    def _parse_services_dict(self, results):
        # SQL results in the form:
        # [
        #   {
        #     'regex': "something",
        #     'url': "something",
        #     'namespace': enum,
        #     'as_id': 0,
        #     'token': "something",
        #     'hs_token': "otherthing",
        #     'id': 0
        #   }
        # ]
        services = {}
        for res in results:
            as_token = res["token"]
            if as_token is None:
                continue
            if as_token not in services:
                # add the service
                services[as_token] = {
                    "id": res["id"],
                    "url": res["url"],
                    "token": as_token,
                    "hs_token": res["hs_token"],
                    "sender": res["sender"],
                    "namespaces": {
                        ApplicationService.NS_USERS: [],
                        ApplicationService.NS_ALIASES: [],
                        ApplicationService.NS_ROOMS: []
                    }
                }
            # add the namespace regex if one exists
            ns_int = res["namespace"]
            if ns_int is None:
                continue
            try:
                services[as_token]["namespaces"][
                    ApplicationService.NS_LIST[ns_int]].append(
                    json.loads(res["regex"])
                )
            except IndexError:
                logger.error("Bad namespace enum '%s'. %s", ns_int, res)
            except JSONDecodeError:
                logger.error("Bad regex object '%s'", res["regex"])

        service_list = []
        for service in services.values():
            service_list.append(ApplicationService(
                token=service["token"],
                url=service["url"],
                namespaces=service["namespaces"],
                hs_token=service["hs_token"],
                sender=service["sender"],
                id=service["id"]
            ))
        return service_list

    def _load_appservice(self, as_info):
        required_string_fields = [
            "url", "as_token", "hs_token", "sender_localpart"
        ]
        for field in required_string_fields:
            if not isinstance(as_info.get(field), basestring):
                raise KeyError("Required string field: '%s'", field)

        localpart = as_info["sender_localpart"]
        if urllib.quote(localpart) != localpart:
            raise ValueError(
                "sender_localpart needs characters which are not URL encoded."
            )
        user = UserID(localpart, self.hostname)
        user_id = user.to_string()

        # namespace checks
        if not isinstance(as_info.get("namespaces"), dict):
            raise KeyError("Requires 'namespaces' object.")
        for ns in ApplicationService.NS_LIST:
            # specific namespaces are optional
            if ns in as_info["namespaces"]:
                # expect a list of dicts with exclusive and regex keys
                for regex_obj in as_info["namespaces"][ns]:
                    if not isinstance(regex_obj, dict):
                        raise ValueError(
                            "Expected namespace entry in %s to be an object,"
                            " but got %s", ns, regex_obj
                        )
                    if not isinstance(regex_obj.get("regex"), basestring):
                        raise ValueError(
                            "Missing/bad type 'regex' key in %s", regex_obj
                        )
                    if not isinstance(regex_obj.get("exclusive"), bool):
                        raise ValueError(
                            "Missing/bad type 'exclusive' key in %s", regex_obj
                        )
        return ApplicationService(
            token=as_info["as_token"],
            url=as_info["url"],
            namespaces=as_info["namespaces"],
            hs_token=as_info["hs_token"],
            sender=user_id,
            id=as_info["as_token"]  # the token is the only unique thing here
        )

    def _populate_appservice_cache(self, config_files):
        """Populates a cache of Application Services from the config files."""
        if not isinstance(config_files, list):
            logger.warning(
                "Expected %s to be a list of AS config files.", config_files
            )
            return

        for config_file in config_files:
            try:
                with open(config_file, 'r') as f:
                    appservice = self._load_appservice(yaml.load(f))
                    logger.info("Loaded application service: %s", appservice)
                    self.services_cache.append(appservice)
            except Exception as e:
                logger.error("Failed to load appservice from '%s'", config_file)
                logger.exception(e)

class ApplicationServiceTransactionStore(SQLBaseStore):

    def __init__(self, hs):
        super(ApplicationServiceTransactionStore, self).__init__(hs)

    @defer.inlineCallbacks
    def get_appservices_by_state(self, state):
        """Get a list of application services based on their state.

        Args:
            state(ApplicationServiceState): The state to filter on.
        Returns:
            A Deferred which resolves to a list of ApplicationServices, which
            may be empty.
        """
        results = yield self._simple_select_list(
            "application_services_state",
            dict(state=state),
            ["as_id"]
        )
        # NB: This assumes this class is linked with ApplicationServiceStore
        as_list = yield self.get_app_services()
        services = []

        for res in results:
            for service in as_list:
                if service.id == res["as_id"]:
                    services.append(service)
        defer.returnValue(services)

    @defer.inlineCallbacks
    def get_appservice_state(self, service):
        """Get the application service state.

        Args:
            service(ApplicationService): The service whose state to set.
        Returns:
            A Deferred which resolves to ApplicationServiceState.
        """
        result = yield self._simple_select_one(
            "application_services_state",
            dict(as_id=service.id),
            ["state"],
            allow_none=True
        )
        if result:
            defer.returnValue(result.get("state"))
            return
        defer.returnValue(None)

    def set_appservice_state(self, service, state):
        """Set the application service state.

        Args:
            service(ApplicationService): The service whose state to set.
            state(ApplicationServiceState): The connectivity state to apply.
        Returns:
            A Deferred which resolves when the state was set successfully.
        """
        return self._simple_upsert(
            "application_services_state",
            dict(as_id=service.id),
            dict(state=state)
        )

    def create_appservice_txn(self, service, events):
        """Atomically creates a new transaction for this application service
        with the given list of events.

        Args:
            service(ApplicationService): The service who the transaction is for.
            events(list<Event>): A list of events to put in the transaction.
        Returns:
            AppServiceTransaction: A new transaction.
        """
        return self.runInteraction(
            "create_appservice_txn",
            self._create_appservice_txn,
            service, events
        )

    def _create_appservice_txn(self, txn, service, events):
        # work out new txn id (highest txn id for this service += 1)
        # The highest id may be the last one sent (in which case it is last_txn)
        # or it may be the highest in the txns list (which are waiting to be/are
        # being sent)
        last_txn_id = self._get_last_txn(txn, service.id)

        txn.execute(
            "SELECT MAX(txn_id) FROM application_services_txns WHERE as_id=?",
            (service.id,)
        )
        highest_txn_id = txn.fetchone()[0]
        if highest_txn_id is None:
            highest_txn_id = 0

        new_txn_id = max(highest_txn_id, last_txn_id) + 1

        # Insert new txn into txn table
        event_ids = json.dumps([e.event_id for e in events])
        txn.execute(
            "INSERT INTO application_services_txns(as_id, txn_id, event_ids) "
            "VALUES(?,?,?)",
            (service.id, new_txn_id, event_ids)
        )
        return AppServiceTransaction(
            service=service, id=new_txn_id, events=events
        )

    def complete_appservice_txn(self, txn_id, service):
        """Completes an application service transaction.

        Args:
            txn_id(str): The transaction ID being completed.
            service(ApplicationService): The application service which was sent
            this transaction.
        Returns:
            A Deferred which resolves if this transaction was stored
            successfully.
        """
        return self.runInteraction(
            "complete_appservice_txn",
            self._complete_appservice_txn,
            txn_id, service
        )

    def _complete_appservice_txn(self, txn, txn_id, service):
        txn_id = int(txn_id)

        # Debugging query: Make sure the txn being completed is EXACTLY +1 from
        # what was there before. If it isn't, we've got problems (e.g. the AS
        # has probably missed some events), so whine loudly but still continue,
        # since it shouldn't fail completion of the transaction.
        last_txn_id = self._get_last_txn(txn, service.id)
        if (last_txn_id + 1) != txn_id:
            logger.error(
                "appservice: Completing a transaction which has an ID > 1 from "
                "the last ID sent to this AS. We've either dropped events or "
                "sent it to the AS out of order. FIX ME. last_txn=%s "
                "completing_txn=%s service_id=%s", last_txn_id, txn_id,
                service.id
            )

        # Set current txn_id for AS to 'txn_id'
        self._simple_upsert_txn(
            txn, "application_services_state", dict(as_id=service.id),
            dict(last_txn=txn_id)
        )

        # Delete txn
        self._simple_delete_txn(
            txn, "application_services_txns",
            dict(txn_id=txn_id, as_id=service.id)
        )

    def get_oldest_unsent_txn(self, service):
        """Get the oldest transaction which has not been sent for this
        service.

        Args:
            service(ApplicationService): The app service to get the oldest txn.
        Returns:
            A Deferred which resolves to an AppServiceTransaction or
            None.
        """
        return self.runInteraction(
            "get_oldest_unsent_appservice_txn",
            self._get_oldest_unsent_txn,
            service
        )

    def _get_oldest_unsent_txn(self, txn, service):
        # Monotonically increasing txn ids, so just select the smallest
        # one in the txns table (we delete them when they are sent)
        txn.execute(
            "SELECT * FROM application_services_txns WHERE as_id=?"
            " ORDER BY txn_id ASC LIMIT 1",
            (service.id,)
        )
        rows = self.cursor_to_dict(txn)
        if not rows:
            return None

        entry = rows[0]

        event_ids = json.loads(entry["event_ids"])
        events = self._get_events_txn(txn, event_ids)

        return AppServiceTransaction(
            service=service, id=entry["txn_id"], events=events
        )

    def _get_last_txn(self, txn, service_id):
        txn.execute(
            "SELECT last_txn FROM application_services_state WHERE as_id=?",
            (service_id,)
        )
        last_txn_id = txn.fetchone()
        if last_txn_id is None or last_txn_id[0] is None:  # no row exists
            return 0
        else:
            return int(last_txn_id[0])  # select 'last_txn' col


import json
import re
import sys
import xmltodict
from functools import reduce

def parse_text(text):
    obj = {}
    for m in re.findall(r'(\{\{pagebanner|\n\=\=(\w+)\=\=\n)(((.|\n)(?!\n\=\=))*)', text):
        label = m[1].lower()
        group = m[2]
        if len(label) < 1:
            if "summary" in obj: continue
            label = "summary"
            group = m[0] + group
        paragraphs = []

        for g in re.findall(r'([^\{]+|\{\{[^\}]+\}\})', group):
            if len(g) < 5: continue

            if '{' in g:
                p = g[2:-2].split('|')
                pk = p[0].strip()
                new_obj = {}

                for s in p[1:]:
                    s = s.split('=')
                    k = s[0].strip()
                    if pk == 'pagebanner' and len(s) == 1:
                        v = k
                        k = 'image'
                    else:
                        v = s[1].strip() if len(s) > 1 else None
                    if v and len(v) > 2:
                        new_obj[k] = v

                if len(new_obj.keys()) > 0:
                    paragraphs.append({pk: new_obj})
            else:
                paragraphs.append(g.strip())

        obj[label] = paragraphs

    return obj

def main(filename):
    file_data = open(filename).read()
    xml_obj = xmltodict.parse(file_data)
    pages = []

    for page in xml_obj['mediawiki']['page']:
        try:
            text = page['revision']['text']['#text']
        except Exception:
            continue
        id = page['id']
        title = page['title']
        text_obj = parse_text(text)

        data_len = len(set([k for k in text_obj.keys() if k not in ['summary']]))
        has_image = False
        has_detailed_summary = False
        if 'summary' in text_obj:
            for s in text_obj['summary']:
                if isinstance(s, str):
                    if len(s) > 150 and len(s.split('.')) > 1:
                        has_detailed_summary = True
                elif 'pagebanner' in s:
                    if 'image' in s['pagebanner']:
                        img = s['pagebanner']['image']
                        if img and 'default' not in img:
                            has_image = True

        if data_len > 0 and has_image and has_detailed_summary:
            pages.append({
                'id': id,
                'title': title,
                'sections': text_obj
            })

    with open('wikivoyage.json', 'w') as f:
        f.write(json.dumps(pages, indent=2))

if __name__ == '__main__':
    """
        data gathered from the latest Wikivoyage data dump
        https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2
    """
    main(sys.argv[-1])


from flask import Flask, request

app = Flask(__name__)

@app.route('/upload', methods = ['POST'])
def upload():
    if request.method == 'POST':
        raw_data = request.data
        print(raw_data)
        return "Data received successfully!"
    return "Data not received successfully"

if __name__ == '__main__':
    app.run()


import logging
from enum import Enum

from cmdb.framework.cmdb_location import CmdbLocation

LOGGER = logging.getLogger(__name__)

class CheckRoutine:
    """
    This class holds checks for the database check routines
    """

    class CheckStatus(Enum):
        """
        Enumeration of the status options for checks
        """
        NOT = 0
        RUNNING = 1
        HAS_UPDATES = 2
        FINISHED = 3

    def __init__(self, dbm):
        self.status: int = CheckRoutine.CheckStatus.NOT
        self.setup_database_manager = dbm

    def get_check_status(self) -> int:
        """
        Returns the CheckStatus of the current CheckRoutine

        Returns:
            (int): Current CheckStatus
        """
        return self.status

    def checker(self):
        LOGGER.info('CHECK ROUTINE: STARTED...')
        self.status = CheckRoutine.CheckStatus.FINISHED

        # check database
        if not self.__is_database_empty():
            if not (self.__check_database_collection_valid() and self.has_updates()):
                LOGGER.warning(
                    'The current database version does not match the valid database version.'
                )
                self.status = CheckRoutine.CheckStatus.HAS_UPDATES
        LOGGER.info('CHECK ROUTINE: FINISHED!')
        return self.status

    def __is_database_empty(self) -> bool:
        """
        Checks if there are any collections in the initialised database

        Returns:
            (bool): Returns bool if there are any collection in the initialised database
        """
        return not self.setup_database_manager.connector.database.list_collection_names()

    def __check_database_collection_valid(self) -> bool:
        LOGGER.info('CHECK ROUTINE: Checking database collection validation')
        from cmdb.framework import __COLLECTIONS__ as FRAMEWORK_CLASSES
        from cmdb.user_management import __COLLECTIONS__ as USER_MANAGEMENT_COLLECTION
        from cmdb.exportd import __COLLECTIONS__ as JOB_MANAGEMENT_COLLECTION
        detected_database = self.setup_database_manager.connector.database
        collection_test = True

        try:
            # framework collections
            for collection in FRAMEWORK_CLASSES:
                collection_test = detected_database.validate_collection(collection.COLLECTION, scandata=True)['valid']

                # setup locations if valid test
                if collection == CmdbLocation and collection_test:
                    root_location: CmdbLocation = self.setup_database_manager.find_one(collection.COLLECTION, 1)

                    if root_location:
                        if self.setup_database_manager.validate_root_location(CmdbLocation.to_data(root_location)):
                            LOGGER.info("CHECK ROUTINE: Root Location valid")
                        else:
                            LOGGER.info("CHECK ROUTINE: Root Location invalalid => Fixing the Issue!")
                            self.setup_database_manager.set_root_location(collection.COLLECTION, create=False)

                    else:
                        LOGGER.info("CHECK ROUTINE: No Root Location => Creating a new Root Location!")
                        self.setup_database_manager.set_root_location(collection.COLLECTION, create=True)

            # user management collections
            for collection in USER_MANAGEMENT_COLLECTION:
                collection_test = detected_database.validate_collection(collection.COLLECTION, scandata=True)['valid']
            # exportdJob management collections
            for collection in JOB_MANAGEMENT_COLLECTION:
                collection_test = detected_database.validate_collection(collection.COLLECTION, scandata=True)[
                    'valid']
        except Exception as ex:
            LOGGER.info(f'CHECK ROUTINE: Database collection validation for "{collection.COLLECTION}" failed. '
                        f'msgerror: {ex}')
            collection_test = False

        LOGGER.info(f'CHECK ROUTINE: Database collection validation status {collection_test}')
        return collection_test

    def has_updates(self) -> bool:
        """
        Check if updates are available

        Returns:
            (bool): True if updates are available else raises an error
        """
        from cmdb.updater import UpdaterModule
        from cmdb.utils.system_reader import SystemSettingsReader, SectionError

        try:
            ssr = SystemSettingsReader(self.setup_database_manager)
            ssr.get_all_values_from_section('updater')

            upd_module = UpdaterModule(ssr)
            new_version = upd_module.get_last_version()['version']
            current_version = upd_module.settings.version

            if new_version > current_version:
                LOGGER.error(
                    f'Please run an update. There are new updates available.'
                    f'Current Updater Version: {current_version} '
                    f'Newest Update Version: {new_version}'
                )
                return False

        except SectionError as err:
            LOGGER.error(err.message)
            return False

        return True


import six
import chainer
import numpy as np
import chainer.links as L
import chainer.functions as F
import nutszebra_chainer
import functools
from collections import defaultdict

class BN_ReLU_Conv(nutszebra_chainer.Model):

    def __init__(self, in_channel, out_channel, filter_size=(3, 3), stride=(1, 1), pad=(1, 1)):
        super(BN_ReLU_Conv, self).__init__(
            conv=L.Convolution2D(in_channel, out_channel, filter_size, stride, pad),
            bn=L.BatchNormalization(in_channel),
        )

    def weight_initialization(self):
        self.conv.W.data = self.weight_relu_initialization(self.conv)
        self.conv.b.data = self.bias_initialization(self.conv, constant=0)

    def __call__(self, x, train=False):
        return self.conv(F.relu(self.bn(x, test=not train)))

    def count_parameters(self):
        return functools.reduce(lambda a, b: a * b, self.conv.W.data.shape)

class Highway_Conv(nutszebra_chainer.Model):

    def __init__(self, in_channel, out_channel, filter_size=(3, 3), stride=(1, 1), pad=(1, 1)):
        stride = stride if type(stride) is not int else (stride, stride)
        super(Highway_Conv, self).__init__()
        modules = []
        modules += [('conv', BN_ReLU_Conv(in_channel, out_channel, filter_size, stride, pad))]
        modules += [('conv_highway_3x3', BN_ReLU_Conv(out_channel, out_channel, filter_size, stride, pad))]
        modules += [('conv_highway_1x1_1', BN_ReLU_Conv(out_channel, out_channel, 1, 1, 0))]
        modules += [('conv_highway_1x1_2', BN_ReLU_Conv(out_channel, out_channel, 1, 1, 0))]
        # register layers
        [self.add_link(*link) for link in modules]
        self.modules = modules
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.filter_size = filter_size
        self.stride = stride
        self.pad = pad

    def weight_initialization(self):
        [link.weight_initialization() for _, link in self.modules]

    def count_parameters(self):
        return int(np.sum([link.count_parameters() for _, link in self.modules]))

    @staticmethod
    def concatenate_zero_pad(x, h_shape, volatile, h_type):
        _, x_channel, _, _ = x.data.shape
        batch, h_channel, h_y, h_x = h_shape
        if x_channel == h_channel:
            return x
        pad = chainer.Variable(np.zeros((batch, h_channel - x_channel, h_y, h_x), dtype=np.float32), volatile=volatile)
        if h_type is not np.ndarray:
            pad.to_gpu()
        return F.concat((x, pad))

    def __call__(self, x, train=False):
        h = self.conv(x, train)
        x = self.concatenate_zero_pad(x, h.data.shape, h.volatile, type(h.data))
        highway = self.conv_highway_3x3(x, train)
        highway = self.conv_highway_1x1_1(highway, train)
        highway = F.sigmoid(self.conv_highway_1x1_2(highway, train))
        self.highway = highway
        return h * highway + x * (1.0 - highway)

class Highway_Fitnet4(nutszebra_chainer.Model):

    def __init__(self, category_num):
        super(Highway_Fitnet4, self).__init__()
        # conv
        modules = []
        modules += [('conv1_1', Highway_Conv(3, 32, 3, 1, 1))]
        modules += [('conv1_2', Highway_Conv(32, 32, 3, 1, 1))]
        modules += [('conv1_3', Highway_Conv(32, 32, 3, 1, 1))]
        modules += [('conv1_4', Highway_Conv(32, 48, 3, 1, 1))]
        modules += [('conv1_5', Highway_Conv(48, 48, 3, 1, 1))]
        modules += [('conv2_1', Highway_Conv(48, 80, 3, 1, 1))]
        modules += [('conv2_2', Highway_Conv(80, 80, 3, 1, 1))]
        modules += [('conv2_3', Highway_Conv(80, 80, 3, 1, 1))]
        modules += [('conv2_4', Highway_Conv(80, 80, 3, 1, 1))]
        modules += [('conv2_5', Highway_Conv(80, 80, 3, 1, 1))]
        modules += [('conv2_6', Highway_Conv(80, 80, 3, 1, 1))]
        modules += [('conv3_1', Highway_Conv(80, 128, 3, 1, 1))]
        modules += [('conv3_2', Highway_Conv(128, 128, 3, 1, 1))]
        modules += [('conv3_3', Highway_Conv(128, 128, 3, 1, 1))]
        modules += [('conv3_4', Highway_Conv(128, 128, 3, 1, 1))]
        modules += [('conv3_5', Highway_Conv(128, 128, 3, 1, 1))]
        modules += [('conv3_6', Highway_Conv(128, 128, 3, 1, 1))]
        modules += [('linear', BN_ReLU_Conv(128, category_num, filter_size=(1, 1), stride=(1, 1), pad=(0, 0)))]
        # register layers
        [self.add_link(*link) for link in modules]
        self.modules = modules
        self.category_num = category_num
        self.block = [5, 6, 6]
        self.name = 'highway_fitnet4_{}'.format(category_num)

    def weight_initialization(self):
        [link.weight_initialization() for _, link in self.modules]

    def count_parameters(self):
        return int(np.sum([link.count_parameters() for _, link in self.modules]))

    def __call__(self, x, train=False):
        for i in six.moves.range(len(self.block)):
            for ii in six.moves.range(1, self.block[i] + 1):
                x = self['conv{}_{}'.format(i + 1, ii)](x, train)
            if i < len(self.block) - 1:
                x = F.max_pooling_2d(x, ksize=(2, 2), stride=(2, 2), pad=(0, 0))
        batch, channels, height, width = x.data.shape
        h = F.reshape(F.average_pooling_2d(x, (height, width)), (batch, channels, 1, 1))
        return F.reshape(self.linear(h, train), (batch, self.category_num))

    def calc_loss(self, y, t):
        loss = F.softmax_cross_entropy(y, t)
        return loss

    def accuracy(self, y, t, xp=np):
        y.to_cpu()
        t.to_cpu()
        indices = np.where((t.data == np.argmax(y.data, axis=1)) == True)[0]
        accuracy = defaultdict(int)
        for i in indices:
            accuracy[t.data[i]] += 1
        indices = np.where((t.data == np.argmax(y.data, axis=1)) == False)[0]
        false_accuracy = defaultdict(int)
        false_y = np.argmax(y.data, axis=1)
        for i in indices:
            false_accuracy[(t.data[i], false_y[i])] += 1
        return accuracy, false_accuracy


from jax import random
import jax.numpy as jnp

def uniform_prior_particles(key, minimums, maximums, N):
    """Produces an array of N particles sampled unformly within a
    multidimensional box defined by minimums and maximums for each parameter.

    Parameters
    ----------
    key : jax.random.PRNGKey
        The pseudo-random number generator key used to seed all
        jax random functions.
    minimums : Vector
        A vector of minimum parameter values
    maximums : Vector
        A vector of maximum parameter values
    N : Int
        The number of particles

    Returns
    -------
    Array
        A set of particles from the desired distribution.
    """
    n_params = len(minimums)
    return random.uniform(key,(n_params,N),
                          minval=jnp.asarray(minimums).reshape(n_params,1),
                          maxval=jnp.asarray(maximums).reshape(n_params,1))


import os
import traceback
from typing import Optional, Union

from Config import experiment_config as cnfg

TEXT_EXTENSION = 'txt'
IMAGE_EXTENSION = 'png'
VIDEO_EXTENSION = 'mp4'
PICKLE_EXTENSION = 'pkl'

def create_subject_output_directory(subject_id: Union[int, str], output_dir: Optional[str] = cnfg.OUTPUT_DIR) -> str:
    """
    Create subject output directory with a unique name with the following format: PATH/TO/OUTPUT_DIR/SXXX
    If argument `output_dir` is not provided, the default output directory from `experiment_config` is used.
    """
    if output_dir is None:
        output_dir = cnfg.OUTPUT_DIR
    if isinstance(subject_id, int):
        subject_id = f"{subject_id:03d}"
    return create_directory(dirname=f"S{subject_id}", parent_dir=output_dir)

def create_directory(dirname: str, parent_dir: str) -> str:
    if not os.path.exists(parent_dir):
        os.makedirs(parent_dir)
    full_path = os.path.join(parent_dir, dirname)
    if not os.path.exists(full_path):
        os.makedirs(full_path)
    return full_path

def get_filename(name: str, extension: Optional[str] = None) -> str:
    split_name = name.split('.')
    if len(split_name) > 2:
        raise ValueError(f'Invalid filename: {name}')
    if len(split_name) == 2:
        if extension is not None:
            raise ValueError(f'File already has an extension: {extension[1]}')
        name, extension = split_name
    if extension is None:
        raise ValueError(f'No extension provided for filename: {name}')
    return f"{name}.{extension}"

def log_and_print(msg: str, log_file: Optional[str] = None) -> None:
    print(msg)
    try:
        if log_file is not None:
            if not log_file.endswith(TEXT_EXTENSION):
                log_file = get_filename(name=log_file, extension=TEXT_EXTENSION)
            with open(log_file, 'a') as f:
                f.write(msg + '\n')
    except Exception as e:
        trace = traceback.format_exc()
        print(f"\tFailed to write to log file: {e}\n\t{trace}\n")


from sunpy.net import Fido, attrs as a, vso
import datetime
from astropy import units as u
import urllib
from bs4 import BeautifulSoup
import os
import time
import numpy as np

#start_time = time.time()

#get utc time rounded down to the previous hour
time_now = datetime.datetime.utcnow().replace(microsecond = 0, second = 0, minute = 0)

def get_xrt(time_now, out_dir):
	if not os.path.exists(out_dir):
		os.mkdir(out_dir)

	base_url = 'http://solar.physics.montana.edu/HINODE/XRT/QL/syn_comp_fits/'

	#find the list of available fits files on webpage
	test = urllib.request.urlopen(base_url)
	soup = BeautifulSoup(test, features="lxml")

	fits_links = []
	for link in soup.findAll('a'):
		if link.get('href')[-4:] == 'fits':
			fits_links.append(base_url  + link.get('href'))

	time_dif = []
	for i in range(len(fits_links)):

		time_dif.append((time_now - datetime.datetime.strptime(fits_links[i][-22:-7], '%Y%m%d_%H%M%S')).total_seconds())

	index = np.where(time_dif == np.min(time_dif))[0][0]
	latest_file = fits_links[index]

	print(latest_file)

	#output path directory
	output_path2 = out_dir + latest_file.split('/')[-1]

	urllib.request.urlretrieve(latest_file, output_path2)
	if os.path.exists(output_path2):
		print('XRT success!')

#end_time = time.time()
#print(end_time - start_time)


import torch
from torch import autocast
from diffusers import StableDiffusionPipeline,StableDiffusionImg2ImgPipeline
from djl_python import Input, Output
import io
from PIL import Image
from io import BytesIO
import requests

textPipe = StableDiffusionPipeline.from_pretrained("/opt/ml/model/stable-diffusion-v1-4",revision="fp16", torch_dtype=torch.float16).to("cuda")
imgPipe = StableDiffusionImg2ImgPipeline.from_pretrained("/opt/ml/model/stable-diffusion-v1-4",revision="fp16", torch_dtype=torch.float16).to("cuda")
#url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
#response = requests.get(url)
def handle(inputs: Input) -> Output:
    global textPipe
    global imgPipe
    if inputs.is_empty():
        return None
    content_type = inputs.get_property("content-type")
    if content_type == "application/json":
        param = inputs.get_as_json()
        prompt =param.get('prompt')
        height =param.get('height')
        width =param.get('width')
        steps =param.get('steps')
        scale =param.get('scale')
        with torch.autocast("cuda"):
            image = textPipe(prompt,guidance_scale=scale, num_inference_steps=steps,height=height, width=width)["sample"][0]
            buf = io.BytesIO()
            image.save(buf, format='PNG')
            byte_im = buf.getvalue()
            return Output().add(byte_im).add_property("content-type","image/png")
    elif content_type is not None and content_type.startswith("text/"):
        prompt =inputs.get_as_string()
        with torch.autocast("cuda"):
            image = textPipe(prompt)["sample"][0]
            buf = io.BytesIO()
            image.save(buf, format='PNG')
            byte_im = buf.getvalue()
            return Output().add(byte_im).add_property("content-type","image/png")
    else:
        init_image = Image.open(BytesIO(inputs.get_as_bytes())).convert("RGB")
        #init_image = Image.open(BytesIO(response.content)).convert("RGB")
        init_image = init_image.resize((768, 512))
        prompt =inputs.get_as_string(key="prompt")
        #prompt = "A fantasy landscape, trending on artstation"
        with torch.autocast("cuda"):
            image = imgPipe(prompt=prompt, init_image=init_image, strength=0.75, guidance_scale=7.5).images[0]
            buf = io.BytesIO()
            image.save(buf, format='PNG')
            byte_im = buf.getvalue()
            return Output().add(byte_im).add_property("content-type","image/png")