#%%
#----データセットの整形
import scipy.io
import numpy as np
from matplotlib import pyplot
import os

#--------------------------------------------------------------
#tsとttsの最も近い値を求めるときに使用
def getNearestValue(list, num):
    """
    概要: リストからある値に最も近い値を返却する関数
    @param list: データ配列
    @param num: 対象値
    @return 対象値に最も近い値
    """

    # リスト要素と対象値の差分を計算し最小値のインデックスを取得
    idx = np.abs(np.asarray(list) - num,dtype="float64").argmin()
    return list[idx]
#--------------------------------------------------------------
#----------------------------------------------------------
subject_number = 1  #被験者数
subject        = 1  #被験者を指定

stride = 0.125       #ストライド(s)
time_window = 4     #時間窓
threshold = 3       #ノイズ除去の閾値
#----------------------------------------------------------
each_data_number = 72 #各runファイルの試行回数
sampling_freq = 256   #サンプリング周波数
ch_number = 3         #チャンネル数
file_number = 3       #各被験者ディレクトリのファイル数
imagination_time = 4  #運動想起時間
label_number = 3      #ラベル数
delete_label = 3 - 1      #削除するラベル-1
start_time = 1.5    #抽出する時間
end_time = 8.5      #抽出する時間

data_magnification = int(((end_time - time_window - start_time) / stride) + 1) #データ増加率
print("data_mag:",data_magnification)

imagination_freq = sampling_freq * imagination_time #運動想起時間→周波数変換time_windowの方がよいか
#stride_freq = int(sampling_freq * stride) #ストライド時間→周波数変換
all_sample_number = int(each_data_number * file_number * subject_number * data_magnification * 2 / 3)
print("asn:",all_sample_number)
#-------------------------------------------------
#print(os.getcwd()) #カレントディレクトリ
#相対パスで全ファイルにアクセス
current_path = os.getcwd() #これでPCによらない
#print("current_path:",current_path)
#-------------------------------------------------

data_ch1 = np.zeros((all_sample_number,imagination_freq),dtype="float64") #ch1,run1
data_ch2 = np.zeros((all_sample_number,imagination_freq),dtype="float64") #ch2,run1
data_ch3 = np.zeros((all_sample_number,imagination_freq),dtype="float64") #ch3,run1

data_label = np.zeros(all_sample_number)

run_cnt = 0
for i in range(subject_number): #各被験者ごとのループ
    #被験者を指定
    i = subject
    print("s"+str(subject))

    #毎回各被験者のデータを読み込む
    data_run1 = (scipy.io.loadmat(current_path + "/s" + str(i) + "/" + "/Run1.mat"))
    data_run2 = (scipy.io.loadmat(current_path + "/s" + str(i) + "/" + "/Run2.mat"))
    data_run3 = (scipy.io.loadmat(current_path + "/s" + str(i) + "/" + "/Run3.mat"))

    data_run1["trial_time_stamps"] = data_run1["trial_time_stamps"].reshape(each_data_number)
    data_run2["trial_time_stamps"] = data_run2["trial_time_stamps"].reshape(each_data_number)
    data_run3["trial_time_stamps"] = data_run3["trial_time_stamps"].reshape(each_data_number)

    data_run1["time_stamps"] = data_run1["time_stamps"].reshape(data_run1["time_stamps"].size)
    data_run2["time_stamps"] = data_run2["time_stamps"].reshape(data_run2["time_stamps"].size)
    data_run3["time_stamps"] = data_run3["time_stamps"].reshape(data_run3["time_stamps"].size)

    starting_imagination_label = np.zeros(file_number,dtype="int64")
    starting_imagination_time  = np.zeros(file_number,dtype="float32")
    label = 0
    for j in range(each_data_number): #ラベル数のループ回数→すべてのtts分を見る
        label_jdg = np.zeros((file_number,label_number),dtype="int8")
        #必要に応じて削除するラベルを決める
        #run1ラベル1か否か
        if data_run1["Y"][0][j] == 1:
            label_jdg[0][0] = 1
        #run1ラベル2か否か
        if data_run1["Y"][0][j] == 2:
            label_jdg[0][1] = 1
        #run1ラベル3か否か
        if data_run1["Y"][0][j] == 3:
            label_jdg[0][2] = 1

        #run2ラベル1か否か
        if data_run2["Y"][0][j] == 1:
            label_jdg[1][0] = 1
        #run2ラベル2か否か
        if data_run2["Y"][0][j] == 2:
            label_jdg[1][1] = 1
        #run2ラベル3か否か
        if data_run2["Y"][0][j] == 3:
            label_jdg[1][2] = 1

        #run3ラベル1か否か
        if data_run3["Y"][0][j] == 1:
            label_jdg[2][0] = 1
        #run3ラベル2か否か
        if data_run3["Y"][0][j] == 2:
            label_jdg[2][1] = 1
        #run3ラベル3か否か
        if data_run3["Y"][0][j] == 3:
            label_jdg[2][2] = 1

        #ttsに最も近いtsを見つける
        starting_imagination_time[0]  = getNearestValue(data_run1["time_stamps"],data_run1["trial_time_stamps"][j])
        starting_imagination_time[1]  = getNearestValue(data_run2["time_stamps"],data_run2["trial_time_stamps"][j])
        starting_imagination_time[2]  = getNearestValue(data_run3["time_stamps"],data_run3["trial_time_stamps"][j])

        starting_imagination_label[0] = np.where(data_run1["time_stamps"]==getNearestValue(data_run1["time_stamps"],data_run1["trial_time_stamps"][j]))[0]
        starting_imagination_label[1] = np.where(data_run2["time_stamps"]==getNearestValue(data_run2["time_stamps"],data_run2["trial_time_stamps"][j]))[0]
        starting_imagination_label[2] = np.where(data_run3["time_stamps"]==getNearestValue(data_run3["time_stamps"],data_run3["trial_time_stamps"][j]))[0]

        #データ増強しながら取り込む
        a = [stride * x for x in range(data_magnification-1,-1,-1)]
        if j % 10 == 0:
            print("j:",j)
        for k in a:
            #stride_label = 0
            k = int(k * sampling_freq)
            #print("k",k)

            if label_jdg[0][delete_label] != 1:
                data_label[run_cnt+label] = data_run1["Y"][0][j]
                for l in range(imagination_freq):
                    data_ch1[run_cnt+label][l] = data_run1["X"][starting_imagination_label[0]-k+l][0]
                    data_ch2[run_cnt+label][l] = data_run1["X"][starting_imagination_label[0]-k+l][1]
                    data_ch3[run_cnt+label][l] = data_run1["X"][starting_imagination_label[0]-k+l][2]
                #stride_label += 1
                label += 1

            if label_jdg[1][delete_label] != 1:
                #stride_label = 0
                data_label[run_cnt+label] = data_run2["Y"][0][j]
                for l in range(imagination_freq):
                    data_ch1[run_cnt+label][l] = data_run2["X"][starting_imagination_label[1]-k+l][0]
                    data_ch2[run_cnt+label][l] = data_run2["X"][starting_imagination_label[1]-k+l][1]
                    data_ch3[run_cnt+label][l] = data_run2["X"][starting_imagination_label[1]-k+l][2]
                #stride_label += 1
                label += 1

            if label_jdg[2][delete_label] != 1:
                #stride_label = 0
                data_label[run_cnt+label] = data_run3["Y"][0][j]
                for l in range(imagination_freq):
                    data_ch1[run_cnt+label][l] = data_run3["X"][starting_imagination_label[2]-k+l][0]
                    data_ch2[run_cnt+label][l] = data_run3["X"][starting_imagination_label[2]-k+l][1]
                    data_ch3[run_cnt+label][l] = data_run3["X"][starting_imagination_label[2]-k+l][2]
                #stride_label += 1
                label += 1
    run_cnt += int(each_data_number * 3 * 2 / 3)

#"""
#グラフ表示
x = np.linspace(0,imagination_time,imagination_freq)
import random
random_label = []
for i in range(5):
    random_label.append(random.randint(0,all_sample_number))
    y = data_ch1[random_label[i]] #ランダムに表示
    pyplot.plot(x,y)
    print("label",random_label[i])
    pyplot.show()
#"""
#"""
#0の要素があるかチェック➡OK:ファイルの名前を変更して対処
print("-----")
for label in range(all_sample_number):
    for j in range(imagination_freq):
        if data_ch1[label][j] == 0:
            print("input_label:",label)
            break
    if data_label[label] == 0:
        print("output_label:",label)
#"""

#%%
#----脳波の前処理
import scipy.signal as signal

filter = [0] * ch_number
tap_number = 61

#"""
#50Hzノッチフィルター
print("50Hz notch filter")
for i in range(all_sample_number):
    filter[0] = signal.firwin(numtaps=tap_number,cutoff=[49,51],fs=sampling_freq)
    filter[1] = signal.firwin(numtaps=tap_number,cutoff=[49,51],fs=sampling_freq)
    filter[2] = signal.firwin(numtaps=tap_number,cutoff=[49,51],fs=sampling_freq)

    data_ch1[i] = signal.lfilter(filter[0],1,data_ch1[i])
    data_ch2[i] = signal.lfilter(filter[1],1,data_ch2[i])
    data_ch3[i] = signal.lfilter(filter[2],1,data_ch3[i])
#"""

#"""
#カットオフ周波数0.5Hzのハイパスフィルター
print("High-pass filter with a cutoff frequency of 0.5Hz")
for i in range(all_sample_number):
    filter[0] = signal.firwin(numtaps=tap_number,cutoff=0.5,fs=sampling_freq)
    filter[1] = signal.firwin(numtaps=tap_number,cutoff=0.5,fs=sampling_freq)
    filter[2] = signal.firwin(numtaps=tap_number,cutoff=0.5,fs=sampling_freq)

    data_ch1[i] = signal.lfilter(filter[0],1,data_ch1[i])
    data_ch2[i] = signal.lfilter(filter[1],1,data_ch2[i])
    data_ch3[i] = signal.lfilter(filter[2],1,data_ch3[i])
#"""

#"""
#5次の0位相バターワースフィルターで2~60Hzのバンドパスフィルター
print("2-60Hz band pass filter")
#------------------------------------バターワースフィルターを要実装
for i in range(all_sample_number):
    filter[0] = signal.firwin(numtaps=tap_number,cutoff=[2,60],fs=sampling_freq)
    filter[1] = signal.firwin(numtaps=tap_number,cutoff=[2,60],fs=sampling_freq)
    filter[2] = signal.firwin(numtaps=tap_number,cutoff=[2,60],fs=sampling_freq)

    data_ch1[i] = signal.lfilter(filter[0],1,data_ch1[i])
    data_ch2[i] = signal.lfilter(filter[1],1,data_ch2[i])
    data_ch3[i] = signal.lfilter(filter[2],1,data_ch3[i])
#"""

#μ+-6σの修正外れ値にクリップ

#"""
#各チャンネルで(xi-μi)/σiで正規化
print("Normalized by (xi-μ) / σi for each channel")
avg_ch1 = np.average(data_ch1)
avg_ch2 = np.average(data_ch2)
avg_ch3 = np.average(data_ch3)

std_ch1 = np.std(data_ch1)
std_ch2 = np.std(data_ch2)
std_ch3 = np.std(data_ch3)

data_ch1 = (data_ch1 - avg_ch1) / std_ch1
data_ch2 = (data_ch2 - avg_ch2) / std_ch2
data_ch3 = (data_ch3 - avg_ch3) / std_ch3
#"""

#-----------------------------------------------------------------------------
#i番目のデータを削除する関数
def delete_data(i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label):
    #print("old:",data_ch1.shape)
    data_ch1 = np.delete(data_ch1,i,0)
    data_ch2 = np.delete(data_ch2,i,0)
    data_ch3 = np.delete(data_ch3,i,0)

    data_label = np.delete(data_label,i,0)

    #print("new:",data_ch1.shape)
    #print("label:",data_label.shape)
    all_sample_number -= 1
    i -= 1
    #print("asn:",all_sample_number)
    #print("i_2:",i)
    return i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label
#-----------------------------------------------------------------------------
#"""
#EOG,EMGアーティファクト除去
print("EOG,EMG artifact removal")
print("asn:",all_sample_number)
tmp = np.zeros(ch_number,dtype="float32")

##max<(x_ne) - min(x_ne)>_N
print("< max(x_ne) - min(x_ne) >_N")
i = 0
while i < all_sample_number:
#for i in range(all_sample_number):
    #i -= 1
    #print("i_4:",i)
    tmp[0] = np.amax(data_ch1[i]) - np.min(data_ch1[i])
    tmp[1] = np.amax(data_ch2[i]) - np.min(data_ch2[i])
    tmp[2] = np.amax(data_ch3[i]) - np.min(data_ch3[i])

    z = np.average(tmp)
    #print("tmp:",tmp)
    #if tmp[0] > threshold or tmp[1] > threshold or tmp[2] > threshold:
    if z > 3:
        i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label = delete_data(i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label)
        #print("i_2:",i)
    if i + 1 == all_sample_number or i + 2 == all_sample_number:
        break
    #print("i_3:",i)
    i = i + 1
print("asn:",all_sample_number)

##<x_ne> - <x_n>
print("<x_ne> - <x_n>")
avg_ch_epoch = np.zeros(ch_number,dtype="float32")
avg_ch       = np.zeros(ch_number,dtype="float32")
while i < all_sample_number:
#for i in range(all_sample_number):
    avg_ch_epoch[0] = np.average(data_ch1[i])
    avg_ch_epoch[1] = np.average(data_ch2[i])
    avg_ch_epoch[2] = np.average(data_ch3[i])

    avg_ch[0] = np.average(data_ch1)
    avg_ch[1] = np.average(data_ch2)
    avg_ch[2] = np.average(data_ch3)

    tmp[0] = np.abs(avg_ch_epoch[0] - avg_ch[0])
    tmp[1] = np.abs(avg_ch_epoch[1] - avg_ch[1])
    tmp[2] = np.abs(avg_ch_epoch[2] - avg_ch[2])

    if tmp[0] > threshold or tmp[1] > threshold or tmp[2] > threshold:
        i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label = delete_data(i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label)

    if i + 1 == all_sample_number:
        break
    i += 1
print("asn:",all_sample_number)

##S^2_x_ne
print("S^2_x_ne")
while i < all_sample_number:
#for i in range(all_sample_number):
    tmp[0] = np.var(data_ch1[i])
    tmp[1] = np.var(data_ch2[i])
    tmp[2] = np.var(data_ch3[i])

    if tmp[0] > threshold or tmp[1] > threshold or tmp[2] > threshold:
        i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label = delete_data(i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label)

    if i + 1 == all_sample_number:
        break
    i += 1
print("asn:",all_sample_number)

##median(d(x_ne) / dt)
print("median(d(x_ne) / dt)")
dt = float(1 / sampling_freq)
while i < all_sample_number:
#for i in range(all_sample_number):
    tmp[0] = np.median(np.diff(data_ch1[i]) / dt)
    tmp[1] = np.median(np.diff(data_ch2[i]) / dt)
    tmp[2] = np.median(np.diff(data_ch3[i]) / dt)
    if tmp[0] > threshold or tmp[1] > threshold or tmp[2] > threshold or tmp[0] < -threshold or tmp[1] < -threshold or tmp[2] < -threshold:
        i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label = delete_data(i,all_sample_number,data_ch1,data_ch2,data_ch3,data_label)

    if i + 1 == all_sample_number:
        break
    i += 1
print("asn:",all_sample_number)

#グラフ描画
#-----------------1セル目と同様にするようにしたほうがよいか
random_label = []
for i in range(5):
    random_label.append(random.randint(0,all_sample_number))
    y = data_ch1[random_label[i]] #ランダムに表示
    pyplot.plot(x,y)
    pyplot.show()

#%%
#----データの整理
from sklearn import model_selection
from keras.utils import np_utils

#input_data = np.zeros((ch_number,all_sample_number,imagination_freq),dtype="float64")
input_data = []
input_data.append(data_ch1)
input_data.append(data_ch2)
input_data.append(data_ch3)
input_data = np.array(input_data)
print("input_data.shape:",input_data.shape)

input_data = input_data.reshape(all_sample_number,ch_number,imagination_freq,1)
classes = 2
test_rate = 0.2

#正解ラベルYを(each_data_number,)へ
print(data_label.shape)
print("input_data.shape:",input_data.shape)

#出力Yのラベルを1,2,3から0,1,2に変更
for i in range(len(data_label)):
    if data_label[i] == 1:
        data_label[i] = 0
    elif data_label[i] == 2:
        data_label[i] = 1
    elif data_label[i] == 3:
        data_label[i] = 2
#print(data_label) #OK

#正解ラベルYをont-hot表現へ
data_label = np_utils.to_categorical(data_label,classes)
#print(data_label) #OK

#以下にモデルの検証の仕方
#https://newtechnologylifestyle.net/%E6%A9%9F%E6%A2%B0%E5%AD%A6%E7%BF%92%E3%80%81%E3%83%87%E3%82%A3%E3%83%BC%E3%83%97%E3%83%A9%E3%83%BC%E3%83%8B%E3%83%B3%E3%82%B0%E3%81%A7%E3%81%AE%E5%AD%A6%E7%BF%92%E3%83%87%E3%83%BC%E3%82%BF%E3%81%A8/

####データの分割####
x_train, x_test, y_train, y_test = model_selection.train_test_split(input_data, data_label, test_size=test_rate)

print(x_train.shape)
print(y_train.shape)
print(x_test.shape)
print(y_test.shape)

#%%
#dCNNモデル構築
from keras.models import Sequential
from keras.layers import Reshape,Dense,Conv2D,Conv3D,MaxPooling2D,MaxPooling3D,Flatten,Activation
from keras.optimizers import Adam
from keras.layers import BatchNormalization,Dropout
from keras.initializers import RandomNormal

dropout_rate = 0.1
model = Sequential()
first_filter_number  = 25
second_filter_number = 25
third_filter_number  = 50
fourth_filter_number = 100
fifth_filter_number  = 200

first_kernel_size  = (1,10)
second_kernel_size = (3,1,25)
third_kernel_size  = (10,25)
fourth_kernel_size = (10,50)
fifth_kernel_size  = (10,100)

activation_function = "relu"
poolsize = (3,1)

model = Sequential()

#1層
model.add(Conv2D(filters=first_filter_number,
            kernel_size=first_kernel_size,
            #padding="same",                #sameじゃないと2層の形状がわからない
            bias_initializer="he_normal",
            kernel_initializer="he_normal",
            input_shape=(ch_number,imagination_freq,1)))
model.add(BatchNormalization())
model.add(Activation(activation_function))

#2層
#print(model.output_shape)

model.add(Reshape((3,model.output_shape[2],first_filter_number,1)))#,input_shape=(ch_number,imagination_freq,first_filter_number)))

#print(model.output_shape)

model.add(Conv3D(filters=second_filter_number,
            kernel_size=second_kernel_size,
            #padding="same",
            bias_initializer="he_normal",
            kernel_initializer="he_normal"))

#print(model.output_shape)

model.add(BatchNormalization())
model.add(Activation(activation_function))
model.add(MaxPooling3D(pool_size=(1,3,1)))
#model.add(Dropout(dropout_rate))

#print(model.output_shape)

#3層
a = model.output_shape[3] * model.output_shape[4]
model.add(Reshape((model.output_shape[2],a,model.output_shape[1])))

#print(model.output_shape)

model.add(Conv2D(filters=third_filter_number,
            kernel_size=third_kernel_size,
            #padding="same",
            bias_initializer="he_normal",
            kernel_initializer="he_normal"))
model.add(BatchNormalization())
model.add(Activation(activation_function))
model.add(MaxPooling2D(pool_size=poolsize))
#model.add(Dropout(dropout_rate))

print(model.output_shape)

#"""
#4層
model.add(Reshape((model.output_shape[1],model.output_shape[3],model.output_shape[2])))
print(model.output_shape)

model.add(Conv2D(filters=fourth_filter_number,
            kernel_size=fourth_kernel_size,
            #padding="same",
            bias_initializer="he_normal",
            kernel_initializer="he_normal"))
model.add(BatchNormalization())
model.add(Activation(activation_function))
model.add(MaxPooling2D(pool_size=poolsize))
#model.add(Dropout(dropout_rate))
#"""
print(model.output_shape)
#"""
#5層
model.add(Reshape((model.output_shape[1],model.output_shape[3],model.output_shape[2])))
print(model.output_shape)

model.add(Conv2D(filters=fifth_filter_number,
            kernel_size=fifth_kernel_size,
            #padding="same",
            bias_initializer="he_normal",
            kernel_initializer="he_normal"))
model.add(BatchNormalization())
model.add(Activation(activation_function))
model.add(MaxPooling2D(pool_size=poolsize))
#model.add(Dropout(dropout_rate))
#"""
print(model.output_shape)

model.add(Flatten())
model.add(Dense(units=classes,activation="softmax"))

model.compile(loss = "binary_crossentropy",
                optimizer = Adam(),
                metrics = ["accuracy"])
model.summary()

#%%
#----学習

#--------------------------------------------
#gpuの必要なメモリしか使わない文
import tensorflow as tf
from keras.backend import tensorflow_backend

config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
session = tf.Session(config=config)
tensorflow_backend.set_session(session)
#--------------------------------------------
#earlystopping
from keras.callbacks import EarlyStopping
early_stopping = EarlyStopping(monitor='val_loss',
                                patience=10,
                                verbose=1,
                                mode='auto'
                                )

history = model.fit(x_train,
                    y_train,
                    epochs=50,
                    batch_size=256,
                    validation_data=(x_test,y_test)
                    #,callbacks=[early_stopping]
                    )

#%%
#----グラフ描画
import matplotlib.pyplot as plt

#Accuracy
#plt.figure(facecolor="azure", edgecolor="coral", linewidth=2)
plt.plot(history.history['acc'])
plt.plot(history.history['val_acc'])
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()

#loss
#plt.figure(facecolor="azure", edgecolor="coral", linewidth=2)
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()

#%%


import random
from typing import Union, Callable, Optional, Dict, List, Tuple

import numpy as np
import torch
from kaggle_environments import evaluate, Environment

from matplotlib import pyplot as plt
from tqdm import tqdm
from src.connectx.environment import convert_state_to_image, ConnectXGymEnv
from src.connectx.opponents import interactive_player
from src.connectx.policy import Policy

WIDTH = 256
HEIGHT = 256

def fix_random(seed: int) -> None:
    """
    Fix all the possible sources of randomness.

    :param seed: the seed to use.
    """

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

    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

def play(env: Environment,
         player1: Union[str, Callable],
         player2: Union[str, Callable],
         width: Optional[int] = WIDTH,
         height: Optional[int] = HEIGHT) -> None:
    """
    Render the game between two agents.

    :param env: Kaggle environment
    :param player1: one of the predefined agents ("random", "negamax") or a custom one (Callable)
    :param player2: one of the predefined agents ("random", "negamax") or a custom one (Callable)
    :param width: board rendering width
    :param height: board rendering width
    """

    env.run([player1, player2])
    env.render(mode='ipython', width=width, height=height)

def interactive_play(env: Environment,
                     player1: Optional[Union[str, Callable]] = None,
                     player2: Optional[Union[str, Callable]] = None,
                     width: Optional[int] = WIDTH,
                     height: Optional[int] = HEIGHT) -> None:
    """
    Interactively play with the computer. The interactive player MUST be passed as None.

    :param env: Kaggle environment
    :param player1: one of the predefined agents ("random", "negamax") or a custom one (Callable)
    :param player2: one of the predefined agents ("random", "negamax") or a custom one (Callable)
    :param width: board rendering width
    :param height: board rendering width
    """

    if player1 is not None and player2 is not None:
        raise ValueError('One of the player must be None to be assigned to the interactive player.')

    env.play([player1, player2], width=width, height=height)

def evaluate_matches(matches: List[Dict[str, int]]):
    """
    Evaluate multiple matches.

    :param matches: List containing dictionaries specifying the matches. Each match defines the 2 players 'ply1' and
    'ply2' and how many times the game is played. A mean of the rewards of each play is displayed for each match.
    """

    for m in matches:
        print(f' {m["ply1"]} vs {m["ply1"]} ({m["eps"]} eps):', np.mean(evaluate('connectx',
                                                                                [m['ply1'], m['ply2']],
                                                                                num_episodes=m['eps'])))

def get_win_percentages(player: Union[str, Callable],
                        opponents: Dict[str, Union[str, Callable]],
                        player_name: Optional[str] = None,
                        config: Optional[Dict] = None,
                        n_rounds_as_1st_player: int = 100,
                        n_rounds_as_2nd_player: int = 100):
    """
    Print win/loss percentages between the player and the list of opponents.

    :param player: Your player
    :param player_name: Your player's name
    :param opponents: List of opponents
    :param config: dict specifying the board ('rows' and 'columns') and the stones to win ('inarow'). Default values are
    {'rows': 6, 'columns': 7, 'inarow': 4}
    :param n_rounds_as_1st_player: number of rounds where agent 1 play as the first player
    :param n_rounds_as_2nd_player: number of rounds where agent 1 play as the second player
    """

    if config is None:
        config = {'rows': 6, 'columns': 7, 'inarow': 4}

    if player_name is None:
        player_name = player.__name__ if callable(player) else player

    for opponent_name, opponent in opponents.items():
        print(f'{player_name} VS {opponent_name}')

        outcomes = evaluate('connectx', [player, opponent], config, [], n_rounds_as_1st_player)
        outcomes += [[b, a] for [a, b] in evaluate('connectx', [opponent, player], config, [], n_rounds_as_2nd_player)]

        # print(outcomes)
        print()
        print(f'{player_name} win percentage: {np.round(outcomes.count([1, -1]) / len(outcomes), 2)} - '
              f'{outcomes.count([1, -1])}/{len(outcomes)}')
        print(f'{opponent_name} win percentage: {np.round(outcomes.count([-1, 1]) / len(outcomes), 2)} - '
              f'{outcomes.count([-1, 1])}/{len(outcomes)}')
        print(f'Draw percentage: {np.round(outcomes.count([0, 0]) / len(outcomes), 2)} - '
              f'{outcomes.count([0, 0])}/{len(outcomes)}')
        print(f'Number of invalid moves by {player_name}: {outcomes.count([None, 0])}')
        print(f'Number of invalid moves by {opponent_name}: {outcomes.count([0, None])}\n')

def record_matches(env: ConnectXGymEnv,
                   policy: Union[Policy, Callable],
                   configuration: dict = None,
                   play_as_first_player: bool = True,
                   num_matches: int = 1,
                   render_env: bool = True,
                   interactive_progress: bool = True,
                   keep_player_colour: bool = True,
                   device: str = 'cpu') -> Tuple[List[torch.Tensor], List[torch.Tensor], List[str]]:
    """
    Record screens and actions performed.

    :param env: the gym environment defining the board size and opponent, used only for testing here
    :param policy: the policy of the agent as a nn.Module implementing predict method or as a generic callable
    :param configuration: game and agent config, default is {'columns': 7, 'rows': 6, 'inarow': 4, 'c_type': None}
    :param play_as_first_player: if True the agent is the first player
    :param num_matches: the number of matches
    :param render_env: if True the environment is rendered
    :param interactive_progress: if True, the user must confirm to continue after each match
    :param keep_player_colour: if True the agent color is maintained between player 1 and player 2. e.g. Your agent
    will always be the red player, otherwise the 1st player will always be the red one
    :param device: the device where the recording occurs, 'cpu', 'gpu' ...

    :return: two list of length num_matches where the first contains the states observed and the second the associate
    actions performed and a list of strings representing the end status of each match
    """

    with torch.no_grad():
        if configuration is None:
            configuration = {'columns': 7, 'rows': 6, 'inarow': 4, 'c_type': None}

        # Results
        action_recording = [torch.Tensor([]) for _ in range(num_matches)]
        state_recording = [torch.Tensor([]) for _ in range(num_matches)]
        results_recording = []

        for m in range(num_matches) if interactive_progress else tqdm(range(num_matches)):
            if interactive_progress:
                print(f'match {m + 1}/{num_matches}')

            # Change first or second player
            if not play_as_first_player and env.opponent is interactive_player:
                print('\n - Game is on ... Press any button -')
            env.set_first(play_as_first_player)
            # Initialize the environment and state# Initialize the environment and state
            if not play_as_first_player and env.opponent is interactive_player:
                print('\n - Table is empty, make your first move! -')
            state = env.reset()
            # Get image and convert to torch tensor
            screen = torch.from_numpy(convert_state_to_image(state=state,
                                                             first_player=play_as_first_player,
                                                             keep_player_colour=keep_player_colour)).to(device)
            done = False
            if render_env:
                env.render(board=np.array(state).reshape((env.rows, env.columns, 1)),
                           mode='rgb_image',
                           render_waiting_time=1,
                           keep_player_colour=keep_player_colour)
            while not done:

                # Get action
                observation = {'board': list(state.ravel()), 'mark': 1 if env.first else 2}

                if isinstance(policy, Policy):
                    action = policy.predict(observation,
                                            configuration)
                elif callable(policy):
                    action = policy(observation,
                                    configuration)
                else:
                    raise ValueError('Type of policy not supported, policy must be either a Policy or a Callable.')

                # Render the board as updated by the agent
                if render_env:
                    # Create the board to render
                    render_state = state.copy()
                    ir = 0
                    for r in render_state:
                        # Spot column to update
                        if r[action] != [0]:
                            break
                        ir += 1
                    if ir - 1 >= 0:
                        render_state[ir - 1, action] = [1] if play_as_first_player else [2]
                    env.render(board=render_state,
                               mode='rgb_image',
                               render_waiting_time=1,
                               keep_player_colour=keep_player_colour)

                # Update env
                state, _, done, info = env.step(action)

                # Update results
                action_recording[m] = torch.cat((action_recording[m], torch.Tensor([action])))
                if len(state_recording[m]) == 0:
                    state_recording[m] = screen
                else:
                    state_recording[m] = torch.cat((state_recording[m], screen), dim=0)

                if done:
                    # Update results
                    results_recording.append(info['end_status'])

                    if interactive_progress:
                        if info['end_status'] == 'victory':
                            print('\nThe trained agent won!\n')
                        elif info['end_status'] == 'lost':
                            print('\nThe trained agent lost!\n')
                        elif info['end_status'] == 'invalid':
                            print('\nThe trained agent performed an invalid action!\n')

                    if render_env:
                        env.render(board=state,
                                   mode='rgb_image',
                                   render_waiting_time=1,
                                   keep_player_colour=keep_player_colour)
                    break

                # Update screen for next iteration
                screen = torch.from_numpy(convert_state_to_image(state=state,
                                                                 first_player=play_as_first_player,
                                                                 keep_player_colour=keep_player_colour)).to(device)

            if interactive_progress:
                input('\n - Game finished! -')
        return state_recording, action_recording, results_recording

def show_recordings(state_recording: torch.Tensor,
                    action_recording: torch.Tensor,
                    render_waiting_time: float = 1) -> None:
    """
    :param state_recording: the states
    :param action_recording: the associated actions
    :param render_waiting_time: in seconds
    """

    image = None
    for i, im in enumerate(state_recording):
        im = im.permute(1, 2, 0)
        if image is None:
            image = plt.imshow(im)
        else:
            image.set_data(im)
        plt.title(f'Frame {i} action performed {action_recording[i]}')
        plt.pause(render_waiting_time)
        plt.draw()


"""
Script for plotting pattern correlations between models and observations
for only land and ocean regions

Author     : Zachary M. Labe
Date       : 11 May 2021
Version    : 1
"""

### Import packages
import sys
import matplotlib.pyplot as plt
import numpy as np
import palettable.cubehelix as cm
import palettable.scientific.sequential as sss
import palettable.cartocolors.qualitative as cc
import cmocean as cmocean
import cmasher as cmr
import calc_Utilities as UT
import scipy.stats as sts

### Plotting defaults
plt.rc('text',usetex=True)
plt.rc('font',**{'family':'sans-serif','sans-serif':['Avant Garde']})

### Paramters
letters = ["a","b","c","d","e","f","g","h","i","j","k","l","m"]
datasetsingle = ['SMILE']
dataset_obs = ['ERA5BE']
monthlychoiceq = ['annual','JFM','AMJ','JAS','OND']
typeOfCorr = ['R','R-DT','R-RMGLO','R-TREND']
variables = ['T2M','P','SLP']
reg_name = 'SMILEGlobe'
level = 'surface'
timeper = 'historical'
option = 8
if timeper == 'historical':
    years = np.arange(1950,2019+1,1)
###############################################################################
###############################################################################
land_only = False
ocean_only = True
###############################################################################
###############################################################################
### Read in data
for vv in range(len(variables)):
    for mo in range(len(monthlychoiceq)):
# for vv in range(1):
#       for mo in range(1):
        variq = variables[vv]
        monthlychoice = monthlychoiceq[mo]
        directorydata = '/Users/zlabe/Documents/Research/ModelComparison/Data/Climatologies/'
        directoryfigure = '/Users/zlabe/Desktop/ModelComparison_v1/Climatologies/patternCorr/%s/%s/' % (variq,option)
        if land_only == True:
            saveData =  monthlychoice + '_LAND_' + variq + '_' + reg_name + '_' + timeper
            typemask = 'LAND'
        elif ocean_only == True:
            saveData =  monthlychoice + '_OCEAN_' + variq + '_' + reg_name + '_' + timeper
            typemask = 'OCEAN'
        else:
            sys.exit()
        print('*Filename == < %s >' % saveData)

        corr = np.load(directorydata + saveData + '_corrs.npz')['arr_0'][:option]
        corrdt = np.load(directorydata + saveData + '_corrsdt.npz')['arr_0'][:option]
        corrglo = np.load(directorydata + saveData + '_corrsglo.npz')['arr_0'][:option]
        corrtrends = np.load(directorydata + saveData + '_corrstrends.npz')['arr_0'][:option]
        modelGCMs = ['CanESM2','MPI','CSIRO-MK3.6','KNMI-ecearth','GFDL-CM3','GFDL-ESM2M','LENS','MMmean'][:option]

        ### Ensemble mean correlations
        correns = np.nanmean(corr,axis=1)
        corrensdt = np.nanmean(corrdt,axis=1)
        corrensglo = np.nanmean(corrglo,axis=1)
        ensALLr = [correns,corrensdt,corrensglo]

        ### Pick highest correlation for each year
        corryr = np.argmax(correns,axis=0)
        corryrdt = np.argmax(corrensdt,axis=0)
        corryrglo = np.argmax(corrensglo,axis=0)
        yrALLr = [corryr,corryrdt,corryrglo]

        ### Counts of data for each year
        uniquecorryr,countcorryr = np.unique(corryr,return_counts=True)
        uniquecorryrdt,countcorryrdt = np.unique(corryrdt,return_counts=True)
        uniquecorryrglo,countcorryrglo = np.unique(corryrglo,return_counts=True)
        uniqueALL = [uniquecorryr,uniquecorryrdt,uniquecorryrglo]
        countALL = [countcorryr,countcorryrdt,countcorryrglo]

        ### Correlation for trends averaged across ensemble members
        correnstrends = np.nanmean(corrtrends,axis=1)

        ###############################################################################
        ###############################################################################
        ###############################################################################
        def adjust_spines(ax, spines):
            for loc, spine in ax.spines.items():
                if loc in spines:
                    spine.set_position(('outward', 5))
                else:
                    spine.set_color('none')
            if 'left' in spines:
                ax.yaxis.set_ticks_position('left')
            else:
                ax.yaxis.set_ticks([])

            if 'bottom' in spines:
                ax.xaxis.set_ticks_position('bottom')
            else:
                    ax.xaxis.set_ticks([])

        for ii in range(len(yrALLr)):
            prediction = yrALLr[ii]
            uni = uniqueALL[ii]
            cou = countALL[ii]

            fig = plt.figure(figsize=(10,5))
            ax = plt.subplot(121)
            adjust_spines(ax, ['left', 'bottom'])
            ax.spines['top'].set_color('none')
            ax.spines['right'].set_color('none')
            ax.spines['left'].set_color('dimgrey')
            ax.spines['bottom'].set_color('dimgrey')
            ax.spines['left'].set_linewidth(2)
            ax.spines['bottom'].set_linewidth(2)
            ax.tick_params('both',length=4,width=2,which='major',color='dimgrey')
            ax.yaxis.grid(zorder=1,color='dimgrey',alpha=0.35)

            x=np.arange(1950,2019+1,1)
            plt.scatter(x,prediction,c=prediction,s=40,clip_on=False,cmap=cc.Antique_8.mpl_colormap,
                        edgecolor='k',linewidth=0.4,zorder=10)

            plt.xticks(np.arange(1950,2021,5),map(str,np.arange(1950,2021,5)),size=6)
            plt.yticks(np.arange(0,12+1,1),modelGCMs,size=6)
            plt.xlim([1950,2020])
            plt.ylim([0,option-1])

            plt.title(r'\textbf{PATTERN CORRELATIONS FOR [ %s ] OF %s}' % (variables[vv],typemask),color='k',fontsize=15)
            plt.xlabel(r'\textbf{highest %s - [climate model data] - %s}' % (typeOfCorr[ii],monthlychoice),
                        color='dimgrey',fontsize=8)

            for te in range(len(uni)):
                plt.text(2022,uni[te],r'\textbf{\#%s}' % cou[te],fontsize=7,color='k',
                          ha='center',va='center')

###############################################################################
###############################################################################
###############################################################################
            ax = plt.subplot(122)
            adjust_spines(ax, ['left', 'bottom'])
            ax.spines['top'].set_color('none')
            ax.spines['right'].set_color('none')
            ax.spines['left'].set_color('dimgrey')
            ax.spines['bottom'].set_color('dimgrey')
            ax.spines['left'].set_linewidth(2)
            ax.spines['bottom'].set_linewidth(0)
            ax.tick_params('y',length=4,width=2,which='major',color='dimgrey')
            ax.tick_params('x',length=0,width=0,which='major',color='dimgrey')

            xbar = np.arange(option)
            rects = plt.bar(xbar,correnstrends, align='center',zorder=2)

            plt.axhline(y=np.max(correnstrends),linestyle='--',linewidth=2,color='k',
                        clip_on=False,zorder=100,dashes=(1,0.3))
            plt.axhline(y=0,linestyle='-',linewidth=2,color='dimgrey',
                        clip_on=False,zorder=1)

            ### Set color
            colorlist = ['darkcyan']*8
            for cccc in range(option):
                rects[cccc].set_color(colorlist[cccc])
                rects[cccc].set_edgecolor(colorlist[cccc])

            plt.xticks(np.arange(0,option+1,1),modelGCMs,size=5.45)
            plt.yticks(np.arange(-1,1.1,0.1),map(str,np.round(np.arange(-1,1.1,0.1),2)),size=6)
            plt.xlim([-0.5,option-1+0.5])
            plt.ylim([-0.1,1])

            plt.xlabel(r'\textbf{highest R - [climate model data] - %s}' % (monthlychoice),color='dimgrey',fontsize=8,labelpad=8)
            plt.title(r'\textbf{...AND FOR LINEAR TRENDS}',color='k',fontsize=15)

            plt.tight_layout()
            plt.savefig(directoryfigure + saveData + '_%s_RawClimateModel_PatternCorrelations_%sclasses.png' % (typeOfCorr[ii],len(corr)),dpi=300)

        if option == 8:
            fig = plt.figure()
            ax = plt.subplot(111)
            adjust_spines(ax, ['left', 'bottom'])
            ax.spines['top'].set_color('none')
            ax.spines['right'].set_color('none')
            ax.spines['left'].set_color('dimgrey')
            ax.spines['bottom'].set_color('dimgrey')
            ax.spines['left'].set_linewidth(2)
            ax.spines['bottom'].set_linewidth(2)
            ax.tick_params('both',length=4,width=2,which='major',color='dimgrey')
            ax.yaxis.grid(zorder=1,color='dimgrey',alpha=0.35)

            color = cmr.infinity(np.linspace(0.00,1,len(modelGCMs)))
            for i,c in zip(range(len(modelGCMs)),color):
                if i == 7:
                    c = 'k'
                else:
                    c = c
                plt.plot(years,correns[i],color=c,linewidth=2.3,
                            label=r'\textbf{%s}' % modelGCMs[i],zorder=11,
                            clip_on=False,alpha=1)

            if variq == 'T2M':
                leg = plt.legend(shadow=False,fontsize=9,loc='upper center',
                              bbox_to_anchor=(0.5,0.2),fancybox=True,ncol=4,frameon=False,
                              handlelength=0,handletextpad=0)
                for line,text in zip(leg.get_lines(), leg.get_texts()):
                    text.set_color(line.get_color())

                plt.xticks(np.arange(1950,2030+1,10),map(str,np.arange(1950,2030+1,10)),size=5.45)
                plt.yticks(np.arange(0.98,1.01,0.005),map(str,np.round(np.arange(0.98,1.01,0.005),3)),size=6)
                plt.xlim([1950,2020])
                plt.ylim([0.98,1.0])
            else:
                leg = plt.legend(shadow=False,fontsize=9,loc='upper center',
                                  bbox_to_anchor=(0.5,0.15),fancybox=True,ncol=4,frameon=False,
                                  handlelength=0,handletextpad=0)
                for line,text in zip(leg.get_lines(), leg.get_texts()):
                    text.set_color(line.get_color())

                plt.xticks(np.arange(1950,2030+1,10),map(str,np.arange(1950,2030+1,10)),size=5.45)
                plt.yticks(np.arange(0,1.01,0.1),map(str,np.round(np.arange(0,1.01,0.1),3)),size=6)
                plt.xlim([1950,2020])
                plt.ylim([0.4,1.0])

            plt.xlabel(r'\textbf{Average R - [climate model data] - %s for %s}' % (monthlychoice,variq),color='dimgrey',fontsize=8,labelpad=8)
            plt.title(r'\textbf{PATTERN CORRELATIONS PER YEAR - %s}' % typemask,color='k',fontsize=15)

            plt.tight_layout()

            plt.savefig(directoryfigure + saveData + '_RawClimateModel_ActualPatternCorrs.png',dpi=300)


'''
Useful functions for plotting with matplotlib
'''

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.colors as clrs
import numpy as np
import scipy.interpolate as interpolate

def get_mpl_colors():
    return plt.rcParams['axes.prop_cycle'].by_key()['color']

def darken(c, f=2):
    '''Darken color by factor. E.g., factor=2 returns color half as bright.'''
    return tuple([v/f for v in clrs.to_rgb(c)])

def lighten(c, f=2):
    '''Lighten color by factor.'''

    return tuple([1 - (1-v)/f for v in clrs.to_rgb(c)])

def greyscale(c):
    (r, g, b) = clrs.to_rgb(c)
    return (r + g + b)/3

def generate_colormap(color_low=None, color_high=None, num_colors=256):

    vals = np.ones((num_colors, 4))

    mpl_colors = get_mpl_colors()
    if color_low is None:
        C0 = clrs.to_rgb(mpl_colors[0])
    else:
        C0 = clrs.to_rgb(color_low)

    if color_high is None:
        C1 = clrs.to_rgb(mpl_colors[1])
    else:
        C1 = clrs.to_rgb(color_high)

    for i in range(3):
        vals[:, i] = np.linspace(C0[i], C1[i], num_colors)

    return mpl.colors.ListedColormap(vals)

def scale_figure(fig, scale):
    ''' Return figure rescaled by the given scale factor.

    Arguments:
        fig     matplotlib Figure object
        scale   Scale factor. Either list (width, height) or scalar.
    '''

    size = fig.get_size_inches()
    try:
        fig.set_size_inches(*[s*z for s, z in zip(scale, size)])
    except TypeError:
        fig.set_size_inches(*[scale*z for z in size])
    return fig

def plot_periodic(ax, x, y, nper=[0.2, 0.2], reset_zero=False, **kwargs):
    '''Plot function, and extend it periodically.

    Arguments:
        ax          matplotlib Axes object
        x,y         x and y data (numpy arrays. x must be evenly spaced)
        nper        Number of periods by which to extend the plot left and
                    right. E.g. nper=[0.5,0.25] will plot the function plus 0.5
                    periods to the left and 0.25 periods to the right. If nper
                    is a scalar, left and right extensions will be the same.
        reset_zero  If True, will set the leftmost x value to zero
        kwargs      Passed directly to ax.plot() function

    returns xp, xy: the periodically-extended plot values.
    '''

    if np.isscalar(nper):
        nper_left = nper
        nper_right = nper
    else:
        nper_left = nper[0]
        nper_right = nper[1]

    dx = np.average(np.diff(x))
    if np.abs(np.std(np.diff(x))/dx) > 1e-10:
        raise ValueError("x array must be evenly spaced")

    x_per = (len(x) + 1)*dx

    N_left = int(nper_left*x_per/dx)+1
    N_right = int(nper_right*x_per/dx)+1

    xp = np.linspace(
        x[0]-N_left*dx,
        x[-1]+N_right*dx,
        N_left + N_right + len(x)
    )

    if reset_zero:
        xp -= xp[0]

    left_pad = y[-(np.arange(N_left) % len(y))-1]
    left_pad = np.flip(left_pad)
    right_pad = y[np.arange(N_right) % len(y)]

    yp = np.concatenate([left_pad, y, right_pad])

    ax.plot(xp, yp, **kwargs)
    ax.set_xlim([xp[0], xp[-1]])

    return xp, yp

def interpolate_2D_grid(x, y, z, Nx, Ny=None, **kwargs):
    '''Simple function to interpolate 2D data assuming x and y form a
    rectangular grid. Interpolates to an Nx by Ny grid.
    If Ny is not given, it defaults as equal to Nx.

    Remaining kwargs are passed to scipy.interpolate.griddata. Also 'rescale'
    default value is changed to True.

    Returns interpolated arrays xi, yi, zi. These are flat 1D arrays of length
    Nx*Ny'''

    if Ny is None:
        Ny = Nx

    X = np.linspace(np.min(x), np.max(x), Nx)
    Y = np.linspace(np.min(y), np.max(y), Ny)

    XX, YY = np.meshgrid(X, Y)

    xi = XX.flatten()
    yi = YY.flatten()

    if 'rescale' not in kwargs:
        kwargs['rescale'] = True

    zi = interpolate.griddata(
        (x, y), z, (xi, yi), **kwargs
    )

    return xi, yi, zi

def plot_2D_color(ax, x, y, c, **kwargs):
    '''Wrapper for matplotlib's tripcolor function.
    Difference is that this  one scales the variables before breaking the grid
    into triangles, to give better results when x and y have different units.

    ax should be a matplotlib Axes object
    Remaining arguments are the same as tripcolor

    '''

    if 'triangles' not in kwargs:
        x_rel = x/np.ptp(x)
        y_rel = y/np.ptp(y)

        kwargs['triangles'] = mpl.tri.Triangulation(x_rel, y_rel).triangles

    return ax.tripcolor(x, y, c, **kwargs)


#!/usr/bin/env python
#  -*- mode: python; indent-tabs-mode: nil; -*- coding: utf-8 -*-

"""
RateLimiter.py

Copyright 2010-2014 by Marcello Perathoner

Distributable under the GNU General Public License Version 3 or newer.

A very basic IP rate limiter tool for cherrypy.

begin;
drop table robots.blocks;
create table robots.blocks (
   ip             cidr      primary key,
   host           text      default null,
   has_info       boolean   default false,
   is_blocked     boolean   default true,
   is_whitelisted boolean   default false,
   created        timestamp default now (),
   expires        timestamp default null,
   types          text[]    default null,
   org            text      default null,
   country        text      default null,
   note           text      default null,
   whois          text      default null,
   user_agents    text[]    default null,
   requests       text[]    default null
);

"""

from __future__ import unicode_literals

import datetime
import re
import logging
import threading
import random
import collections
import itertools

import cherrypy
import requests
import netaddr

from libgutenberg.GutenbergDatabase import DatabaseError
import BaseSearcher
import Page
import asyncdns
import ipinfo

# no. of requests to keep in LIFO queue for each IP
KEEP_REQUESTS = 10

RE_SPOOFED_BOTS = re.compile (
    'Googlebot|msnbot|bingbot|Yahoo! Slurp|YandexBot|Baiduspider|ia_archiver')

TRUSTED_XFF_URL = "https://meta.wikimedia.org/w/extensions/TrustedXFF/trusted-hosts.txt"

def log (ip, msg, context = 'RATELIMITER', severity = logging.INFO):
    """ Log with IP. """

    a = []
    if ip:
        a.append (str (ip))
    if msg:
        a.append (str (msg))

    cherrypy.log (' '.join (a), context = context, severity = severity)

def to_cidr (ip):
    """ Build cidr out of ip. """

    if ip and ip != 'unknown':
        if ':' in ip:
            return str (ip) + '/128'
        return str (ip) + '/32'
    return None

# See: https://docs.python.org/3/library/itertools.html
def unique_everseen (iterable):
    "List unique elements, preserving order. Remember all elements ever seen."
    # unique_everseen('AAAABBBCCDAABBB') --> A B C D
    # unique_everseen('ABBCcAD', str.lower) --> A B C D
    seen = set ()
    seen_add = seen.add
    for element in itertools.filterfalse (seen.__contains__, iterable):
        seen_add (element)
        yield element

class IPStruct (object):
    """ A container for a client and proxy address. """

    def __init__ (self):
                                    # ips are always of type string
        self.client_ip = None       # ip of client (or of untrusted proxy)
        self.proxy_ip = None        # ip of proxy or None
        self.proxy_type = None      # None, 'anonymous', 'untrusted', 'trusted', 'private'
        self.whitelisted = False    # if whitelisted ip (eg. Googlebot)
        self.ipinfo = None          # more (dns, whois) info about this ip

    def from_request (self, request, rl):
        """ Init from remote IP and X-Forwarded-For header. """

        xffh = request.headers.get ('X-Forwarded-For') or "%s, %s" % (
            cherrypy.request.remote.ip, cherrypy.request.remote.ip)

        log ('X-FORWARDED-FOR', xffh)

        # split string header into array of IPs
        xff = [x.strip () for x in xffh.split (',')]

        # strip ibiblio haproxy IP
        xff = [x for x in xff if not x.startswith ("2610:28:3091:3001:")]

        xff = list (reversed (list (unique_everseen (xff))))
        if not xff:
            xff = [cherrypy.request.remote.ip]

        try:
            nip = netaddr.ip.IPAddress (xff[0])
        except netaddr.AddrFormatError:
            raise cherrypy.HTTPError (400, 'Bad Request. Invalid IP: %s.' % xff[0])

        self.whitelisted = nip in rl.whitelist

        # the request reaches the app server either this way:
        #
        #   user_proxy* haproxy varnish mod_proxy app_server
        #
        # or this way:
        #
        #   user_proxy* haproxy varnish php_curl app_server
        #
        # so the layout of xff is either:
        #
        # xff[0] address of user    (added by curl)
        # xff[1] address of haproxy (added by varnish)
        # xff[2] address of user    (added by haproxy)
        # xff[3] address of user    (added by user_proxy)
        #
        # or
        #
        # xff[0] address of user    (added by mod_proxy)
        # xff[1] address of haproxy (added by varnish)
        # xff[2] address of user    (added by haproxy)
        # xff[3] address of user    (added by user_proxy)

        self.client_ip = xff[0]
        offset = 1
        if len (xff) >= 2 and xff[0] == xff[1]:
            # skip identic ip
            offset = 2
        if len (xff) > offset:
            # some kind of proxy detected
            self.client_ip = xff[offset]
            self.proxy_ip  = xff[0]

            if self.client_ip == 'unknown':
                self.proxy_type = 'anonymous'
            else:
                try:
                    if netaddr.ip.IPAddress (self.client_ip).is_private ():
                        self.proxy_type = 'private'
                except netaddr.AddrFormatError:
                    raise cherrypy.HTTPError (400, 'Bad Request. Invalid IP: %s.' % self.client_ip)

            if self.proxy_type is None:
                # Check if this request came through a trusted proxy.
                # This check is expensive, so do it only if needed.
                if nip in rl.trusted_xff_list:
                    self.proxy_type = 'trusted'
                else:
                    self.proxy_type = 'untrusted'

    def get_cache_key (self):
        """ Make a key suitable for caching an ipsession.

        This key will distinguish users.

        """
        if self.proxy_type == 'trusted':
            return "%s-%s" % (self.proxy_ip, self.client_ip)
        else:
            return self.proxy_ip or self.client_ip

    def get_ip_to_block (self):
        """ Get the IP to block.

        Because we use apache 'deny from' to block, we need
        the proxy address, not the address behind the proxy.

        """
        return self.proxy_ip or self.client_ip

    def sort_key (self):
        """ Return a key for sorting of IPStructs. """
        return netaddr.ip.IPAddress (self.get_ip_to_block ()).sort_key ()

    def __str__ (self):
        """ Return a human-readable representation. """
        if self.proxy_ip:
            return "%s via %s (%s)" % (self.client_ip, self.proxy_ip, self.proxy_type)
        return self.client_ip

class RatelimiterSession (dict):
    """ A session maintained for every remote IP. """

    def __init__ (self, rl, ips):
        super (RatelimiterSession, self).__init__ ()

        self.ips = ips
        self.sessions = set ()
        now = datetime.datetime.now ()

        # pylint: disable=too-many-function-args
        self.update (
            {
                'blocked':     0,                  # 1 = temp, 2 = perm, 9 = written to db
                'expires':     now,                # expiration of this session
                'last_hit':    now,                # time of last hit
                'hits':        0,                  # all hits
                'rhits':       0,                  # ip score
                'rhits_max':   rl.rhits_max,       # next captcha appears at this score
                'captchas':    0,                  # no. of captchas presented
                'rearms':      0,                  # how many times rearmed (thru captcha)
                'dhits':       0,                  # denied hits
                'user_agents': set (),             # all user agents seen from this ip
                'signatures':  set (),             # all http signatures seen from this ip
                'requests':    collections.deque ([], KEEP_REQUESTS), # last n requests
                'referrers':   collections.deque ([], KEEP_REQUESTS), # last n referrers
                'categories':  collections.deque ([], KEEP_REQUESTS), # last n categories
                'headers':     [],                 # request headers
                'active':      True,               # for stats page
                'css_ok':      None,               # flag for stats page
                'expire_days': 1,                  # expiration of block
            }
        )

        self.set_expiration (rl.expiration)

    def append_trace (self, rl, ua, signature):
        """ Append trace of last X requests. """

        headers = cherrypy.request.headers

        session = cherrypy.session

        ipset = session.get ('ips', set ())
        ipset.add (self.ips.client_ip)
        session['ips'] = ipset

        uaset = session.get ('user_agents', set ())
        uaset.add (ua)
        session['user_agents'] = uaset

        self['signatures'].add (signature)

        request_str = headers.get ('X-Request')
        if request_str is None:
            request_str = cherrypy.request.path_info
            if cherrypy.request.query_string:
                request_str += '?' + cherrypy.request.query_string
        request_str = request_str[:100]
        self['requests'].append (request_str)
        self['referrers'].append (headers.get ('Referer'))

        # categorize request
        for cat in rl.categories:
            if request_str.startswith (cat):
                self['categories'].append (cat)
                break

    def set_add_with_trigger (self, set_, value, trigger_size):
        """ Insert value into set and trigger on size. """

        if value in set_:
            return False
        set_.add (value)
        return len (set_) == trigger_size

    def calc_block (self, rl, ua):
        """ Calculate the penalties for this hit.

        Penalties make rhits grow faster, thus captchas come earlier.

        """

        headers = cherrypy.request.headers
        session = cherrypy.session
        path = cherrypy.request.path_info
        now = datetime.datetime.now ()

        self['rhits'] += 1

        # penalties
        if self['last_hit'] + datetime.timedelta (seconds = rl.min_interval) > now:
            self['rhits'] += rl.rate_malus

        # new session?
        if 'captchas' not in session:
            self['rhits'] += rl.session_malus
            session['captchas'] = 0
        elif session['captchas'] == 0:
            self['rhits'] += rl.no_captcha_malus

        if ('.opds' not in path) and ('.stanza' not in path):
            if self.has_few ('referrers'):
                self['rhits'] += rl.referrer_malus

            if self.has_few ('categories'):
                self['rhits'] += rl.category_malus

            if 'Accept-Encoding' not in headers:
                self['rhits'] += rl.encoding_malus

            if self['css_ok'] == False: # 3-state logic
                self['rhits'] += rl.css_malus

        penalty = headers.get ('X-Penalize-Me')
        if penalty is not None:
            self['rhits'] += max (int (penalty), 1)

        if 'X-Error-404' in headers:
            self['rhits'] += rl.e404_malus

        # soft block to captcha
        if self['blocked'] == 0:
            if self['rhits'] >= self['rhits_max']:
                self['blocked'] = 1

            if self.set_add_with_trigger (self['user_agents'], ua, 5):
                if len (self['signatures']) <= 3:
                    self.log ('Captcha on too few http header signatures')
                    self['blocked'] = 1

        if self['blocked'] == 1:
            if self['rearms'] > rl.rearms_max:
                self.log ('Blocked on rearms_max')
                self['blocked'] = 2

        # hard block
        if self['blocked'] < 2:
            if self['captchas'] > rl.captchas_max:
                self.log ('Blocked on captchas_max')
                self['blocked'] = 2

            penalty = headers.get ('X-Block-Me')
            if penalty is not None:
                self.log ('Blocked on X-Block-Me')
                self['blocked'] = 2
                self['expire_days'] = max (int (penalty), 1)

            if len (session['user_agents']) > rl.ua_limit:
                # spoofed random user-agents
                self.log ('Blocked on session ua_limit')
                self['blocked'] = 2

            if len (session['ips']) > rl.ip_limit:
                # session from multiple ips (probably proxies)
                self.log ('Blocked on session ip_limit')
                self['blocked'] = 2

            if (ua is not None) and (RE_SPOOFED_BOTS.search (ua) is not None):
                # spoofed user-agent because not from whitelisted ip
                self.log ('Blocked on spoofed user agent')
                self['blocked'] = 2

            # if self.ips.proxy_type == 'anonymous':
            #     self.log ('Blocked on anon_proxy')
            #     self['blocked'] = 2

    def extend_expiration (self, timeout):
        """ Extend expiration time to now + timeout. """

        new_expires = datetime.datetime.now () + datetime.timedelta (seconds = timeout)
        self['expires'] = max (new_expires, self['expires'])

    def set_expiration (self, timeout):
        """ Set expiration time to now + timeout. """

        new_expires = datetime.datetime.now () + datetime.timedelta (seconds = timeout)
        self['expires'] = new_expires

    def has_few (self, what):
        """ IPs that always use the same referrer or always hit the
        same category are most likely bots. """

        return self['hits'] > KEEP_REQUESTS and len (set (list (self[what]))) < 2

    def get_tags (self):
        """ Get all tags from ipinfo and add oir own. """
        tags = list (self.ips.ipinfo.tags)
        if self['css_ok'] == False: # 3-state
            tags.append ('NO_CSS')
        return sorted (tags)

    def captcha_answer (self, response):
        """ Check answer to captcha. """

        if response.is_valid:
            # session.captchas NOT ipsession.captchas
            captchas = cherrypy.session.get ('captchas', 0)
            cherrypy.session['captchas'] = captchas + 1
            self.rearm ()
            self.log ('Correct answer', 'CAPTCHA')
        else:
            self.log ('Wrong answer', 'CAPTCHA')

    def rearm (self):
        """ Rest counters. """

        rl = cherrypy.tools.rate_limiter # pylint: disable=E1101

        if self['blocked'] < 2 and self['rearms'] <= rl.rearms_max:
            self.log ('Rearmed')
            self['rearms'] += 1
            self['blocked'] = 0
            self['captchas'] = 0
            # recharge, but don't add up multiple recharges
            self['rhits_max'] = min (self['rhits'] + rl.rhits_max,
                                     self['rhits_max'] + rl.rhits_max)
            self.set_expiration (rl.expiration)

    def block_me (self):
        """ Immediately block this IP. """

        self.log ('Blocked on block_me')
        self['blocked'] = 2

    def unblock_me (self):
        """ Unblock this IP. """

        self.rearm ()

    def log (self, msg, context = 'RATELIMITER'):
        """ Log with IP. """

        log (str (self.ips), str (msg), context)

    def log_state (self):
        """ Log the session state. """

        tup = (
            self['blocked'],
            self['hits'],
            self['rhits'],
            self['rhits_max'],
            self['requests'][-1],
            )
        self.log ("%d %d %d/%d %s" % tup)

        if 'is-catalog-maintainer' in cherrypy.request.cookie:
            cherrypy.response.headers['X-Ratelimiter-Info'] = "b:%d h:%d %d/%d %s" % tup

class RateLimiterTool (cherrypy.Tool):
    """ Limit rate of access per IP. """

    lock = threading.Lock ()
    dblock = threading.Lock () # database accesses
    cache = {}

    # http header fields that should change if user-agent changes,
    # but bots that spoof user-agent won't bother to randomize

    http_headers_no_random = set (
        """Accept Accept-Language Accept-Charset Accept-Encoding
        Cache-Control Connection Dnt Keep-Alive Pragma"""
        .split ())

    http_headers_ignore = set (
        """X-Forwarded-For X-Forwarded-Server X-Forwarded-Host
        X-Request X-Varnish X-Block-Me X-Penalize-Me Remote-Addr
        Cookie"""
        .split ())

    def __init__ (self):
        """ Initialize the Rate Limiter.

        Needs to be on before_request_body on priority > 30 because
        tools.proxy (on prio 30) needs a chance to rewrite the IP.

        """

        super (RateLimiterTool, self).__init__ ('before_request_body', None)
        self.callable = self.limit
        self._priority = 60

        conf = cherrypy.config.get

        # ipsession expiration in seconds
        self.expiration       = conf ('tools.rate_limiter.expiration',      30 * 60)

        # cleanup thread frequency
        self.frequency        = conf ('tools.rate_limiter.frequency',       60)

        # when rhits are reached a captcha must be solved
        self.rhits_max        = conf ('tools.rate_limiter.rhits_max',       500)

        # this many captchas are presented before hard block
        self.captchas_max     = conf ('tools.rate_limiter.captchas_max',    10)

        # how many times rearm thru captcha works
        self.rearms_max       = conf ('tools.rate_limiter.rearms_max',      20)

        # more than this many user agents on *session* will get ip blocked
        self.ua_limit         = conf ('tools.rate_limiter.ua_limit',         4)

        # more than this many ips on *session* will get ip blocked
        self.ip_limit         = conf ('tools.rate_limiter.ip_limit',        20)

        # leave this much time between requests or get rate_malus
        self.min_interval     = conf ('tools.rate_limiter.min_interval',     1)

        self.session_malus    = conf ('tools.rate_limiter.session_malus',   20)
        self.no_captcha_malus = conf ('tools.rate_limiter.no_captcha_malus', 1)
        self.e404_malus       = conf ('tools.rate_limiter.e404_malus',     200)
        self.rate_malus       = conf ('tools.rate_limiter.rate_malus',      10)
        self.encoding_malus   = conf ('tools.rate_limiter.encoding_malus',   5)
        self.referrer_malus   = conf ('tools.rate_limiter.referrer_malus',   5)
        self.category_malus   = conf ('tools.rate_limiter.category_malus',   0)
        self.css_malus        = conf ('tools.rate_limiter.css_malus',        0)

        self.whitelist        = self.load_whitelist (conf)
        self.trusted_xff_list = self.load_trusted_xff_list (conf)

        self.users = 0           # no. of active users
        self.hits = 0            # no. of hits
        self.whitelist_hits = 0  # no. of hits from whitelisted ips
        self.denied_hits = 0     # no. of denied accesses

        self.hit_accumulator = 0
        self.whitelist_hit_accumulator = 0
        self.denied_hit_accumulator = 0

        self.categories = '/ebooks/search/ /ebooks/ /files/ /cache/'.split ()

        # reasons
        # suggest: ajax call, user cannot see the captcha
        # captcha: user must see captcha even if soft-blocked
        # stats: don't block the sysop
        # re must match whole path !
        self.whitelisted_paths = re.compile (r'/ebooks/suggest/|/stats/|/w/captcha/.*')

    @staticmethod
    def load_whitelist (conf):
        """ Load IP whitelist. """

        wl = netaddr.ip.sets.IPSet ()
        for cidr in conf ('tools.rate_limiter.whitelist', '').splitlines ():
            cidr = cidr.partition ('#')[0]
            cidr = cidr.strip ()
            if cidr:
                wl.add (netaddr.ip.IPNetwork (cidr))
                # print cidr
        return wl

    @staticmethod
    def load_trusted_xff_list (conf):
        """ Load a list of proxies whose X-Forwarded-For header we trust. """

        tl = netaddr.ip.sets.IPSet ()
        try:
            url = conf ('tools.rate_limiter.trusted_xff_url')
            if url:
                log (None, "Requesting trusted xff list from: %s ..." % url)
                trusted_xff_list = requests.get (url)
                log (None, "Scanning trusted xff list ...")
                cidrs = []
                to_resolve = {}
                skip = True # horrible hack to skip AOL proxies in list
                for line in trusted_xff_list.text.splitlines ():
                    if line.startswith ('# British Columbia school network'):
                        skip = False
                    if skip:
                        continue
                    line = line.partition ('#')[0]
                    line = line.strip ()
                    if line:
                        try:
                            cidr = netaddr.ip.IPNetwork (line)
                            cidrs.append (cidr)
                        except (ValueError, netaddr.core.AddrFormatError):
                            # not an IP, maybe a hostname we can resolve
                            to_resolve[line] = None
                trusted_xff_list.close ()
                tl.update (cidrs)
                if to_resolve:
                    log (None, "Resolving hosts in xff list ...")
                    asyncdns.bulk_query (to_resolve, asyncdns.adns.rr.A, 5)
                    tl.update ( (ip[0] for ip in to_resolve.values () if ip is not None) )

        except Exception as what:
            log (None, what)

        for cidr in conf ('tools.rate_limiter.trusted_xff_list', '').splitlines ():
            cidr = cidr.partition ('#')[0]
            cidr = cidr.strip ()
            if cidr:
                tl.add (netaddr.ip.IPNetwork (cidr))
                # print cidr

        log (None, "Loaded {count} trusted xff ips.".format (count = len (tl)))
        return tl

    def e404 (self):
        """ 404 penalty when hitting non-existing ebook nos.
        for people who write naive scraper scripts. """
        cherrypy.ipsession['rhits'] += self.e404_malus

    def block (self, ip):
        """ Block IP. """

        ipsession = self.cache.get (ip)
        if ipsession is not None:
            ipsession.block_me ()
            return True
        return False

    def unblock (self, ip):
        """ Unblock IP. """

        ipsession = self.cache.get (ip)
        if ipsession is None:
            return False

        ipsession.unblock_me ()

        conn = cherrypy.engine.pool.connect ()
        c  = conn.cursor ()
        with self.dblock:
            try:
                c.execute ('begin')
                c.execute ('set transaction isolation level serializable')
                query = 'UPDATE robots.blocks SET expires = now () WHERE ip = %(ip)s'
                params = { 'ip': to_cidr (ip) }
                c.execute (query, params)
                c.execute ('commit')
            except DatabaseError as what:
                cherrypy.log ("SQL Error: %s\n" % what,
                              context = 'RATELIMITER', severity = logging.ERROR)
                cherrypy.log ("Query was: %s\n" % c.mogrify (query, params),
                              context = 'RATELIMITER', severity = logging.ERROR)
                c.execute ('rollback')

        return True

    @staticmethod
    def get_challenge ():
        """ Generate a challenge. Called from genshi templates. """

        # currently the challenge gets applied as xml id
        challenge = cherrypy.session['challenge'] = 'id' + str (random.randint (1000, 9999))
        return challenge

    @staticmethod
    def check_challenge ():
        """ User agent handles cookies, JS and CSS if answer equals challenge. """

        # Note that we cannot insert a challenge into every html page,
        # only in those the app server generates. Thus we may get back
        # the answer cookie for one challenge more than once.

        s = cherrypy.session
        if 'challenge' not in s:
            return None # no challenge emitted yet
        cookies = cherrypy.request.cookie
        if 'bonus' not in cookies:
            return False

        challenge = s['challenge']
        answer = cookies['bonus'].value

        if 'is-catalog-maintainer' in cookies:
            # send some debug info
            cherrypy.response.headers['X-CSS-Bonus'] = (
                "Challenge: %s - Response: %s" % (challenge, answer))

        return answer == challenge

    def _get_ip_session (self):
        """ Retrieve or create an ipsession for this request. """

        ips = IPStruct ()
        ips.from_request (cherrypy.request, self)

        # get or set ipsession object
        cache_key = ips.get_cache_key ()
        try:
            ipsession = self.cache[cache_key]
        except KeyError:
            ipsession = RatelimiterSession (self, ips)
            # To guard against `list (dict.items ())` failing in
            # `self.reset ()` because of 'RuntimeError: dictionary
            # changed size during iteration'
            with self.lock:
                self.cache[cache_key] = ipsession

        return ipsession

    def _get_headers_array (self, headers):
        """ All request headers in one array of string. """

        a = []
        for k, v in headers.items ():
            a.append (k + ': ' + v)
        return a

    def _get_http_signature (self, headers):
        """ Get a signature of the header fields.

        Rationale: to detect randomly spoofed user-agents we
        check if the signature stays the same over a range of
        (possibly spoofed) user-agents.

        """

        a = []
        for k, v in headers.items ():
            if k in self.http_headers_ignore or k.startswith ('X-'):
                continue
            if k in self.http_headers_no_random:
                a.append (k + ': ' + v)
            else:
                a.append (k)
        # FIXME: we have no way of getting the header order
        # coordinate with ibiblio varnish team?
        # FIXME: just return a hash after debugging this
        return '|'.join (sorted (a))

    def limit (self, **dummy_kwargs):
        """ Do accounting for ip and return 403 if ip exceeded quota. """

        ipsession = self._get_ip_session ()

        cherrypy.serving.ipsession = ipsession
        if not hasattr (cherrypy, "ipsession"):
            cherrypy.ipsession = cherrypy._ThreadLocalProxy ('ipsession') # pylint: disable=W0212

        ipsession.sessions.add (cherrypy.session.id)

        headers = cherrypy.request.headers

        ua = headers.get ('User-Agent')
        signature = self._get_http_signature (headers)

        whitelisted = (ipsession.ips.whitelisted or
                       self.whitelisted_paths.match (cherrypy.request.path_info))

        with self.lock:
            ipsession['active'] = True
            ipsession.extend_expiration (self.expiration)
            ipsession.append_trace (self, ua, signature)
            ipsession['css_ok'] = self.check_challenge () # store value for stats page

            if not whitelisted:
                ipsession.calc_block (self, ua)

            ipsession['last_hit'] = datetime.datetime.now () # set this after calc_block
            ipsession.log_state ()

            # allow hit
            if whitelisted:
                ipsession.log ('whitelisted')
                cherrypy.response.headers['X-Whitelisted'] = 'ratelimiter'
                ipsession['hits'] += 1
                self.whitelist_hit_accumulator += 1
                return

            if ipsession['blocked'] == 0:
                ipsession['hits'] += 1
                self.hit_accumulator += 1
                return

            # soft or hard deny
            self.denied_hit_accumulator += 1
            cherrypy.lib.caching.expires (0, True)

            if ipsession['blocked'] == 1:
                ipsession['captchas'] += 1
                ipsession.extend_expiration (10 * self.expiration) # watch 'em longer
                ipsession.log ('Redirected to captcha')
                cherrypy.response.headers['X-Ratelimiter-Denied'] = "503"
                raise cherrypy.HTTPRedirect ('//www.gutenberg.org/w/captcha/question/')
            else:
                ipsession['dhits'] += 1
                ipsession.set_expiration (self.expiration)
                if not ipsession['headers']:
                    ipsession['headers'] = self._get_headers_array (headers)
                ipsession.log ('Blocked 403')
                cherrypy.response.headers['X-Ratelimiter-Denied'] = "403"
                raise cherrypy.HTTPError (403, 'Forbidden')

    def reset (self):
        """ Reset the hit counters. """

        try:
            users = 0
            now = datetime.datetime.now ()
            active_window = now - datetime.timedelta (seconds = 60)

            # To guard against `list (self.cache.items ())` failing because
            # of 'RuntimeError: dictionary changed size during iteration'
            with self.lock:
                sessions = list (self.cache.items ())

            # We have to lock self.cache while deleting items, so to keep
            # the lock short, we first compute a list of keys to delete.
            expired_keys = []
            for key, data in sessions:
                if data['expires'] < now:
                    expired_keys.append (key)
                if data['last_hit'] < active_window:
                    data['active'] = False
                if data['active']:
                    users += 1

            with self.lock:
                for key in expired_keys:
                    try:
                        del self.cache[key]
                    except KeyError:
                        cherrypy.log ('KeyError in reset ()', context = 'RATELIMITER',
                                      severity = logging.ERROR)
                self.users = users
                self.hits = self.hit_accumulator
                self.whitelist_hits = self.whitelist_hit_accumulator
                self.denied_hits = self.denied_hit_accumulator

                self.hit_accumulator = 0
                self.whitelist_hit_accumulator = 0
                self.denied_hit_accumulator = 0

            BaseSearcher.formats_acc.reset ()
            BaseSearcher.formats_sum_acc.reset ()

            cherrypy.log ('Run reset ()', context = 'RATELIMITER', severity = logging.INFO)

        except Exception as what:
            # make sure background thread never stops
            cherrypy.log ("Error in reset (): %s" % what,
                          context = 'RATELIMITER', severity = logging.ERROR)

    def debug_reset_all (self):
        """ Reset all hit counters unconditionally.

        Used only by the test suite.

        """
        for key in list (self.cache.keys ()):
            try:
                del self.cache[key]
            except KeyError:
                pass
        cherrypy.log ('Run debug_reset_all ()', context = 'RATELIMITER', severity = logging.INFO)

    @staticmethod
    def get_ipinfo_for (sessions):
        """ Retrieve IP info for all sessions.

        Retrieves info asynchronously in parallel for all sessions.

        """

        adns = asyncdns.AsyncDNS (asyncdns.RESOLV_CONF)
        for ipsession in sessions:
            try:
                if ipsession.ips.ipinfo is None:
                    ip = ipsession.ips.get_ip_to_block ()
                    ipsession.ips.ipinfo = ipinfo.IPInfo (adns, ip)
            except Exception as what:
                cherrypy.log ("Error in get_ipinfo_for (): %s" % what,
                              context = 'RATELIMITER', severity = logging.ERROR)
        adns.wait ()
        adns.cancel ()

    @staticmethod
    def ua_decode (ua):
        """ Try if we can upgrade the encoding to utf-8.

        'headers' items are already decoded to unicode by
        cherrypy, but I found that many foreign (chinese) user
        agents use but fail to specify the utf-8 encoding, so
        that cherrypy wrongly decodes from iso-8859-1.
        """

        if ua is None:
            return 'NO_UA'
        try:
            return ua.encode ('iso-8859-1').decode ('utf-8')
        except UnicodeError:
            return ua

    def to_database (self):
        """ Write IPs of most blatant offenders to database.

        In get_ip_info_for we collect information from many services
        which may be slow or down, so we run in an extra thread to
        avoid delaying the reset thread.

        """

        with self.dblock:  # make sure we don't reenter if we take too long
            try:
                blocked_sessions = [s for s in self.cache.values ()
                                    if s is not None and s['blocked'] == 2]
                self.get_ipinfo_for (blocked_sessions)

                cursor = cherrypy.engine.pool.connect ().cursor ()
                cursor.execute ('set transaction isolation level serializable')

                for ipsession in blocked_sessions:
                    try:
                        ips = ipsession.ips
                        info = ips.ipinfo

                        expire_days = ipsession['expire_days']
                        if info.is_blackhat ():
                            if info.is_dialup ():
                                # at least 2 days
                                expire_days = max (2, expire_days)
                            else:
                                # at least 7 days
                                expire_days = max (7, expire_days)

                        params = {
                            'ip':          to_cidr (ips.get_ip_to_block ()),
                            'client_ip':   to_cidr (ips.client_ip),
                            'proxy_ip':    to_cidr (ips.proxy_ip),
                            'proxy_type':  ips.proxy_type,
                            'host':        info.hostname,
                            'expires':     (datetime.datetime.now () +
                                            datetime.timedelta (days = expire_days)),
                            'types':       ipsession.get_tags (),
                            'whois':       ' - '.join (filter (None, info.whois.values ())),
                            'user_agents': sorted ([self.ua_decode (ua)
                                                    for ua in ipsession['user_agents']]),
                            'requests':    list (ipsession['requests']),
                            'headers':     [self.ua_decode (h) for h in ipsession['headers']],
                            'count':       1,
                            'hits':        ipsession['hits'],
                            'country':     None,
                            'org':         None,
                            'cidr':        None,
                            'asn':         None,
                            }
                        params.update (info.whois)

                        cursor.execute ('begin')
                        query = """select count, hits from robots.blocks
                                   where ip = %(ip)s and not is_whitelisted"""
                        cursor.execute (query, params)
                        rows = cursor.fetchall ()
                        if len (rows) > 0:
                            count = rows[0][0]
                            hits = rows[0][1]
                            params['count'] = count + 1
                            params['hits']  = hits + ipsession['hits']
                            params['expires'] = (
                                datetime.datetime.now () +
                                datetime.timedelta (
                                    days = expire_days * (2 ** min (count, 10))))
                            query = """update robots.blocks set
                                       count = %(count)s,
                                       expires = %(expires)s,
                                       types = %(types)s,
                                       user_agents = %(user_agents)s,
                                       requests = %(requests)s,
                                       headers = %(headers)s,
                                       hits = %(hits)s
                                       where ip = %(ip)s"""
                        else:
                            query = """insert into robots.blocks
                                       (ip, client_ip, proxy_ip, proxy_type,
                                        host, expires, types, country, org, cidr, asn,
                                        whois, user_agents, requests, headers, hits)
                                       values
                                       (%(ip)s, %(client_ip)s, %(proxy_ip)s, %(proxy_type)s,
                                        %(host)s, %(expires)s, %(types)s,
                                        %(country)s, %(org)s, %(cidr)s, %(asn)s,
                                        %(whois)s, %(user_agents)s, %(requests)s,
                                        %(headers)s, %(hits)s)"""
                        cursor.execute (query, params)
                        cursor.execute ('commit')
                        ipsession['blocked'] = 9
                        # keep on stats page for 5 minutes
                        ipsession.set_expiration (5 * 60)
                        cherrypy.log ("BLOCKED: %s\n" % cursor.mogrify (query, params),
                                      context = 'RATELIMITER', severity = logging.INFO)
                    except DatabaseError as what:
                        cherrypy.log ("SQL Error: %s\n" % what,
                                      context = 'RATELIMITER', severity = logging.ERROR)
                        cherrypy.log ("Query was: %s\n" % cursor.mogrify (query, params),
                                      context = 'RATELIMITER', severity = logging.ERROR)
                        cursor.execute ('rollback')
                    except Exception as what:
                        cherrypy.log ("Error in to_database (): %s" % what,
                                      context = 'RATELIMITER', severity = logging.ERROR)
                        cursor.execute ('rollback')

                cherrypy.log ('Run to_database ()',
                              context = 'RATELIMITER', severity = logging.INFO)

            except Exception as what:
                # make sure background thread never stops
                cherrypy.log ("Error in to_database (): %s" % what,
                              context = 'RATELIMITER', severity = logging.ERROR)

class BlockPage (Page.Page):
    """ Block IP. """

    def index (self, **kwargs): # pylint: disable=R0201
        """ Block IP. """

        if cherrypy.tools.rate_limiter.block (kwargs['ip']):  # pylint: disable=E1101
            return "<p>blocked</p>"
        return "<p>not found</p>"

class UnblockPage (Page.Page):
    """ Unblock IP. """

    def index (self, **kwargs): # pylint: disable=R0201
        """ Unblock IP. """

        if cherrypy.tools.rate_limiter.unblock (kwargs['ip']):  # pylint: disable=E1101
            return "<p>unblocked</p>"
        return "<p>not found</p>"

class TracebackPage (Page.Page):
    """ Produce a traceback in the logs. """

    def index (self, **dummy_kwargs): # pylint: disable=R0201
        """ Traceback. """

        raise Exception

class RateLimiterResetPlugin (cherrypy.process.plugins.Monitor):
    """ Plugin to start the counter reset thread.

    We cannot start any threads before daemonizing,
    so we must start the timer thread by this plugin.

    """

    def __init__ (self, bus):
        # frequency of reset () calls in seconds
        frequency = cherrypy.config.get ('tools.rate_limiter.frequency', 60)
        # pylint: disable=E1101
        super (RateLimiterResetPlugin, self).__init__ (
            bus, cherrypy.tools.rate_limiter.reset, frequency)
        self.name = 'rate_limiter_reset'

class RateLimiterDatabasePlugin (cherrypy.process.plugins.Monitor):
    """ Plugin to start the database writer thread.

    We cannot start any threads before daemonizing,
    so we must start the timer thread by this plugin.

    """

    def __init__ (self, bus):
        # frequency of database calls in seconds
        frequency = cherrypy.config.get ('tools.rate_limiter_database.frequency', 60)
        # pylint: disable=E1101
        super (RateLimiterDatabasePlugin, self).__init__ (
            bus, cherrypy.tools.rate_limiter.to_database, frequency)
        self.name = 'rate_limiter_database'

cherrypy.process.plugins.RateLimiterReset    = RateLimiterResetPlugin
cherrypy.process.plugins.RateLimiterDatabase = RateLimiterDatabasePlugin


import pytest

from helpers.cluster import ClickHouseCluster
from helpers.network import PartitionManager

test_recover_staled_replica_run = 1

cluster = ClickHouseCluster(__file__)

node = cluster.add_instance(
    "node",
    main_configs=["configs/enable_keeper_map.xml"],
    with_zookeeper=True,
    stay_alive=True,
)

@pytest.fixture(scope="module")
def started_cluster():
    try:
        cluster.start()
        yield cluster

    finally:
        cluster.shutdown()

def get_genuine_zk():
    return cluster.get_kazoo_client("zoo1")

def remove_children(client, path):
    children = client.get_children(path)

    for child in children:
        child_path = f"{path}/{child}"
        remove_children(client, child_path)
        client.delete(child_path)

def test_keeper_map_without_zk(started_cluster):
    def assert_keeper_exception_after_partition(query):
        with PartitionManager() as pm:
            pm.drop_instance_zk_connections(node)
            error = node.query_and_get_error(query)
            assert "Coordination::Exception" in error

    assert_keeper_exception_after_partition(
        "CREATE TABLE test_keeper_map_without_zk (key UInt64, value UInt64) ENGINE = KeeperMap('/test_without_zk') PRIMARY KEY(key);"
    )

    node.query(
        "CREATE TABLE test_keeper_map_without_zk (key UInt64, value UInt64) ENGINE = KeeperMap('/test_without_zk') PRIMARY KEY(key);"
    )

    assert_keeper_exception_after_partition(
        "INSERT INTO test_keeper_map_without_zk VALUES (1, 11)"
    )
    node.query("INSERT INTO test_keeper_map_without_zk VALUES (1, 11)")

    assert_keeper_exception_after_partition("SELECT * FROM test_keeper_map_without_zk")
    node.query("SELECT * FROM test_keeper_map_without_zk")

    with PartitionManager() as pm:
        pm.drop_instance_zk_connections(node)
        node.restart_clickhouse(60)
        error = node.query_and_get_error("SELECT * FROM test_keeper_map_without_zk")
        assert "Failed to activate table because of connection issues" in error

    node.query("SELECT * FROM test_keeper_map_without_zk")

    client = get_genuine_zk()
    remove_children(client, "/test_keeper_map/test_without_zk")
    node.restart_clickhouse(60)
    error = node.query_and_get_error("SELECT * FROM test_keeper_map_without_zk")
    assert "Failed to activate table because of invalid metadata in ZooKeeper" in error

    node.query("DETACH TABLE test_keeper_map_without_zk")

    client.stop()


######################################################################
#
# File: b2sdk/transfer/emerge/emerger.py
#
# Copyright 2020 Backblaze Inc. All Rights Reserved.
#
# License https://www.backblaze.com/using_b2_code.html
#
######################################################################

import logging

from b2sdk.utils import B2TraceMetaAbstract

from b2sdk.transfer.emerge.executor import EmergeExecutor
from b2sdk.transfer.emerge.planner.planner import EmergePlanner

logger = logging.getLogger(__name__)

class Emerger(metaclass=B2TraceMetaAbstract):
    """
    Handle complex actions around multi source copy/uploads.

    This class can be used to build advanced copy workflows like incremental upload.

    It creates a emerge plan and pass it to emerge executor - all complex logic
    is actually implemented in :class:`b2sdk.transfer.emerge.planner.planner.EmergePlanner`
    and :class:`b2sdk.transfer.emerge.executor.EmergeExecutor`
    """

    DEFAULT_STREAMING_MAX_QUEUE_SIZE = 100

    def __init__(self, services):
        """
        :param b2sdk.v1.Services services:
        """
        self.services = services
        self.emerge_executor = EmergeExecutor(services)

    def emerge(
        self,
        bucket_id,
        write_intents,
        file_name,
        content_type,
        file_info,
        progress_listener,
        recommended_upload_part_size=None,
        continue_large_file_id=None,
    ):
        """
        Create a new file (object in the cloud, really) from an iterable (list, tuple etc) of write intents.

        :param str bucket_id: a bucket ID
        :param write_intents: write intents to process to create a file
        :type write_intents: List[b2sdk.v1.WriteIntent]
        :param str file_name: the file name of the new B2 file
        :param str,None content_type: the MIME type or ``None`` to determine automatically
        :param dict,None file_info: a file info to store with the file or ``None`` to not store anything
        :param b2sdk.v1.AbstractProgressListener progress_listener: a progress listener object to use

        """
        planner = self.get_emerge_planner(recommended_upload_part_size=recommended_upload_part_size)
        emerge_plan = planner.get_emerge_plan(write_intents)
        return self.emerge_executor.execute_emerge_plan(
            emerge_plan,
            bucket_id,
            file_name,
            content_type,
            file_info,
            progress_listener,
            continue_large_file_id=continue_large_file_id,
        )

    def emerge_stream(
        self,
        bucket_id,
        write_intent_iterator,
        file_name,
        content_type,
        file_info,
        progress_listener,
        recommended_upload_part_size=None,
        continue_large_file_id=None,
        max_queue_size=DEFAULT_STREAMING_MAX_QUEUE_SIZE,
    ):
        """
        Create a new file (object in the cloud, really) from a stream of write intents.

        :param str bucket_id: a bucket ID
        :param write_intents: iterator of :class:`~b2sdk.v1.WriteIntent`
        :param str file_name: the file name of the new B2 file
        :param str,None content_type: the MIME type or ``None`` to determine automatically
        :param dict,None file_info: a file info to store with the file or ``None`` to not store anything
        :param b2sdk.v1.AbstractProgressListener progress_listener: a progress listener object to use

        """
        planner = self.get_emerge_planner(recommended_upload_part_size=recommended_upload_part_size)
        emerge_plan = planner.get_streaming_emerge_plan(write_intent_iterator)
        return self.emerge_executor.execute_emerge_plan(
            emerge_plan,
            bucket_id,
            file_name,
            content_type,
            file_info,
            progress_listener,
            continue_large_file_id=continue_large_file_id,
            max_queue_size=max_queue_size,
        )

    def get_emerge_planner(self, recommended_upload_part_size=None):
        return EmergePlanner.from_account_info(
            self.services.session.account_info,
            recommended_upload_part_size=recommended_upload_part_size,
        )


import pygame
from pygame.sprite import Group

from settings import Settings
from ship import Ship
from alien import Alien

import game_functions as gf

def run_game():
    # 初始化游戏并创建一个屏幕对象
    pygame.init()
    ai_settings = Settings()
    screen = pygame.display.set_mode((ai_settings.screen_width,ai_settings.screen_height))
    pygame.display.set_caption(ai_settings.caption)
    bg_color = ai_settings.bgcolor

    # 创建一艘飞船
    my_ship = Ship(ai_settings,screen)
    # 创建一个储存子弹的编组
    bullets = Group()
    # 创建一个外星人
    alien = Alien(ai_settings,screen)

    # 开始游戏主循环
    while True:

        # 监视键盘与鼠标事件
        gf.check_events(ai_settings,screen,my_ship,bullets)
        # 更新屏幕
        my_ship.update()

        #更新子弹位置
        gf.update_bullets(bullets)
        # 每次循环重绘屏幕以及让最近绘制的屏幕可见
        gf.update_screen(ai_settings,screen,my_ship,alien,bullets)

run_game()


import rubik

"""
invariant: Lpath stores the path to the overlap and Rpath stores the inverse moves to end
"""
def construct_path(overlap, Ldict, Rdict):

    Lpath = []
    Rpath = []
    Lcurr = overlap
    Rcurr = overlap
    # starting with overlap, prepend moves to the left list until the start is found
    while True:
        if Ldict[Lcurr][0] == 's':
            break
        Lpath.insert(0, Ldict[Lcurr][0])
        Lcurr = Ldict[Lcurr][1]
    # starting with overlap, append inverse moves to the right list until the end is found
    while True:
        if Rdict[Rcurr][0] == 'e':
            break
        Rpath.append(rubik.perm_inverse(Rdict[Rcurr][0]))
        Rcurr = Rdict[Rcurr][1]
    # return right list appended to left list for the complete moveset
    return Lpath + Rpath

def shortest_path(start, end):
    # dict:
        # key = current configuration
        # value = (move, parent's configuration)
    Ldict = { start : ("s", None) }
    LQueue = [start]
    Lcount = 1
    Rdict = { end : ("e", None) }
    RQueue = [end]
    Rcount = 1
    overlap = tuple()
    depth = 0
    '''
    invariant: both queues will always store the configurations of the previous depth's permuations

    invariant: if a configuration that exists in a queue, as well as the oposite side's dictionary,
    then there must exist a path from start to finish of length depth
    '''
    '''
    Initialization: LQueue stores the start config and RQueue Stores the end config

    Mainenance: All permuations of the current configuration are stored in their respective dictionaries,
    and checked against the oposite diction to check for overlaps. Unique permuations are stored in their
    respective queues to create new permutations at the next depth if no overlap is found

    Termination: an overlap is found, meaning there exists a path from start to finish of length depth,
    or depth 7 of the BFS is reached meaning no path exists from start to finish.
    '''
    while True:
        # if depth > 7 then the cube has an invalid configuration and there is no solution
        if depth > 7:
            return None
        newLCount = 0
        # find all permuations of all the configurations in the left queue
        for i in range(Lcount):
            curr = LQueue.pop(0)
            # if an overlap is found
            if curr in Rdict:
                overlap = curr
                return construct_path(overlap, Ldict, Rdict)
            else:
                # add all permuations of the current configuration to the queue
                for move in rubik.quarter_twists:
                    perm = rubik.perm_apply(move, curr)
                    # add all the permuations of the current configuration to the Ldict if not already in there
                    if perm not in Ldict:
                        Ldict[perm] = (move, curr)
                        LQueue.append(perm)
                        newLCount+=1
        Lcount = newLCount
        newRCount = 0
        for i in range(Rcount):
            curr = RQueue.pop(0)
            # if an overlap is found
            if curr in Ldict:
                overlap = curr
                return construct_path(overlap, Ldict, Rdict)
            else:
                # add all permuations of the current configuration to the queue
                for move in rubik.quarter_twists:
                    perm = rubik.perm_apply(move, curr)
                    # add all the permuations of the current configuration to the Ldict if not already in there
                    if perm not in Rdict:
                        Rdict[perm] = (move, curr)
                        RQueue.append(perm)
                        newRCount+=1
        Rcount = newRCount
        depth+=1
    # end while
    return None