# -*- coding: utf-8 -*-
"""
Created on Wed Mar 10 19:08:44 2021

@author: 11911627 Tan Sixu
"""

dic1={2:1,1:2}
dic2={1:2, 2:1}

print(dic1==dic2)  #orders not matter in dics

a=dic1.keys()
print(type(a))    #new data type: dickeys

b=list(dic1.keys())  #transform to list
print(b)

c=dic1.values()    #same as keys()

d=dic1.items()
print(d)          #datatype: dict_items

e=list(d)
print(e)         #tranform to tuple in list
# =============================================================================
# True
# <class 'dict_keys'>
# [2, 1]
# dict_items([(2, 1), (1, 2)])
# [(2, 1), (1, 2)]
# =============================================================================

#%%
dic=dict(one=1, two=2, three=3)   #another way to creat a dict
print(dic)

l=len(dic)
print (l)             #count items

dic3=sorted(dic)      #sorted return a list of keys, not change the original one
print(dic3)         #sort based on dictionary order of keys

dic3=sorted(dic, reverse=True)      #reverse=True, get a reverse sort
print(dic3)

print(sorted(dic.items()))       #sort based on keys, display with values

print ('one' in dic)              #"in" method of dic
print('one' in dic.keys())
print('one' in dic.items())

# =============================================================================
# {'one': 1, 'two': 2, 'three': 3}
# 3
# ['one', 'three', 'two']
# ['two', 'three', 'one']
# [('one', 1), ('three', 3), ('two', 2)]
# True
# True
# False
# =============================================================================

#%%
numbers={'one':1, 'two':2, 'three':3}
nf=numbers.fromkeys(numbers)          #fromkeys, create a new dict with same keys
print(nf)

nf=numbers.fromkeys(numbers, "num")   #can assign a value to every key(same as numebrs) in nf
print(nf)

print(numbers.get('one'))    #get function is safe to get values, even key not exist
print(numbers.get('four'))

#add and remove
numbers.update({'four':4})
numbers['five']=5
print(numbers)

numbers.pop('five')
print(numbers)

numbers.popitem()      #remove the last one and return the tuple
print(numbers)

# numbers.clear()       #clear all, keep name
# print(numbers)

del numbers['three']
print (numbers)          #del one element

del numbers
# print(numbers)        #del space

# =============================================================================
# {'one': None, 'two': None, 'three': None}
# {'one': 'num', 'two': 'num', 'three': 'num'}
# 1
# None
# {'one': 1, 'two': 2, 'three': 3, 'four': 4, 'five': 5}
# {'one': 1, 'two': 2, 'three': 3, 'four': 4}
# {'one': 1, 'two': 2, 'three': 3}
# {'one': 1, 'two': 2}
# =============================================================================

#%% excercise
chance=int(input("counts?:"))
info={}
for i in range(chance):
    name=input("name:")
    height=int(input("height:"))
    info[name]=height
# print(info)

for i, j in info.items():
    print(f"name: {i}, height:{j}")

#%% exercise
n=int(input("your n:"))
d1={}
for i in range(1,n):
    d1[i]=i**2

print(d1)

#%%excercise
Dynasties = {"tang": 618, "song": 960, "yuan":
1279, "ming": 1368, "qing": 1644}

indy=input("your input:")
print(indy in Dynasties)
s=sum(list(Dynasties.values()))
print(f"The sum of values is {s}")

#%%excercise
# =============================================================================
# Write a program to count the number of characters in a word
# using dictionary. Display the character and their appeared
# times in alphabetical order:
# pneumonoultramicroscopicsilicovolcanoconiosis
#
# =============================================================================

s='pneumonoultramicroscopicsilicovolcanoconiosis'
dic={}
# for c in s:
#     if c in dic:
#         dic[c]+=1
#     else:
#         dic[c]=1

#O(n) method, common manipulation of counting letters
letter=[chr(i) for i in range(97,126)]
zero=[0 for i in range(26)]
dic=dict(zip(letter,zero))
for c in s:
    dic[c]+=1

ss=sorted(dic.items())
for i in ss:
    if i[1]:
        print(f"{i[0]}:{i[1]}")
# =============================================================================
#
# note here chr() not chr[]!
# and sorted function return a new dict, not changing the original one.
# zip()funt can generate a tuple used in dict
# =============================================================================

#%%exercise
# =============================================================================
# Write a program to calculate the number of digits, uppercase
# and lowercase letters of the following sentence:
# I am Time, the great destroyer of the world -
# Bhagavad Gita 11.32
# =============================================================================
s1='I am Time, the great destroyer of the world -Bhagavad Gita 11.32'
dic={"upper":0, "lower":0, "digits":0}
for c in s1:
    if c.isdigit():
        dic["digits"]+=1
    elif c.isupper():
        dic["upper"]+=1
    elif c.islower():
        dic["lower"]+=1

for i,j in dic.items():
    print(f"{i}:{j}")

#%%exercise
# =============================================================================
# Write a program to print the following student's name, ID, and
# average score:
# student_details = {"name": "John", "ID": "100101",
# "marks": {"python": 95, "java": 90, "C": 85}}
# dict in a dict
# =============================================================================
student_details = {"name": "John", "ID": "100101",
"marks": {"python": 95, "java": 90, "C": 85}}
for attribution, info in student_details.items():
    if type(info)==dict:
        print(f"{attribution}:")
        for i, j in info.items():
            print(f"{i}:{j}")
    else:print(f"{attribution}:{info}")


import math
from selenium import webdriver
from config.options import options
from bs4 import BeautifulSoup

def get_total_pages(soup, limit):
    try:
        job_count_response = soup.find(
            "div", class_="jobsearch-JobCountAndSortPane-jobCount"
        ).text
        job_total = int(job_count_response.split()[1][:-1].replace(",", ""))
        page_total = math.ceil(job_total / limit)
        if page_total > 10:
            page_total = 10
        print("\n", f"{format(job_total, ',')}개를 발견하였습니다.")
        return page_total
    except AttributeError:
        print("\n", "검색 결과를 찾지 못했습니다.")
        return 0

def parse_job_data(soup):
    results = []

    job_list = soup.find("ul", class_="jobsearch-ResultsList")
    jobs = job_list.find_all("li", recursive=False)
    for job in jobs:
        zone = job.find("div", class_="mosaic-zone")
        if zone is not None:
            continue
        anchor = job.select_one("h2 a", class_="jobTitle")
        position = job.find("h2", class_="jobTitle")
        link = anchor["href"]
        company = job.find("span", class_="companyName")
        location = job.find("div", class_="companyLocation")
        job_data = {
            "link": f"https://kr.indeed.com{link}",
            "company": company.string.replace(",", " "),
            "location": location.string.replace(",", " "),
            "position": position.string.replace(",", " "),
        }
        results.append(job_data)
    return results

def extract_indeed_jobs(keyword):

    base_url = "https://kr.indeed.com/jobs"
    limit = 50
    results = []

    browser = webdriver.Chrome(options=options)
    browser.get(f"{base_url}?q={keyword}")
    print("\n", f"Indeed에서 {keyword} 직업 검색을 시작합니다..", "\n")
    soup = BeautifulSoup(browser.page_source, "html.parser")

    page_total = get_total_pages(soup, limit)

    if page_total == 0:
        return []

    for page in range(page_total):
        request_url = f"{base_url}?q={keyword}&limit={limit}&start={page*limit}"
        print(f"요청중... {request_url}")

        browser.get(request_url)
        soup = BeautifulSoup(browser.page_source, "html.parser")
        result = parse_job_data(soup)
        results += result

    return results


import sys
from PyQt5.QtWidgets import QApplication, QWidget
from PyQt5.QtGui import QIcon
from PyQt5 import uic

#load both ui file
uifile_1 = 'Main.ui'
form_1, base_1 = uic.loadUiType(uifile_1)

uifile_2 = 'Station_1.ui'
form_2, base_2 = uic.loadUiType(uifile_2)

class Example(base_2, form_2):
    def __init__(self):
        super(base_2,self).__init__()
        self.setupUi(self)
        #self.pushButton14.clicked.connect(self.change)

    def change(self):
        self.main = MainPage()
        self.main.show()
        self.close()

class MainPage(base_1, form_1):
    def __init__(self):
        super(base_1, self).__init__()
        self.setupUi(self)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = Example()
    ex.show()
    sys.exit(app.exec_())


import numpy as np
from .tanh_layer import tanh
from .relu_layer import ReLU
from .prelu_layer import pReLU
from .leakyrelu_layer import LeakyReLU
from .sigmoid_layer import Sigmoid
from .out_layer import outactivation
from .conv_layer import convolution
from .SPP_layer import SPP
from .max_pool_layer import maxpool
from .flatten_layer import flatten
from TSML.TSPreprocessandCost import *

activ_functs={}
activ_functs['tanh']=tanh
activ_functs['relu']=ReLU
activ_functs['prelu']=pReLU
activ_functs['leakyrelu']=LeakyReLU
activ_functs['sigmoid']=Sigmoid
activ_functs['outact']=outactivation
activ_functs['conv']=convolution
activ_functs['spp']=SPP
activ_functs['maxpool']=maxpool
activ_functs['flatten']=flatten

class NeuralNet:
    def __init__(self, nodes, indims, activations, filtersize=[], cost=False, task='Classification', solver="Basic", seed=False):
        self.nodes=nodes
        self.layers=len(nodes)
        self.indims=indims
        activations.append('outact')
        self.activations=activations
        self.task=task.lower()
        self.t=0
        if not cost and task.lower()=='classification':
            if nodes[-1]==1:
                self.cost=BCEC
            else:
                self.cost=GCEC
        elif not cost and task.lower()=='regression':
            if nodes[-1]==1:
                self.cost=SSE
            else:
                self.cost=FSSE
        if self.task !='regression' and self.task !='classification':
            print("Invalid task, use 'Regression' or 'Classification'")
            return None
        solvers=['momentum','nesterov','nesterovrms','nesterovada','momentumada','momentumrms','rmsprop','adagrad','basic', 'adam']

        self.y_classes=np.array([])
        self.Xmean=np.array([])
        self.seed=seed
        self.solver=solver.lower()
        if self.solver not in solvers:
            print("Possible solvers are " +', '.join(solvers))
            return None
        self.layer_classes=[]
        nodes=[indims]+nodes
        f=0
        for i in range(self.layers-1):
            if activations[i].lower()=='conv':
                self.layer_classes.append(activ_functs[activations[i].lower()](nodes[i+1],filtersize[f],nodes[i],solver=self.solver,seed=self.seed))
                f+=1
            else:
                self.layer_classes.append(activ_functs[activations[i].lower()](nodes[i+1],nodes[i],solver=self.solver,seed=self.seed))

        self.layer_classes.append(activ_functs[activations[self.layers-1]](nodes[self.layers],nodes[self.layers-1],solver=self.solver,
                                    seed=self.seed, task=self.task))

    def forward(self, X, p, normalize=True):

        if not self.dropout:
            self.predict(X,False)
        else:
            if not self.Xmean.size:
                return("Please train before predicting." )

            elif normalize:
                X=norm(X,self.Xmean,self.Xstd)
            A=X
            if 'nesterov' in self.solver:
                for i in range(self.layers):
                    A=self.layer_classes[i].nesterov_forwardprop_dropout(A,self.mu,p[i])
            else:
                for i in range(self.layers):
                    A=self.layer_classes[i].forwardprop_dropout(A,p[i])

            if self.task=='classification':
                self.probabilities=A
            else:
                self.normpredictions=A

    def predict(self, X, normalize=True):
        if not self.Xmean.size:
            return("Please train before predicting." )

        elif normalize:
            X=norm(X,self.Xmean,self.Xstd)
        A=X
        if 'nesterov' in self.solver:
            for i in range(self.layers):
                A=self.layer_classes[i].nesterov_forwardprop(A,self.mu)
        else:
            for i in range(self.layers):
                A=self.layer_classes[i].forwardprop(A)

        if self.task=='classification':
            self.probabilities=A
            if self.convolutional==False:
                self.ohe_predictions=np.eye(self.probabilities.shape[1])[np.argmax(self.probabilities, axis=1)]

                if self.nodes[-1]>1 and self.y_classes.size:
                     self.predictions=self.y_classes[np.argmax(self.probabilities, axis=1)].reshape(-1,1)

                elif self.nodes[-1]>1:
                    self.predictions=self.ohe_predictions

                else:
                    self.bin_predictions=np.rint(self.probabilities)
                    self.predictions=np.where(self.bin_predictions==1,self.positive_class,self.negative_class).reshape(-1,1)

        else:
            self.normpredictions=A
            if convolutional==False:
                self.predictions=unnorm(self.normpredictions,self.ymean,self.ystd)

    def weight_initialization(self):
        self.t=0
        for i in range(self.layers-1,-1,-1):
            self.layer_classes[i].weightinit(self.seed)

    def weight_update(self,eta,y):
        D=y
        if self.solver=='nesterov':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].nesterov_backprop(eta,self.mu,D,self.lam1,self.lam2)

        elif self.solver=='basic':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].backprop(eta,D,self.lam1,self.lam2)

        elif self.solver=='momentum':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].momentum_backprop(eta,self.mu,D,self.lam1,self.lam2)

        elif self.solver=='adagrad':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].ada_backprop(eta,D,self.nesterov,self.lam1,self.lam2)

        elif self.solver=='rmsprop':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].RMS_prop(eta,self.gamma,D,self.nesterov,self.lam1,self.lam2)

        elif self.solver=='adam':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].adamoptimizer(eta,self.mu,self.gamma,self.t,D,self.lam1,self.lam2)

        elif self.solver=='momentumrms':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].momentum_RMS_prop(eta,self.mu,self.gamma,D,self.lam1,self.lam2)

        elif self.solver=='momentumada':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].momentum_ada_backprop(eta,self.mu,D,self.lam1,self.lam2)

        elif self.solver=='nesterovrms':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].nesterov_RMS_prop(eta,self.mu,D,self.lam1,self.lam2)

        elif self.solver=='nesterovada':
            for i in range(self.layers-1,-1,-1):
                D=self.layer_classes[i].nesterov_ada_backprop(eta,self.mu,D,self.lam1,self.lam2)

    def train(self,X,y,epochs,eta,mu=0.1,gamma=.9,lam1=0,lam2=0,decay=False,k=0,T=1,batch_size=0,error_calc=True,reinitialize=False,dropout=False,
                p=[],X_val=np.array([]),y_val=np.array([]), convolutional=False):
        if reinitialize:
            self.weight_initialization()
        if batch_size==0:
            batch_size=X.shape[0]
        self.valerror=[]
        self.error=[]
        self.Xmean=X.mean(0)
        self.Xstd=X.std(0)
        self.gamma=gamma
        self.mu=mu
        self.nesterov=False
        self.lam1=lam1
        self.lam2=lam2
        self.convolutional=convolutional
        eta0=eta
        if len(p)==1:
            p=p*(self.layers)
        self.dropout=dropout
        if convolutional==False:
            X=norm(X,self.Xmean,self.Xstd)

            if X_val.size:
                X_val=norm(X_val,self.Xmean,self.Xstd)
            if (self.task=="classification" and y.shape[1]<self.nodes[-1]):
                self.y_classes=np.unique(y)
                y=mat_ohe(y,[0])
                if y_val.size:
                    y_val=mat_ohe(y_val,[0])

            elif (self.task=="classification" and y.shape[1]==self.nodes[-1]):
                self.positive_class=np.unique(y)[1]
                self.negative_class=np.unique(y)[0]
                y=np.where(y==self.positive_class,1,0)
                if y_val.size:
                    y_val=np.where(y_val==self.positive_class,1,0)

            elif self.task=="regression":
                self.ymean=y.mean(0)
                self.ystd=y.std(0)
                y=norm(y,self.ymean,self.ystd)

                if y_val.size:
                    y_val=norm(y_val,self.ymean,self.ystd)

        batches=np.ceil(X.shape[0]/batch_size).astype(int)
        for i in range(epochs):
            if convolutional==False:
                train=np.hstack((X,y))
                np.random.shuffle(train)
                train=np.array_split(train,batches)
            else:
                train=[0]
            for a in train:
                if convolutional==False:
                    X_temp=np.array(a[:,:-self.nodes[-1]])
                    y_temp=np.array(a[:,-self.nodes[-1]:])
                else:
                    X_temp=X
                    y_temp=y

                if self.solver == 'nesterovada':
                    self.solver ='adagrad'
                    self.nesterov=True
                    self.forward(X_temp,p,False)
                    self.weight_update(eta,y_temp)
                    self.solver='nesterovada'
                    self.forward(X_temp,p,False)
                    self.weight_update(eta,y_temp)
                elif self.solver == 'nesterovrms':
                    self.solver ='rmsprop'
                    self.nesterov=True
                    self.forward(X_temp,p,False)
                    self.weight_update(eta,y_temp)
                    self.solver='nesterovrms'
                    self.forward(X_temp,p,False)
                    self.weight_update(eta,y_temp)
                else:
                    self.forward(X_temp,p,False)
                    self.weight_update(eta,y_temp)

                if error_calc and self.task=='classification':
                    self.forward(X_temp,p,False)
                    self.error.append(self.cost(y_temp,self.probabilities))

                elif error_calc:
                    self.forward(X_temp,p,False)
                    self.error.append(self.cost(y_temp,self.normpredictions))
                self.t += 1

                if decay=="Scheduled":
                    eta=eta0*k**(self.t/T)
                elif decay=="Inverse":
                    eta=eta0/(k*self.t+1)
                elif decay=="Exponential":
                    eta=eta0*np.exp(-k*self.t)

                if convolutional==False:
                    if self.task=='classification':
                        self.predict(X,False)
                        if self.nodes[-1]>1:
                            acc=(self.ohe_predictions==y).mean()

                        elif self.nodes[-1]==1:
                            acc=(self.bin_predictions==y).mean()
                        if acc==1.0:
                            return acc
            if self.task=='classification':
                if y_val.size:
                    self.predict(X_val,False)
                    self.valerror.append(self.cost(y_val,self.probabilities))

            else:
                if y_val.size:
                    self.predict(X_val,False)
                    self.valerror.append(self.cost(y_val,self.normpredictions))


import dose
import uuid
from mongokit import Document
from app.DAO import mongo, configParser

__author__ = 'Luka Strizic'

class Institution:
    def __init__(self, document):
        assert(isinstance(document, InstitutionDocument))
        self.document = document

    def __getitem__(self, item):
        ret = None
        try:
            ret = self.document[str(item)]
        except:
            pass
        return ret

    def save(self):
        self.document.save()

    def setCountry(self, country, save=False):
        self.document['country'] = country
        if(save): self.save()
        return

    def setCity(self, city, save=False):
        self.document['city'] = city
        if(save): self.save()
        return

    def setAddress(self, address, save=False):
        self.document['address'] = address
        if(save): self.save()
        return

    def setName(self, name, save=False):
        self.document['name'] = name
        if(save): self.save()
        return

    def activate(self, save = False):
        self.document['active'] = True
        if(save): self.save()
        return

    def deactivate(self, save = False):
        self.document['active'] = False
        if(save): self.save()
        return

    def setAplusLow(self, low, save = False):
        self.document['A+low'] = low
        if(save): self.save()
        return

    def setAminusLow(self, low, save = False):
        self.document['A-low'] = low
        if(save): self.save()
        return

    def setBplusLow(self, low, save = False):
        self.document['B+low'] = low
        if(save): self.save()
        return

    def setBminusLow(self, low, save = False):
        self.document['B-low'] = low
        if(save): self.save()
        return

    def setABplusLow(self, low, save = False):
        self.document['AB+low'] = low
        if(save): self.save()
        return

    def setABminusLow(self, low, save = False):
        self.document['AB-low'] = low
        if(save): self.save()
        return

    def set0plusLow(self, low, save = False):
        self.document['0+low'] = low
        if(save): self.save()
        return

    def set0minusLow(self, low, save = False):
        self.document['0-low'] = low
        if(save): self.save()
        return

    def getAplus(self):
        return mongo.DoseDocument.find({'institutionID' : self.document.id,
                                        'AB0' : 'A', 'RH' : '+',
                                        'status' : 'available'}).count()

    def getAminus(self):
        return mongo.DoseDocument.find({'institutionID' : self.document.id,
                                        'AB0' : 'A',
                                        'RH' : '-',
                                        'status' : 'available'}).count()

    def getBplus(self):
        return mongo.DoseDocument.find({'institutionID' : self.document.id,
                                        'AB0' : 'B',
                                        'RH' : '+',
                                        'status' : 'available'}).count()

    def getBminus(self):
        return mongo.DoseDocument.find({'institutionID' : self.document.id,
                                        'AB0' : 'B',
                                        'RH' : '-',
                                        'status' : 'available'}).count()

    def getABplus(self):
        return mongo.DoseDocument.find({'institutionID' : self.document.id,
                                        'AB0' : 'AB',
                                        'RH' : '+',
                                        'status' : 'available'}).count()

    def getABminus(self):
        return mongo.DoseDocument.find({'institutionID' : self.document.id,
                                        'AB0' : 'AB',
                                        'RH' : '-',
                                        'status' : 'available'}).count()

    def get0plus(self):
        return mongo.DoseDocument.find({'institutionID' : self.document.id,
                                        'AB0' : '0',
                                        'RH' : '+',
                                        'status' : 'available'}).count()

    def get0minus(self):
        return mongo.DoseDocument.find({'institutionID' : self.document.id,
                                        'AB0' : '0',
                                        'RH' : '-',
                                        'status' : 'available'}).count()

    def warning(self):
        ret =  self.getAplus() <= self.document['A+low'] or \
                self.getAminus() <= self.document['A-low'] or \
                self.getBplus() <= self.document['B+low'] or \
                self.getBminus() <= self.document['B-low'] or \
                self.getABplus() <= self.document['AB+low'] or \
                self.getABminus() <= self.document['AB-low'] or \
                self.get0plus() <= self.document['0+low'] or \
                self.get0minus() <= self.document['0-low']
        return ret

    def get_string_presentation(self):
        return self["name"]+(", "+self["city"]) if self["city"] else""

def create_institution_from_request(req):
    doc = mongo.InstitutionDocument()
    doc['name'] = req['name']
    doc['city'] = req['city']
    doc['address'] = req['address']
    doc['country'] = req['country']
    return Institution(doc)

def get_by_id(id):
    doc = mongo.InstitutionDocument.find_one({'id':id})
    if (doc):
        return Institution(doc)
    return None

def get_all_institutions_cursor():
    return mongo.InstitutionDocument.find()

def get_all_institutions_array():
    ret =[]
    for inst in get_all_institutions_cursor():
        ret.append(Institution(inst))
    return ret

def uuidStr():
    return str(uuid.uuid4())

@mongo.register
class InstitutionDocument(Document):
    __database__ = configParser.get("Mongo","DBname")
    __collection__ = 'institutions'

    use_schemaless = True
    use_dot_notation = True

    structure = {
        'country' : basestring,
        'city' : basestring,
        'address' : basestring,
        'name' : basestring,
        'id' : basestring,
        'A+low' : float,
        'A-low' : float,
        'B+low' : float,
        'B-low' : float,
        'AB+low' : float,
        'AB-low' : float,
        '0+low' : float,
        '0-low' : float,
        'active' : bool
    }

    default_values = {
        'id' : uuidStr,
        'A+low' : 1.0,
        'A-low' : 1.0,
        'B+low' : 1.0,
        'B-low' : 1.0,
        'AB+low' : 1.0,
        'AB-low' : 1.0,
        '0+low' : 1.0,
        '0-low' : 1.0,
        'active' : True
    }


import datetime
import os
import threading
import time
import unittest
from typing import cast
from collections import namedtuple

from pyspark import SparkContext, SparkConf
from pyspark.sql import Row, SparkSession
from pyspark.sql.functions import rand, udf, assert_true, lit
from pyspark.sql.types import (
    StructType,
    StringType,
    IntegerType,
    LongType,
    FloatType,
    DoubleType,
    DecimalType,
    DateType,
    TimestampType,
    TimestampNTZType,
    BinaryType,
    StructField,
    ArrayType,
    MapType,
    NullType,
    DayTimeIntervalType,
)
from pyspark.testing.sqlutils import (
    ReusedSQLTestCase,
    have_pandas,
    have_pyarrow,
    pandas_requirement_message,
    pyarrow_requirement_message,
    ExamplePoint,
    ExamplePointUDT,
)
from pyspark.testing.utils import QuietTest
from pyspark.errors import ArithmeticException, PySparkTypeError, UnsupportedOperationException

if have_pandas:
    import pandas as pd
    from pandas.testing import assert_frame_equal

if have_pyarrow:
    import pyarrow as pa  # noqa: F401

class ArrowTestsMixin:
    @classmethod
    def setUpClass(cls):
        from datetime import date, datetime
        from decimal import Decimal

        super().setUpClass()
        cls.warnings_lock = threading.Lock()

        # Synchronize default timezone between Python and Java
        cls.tz_prev = os.environ.get("TZ", None)  # save current tz if set
        tz = "America/Los_Angeles"
        os.environ["TZ"] = tz
        time.tzset()

        cls.spark.conf.set("spark.sql.session.timeZone", tz)

        # Test fallback
        cls.spark.conf.set("spark.sql.execution.arrow.enabled", "false")
        assert cls.spark.conf.get("spark.sql.execution.arrow.pyspark.enabled") == "false"
        cls.spark.conf.set("spark.sql.execution.arrow.enabled", "true")
        assert cls.spark.conf.get("spark.sql.execution.arrow.pyspark.enabled") == "true"

        cls.spark.conf.set("spark.sql.execution.arrow.fallback.enabled", "true")
        assert cls.spark.conf.get("spark.sql.execution.arrow.pyspark.fallback.enabled") == "true"
        cls.spark.conf.set("spark.sql.execution.arrow.fallback.enabled", "false")
        assert cls.spark.conf.get("spark.sql.execution.arrow.pyspark.fallback.enabled") == "false"

        # Enable Arrow optimization in this tests.
        cls.spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
        # Disable fallback by default to easily detect the failures.
        cls.spark.conf.set("spark.sql.execution.arrow.pyspark.fallback.enabled", "false")

        cls.schema_wo_null = StructType(
            [
                StructField("1_str_t", StringType(), True),
                StructField("2_int_t", IntegerType(), True),
                StructField("3_long_t", LongType(), True),
                StructField("4_float_t", FloatType(), True),
                StructField("5_double_t", DoubleType(), True),
                StructField("6_decimal_t", DecimalType(38, 18), True),
                StructField("7_date_t", DateType(), True),
                StructField("8_timestamp_t", TimestampType(), True),
                StructField("9_binary_t", BinaryType(), True),
            ]
        )
        cls.schema = cls.schema_wo_null.add("10_null_t", NullType(), True)
        cls.data_wo_null = [
            (
                "a",
                1,
                10,
                0.2,
                2.0,
                Decimal("2.0"),
                date(1969, 1, 1),
                datetime(1969, 1, 1, 1, 1, 1),
                bytearray(b"a"),
            ),
            (
                "b",
                2,
                20,
                0.4,
                4.0,
                Decimal("4.0"),
                date(2012, 2, 2),
                datetime(2012, 2, 2, 2, 2, 2),
                bytearray(b"bb"),
            ),
            (
                "c",
                3,
                30,
                0.8,
                6.0,
                Decimal("6.0"),
                date(2100, 3, 3),
                datetime(2100, 3, 3, 3, 3, 3),
                bytearray(b"ccc"),
            ),
            (
                "d",
                4,
                40,
                1.0,
                8.0,
                Decimal("8.0"),
                date(2262, 4, 12),
                datetime(2262, 3, 3, 3, 3, 3),
                bytearray(b"dddd"),
            ),
        ]
        cls.data = [tuple(list(d) + [None]) for d in cls.data_wo_null]

    @classmethod
    def tearDownClass(cls):
        del os.environ["TZ"]
        if cls.tz_prev is not None:
            os.environ["TZ"] = cls.tz_prev
        time.tzset()
        super().tearDownClass()

    def create_pandas_data_frame(self):
        import numpy as np

        data_dict = {}
        for j, name in enumerate(self.schema.names):
            data_dict[name] = [self.data[i][j] for i in range(len(self.data))]
        # need to convert these to numpy types first
        data_dict["2_int_t"] = np.int32(data_dict["2_int_t"])
        data_dict["4_float_t"] = np.float32(data_dict["4_float_t"])
        return pd.DataFrame(data=data_dict)

    @property
    def create_np_arrs(self):
        import numpy as np

        int_dtypes = ["int8", "int16", "int32", "int64"]
        float_dtypes = ["float32", "float64"]
        return (
            [np.array([1, 2]).astype(t) for t in int_dtypes]
            + [np.array([0.1, 0.2]).astype(t) for t in float_dtypes]
            + [np.array([[1], [2]]).astype(t) for t in int_dtypes]
            + [np.array([[0.1], [0.2]]).astype(t) for t in float_dtypes]
            + [np.array([[1, 1, 1], [2, 2, 2]]).astype(t) for t in int_dtypes]
            + [np.array([[0.1, 0.1, 0.1], [0.2, 0.2, 0.2]]).astype(t) for t in float_dtypes]
        )

    def test_toPandas_empty_df_arrow_enabled(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_toPandas_empty_df_arrow_enabled(arrow_enabled)

    def check_toPandas_empty_df_arrow_enabled(self, arrow_enabled):
        # SPARK-30537 test that toPandas() on an empty dataframe has the correct dtypes
        # when arrow is enabled
        from datetime import date
        from decimal import Decimal

        schema = StructType(
            [
                StructField("a", StringType(), True),
                StructField("a", IntegerType(), True),
                StructField("c", TimestampType(), True),
                StructField("d", NullType(), True),
                StructField("e", LongType(), True),
                StructField("f", FloatType(), True),
                StructField("g", DateType(), True),
                StructField("h", BinaryType(), True),
                StructField("i", DecimalType(38, 18), True),
                StructField("k", TimestampNTZType(), True),
                StructField("L", DayTimeIntervalType(0, 3), True),
            ]
        )
        df = self.spark.createDataFrame([], schema=schema)
        non_empty_df = self.spark.createDataFrame(
            [
                (
                    "a",
                    1,
                    datetime.datetime(1969, 1, 1, 1, 1, 1),
                    None,
                    10,
                    0.2,
                    date(1969, 1, 1),
                    bytearray(b"a"),
                    Decimal("2.0"),
                    datetime.datetime(1969, 1, 1, 1, 1, 1),
                    datetime.timedelta(microseconds=123),
                )
            ],
            schema=schema,
        )

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            pdf = df.toPandas()
            pdf_non_empty = non_empty_df.toPandas()
        self.assertTrue(pdf.dtypes.equals(pdf_non_empty.dtypes))

    def test_null_conversion(self):
        df_null = self.spark.createDataFrame(
            [tuple([None for _ in range(len(self.data_wo_null[0]))])] + self.data_wo_null
        )
        pdf = df_null.toPandas()
        null_counts = pdf.isnull().sum().tolist()
        self.assertTrue(all([c == 1 for c in null_counts]))

    def _toPandas_arrow_toggle(self, df):
        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}):
            pdf = df.toPandas()

        pdf_arrow = df.toPandas()

        return pdf, pdf_arrow

    def test_toPandas_arrow_toggle(self):
        df = self.spark.createDataFrame(self.data, schema=self.schema)
        pdf, pdf_arrow = self._toPandas_arrow_toggle(df)
        expected = self.create_pandas_data_frame()
        assert_frame_equal(expected, pdf)
        assert_frame_equal(expected, pdf_arrow)

    def test_create_data_frame_to_pandas_timestamp_ntz(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_create_data_frame_to_pandas_timestamp_ntz(arrow_enabled)

    def check_create_data_frame_to_pandas_timestamp_ntz(self, arrow_enabled):
        # SPARK-36626: Test TimestampNTZ in createDataFrame and toPandas
        with self.sql_conf({"spark.sql.session.timeZone": "America/Los_Angeles"}):
            origin = pd.DataFrame({"a": [datetime.datetime(2012, 2, 2, 2, 2, 2)]})
            df = self.spark.createDataFrame(
                origin, schema=StructType([StructField("a", TimestampNTZType(), True)])
            )
            df.selectExpr("assert_true('2012-02-02 02:02:02' == CAST(a AS STRING))").collect()

            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                pdf = df.toPandas()
            assert_frame_equal(origin, pdf)

    def test_create_data_frame_to_pandas_day_time_internal(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_create_data_frame_to_pandas_day_time_internal(arrow_enabled)

    def check_create_data_frame_to_pandas_day_time_internal(self, arrow_enabled):
        # SPARK-37279: Test DayTimeInterval in createDataFrame and toPandas
        origin = pd.DataFrame({"a": [datetime.timedelta(microseconds=123)]})
        df = self.spark.createDataFrame(origin)
        df.select(
            assert_true(lit("INTERVAL '0 00:00:00.000123' DAY TO SECOND") == df.a.cast("string"))
        ).collect()

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            pdf = df.toPandas()
        assert_frame_equal(origin, pdf)

    def test_toPandas_respect_session_timezone(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_toPandas_respect_session_timezone(arrow_enabled)

    def check_toPandas_respect_session_timezone(self, arrow_enabled):
        df = self.spark.createDataFrame(self.data, schema=self.schema)

        timezone = "America/Los_Angeles"
        with self.sql_conf({"spark.sql.session.timeZone": timezone}):
            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                pdf_la = df.toPandas()

        timezone = "America/New_York"
        with self.sql_conf({"spark.sql.session.timeZone": timezone}):
            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                pdf_ny = df.toPandas()

            self.assertFalse(pdf_ny.equals(pdf_la))

            from pyspark.sql.pandas.types import _check_series_convert_timestamps_local_tz

            pdf_la_corrected = pdf_la.copy()
            for field in self.schema:
                if isinstance(field.dataType, TimestampType):
                    pdf_la_corrected[field.name] = _check_series_convert_timestamps_local_tz(
                        pdf_la_corrected[field.name], timezone
                    )
            assert_frame_equal(pdf_ny, pdf_la_corrected)

    def test_pandas_round_trip(self):
        pdf = self.create_pandas_data_frame()
        df = self.spark.createDataFrame(self.data, schema=self.schema)
        pdf_arrow = df.toPandas()
        assert_frame_equal(pdf_arrow, pdf)

    def test_pandas_self_destruct(self):
        import pyarrow as pa

        rows = 2**10
        cols = 4
        expected_bytes = rows * cols * 8
        df = self.spark.range(0, rows).select(*[rand() for _ in range(cols)])
        # Test the self_destruct behavior by testing _collect_as_arrow directly
        allocation_before = pa.total_allocated_bytes()
        batches = df._collect_as_arrow(split_batches=True)
        table = pa.Table.from_batches(batches)
        del batches
        pdf_split = table.to_pandas(self_destruct=True, split_blocks=True, use_threads=False)
        allocation_after = pa.total_allocated_bytes()
        difference = allocation_after - allocation_before
        # Should be around 1x the data size (table should not hold on to any memory)
        self.assertGreaterEqual(difference, 0.9 * expected_bytes)
        self.assertLessEqual(difference, 1.1 * expected_bytes)

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.selfDestruct.enabled": False}):
            no_self_destruct_pdf = df.toPandas()
            # Note while memory usage is 2x data size here (both table and pdf hold on to
            # memory), in this case Arrow still only tracks 1x worth of memory (since the
            # batches are not allocated by Arrow in this case), so we can't make any
            # assertions here

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.selfDestruct.enabled": True}):
            self_destruct_pdf = df.toPandas()

        assert_frame_equal(pdf_split, no_self_destruct_pdf)
        assert_frame_equal(pdf_split, self_destruct_pdf)

    def test_filtered_frame(self):
        df = self.spark.range(3).toDF("i")
        pdf = df.filter("i < 0").toPandas()
        self.assertEqual(len(pdf.columns), 1)
        self.assertEqual(pdf.columns[0], "i")
        self.assertTrue(pdf.empty)

    def test_no_partition_frame(self):
        schema = StructType([StructField("field1", StringType(), True)])
        df = self.spark.createDataFrame(self.sc.emptyRDD(), schema)
        pdf = df.toPandas()
        self.assertEqual(len(pdf.columns), 1)
        self.assertEqual(pdf.columns[0], "field1")
        self.assertTrue(pdf.empty)

    def test_propagates_spark_exception(self):
        with QuietTest(self.sc):
            self.check_propagates_spark_exception()

    def check_propagates_spark_exception(self):
        df = self.spark.range(3).toDF("i")

        def raise_exception():
            raise RuntimeError("My error")

        exception_udf = udf(raise_exception, IntegerType())
        df = df.withColumn("error", exception_udf())

        with self.assertRaisesRegex(Exception, "My error"):
            df.toPandas()

    def _createDataFrame_toggle(self, data, schema=None):
        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}):
            df_no_arrow = self.spark.createDataFrame(data, schema=schema)

        df_arrow = self.spark.createDataFrame(data, schema=schema)

        return df_no_arrow, df_arrow

    def test_createDataFrame_toggle(self):
        pdf = self.create_pandas_data_frame()
        df_no_arrow, df_arrow = self._createDataFrame_toggle(pdf, schema=self.schema)
        self.assertEqual(df_no_arrow.collect(), df_arrow.collect())

    def test_createDataFrame_respect_session_timezone(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_createDataFrame_respect_session_timezone(arrow_enabled)

    def check_createDataFrame_respect_session_timezone(self, arrow_enabled):
        from datetime import timedelta

        pdf = self.create_pandas_data_frame()
        timezone = "America/Los_Angeles"
        with self.sql_conf({"spark.sql.session.timeZone": timezone}):
            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                df_la = self.spark.createDataFrame(pdf, schema=self.schema)
            result_la = df_la.collect()

        timezone = "America/New_York"
        with self.sql_conf({"spark.sql.session.timeZone": timezone}):
            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                df_ny = self.spark.createDataFrame(pdf, schema=self.schema)
            result_ny = df_ny.collect()

            self.assertNotEqual(result_ny, result_la)

            # Correct result_la by adjusting 3 hours difference between Los Angeles and New York
            result_la_corrected = [
                Row(
                    **{
                        k: v - timedelta(hours=3) if k == "8_timestamp_t" else v
                        for k, v in row.asDict().items()
                    }
                )
                for row in result_la
            ]
            self.assertEqual(result_ny, result_la_corrected)

    def test_createDataFrame_with_schema(self):
        pdf = self.create_pandas_data_frame()
        df = self.spark.createDataFrame(pdf, schema=self.schema)
        self.assertEqual(self.schema, df.schema)
        pdf_arrow = df.toPandas()
        assert_frame_equal(pdf_arrow, pdf)

    def test_createDataFrame_with_incorrect_schema(self):
        with QuietTest(self.sc):
            self.check_createDataFrame_with_incorrect_schema()

    def check_createDataFrame_with_incorrect_schema(self):
        pdf = self.create_pandas_data_frame()
        fields = list(self.schema)
        fields[5], fields[6] = fields[6], fields[5]  # swap decimal with date
        wrong_schema = StructType(fields)
        with self.sql_conf({"spark.sql.execution.pandas.convertToArrowArraySafely": False}):
            with self.assertRaises(Exception) as context:
                self.spark.createDataFrame(pdf, schema=wrong_schema)

            # the exception provides us with the column that is incorrect
            exception = context.exception
            self.assertTrue(hasattr(exception, "args"))
            self.assertEqual(len(exception.args), 1)
            self.assertRegex(
                exception.args[0],
                "with name '7_date_t' " "to Arrow Array \\(decimal128\\(38, 18\\)\\)",
            )

            # the inner exception provides us with the incorrect types
            exception = exception.__context__
            self.assertTrue(hasattr(exception, "args"))
            self.assertEqual(len(exception.args), 1)
            self.assertRegex(exception.args[0], "[D|d]ecimal.*got.*date")

    def test_createDataFrame_with_names(self):
        pdf = self.create_pandas_data_frame()
        new_names = list(map(str, range(len(self.schema.fieldNames()))))
        # Test that schema as a list of column names gets applied
        df = self.spark.createDataFrame(pdf, schema=list(new_names))
        self.assertEqual(df.schema.fieldNames(), new_names)
        # Test that schema as tuple of column names gets applied
        df = self.spark.createDataFrame(pdf, schema=tuple(new_names))
        self.assertEqual(df.schema.fieldNames(), new_names)

    def test_createDataFrame_column_name_encoding(self):
        pdf = pd.DataFrame({"a": [1]})
        columns = self.spark.createDataFrame(pdf).columns
        self.assertTrue(isinstance(columns[0], str))
        self.assertEqual(columns[0], "a")
        columns = self.spark.createDataFrame(pdf, ["b"]).columns
        self.assertTrue(isinstance(columns[0], str))
        self.assertEqual(columns[0], "b")

    def test_createDataFrame_with_single_data_type(self):
        with QuietTest(self.sc):
            self.check_createDataFrame_with_single_data_type()

    def check_createDataFrame_with_single_data_type(self):
        for schema in ["int", IntegerType()]:
            with self.subTest(schema=schema):
                with self.assertRaises(PySparkTypeError) as pe:
                    self.spark.createDataFrame(pd.DataFrame({"a": [1]}), schema=schema).collect()

                self.check_error(
                    exception=pe.exception,
                    error_class="UNSUPPORTED_DATA_TYPE_FOR_ARROW",
                    message_parameters={"data_type": "IntegerType()"},
                )

    def test_createDataFrame_does_not_modify_input(self):
        # Some series get converted for Spark to consume, this makes sure input is unchanged
        pdf = self.create_pandas_data_frame()
        # Use a nanosecond value to make sure it is not truncated
        pdf.iloc[0, 7] = pd.Timestamp(1)
        # Integers with nulls will get NaNs filled with 0 and will be casted
        pdf.iloc[1, 1] = None
        pdf_copy = pdf.copy(deep=True)
        self.spark.createDataFrame(pdf, schema=self.schema)
        self.assertTrue(pdf.equals(pdf_copy))

    def test_schema_conversion_roundtrip(self):
        from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema

        arrow_schema = to_arrow_schema(self.schema)
        schema_rt = from_arrow_schema(arrow_schema, prefer_timestamp_ntz=True)
        self.assertEqual(self.schema, schema_rt)

    def test_createDataFrame_with_ndarray(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_createDataFrame_with_ndarray(arrow_enabled)

    def check_createDataFrame_with_ndarray(self, arrow_enabled):
        import numpy as np

        dtypes = ["tinyint", "smallint", "int", "bigint", "float", "double"]
        expected_dtypes = (
            [[("value", t)] for t in dtypes]
            + [[("value", t)] for t in dtypes]
            + [[("_1", t), ("_2", t), ("_3", t)] for t in dtypes]
        )
        arrs = self.create_np_arrs

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            for arr, dtypes in zip(arrs, expected_dtypes):
                df = self.spark.createDataFrame(arr)
                self.assertEqual(df.dtypes, dtypes)
                np.array_equal(np.array(df.collect()), arr)

            with self.assertRaisesRegex(
                ValueError, "NumPy array input should be of 1 or 2 dimensions"
            ):
                self.spark.createDataFrame(np.array(0))

    def test_createDataFrame_with_array_type(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_createDataFrame_with_array_type(arrow_enabled)

    def check_createDataFrame_with_array_type(self, arrow_enabled):
        pdf = pd.DataFrame({"a": [[1, 2], [3, 4]], "b": [["x", "y"], ["y", "z"]]})
        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            df = self.spark.createDataFrame(pdf)
        result = df.collect()
        expected = [tuple(list(e) for e in rec) for rec in pdf.to_records(index=False)]
        for r in range(len(expected)):
            for e in range(len(expected[r])):
                self.assertTrue(expected[r][e] == result[r][e])

    def test_toPandas_with_array_type(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_toPandas_with_array_type(arrow_enabled)

    def check_toPandas_with_array_type(self, arrow_enabled):
        expected = [([1, 2], ["x", "y"]), ([3, 4], ["y", "z"])]
        array_schema = StructType(
            [StructField("a", ArrayType(IntegerType())), StructField("b", ArrayType(StringType()))]
        )
        df = self.spark.createDataFrame(expected, schema=array_schema)
        pdf = df.toPandas()
        result = [tuple(list(e) for e in rec) for rec in pdf.to_records(index=False)]
        for r in range(len(expected)):
            for e in range(len(expected[r])):
                self.assertTrue(expected[r][e] == result[r][e])

    def test_createDataFrame_with_map_type(self):
        with QuietTest(self.sc):
            for arrow_enabled in [True, False]:
                with self.subTest(arrow_enabled=arrow_enabled):
                    self.check_createDataFrame_with_map_type(arrow_enabled)

    def check_createDataFrame_with_map_type(self, arrow_enabled):
        map_data = [{"a": 1}, {"b": 2, "c": 3}, {}, None, {"d": None}]

        pdf = pd.DataFrame({"id": [0, 1, 2, 3, 4], "m": map_data})
        for schema in (
            "id long, m map<string, long>",
            StructType().add("id", LongType()).add("m", MapType(StringType(), LongType())),
        ):
            with self.subTest(schema=schema):
                with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                    df = self.spark.createDataFrame(pdf, schema=schema)

                    result = df.collect()

                    for row in result:
                        i, m = row
                        self.assertEqual(m, map_data[i])

    def test_createDataFrame_with_struct_type(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_createDataFrame_with_struct_type(arrow_enabled)

    def check_createDataFrame_with_struct_type(self, arrow_enabled):
        pdf = pd.DataFrame(
            {"a": [Row(1, "a"), Row(2, "b")], "b": [{"s": 3, "t": "x"}, {"s": 4, "t": "y"}]}
        )
        for schema in (
            "a struct<x int, y string>, b struct<s int, t string>",
            StructType()
            .add("a", StructType().add("x", LongType()).add("y", StringType()))
            .add("b", StructType().add("s", LongType()).add("t", StringType())),
        ):
            with self.subTest(schema=schema):
                with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                    df = self.spark.createDataFrame(pdf, schema)
                result = df.collect()
                expected = [(rec[0], Row(**rec[1])) for rec in pdf.to_records(index=False)]
                for r in range(len(expected)):
                    for e in range(len(expected[r])):
                        self.assertTrue(
                            expected[r][e] == result[r][e], f"{expected[r][e]} == {result[r][e]}"
                        )

    def test_createDataFrame_with_string_dtype(self):
        # SPARK-34521: spark.createDataFrame does not support Pandas StringDtype extension type
        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}):
            data = [["abc"], ["def"], [None], ["ghi"], [None]]
            pandas_df = pd.DataFrame(data, columns=["col"], dtype="string")
            schema = StructType([StructField("col", StringType(), True)])
            df = self.spark.createDataFrame(pandas_df, schema=schema)

            # dtypes won't match. Pandas has two different ways to store string columns:
            # using ndarray (when dtype isn't specified) or using a StringArray when dtype="string".
            # When calling dataframe#toPandas() it will use the ndarray version.
            # Changing that to use a StringArray would be backwards incompatible.
            assert_frame_equal(pandas_df, df.toPandas(), check_dtype=False)

    def test_createDataFrame_with_int64(self):
        # SPARK-34521: spark.createDataFrame does not support Pandas StringDtype extension type
        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}):
            pandas_df = pd.DataFrame({"col": [1, 2, 3, None]}, dtype="Int64")
            df = self.spark.createDataFrame(pandas_df)
            assert_frame_equal(pandas_df, df.toPandas(), check_dtype=False)

    def test_toPandas_with_map_type(self):
        with QuietTest(self.sc):
            for arrow_enabled in [True, False]:
                with self.subTest(arrow_enabled=arrow_enabled):
                    self.check_toPandas_with_map_type(arrow_enabled)

    def check_toPandas_with_map_type(self, arrow_enabled):
        origin = pd.DataFrame(
            {"id": [0, 1, 2, 3], "m": [{}, {"a": 1}, {"a": 1, "b": 2}, {"a": 1, "b": 2, "c": 3}]}
        )

        for schema in [
            "id long, m map<string, long>",
            StructType().add("id", LongType()).add("m", MapType(StringType(), LongType())),
        ]:
            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}):
                df = self.spark.createDataFrame(origin, schema=schema)

            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                pdf = df.toPandas()
                assert_frame_equal(origin, pdf)

    def test_toPandas_with_map_type_nulls(self):
        with QuietTest(self.sc):
            for arrow_enabled in [True, False]:
                with self.subTest(arrow_enabled=arrow_enabled):
                    self.check_toPandas_with_map_type_nulls(arrow_enabled)

    def check_toPandas_with_map_type_nulls(self, arrow_enabled):
        origin = pd.DataFrame(
            {"id": [0, 1, 2, 3, 4], "m": [{"a": 1}, {"b": 2, "c": 3}, {}, None, {"d": None}]}
        )

        for schema in [
            "id long, m map<string, long>",
            StructType().add("id", LongType()).add("m", MapType(StringType(), LongType())),
        ]:
            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}):
                df = self.spark.createDataFrame(origin, schema=schema)

            with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
                pdf = df.toPandas()
                assert_frame_equal(origin, pdf)

    def test_createDataFrame_with_int_col_names(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_createDataFrame_with_int_col_names(arrow_enabled)

    def check_createDataFrame_with_int_col_names(self, arrow_enabled):
        import numpy as np

        pdf = pd.DataFrame(np.random.rand(4, 2))
        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            df = self.spark.createDataFrame(pdf)
        pdf_col_names = [str(c) for c in pdf.columns]
        self.assertEqual(pdf_col_names, df.columns)

    # Regression test for SPARK-23314
    def test_timestamp_dst(self):
        # Daylight saving time for Los Angeles for 2015 is Sun, Nov 1 at 2:00 am
        dt = [
            datetime.datetime(2015, 11, 1, 0, 30),
            datetime.datetime(2015, 11, 1, 1, 30),
            datetime.datetime(2015, 11, 1, 2, 30),
        ]
        pdf = pd.DataFrame({"time": dt})

        df_from_python = self.spark.createDataFrame(dt, "timestamp").toDF("time")
        df_from_pandas = self.spark.createDataFrame(pdf)

        assert_frame_equal(pdf, df_from_python.toPandas())
        assert_frame_equal(pdf, df_from_pandas.toPandas())

    # Regression test for SPARK-28003
    def test_timestamp_nat(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_timestamp_nat(arrow_enabled)

    def check_timestamp_nat(self, arrow_enabled):
        dt = [pd.NaT, pd.Timestamp("2019-06-11"), None] * 100
        pdf = pd.DataFrame({"time": dt})
        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            df = self.spark.createDataFrame(pdf)

        assert_frame_equal(pdf, df.toPandas())

    def test_toPandas_batch_order(self):
        def delay_first_part(partition_index, iterator):
            if partition_index == 0:
                time.sleep(0.1)
            return iterator

        # Collects Arrow RecordBatches out of order in driver JVM then re-orders in Python
        def run_test(num_records, num_parts, max_records, use_delay=False):
            df = self.spark.range(num_records, numPartitions=num_parts).toDF("a")
            if use_delay:
                df = df.rdd.mapPartitionsWithIndex(delay_first_part).toDF()
            with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": max_records}):
                pdf, pdf_arrow = self._toPandas_arrow_toggle(df)
                assert_frame_equal(pdf, pdf_arrow)

        cases = [
            (1024, 512, 2),  # Use large num partitions for more likely collecting out of order
            (64, 8, 2, True),  # Use delay in first partition to force collecting out of order
            (64, 64, 1),  # Test single batch per partition
            (64, 1, 64),  # Test single partition, single batch
            (64, 1, 8),  # Test single partition, multiple batches
            (30, 7, 2),  # Test different sized partitions
        ]

        for case in cases:
            run_test(*case)

    def test_createDataFrame_with_category_type(self):
        pdf = pd.DataFrame({"A": ["a", "b", "c", "a"]})
        pdf["B"] = pdf["A"].astype("category")
        category_first_element = dict(enumerate(pdf["B"].cat.categories))[0]

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}):
            arrow_df = self.spark.createDataFrame(pdf)
            arrow_type = arrow_df.dtypes[1][1]
            result_arrow = arrow_df.toPandas()
            arrow_first_category_element = result_arrow["B"][0]

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}):
            df = self.spark.createDataFrame(pdf)
            spark_type = df.dtypes[1][1]
            result_spark = df.toPandas()
            spark_first_category_element = result_spark["B"][0]

        assert_frame_equal(result_spark, result_arrow)

        # ensure original category elements are string
        self.assertIsInstance(category_first_element, str)
        # spark data frame and arrow execution mode enabled data frame type must match pandas
        self.assertEqual(spark_type, "string")
        self.assertEqual(arrow_type, "string")
        self.assertIsInstance(arrow_first_category_element, str)
        self.assertIsInstance(spark_first_category_element, str)

    def test_createDataFrame_with_float_index(self):
        # SPARK-32098: float index should not produce duplicated or truncated Spark DataFrame
        self.assertEqual(
            self.spark.createDataFrame(pd.DataFrame({"a": [1, 2, 3]}, index=[2.0, 3.0, 4.0]))
            .distinct()
            .count(),
            3,
        )

    def test_no_partition_toPandas(self):
        # SPARK-32301: toPandas should work from a Spark DataFrame with no partitions
        # Forward-ported from SPARK-32300.
        pdf = self.spark.sparkContext.emptyRDD().toDF("col1 int").toPandas()
        self.assertEqual(len(pdf), 0)
        self.assertEqual(list(pdf.columns), ["col1"])

    def test_createDataFrame_empty_partition(self):
        pdf = pd.DataFrame({"c1": [1], "c2": ["string"]})
        df = self.spark.createDataFrame(pdf)
        self.assertEqual([Row(c1=1, c2="string")], df.collect())
        self.assertGreater(self.spark.sparkContext.defaultParallelism, len(pdf))

    def test_toPandas_error(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_toPandas_error(arrow_enabled)

    def check_toPandas_error(self, arrow_enabled):
        with self.sql_conf(
            {
                "spark.sql.ansi.enabled": True,
                "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled,
            }
        ):
            with self.assertRaises(ArithmeticException):
                self.spark.sql("select 1/0").toPandas()

    def test_toPandas_duplicate_field_names(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_toPandas_duplicate_field_names(arrow_enabled)

    def check_toPandas_duplicate_field_names(self, arrow_enabled):
        data = [Row(Row("a", 1), Row(2, 3, "b", 4, "c")), Row(Row("x", 6), Row(7, 8, "y", 9, "z"))]
        schema = (
            StructType()
            .add("struct", StructType().add("x", StringType()).add("x", IntegerType()))
            .add(
                "struct",
                StructType()
                .add("a", IntegerType())
                .add("x", IntegerType())
                .add("x", StringType())
                .add("y", IntegerType())
                .add("y", StringType()),
            )
        )
        for struct_in_pandas in ["legacy", "row", "dict"]:
            df = self.spark.createDataFrame(data, schema=schema)

            with self.subTest(struct_in_pandas=struct_in_pandas):
                with self.sql_conf(
                    {
                        "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled,
                        "spark.sql.execution.pandas.structHandlingMode": struct_in_pandas,
                    }
                ):
                    if arrow_enabled and struct_in_pandas == "legacy":
                        with self.assertRaisesRegexp(
                            UnsupportedOperationException, "DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT"
                        ):
                            df.toPandas()
                    else:
                        if struct_in_pandas == "dict":
                            expected = pd.DataFrame(
                                [
                                    [
                                        {"x_0": "a", "x_1": 1},
                                        {"a": 2, "x_0": 3, "x_1": "b", "y_0": 4, "y_1": "c"},
                                    ],
                                    [
                                        {"x_0": "x", "x_1": 6},
                                        {"a": 7, "x_0": 8, "x_1": "y", "y_0": 9, "y_1": "z"},
                                    ],
                                ],
                                columns=schema.names,
                            )
                        else:
                            expected = pd.DataFrame.from_records(data, columns=schema.names)
                        assert_frame_equal(df.toPandas(), expected)

    def test_createDataFrame_duplicate_field_names(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_createDataFrame_duplicate_field_names(arrow_enabled)

    def check_createDataFrame_duplicate_field_names(self, arrow_enabled):
        schema = (
            StructType()
            .add("struct", StructType().add("x", StringType()).add("x", IntegerType()))
            .add(
                "struct",
                StructType()
                .add("a", IntegerType())
                .add("x", IntegerType())
                .add("x", StringType())
                .add("y", IntegerType())
                .add("y", StringType()),
            )
        )

        data = [Row(Row("a", 1), Row(2, 3, "b", 4, "c")), Row(Row("x", 6), Row(7, 8, "y", 9, "z"))]
        pdf = pd.DataFrame.from_records(data, columns=schema.names)

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            df = self.spark.createDataFrame(pdf, schema)

        self.assertEqual(df.collect(), data)

    def test_toPandas_empty_columns(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_toPandas_empty_columns(arrow_enabled)

    def check_toPandas_empty_columns(self, arrow_enabled):
        df = self.spark.range(2).select([])

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            assert_frame_equal(df.toPandas(), pd.DataFrame(columns=[], index=range(2)))

    def test_createDataFrame_nested_timestamp(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_createDataFrame_nested_timestamp(arrow_enabled)

    def check_createDataFrame_nested_timestamp(self, arrow_enabled):
        schema = (
            StructType()
            .add("ts", TimestampType())
            .add("ts_ntz", TimestampNTZType())
            .add(
                "struct", StructType().add("ts", TimestampType()).add("ts_ntz", TimestampNTZType())
            )
            .add("array", ArrayType(TimestampType()))
            .add("array_ntz", ArrayType(TimestampNTZType()))
            .add("map", MapType(StringType(), TimestampType()))
            .add("map_ntz", MapType(StringType(), TimestampNTZType()))
        )
        data = [
            Row(
                datetime.datetime(2023, 1, 1, 0, 0, 0),
                datetime.datetime(2023, 1, 1, 0, 0, 0),
                Row(
                    datetime.datetime(2023, 1, 1, 0, 0, 0),
                    datetime.datetime(2023, 1, 1, 0, 0, 0),
                ),
                [datetime.datetime(2023, 1, 1, 0, 0, 0)],
                [datetime.datetime(2023, 1, 1, 0, 0, 0)],
                dict(ts=datetime.datetime(2023, 1, 1, 0, 0, 0)),
                dict(ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0)),
            )
        ]
        pdf = pd.DataFrame.from_records(data, columns=schema.names)

        with self.sql_conf(
            {
                "spark.sql.session.timeZone": "America/New_York",
                "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled,
            }
        ):
            df = self.spark.createDataFrame(pdf, schema)

        expected = Row(
            ts=datetime.datetime(2022, 12, 31, 21, 0, 0),
            ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0),
            struct=Row(
                ts=datetime.datetime(2022, 12, 31, 21, 0, 0),
                ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0),
            ),
            array=[datetime.datetime(2022, 12, 31, 21, 0, 0)],
            array_ntz=[datetime.datetime(2023, 1, 1, 0, 0, 0)],
            map=dict(ts=datetime.datetime(2022, 12, 31, 21, 0, 0)),
            map_ntz=dict(ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0)),
        )

        self.assertEqual(df.first(), expected)

    def test_toPandas_nested_timestamp(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_toPandas_nested_timestamp(arrow_enabled)

    def check_toPandas_nested_timestamp(self, arrow_enabled):
        schema = (
            StructType()
            .add("ts", TimestampType())
            .add("ts_ntz", TimestampNTZType())
            .add(
                "struct", StructType().add("ts", TimestampType()).add("ts_ntz", TimestampNTZType())
            )
            .add("array", ArrayType(TimestampType()))
            .add("array_ntz", ArrayType(TimestampNTZType()))
            .add("map", MapType(StringType(), TimestampType()))
            .add("map_ntz", MapType(StringType(), TimestampNTZType()))
        )
        data = [
            Row(
                datetime.datetime(2023, 1, 1, 0, 0, 0),
                datetime.datetime(2023, 1, 1, 0, 0, 0),
                Row(
                    datetime.datetime(2023, 1, 1, 0, 0, 0),
                    datetime.datetime(2023, 1, 1, 0, 0, 0),
                ),
                [datetime.datetime(2023, 1, 1, 0, 0, 0)],
                [datetime.datetime(2023, 1, 1, 0, 0, 0)],
                dict(ts=datetime.datetime(2023, 1, 1, 0, 0, 0)),
                dict(ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0)),
            )
        ]
        df = self.spark.createDataFrame(data, schema)

        with self.sql_conf(
            {
                "spark.sql.session.timeZone": "America/New_York",
                "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled,
                "spark.sql.execution.pandas.structHandlingMode": "row",
            }
        ):
            pdf = df.toPandas()

        expected = pd.DataFrame(
            {
                "ts": [datetime.datetime(2023, 1, 1, 3, 0, 0)],
                "ts_ntz": [datetime.datetime(2023, 1, 1, 0, 0, 0)],
                "struct": [
                    Row(
                        datetime.datetime(2023, 1, 1, 3, 0, 0),
                        datetime.datetime(2023, 1, 1, 0, 0, 0),
                    )
                ],
                "array": [[datetime.datetime(2023, 1, 1, 3, 0, 0)]],
                "array_ntz": [[datetime.datetime(2023, 1, 1, 0, 0, 0)]],
                "map": [dict(ts=datetime.datetime(2023, 1, 1, 3, 0, 0))],
                "map_ntz": [dict(ts_ntz=datetime.datetime(2023, 1, 1, 0, 0, 0))],
            }
        )

        assert_frame_equal(pdf, expected)

    def test_createDataFrame_udt(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_createDataFrame_udt(arrow_enabled)

    def check_createDataFrame_udt(self, arrow_enabled):
        schema = (
            StructType()
            .add("point", ExamplePointUDT())
            .add("struct", StructType().add("point", ExamplePointUDT()))
            .add("array", ArrayType(ExamplePointUDT()))
            .add("map", MapType(StringType(), ExamplePointUDT()))
        )
        data = [
            Row(
                ExamplePoint(1.0, 2.0),
                Row(ExamplePoint(3.0, 4.0)),
                [ExamplePoint(5.0, 6.0)],
                dict(point=ExamplePoint(7.0, 8.0)),
            )
        ]
        pdf = pd.DataFrame.from_records(data, columns=schema.names)

        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": arrow_enabled}):
            df = self.spark.createDataFrame(pdf, schema)

        self.assertEqual(df.collect(), data)

    def test_toPandas_udt(self):
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_toPandas_udt(arrow_enabled)

    def check_toPandas_udt(self, arrow_enabled):
        schema = (
            StructType()
            .add("point", ExamplePointUDT())
            .add("struct", StructType().add("point", ExamplePointUDT()))
            .add("array", ArrayType(ExamplePointUDT()))
            .add("map", MapType(StringType(), ExamplePointUDT()))
        )
        data = [
            Row(
                ExamplePoint(1.0, 2.0),
                Row(ExamplePoint(3.0, 4.0)),
                [ExamplePoint(5.0, 6.0)],
                dict(point=ExamplePoint(7.0, 8.0)),
            )
        ]
        df = self.spark.createDataFrame(data, schema)

        with self.sql_conf(
            {
                "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled,
                "spark.sql.execution.pandas.structHandlingMode": "row",
            }
        ):
            pdf = df.toPandas()

        expected = pd.DataFrame.from_records(data, columns=schema.names)

        assert_frame_equal(pdf, expected)

    def test_create_dataframe_namedtuples(self):
        # SPARK-44980: Inherited namedtuples in createDataFrame
        for arrow_enabled in [True, False]:
            with self.subTest(arrow_enabled=arrow_enabled):
                self.check_create_dataframe_namedtuples(arrow_enabled)

    def check_create_dataframe_namedtuples(self, arrow_enabled):
        MyTuple = namedtuple("MyTuple", ["a", "b", "c"])

        class MyInheritedTuple(MyTuple):
            pass

        with self.sql_conf(
            {
                "spark.sql.execution.arrow.pyspark.enabled": arrow_enabled,
            }
        ):
            df = self.spark.createDataFrame([MyInheritedTuple(1, 2, 3)])
            self.assertEqual(df.first(), Row(a=1, b=2, c=3))

            df = self.spark.createDataFrame([MyInheritedTuple(1, 2, MyInheritedTuple(1, 2, 3))])
            self.assertEqual(df.first(), Row(a=1, b=2, c=Row(a=1, b=2, c=3)))

@unittest.skipIf(
    not have_pandas or not have_pyarrow,
    cast(str, pandas_requirement_message or pyarrow_requirement_message),
)
class ArrowTests(ArrowTestsMixin, ReusedSQLTestCase):
    pass

@unittest.skipIf(
    not have_pandas or not have_pyarrow,
    cast(str, pandas_requirement_message or pyarrow_requirement_message),
)
class MaxResultArrowTests(unittest.TestCase):
    # These tests are separate as 'spark.driver.maxResultSize' configuration
    # is a static configuration to Spark context.

    @classmethod
    def setUpClass(cls):
        cls.spark = SparkSession(
            SparkContext(
                "local[4]", cls.__name__, conf=SparkConf().set("spark.driver.maxResultSize", "10k")
            )
        )

        # Explicitly enable Arrow and disable fallback.
        cls.spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
        cls.spark.conf.set("spark.sql.execution.arrow.pyspark.fallback.enabled", "false")

    @classmethod
    def tearDownClass(cls):
        if hasattr(cls, "spark"):
            cls.spark.stop()

    def test_exception_by_max_results(self):
        with self.assertRaisesRegex(Exception, "is bigger than"):
            self.spark.range(0, 10000, 1, 100).toPandas()

class EncryptionArrowTests(ArrowTests):
    @classmethod
    def conf(cls):
        return super(EncryptionArrowTests, cls).conf().set("spark.io.encryption.enabled", "true")

class RDDBasedArrowTests(ArrowTests):
    @classmethod
    def conf(cls):
        return (
            super(RDDBasedArrowTests, cls)
            .conf()
            .set("spark.sql.execution.arrow.localRelationThreshold", "0")
            # to test multiple partitions
            .set("spark.sql.execution.arrow.maxRecordsPerBatch", "2")
        )

if __name__ == "__main__":
    from pyspark.sql.tests.test_arrow import *  # noqa: F401

    try:
        import xmlrunner  # type: ignore

        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
    except ImportError:
        testRunner = None
    unittest.main(testRunner=testRunner, verbosity=2)


from flask import render_template, request, redirect, url_for, jsonify
from application import service1
import random, requests, json

f_f_name = ["Yarsubelle", "Kobrebella", "Voddaebella", "Katdrouna", "Ulmaengrid",
            "Thraluhilde", "Ornetrude", "Huvribella", "Fosgribella", "Nekonelyn"]

m_f_name = ["Gormon", "Halamar", "Luvon", "Vorvorick", "Skarseat",
          "Thaznat", "Ylra", "Guka", "Elberar", "Joilbe"]

l_name = ["The Magnificent", "The Protective", "The Hungry", "Boulderstone", "Copperguard",
          "Understone", "Mudsunder", "Grumblehelm", "Snowbane", "Flintgrip"]

name = {"firstname":"none", "lastname":"none"}

@service1.route('/name', methods=['GET', 'POST'])
def random_name():
    if request.method == 'POST':
        data = request.args.get("gender")
        if data == 'Female':
            name["firstname"] = f_f_name[random.randint(0, 9)]
            name["lastname"] = l_name[random.randint(0, 9)]
            return name
        elif data == 'Male':
            name["firstname"] = m_f_name[random.randint(0, 9)]
            name["lastname"] = l_name[random.randint(0, 9)]
            return name
        return {"firstname":"failure"}


import torch
import numpy as np
from dm_control import suite
from collections import deque
from PIL import Image
import random
import gym
import os
from time import sleep, time
from copy import deepcopy

from replay_memory import ReplayMemory
from models import ActorNet, CriticNet
from utils import soft_update, get_obs, calc_priority, invertical_vf

def actor_process(actor_id):
    actor = Actor(actor_id)
    actor.run()

class Actor:
    def __init__(self, actor_id):
        self.env = suite.load(domain_name="walker", task_name="run")
        self.action_size = self.env.action_spec().shape[0]
        self.obs_size = get_obs(self.env.reset().observation).shape[1]

        self.actor_id = actor_id
        self.burn_in_length = 20 # 40-80
        self.learning_length = 40
        self.sequence_length = self.burn_in_length + self.learning_length
        self.n_step = 5
        self.sequence = []
        self.recurrent_state = []
        self.priority = []
        self.td_loss = deque(maxlen=self.learning_length)
        self.memory_sequence_size = 1000
        self.memory = ReplayMemory(memory_sequence_size=self.memory_sequence_size)
        self.memory_save_interval = 3

        self.gamma = 0.997
        self.actor_parameter_update_interval = 500
        self.model_path = './model_data/'
        self.actor = ActorNet(self.obs_size, self.action_size, cuda_id=self.actor_id%2+1).cuda(self.actor_id%2+1).eval()
        self.target_actor = deepcopy(self.actor)
        self.critic = CriticNet(self.obs_size, self.action_size, cuda_id=self.actor_id%2+1).cuda(self.actor_id%2+1).eval()
        self.target_critic = deepcopy(self.critic)
        self.load_model()
        self.epsilon = 1
        self.last_obs = None

    def load_model(self):
        if os.path.isfile(self.model_path + 'model.pt'):
            while True:
                try:
                    # TODO: Delete
                    self.actor = ActorNet(self.obs_size, self.action_size, self.actor_id%2+1).cuda().eval()
                    self.target_actor = deepcopy(self.actor)
                    self.critic = CriticNet(self.obs_size, self.action_size, self.actor_id%2+1).cuda().eval()
                    self.target_critic = deepcopy(self.critic)
                    #model_dict = torch.load(self.model_path + 'model.pt', map_location={'cuda:0':'cuda:{}'.format(self.actor_id%2+1)})
                    model_dict = torch.load(self.model_path + 'model.pt')
                    self.actor.load_state_dict(model_dict['actor'])
                    self.target_actor.load_state_dict(model_dict['target_actor'])
                    self.critic.load_state_dict(model_dict['critic'])
                    self.target_critic.load_state_dict(model_dict['target_critic'])
                    self.actor.cuda(self.actor_id%2+1)
                    self.target_actor.cuda(self.actor_id%2+1)
                    self.critic.cuda(self.actor_id%2+1)
                    self.target_critic.cuda(self.actor_id%2+1)
                except:
                    sleep(np.random.rand() * 5 + 2)
                else:
                    break

    def calc_nstep_reward(self):
        for i in range(len(self.sequence) - self.n_step):
            self.sequence[i][2][0] = sum([self.sequence[i+j][2][0] * (self.gamma ** j) for j in range(self.n_step)])

    def calc_priorities(self):
        self.actor.reset_state()
        self.critic.reset_state()
        self.target_actor.reset_state()
        self.target_critic.reset_state()
        self.td_loss = deque(maxlen=self.learning_length)
        self.priority = []

        for i in range(self.n_step):
            next_obs = torch.from_numpy(self.sequence[i][0]).cuda(self.actor_id%2+1).unsqueeze(0)
            next_action = self.target_actor(next_obs)
            next_q_value = self.target_critic(next_obs, next_action).detach().cpu().numpy()

        for i in range(len(self.sequence) - self.n_step):
            obs = torch.from_numpy(self.sequence[i][0]).cuda(self.actor_id%2+1).unsqueeze(0)
            action = torch.from_numpy(self.sequence[i][1]).cuda(self.actor_id%2+1).unsqueeze(0)
            next_obs = torch.from_numpy(self.sequence[i + self.n_step][0]).cuda(self.actor_id%2+1).unsqueeze(0)
            next_action = self.target_actor(next_obs)

            q_value = self.critic(obs, action).detach().cpu().numpy()
            reward = self.sequence[i][2][0]
            terminal = self.sequence[i + self.n_step - 1][3][0]
            next_q_value = self.target_critic(next_obs, next_action).detach().cpu().numpy()

            if i >= self.burn_in_length:
                target_q_value = (reward + (self.gamma ** self.n_step) * (1.-terminal) * next_q_value)
                target_q_value = invertical_vf(torch.tensor(target_q_value).cuda(self.actor_id%2+1)).detach().cpu().numpy()
                self.td_loss.append((q_value - target_q_value).mean())
            if i >= self.sequence_length:
                self.priority.append(calc_priority(np.array(list(self.td_loss), dtype=np.float32) ** 2.))

    def run(self):
        episode = 0
        step = 0
        reward_sum = 0

        while True:
            time_step = self.env.reset()
            obs = get_obs(time_step.observation)
            self.actor.reset_state()
            self.critic.reset_state()
            self.target_actor.reset_state()
            self.target_critic.reset_state()
            self.sequence = []
            self.recurrent_state = []
            self.priority = []
            self.td_loss.clear()
            last_obs = None
            episode_step = 0
            done = False
            if self.actor_id == 0 and episode != 0:
                print('episode:', episode,'step:', step, 'reward:', reward_sum)
            episode += 1
            reward_sum = 0

            while not time_step.last():

                # get recurrent state
                actor_hx, actor_cx = self.actor.get_state()
                target_actor_hx, target_actor_cx = self.target_actor.get_state()
                critic_hx, critic_cx = self.critic.get_state()
                target_critic_hx, target_critic_cx = self.target_critic.get_state()

                action = self.actor(torch.from_numpy(obs).cuda(self.actor_id%2+1))
                target_action = self.target_actor(torch.from_numpy(obs).cuda(self.actor_id%2+1))
                _ = self.critic(torch.from_numpy(obs).cuda(self.actor_id%2+1), action)
                _ = self.target_critic(torch.from_numpy(obs).cuda(self.actor_id%2+1), target_action)

                action = action.detach().cpu().numpy()[0]
                action += np.random.normal(0, 0.3, (self.action_size))
                action = np.clip(action, -1, 1)

                reward = 0.
                sleep(0.01)
                for i in range(4):
                    time_step = self.env.step(action)
                    next_obs = get_obs(time_step.observation)
                    reward += time_step.reward
                    if time_step.last():
                        break

                reward_sum += reward
                step += 1
                episode_step += 1
                terminal = 1. if time_step.last() else 0.
                self.sequence.append((obs[0], action, [reward], [terminal]))
                obs = next_obs.copy()

                self.recurrent_state.append([[actor_hx[0], actor_cx[0]], [target_actor_hx[0], target_actor_cx[0]],
                                                [critic_hx[0], critic_cx[0]], [target_critic_hx[0], target_critic_cx[0]]])

                if step % self.actor_parameter_update_interval == 0:
                    self.load_model()

            if len(self.sequence) >= self.sequence_length:
                self.sequence.extend([(np.zeros((self.obs_size), dtype=np.float32), np.zeros((self.action_size), dtype=np.float32), [0.], [1.]) for i in range(self.n_step)])
                self.calc_nstep_reward()
                self.calc_priorities()
                self.memory.add(self.sequence, self.recurrent_state, self.priority)

            if len(self.memory.memory) > self.memory_save_interval:
                self.memory.save(self.actor_id)


commands = {

    '!kajfklajflksajlkfjlksjf': {
        'limit': 200,
        'argc': 1,
        'return': 'command',
        'space_case': True,
        'ul': 'mod',
        'usage': "!report [insert bug report text here]"

    },

    '!asndan,mna,n,awmn,man,m': {
        'limit': 15,
        'return': 'There is a README for me on my profile',
        'usage': '!help'

    },

}

def initalizeCommands(config):
    for channel in config['channels']:
        for command in commands:
            commands[command][channel] = {}
            commands[command][channel]['last_used'] = 0


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import re
import logging
import argparse

logging.basicConfig(level=logging.INFO)

parser = argparse.ArgumentParser(\
            description='TRSQ-8 assembler (verilog version)')
parser.add_argument('path')
args = parser.parse_args()

def get_inst(line):
  return line.split()[0]

def assembler(line):
    inst = line[0]
    logging.debug(inst)
    bit = ''
    if (inst == 'ADD'):
        bit = '0100000' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'SUB'):
        bit = '0100001' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'AND'):
        bit = '0100111' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'OR'):
        bit = '0101000' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'NOT'):
        bit = '0101001' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'XOR'):
        bit = '0101011' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'BTC'):
        bit = '1000' + format(int(line[1]), 'b').zfill(3) \
                     + format(int(line[2]), 'b').zfill(8)
    elif (inst == 'BTS'):
        bit = '1001' + format(int(line[1]), 'b').zfill(3) \
                     + format(int(line[2]), 'b').zfill(8)
    elif (inst == 'ST'):
        bit = '0101100' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'LD'):
        bit = '0101101' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'LDL'):
        bit = '0101110' + format(int(line[1]), 'b').zfill(8)
    elif (inst == 'STR'):
        bit = '011110000000000'
    elif (inst == 'LDR'):
        bit = '011110100000000'
    elif (inst == 'SKZ'):
        bit = '000010100000000'
    elif (inst == 'SKC'):
        bit = '000011000000000'
    elif (inst == 'NOP'):
        bit = '000000000000000'
    elif (inst == 'HALT'):
        bit = '000000100000000'
    elif (inst == 'GOTO'):
        bit = '11' + format(int(line[1]), 'b').zfill(13)
    elif (inst == 'RETURN'):
        bit = '000001000000000'
    else:
        logging.error('err: ' + inst)

    logging.debug(bit)
    return bit

if __name__ == '__main__':
    # Open the file and split by return
    lines_raw = []
    with open(args.path, 'r', encoding='utf-8-sig') as text:
        lines_raw = text.read()
        lines_raw = lines_raw.split('\n')

    logging.debug('------- lines_raw -----------')
    logging.debug(lines_raw)
    logging.debug(len(lines_raw))

    # Delete comments
    lines = []
    for line in lines_raw:
        i = line.find('#')
        if (i > -1):
            line = line[0:i]
        lines.append(line)
    del(lines_raw)
    logging.debug('------- lines -----------')
    logging.debug(lines)
    logging.debug(len(lines))

    # Remove blank line
    asm = []
    for line in lines:
        i = re.search(r'\S', line)
        if (i != None):
            asm.append(line)
    del(lines)
    logging.debug('------- asm -----------')
    logging.debug(asm)
    logging.debug(len(asm))

    # TODO:
    # Delete white-space from ':' to next character
    # asm_shaped = []
    # tmp = ""
    # flag = False
    # for line in asm:
    #     if (flag = True):
    #         line = tmp + line
    #         flag = False

    #     # search characters from ':' to EOL
    #     i = line.find(':')
    #     j = re.search(r'\S', line[i+1:])
    #     if (j == None):
    #         tmp = line[:i]
    #         flag = True

    labels = {}
    lines2 = []
    # Create labels dict.
    label_cache = []
    for line in asm:
      label_m = re.search(r'(\w+):', line)
      if label_m:
        l_name = label_m.group(1)
        label_cache.append(l_name)
        continue
      for l_name in label_cache:
        labels[l_name] = len(lines2)
      label_cache = []
      lines2.append(line)

    logging.debug('------- labels -----------')
    logging.debug(labels)

    # Delete Labels
    asm_nolabels = []
    for line in asm:
        i = line.find(':')
        if (i > -1):
            line = line[i+1:]
        asm_nolabels.append(line)
    del(asm)
    logging.debug('------- asm_nolabels -----------')
    logging.debug(asm_nolabels)
    logging.debug(len(asm_nolabels))

    asm_list = []

    # Replace Labels to each line number
    for line in lines2:
      if get_inst(line) == 'GOTO':
        l_name = line.split()[1]
        ln = labels[l_name]
        asm_list.append(line.replace(l_name, str(ln)))
      else:
        asm_list.append(line)

    # Split lines by white-spaces
    asm_split = []
    for line in asm_list:
        line = line.split()
        asm_split.append(line)
    logging.debug('------- asm_split -----------')
    logging.debug(asm_split)
    logging.debug(len(asm_split))

    logging.debug('------- Start encode -----------')
    bit = []
    for line in asm_split:
        assemble = assembler(line)
        if (assemble != ''):
            bit.append(assemble)

    logging.debug('------- Finish encode -----------')
    # write bin file
    with open('./prom.bin', 'w', encoding='utf-8-sig') as text:
        for line in bit:
            text.write(line)
            text.write('\n')

    # write VHDL entity file
    with open('./prom.v', 'w', encoding='utf-8-sig') as text:
        hdl_text_0 = [
                'module prom (\n',
                '    input CLK_ip,\n',
                '    input [12:0] ADDR_ip,\n',
                '    output [14:0] DATA_op);\n',
                '\n',
                '    assign DATA_op = \n'
                ]

        hdl_text_1 = [
                '                           15\'b000000000000000;\n',
                'endmodule\n'
                ]

        text.writelines('// ' + args.path + '\n')
        text.writelines(hdl_text_0)

        i = 0
        for line in bit:
            if (line != ''):
                # TODO
                # comment out the original code
                text.write(
                        '        ADDR_ip==13\'d' + str(i) + ' ? ' + \
                                '15\'b' + line + ': // ' + '' + '\n'
                        )
                i += 1

        text.writelines(hdl_text_1);


import pandas as pd
import matplotlib.pyplot as plt
import datetime as dt
import numpy as np
import boto3
import matplotlib.dates as mdates
from pandas.plotting import register_matplotlib_converters
from matplotlib import cm

register_matplotlib_converters()

def lambda_handler(event, context):

    # Read in arguments / event data
    data_uid = event['data_uid']
    user_name = event['user_name']
    image_uid = event['image_uid']

    # Download the data set from S3
    data_file_name = f'{data_uid}-channels.csv.gz'
    aws_s3_bucket_prefix = 'deepfake-discord-bot'
    s3 = boto3.resource('s3')
    s3.Bucket(aws_s3_bucket_prefix) \
        .download_file(data_file_name, '/tmp/' + data_file_name)

    # Make the plots
    activity_file = time_series_chart(data_uid, image_uid, user_name)

    channels_file = channels_chart(data_uid, image_uid, user_name)

    # Upload to S3
    s3.Object(aws_s3_bucket_prefix, activity_file).upload_file(f'/tmp/{activity_file}')
    s3.Object(aws_s3_bucket_prefix, channels_file).upload_file(f'/tmp/{channels_file}')

    return {
        'statusCode': 200,
        'image_file_names': [activity_file, channels_file]
    }

def day_filler(dates, counts):
    """Returns completed lists of dates and message counts with 0's added to every day with no messages"""
    filled_dates = []
    filled_counts = []
    for n in range(int((dates.max() - dates.min()).days) + 1):
        single_date = dates.min() + dt.timedelta(n)
        filled_dates.append(single_date)
        if single_date in dates:
            i = np.where(dates == single_date)
            filled_counts.append(counts[i])
        else:
            filled_counts.append(0)

    return filled_dates, filled_counts

def auto_time_scale(td):
    """Used to format the x-axis based on the length of time to be plotted"""
    if td.days > 365:
        date_format = mdates.DateFormatter('%Y')
        major_tick = mdates.YearLocator()
    elif td.days > 65:
        date_format = mdates.DateFormatter('%b-%Y')
        major_tick = mdates.MonthLocator()
    elif td.days > 15:
        date_format = mdates.DateFormatter('%d %b-%Y')
        major_tick = mdates.DayLocator(interval=7)
    else:
        date_format = mdates.DateFormatter('%d %b-%Y')
        major_tick = mdates.DayLocator()

    return date_format, major_tick

def time_series_chart(data_id, image_uid, user_name):
    """Plots a user's activity over time. I.e. number of messages vs. date"""

    # Open our file from S3 and read in the
    data_file_name = f'/tmp/{data_id}-channels.csv.gz'
    try:
        df = pd.read_csv(data_file_name, compression='gzip', encoding='utf-8')
    except FileNotFoundError:
        df = pd.read_csv('.' + data_file_name, compression='gzip', encoding='utf-8')

    # Some data transformations
    df['datetime'] = df['timestamp'].apply(lambda t: dt.datetime.fromtimestamp(t))
    df['date'] = df['datetime'].apply(lambda t: t.date())

    ch_gb = df.groupby('date')
    message_dates = ch_gb['date'].count().index.values
    message_counts = ch_gb['timestamp'].count().values

    filled_dates, filled_counts = day_filler(message_dates, message_counts)

    # Make the time series plots
    fig, ax = plt.subplots()
    fig.set_figheight(7)
    fig.set_figwidth(10)

    ax.plot(filled_dates, filled_counts)
    ax.set_title(f'{user_name}\'s Activity')
    ax.set_ylabel('# messages')

    message_date_range = df['datetime'].max() - df['datetime'].min()

    if len(filled_dates) > 1:
        auto_format, auto_tick = auto_time_scale(message_date_range)
        ax.xaxis.set_major_formatter(auto_format)
        ax.xaxis.set_major_locator(auto_tick)

    ax.grid()
    for tick in ax.get_xticklabels():
        tick.set_rotation(45)

    file_name = f'{image_uid}-activity.png'
    try:
        fig.savefig(f'/tmp/{file_name}')
    except FileNotFoundError:
        fig.savefig(f'./tmp/{file_name}')

    return file_name

def channels_chart(data_id, image_uid, user_name):
    """Plots a user's most active channels"""

    # Open our file from S3 and read in the
    data_file_name = f'/tmp/{data_id}-channels.csv.gz'
    df = pd.read_csv(data_file_name, compression='gzip', encoding='utf-8')

    # Some data transformations
    df['datetime'] = df['timestamp'].apply(lambda t: dt.datetime.fromtimestamp(t))
    df['day'] = df['datetime'].apply(lambda t: t.strftime("%A"))
    df['date'] = df['datetime'].apply(lambda t: t.date())

    ch_gb = df.groupby('channel')

    pie_labels = ch_gb['timestamp'].count().index.values
    pie_values = ch_gb['timestamp'].count().values

    # Make the channels pie chart
    fig, ax = plt.subplots()
    color_scale = cm.Blues(np.flip(np.arange(pie_labels.size)) / pie_labels.size)
    ax.pie(pie_values, labels=pie_labels, autopct='%1.1f%%', explode=[0.05]*pie_labels.size, colors=color_scale)
    ax.set_title(f'{user_name}\'s Favorite Channels')

    file_name_channels = f'{image_uid}-pie-chart-channels.png'
    fig.savefig(f'/tmp/{file_name_channels}')

    return file_name_channels


from sqlalchemy import Table, MetaData, Column, Integer, Float, String, Date,  ForeignKey
from sqlalchemy.orm import mapper, relationship

from domain import core

metadata = MetaData()

group_table = Table(
    'group', metadata,
    Column('group_id', String(255), primary_key = True),
    Column('name', String(255)),
    Column('parent_reference', String(255)),
    Column('individual_debt_limit', Float),
    Column('individual_credit_limit', Float),
)

member_table = Table(
    'member', metadata,
    Column('member_id', String(255), primary_key = True),
    Column('name', String(255)),
    Column('group', String(255)),
)

tx_table = Table(
    'transaction_', metadata,
    Column('id', Integer, primary_key=True, autoincrement=True),
    Column('group_id', String(255)),
    Column('from_member', String(255)),
    Column('to_member', String(255)),
    Column('amount', Float),
)

def start_mappers():
    tx_mapper = mapper(core.Transaction, tx_table)
    member_mapper = mapper(core.Group, group_table)
    group_mapper = mapper(core.Member, member_table)