"""
This is the main file for the code.
1. This file uses code from other .py files to fit and forecast data.
3. A documentation pdf along with readme.txt file is submitted for reference.
"""

#Importing necessary Libraries
import scipy.signal as signal
import time
import math
import numpy as np
import pandas as pd
from scipy.optimize import fmin_tnc
from scipy import stats
from pylab import *
from scipy.optimize import curve_fit
from multiprocessing import Pool, cpu_count
from contextlib import closing
import pickle

#Importing other .py files
import prepare_data
import initial_estimate_of_time_to_peak as init_peak
import outlier_analysis
import identify_number_and_volume_of_peaks as find_all_peaks
import decline_functions as dec_func
import signal_features

def fit_and_forecast(well_data):
	"""
	Description: This function takes in well_data.
	Input: df
	Output: Preprocessed data for oil and gas for all wells.
	"""

	#Saving the well_name of our well_data into a variable.
	well_name = "Well_170245_gas"
	print("Processing",well_name)

	#Converting pandas series to a numpy array for much faster processing without overhead.
	well_data = np.array(well_data)

	#Creating a time array for our well_data where t1 and y1 are time and welldata respectively.
	t1 = np.array(list(range(len(well_data)))) + 1
	y1 = well_data

	#Remove downtime outliers.
	y2 = y1[np.nonzero(y1)]
	t2 = np.array(list(range(len(y2)))) + 1

	#features = signal_features.reduce_signal_to_basic_features(y2)

	#Outlier analysis using data distribution --> Removes both global and local outliers.
	y2,outliers = np.array(outlier_analysis.remove_outliers(y2))

	#Identify initial time to peak
	max_peak = init_peak.indentify_initial_peak(y2)

	if math.isnan(max_peak):
		index_max_peak = 0
	else:
		index_max_peak = list(y2).index(max_peak)

	#Seperate decline curves and remove operational outlier peaks.
	y3 = y2[index_max_peak:]
	peaks = find_all_peaks.find_all_peaks(y3,max_peak,index_max_peak)
	list_of_subdeclines,all_peaks = find_all_peaks.get_seperate_declines(y2,peaks,index_max_peak)
	find_indexes_for_these_values = y2[all_peaks]
	indexes_for_prediction = [list(y1).index(each_value) for each_value in find_indexes_for_these_values]

	#Initialize lists to store models.
	list_of_all_fits_exponential = []
	list_of_all_fits_harmonic = []
	list_of_all_fits_hyperbolic = []
	list_of_all_fits_hyp2exp = []

	#Check if subdeclines are less than length 7 as per R code.
	length_of_first_sub_decline = len(list_of_subdeclines[0])
	length_of_last_sub_decline = len(list_of_subdeclines[-1])

	if length_of_last_sub_decline <= 7:
		list_of_subdeclines[-2] = np.concatenate((list_of_subdeclines[-2],list_of_subdeclines[-1]))
		list_of_subdeclines = list_of_subdeclines[:-1]
		indexes_for_prediction = indexes_for_prediction[:-1]

	if length_of_first_sub_decline <= 7:
		list_of_subdeclines[1] = np.concatenate((list_of_subdeclines[0],list_of_subdeclines[1]))
		list_of_subdeclines = list_of_subdeclines[1:]
		indexes_for_prediction = indexes_for_prediction[1:]

	#Model traininig to fit the data for several decline curves.
	for index,each_decline in enumerate(list_of_subdeclines):
		if index == 0:
			initial_guess_for_q_i = 200
		elif index == 0 and len(list_of_subdeclines)==1:
			initial_guess_for_q_i = np.mean(each_decline[:3])
		else:
			initial_guess_for_q_i = np.mean(each_decline[:3])
		t = np.array(list(range(len(each_decline)))) + 1
		try:
			popt_1, pcov_1 = curve_fit(dec_func.exponential_decline, t,each_decline,p0=(initial_guess_for_q_i,0.01), bounds = ([0,0.0002738226],[initial_guess_for_q_i*2, 1]),method = "trf")
			list_of_all_fits_exponential.append(popt_1)
		except Exception as e:
			list_of_all_fits_exponential.append(["No_Model"])
			print("Error: "+str(e))
		try:
			popt_2, pcov_2 = curve_fit(dec_func.harmonic_decline, t,each_decline,p0=(initial_guess_for_q_i,0.01), bounds = ([0,0.0002738226],[initial_guess_for_q_i*2, 1]),method = "trf")
			list_of_all_fits_harmonic.append(popt_2)
		except Exception as e:
			list_of_all_fits_harmonic.append(["No_Model"])
			print("Error: "+str(e))
		try:
			popt_3, pcov_3 = curve_fit(dec_func.hyperbolic_decline, t,each_decline,p0=(initial_guess_for_q_i,0.01,0.7), bounds = ([0,0.0002738226,0.6],[initial_guess_for_q_i*2, 1, 4]),method = "trf" )
			list_of_all_fits_hyperbolic.append(popt_3)
		except Exception as e:
			list_of_all_fits_hyperbolic.append(["No_Model"])
			print("Error: "+str(e))
		try:
			popt_4, pcov_4 = curve_fit(dec_func.hyp2exp, t,each_decline,p0=(initial_guess_for_q_i,0.01,0.7,0.1), bounds = ([0,0.0002738226,0.6,0],[initial_guess_for_q_i*2, 1, 4,0.5]),method = "trf")
			list_of_all_fits_hyp2exp.append(popt_4)
		except Exception as e:
			list_of_all_fits_hyp2exp.append(["No_Model"])
			print("Error: "+str(e))

	list_of_prediction_declines = np.split(y1,indexes_for_prediction)
	forecast_from = len(list_of_prediction_declines[-1])+ 1
	forecast_until = forecast_from + 1100
	t_forecast = np.array(list(range(forecast_from,forecast_until+1)))

	#Check how the trained model fits the training data
	y_pred_exponential = []
	for index,each_model in enumerate(list_of_all_fits_exponential):
		current_length = len(list_of_prediction_declines[index])
		t = np.array(list(range(current_length))) + 1
		if index != 0 and "No_Model" not in list(each_model):
			y_pred_exponential = y_pred_exponential + list(dec_func.exponential_decline(t,each_model[0],each_model[1]))
		elif index == 0 and len(list_of_subdeclines)==1:
			y_pred_exponential = y_pred_exponential + list(dec_func.exponential_decline(t,each_model[0],each_model[1]))
		else:
			y_pred_exponential = y_pred_exponential + [float('nan')]*current_length

	y_pred_harmonic = []
	for index,each_model in enumerate(list_of_all_fits_harmonic):
		current_length = len(list_of_prediction_declines[index])
		t = np.array(list(range(current_length))) + 1
		if index!=0 and "No_Model" not in list(each_model):
			y_pred_harmonic = y_pred_harmonic + list(dec_func.harmonic_decline(t,each_model[0],each_model[1]))
		elif index == 0 and len(list_of_subdeclines)==1:
			y_pred_harmonic = y_pred_harmonic + list(dec_func.harmonic_decline(t,each_model[0],each_model[1]))
		else:
			y_pred_harmonic = y_pred_harmonic + [float('nan')]*current_length

	y_pred_hyperbolic = []
	for index,each_model in enumerate(list_of_all_fits_hyperbolic):
		current_length = len(list_of_prediction_declines[index])
		t = np.array(list(range(current_length))) + 1
		if index!=0 and "No_Model" not in list(each_model):
			y_pred_hyperbolic = y_pred_hyperbolic + list(dec_func.hyperbolic_decline(t,each_model[0],each_model[1],each_model[2]))
		elif index == 0 and len(list_of_subdeclines)==1:
			y_pred_hyperbolic = y_pred_hyperbolic + list(dec_func.hyperbolic_decline(t,each_model[0],each_model[1],each_model[2]))
		else:
			y_pred_hyperbolic = y_pred_hyperbolic + [float('nan')]*current_length

	y_pred_hyp2exp = []
	for index,each_model in enumerate(list_of_all_fits_hyp2exp):
		current_length = len(list_of_prediction_declines[index])
		t = np.array(list(range(current_length))) + 1
		if index!=0 and "No_Model" not in list(each_model):
			y_pred_hyp2exp = y_pred_hyp2exp + list(dec_func.hyp2exp(t,each_model[0],each_model[1],each_model[2],each_model[3]))
		elif index == 0 and len(list_of_subdeclines)==1:
			y_pred_hyp2exp = y_pred_hyp2exp + list(dec_func.hyp2exp(t,each_model[0],each_model[1],each_model[2],each_model[3]))
		else:
			y_pred_hyp2exp = y_pred_hyp2exp + [float('nan')]*current_length

	#Forecast for three years in the future based on recent decline.
	if "No_Model" not in list(list_of_all_fits_exponential[-1]):
		exponential_forecast = dec_func.exponential_decline(t_forecast,list_of_all_fits_exponential[-1][0],list_of_all_fits_exponential[-1][1])
	else:
		exponential_forecast = "No_Forecast"

	if "No_Model" not in list(list_of_all_fits_harmonic[-1]):
		harmonic_forecast = dec_func.harmonic_decline(t_forecast,list_of_all_fits_harmonic[-1][0],list_of_all_fits_harmonic[-1][1])
	else:
		harmonic_forecast = "No_Forecast"

	if "No_Model" not in list(list_of_all_fits_hyperbolic[-1]):
		hyperbolic_forecast = dec_func.hyperbolic_decline(t_forecast,list_of_all_fits_hyperbolic[-1][0],list_of_all_fits_hyperbolic[-1][1],list_of_all_fits_hyperbolic[-1][2])
	else:
		hyperbolic_forecast = "No_Forecast"

	if "No_Model" not in list(list_of_all_fits_hyp2exp[-1]):
		hyp2exp_forecast = dec_func.hyp2exp(t_forecast,list_of_all_fits_hyp2exp[-1][0],list_of_all_fits_hyp2exp[-1][1],list_of_all_fits_hyp2exp[-1][2],list_of_all_fits_hyp2exp[-1][3])
	else:
		hyp2exp_forecast = "No_Forecast"

	list_save = []
	list_save.append(y1)
	list_save.append(y_pred_exponential)
	list_save.append(y_pred_harmonic)
	list_save.append(y_pred_hyperbolic)
	list_save.append(y_pred_hyp2exp)
	list_save.append(exponential_forecast)
	list_save.append(harmonic_forecast)
	list_save.append(hyperbolic_forecast)
	list_save.append(hyp2exp_forecast)
	list_save.append(len(list_of_subdeclines[-1]))
	list_save.append(outliers)

	return list_save


# -*- coding: utf-8 -*-
# @Time    : 2020/3/6 15:35
# @Author  : Deng Wenxing
# @Email   : dengwenxingae86@163.com
# @File    : call_msg.py
# @Software: PyCharm
# @content : 节点结构信息，关系依赖信息

from collections import OrderedDict

def get_table_index(tablename):
    table_index = {
        'vertex_person': 0,	# 人 40亿空间
        'vertex_imsi': 4000000000,	# IMSI 20亿
        'vertex_imei': 5000000000,	# IMEI 10亿
        'vertex_autolpn': 6000000000,	# 车牌号码' 10亿
        'vertex_vehicle': 7000000000,	# 机车（车架号）' 10亿
        'vertex_case': 8000000000,	# 案件' 10亿
        'vertex_jail': 9000000000,	# 所' 10亿
        'vertex_hukou': 10000000000,	# 户口'  10亿
        'vertex_internetbar': 11000000000,	# 网吧' 10亿
        'vertex_ip': 12000000000,	# IP地址' 10亿
        'vertex_hotel': 13000000000,	# 宾馆' 10亿
        'vertex_airline': 14000000000,	# 航班班次' 10亿
        'vertex_ssaccount': 15000000000,	# 社保账号' 10亿
        'vertex_trainline': 17000000000,	# 火车班次' 10亿
        'vertex_package': 21000000000,	# 快递(快递单号）' 10亿
        'vertex_certificate': 25000000000,	# 证件' 10亿
        'vertex_school': 27000000000,	# 学校' 10亿
        'vertex_college': 28000000000,	# 学院' 10亿
        'vertex_company': 29000000000,	# 公司' 10亿
        'vertex_vbn': 30000000000,	# 宽带账号' 10亿
        'vertex_phonenumber': 31000000000,	# 电话 20亿空间
    }
    return table_index[tablename]

# 所有表字段配置
vertex_table_info = OrderedDict()
vertex_table_info['vertex_person'] = ['zjhm', 'zjlx', 'gj', 'xm', 'ywxm', 'zym', 'xb', 'csrq', 'mz', 'jg',
                                      'whcd', 'hyzk', 'zzmm', 'hkszdxz', 'sjjzxz']
vertex_table_info['vertex_phonenumber'] = ['phone', 'xm','province', 'city']
vertex_table_info['vertex_imsi'] = ['imsi']
vertex_table_info['vertex_imei'] = ['imei']
vertex_table_info['vertex_autolpn'] = ['autolpn','hpzl']
vertex_table_info['vertex_vehicle'] = ['vehicle', 'fdjh', 'cllx', 'csys']
vertex_table_info['vertex_case'] = ['asjbh', 'ajmc', 'asjfskssj', 'asjfsjssj', 'asjly', 'ajlb', 'fxasjsj',
                                    'fxasjdd_dzmc', 'jyaq', 'asjswry_rs', 'asjssry_rs', 'ssjzrmby', 'slsj',
                                    'sldw_gajgjgdm', 'sldw_gajgmc', 'larq']
vertex_table_info['vertex_jail'] = ['unitid', 'unitname', 'unitname_s', 'unitkindid', 'dwdz']
vertex_table_info['vertex_hukou'] = ['hukouid', 'hukou']
vertex_table_info['vertex_internetbar'] = ['siteid', 'regtime', 'serial', 'sitetype', 'businessstate', 'title',
                                           'address', 'sp', 'servercount', 'terminalcount', 'nettype', 'employeecount',
                                           'ownercertnumber']
vertex_table_info['vertex_ip'] = ['ip', 'tablename']
vertex_table_info['vertex_hotel'] = ['lgdm', 'qiyemc', 'yyzz', 'xxdz']
vertex_table_info['vertex_airline'] = ['airlineid', 'hbh', 'hbrq', 'yjqfsj', 'yjddsj', 'sfd', 'mdd']
vertex_table_info['vertex_ssaccount'] = ['sbkh', 'tablename']
vertex_table_info['vertex_trainline'] = ['trianid', 'cc', 'fcrq']
vertex_table_info['vertex_email'] = ['email', 'nick', 'tablename']
vertex_table_info['vertex_package'] = ['packageid', 'unitcode', 'orderid', 'accepttime', 'inputtime', 'realarrivetime',
                                       'acceptoperatorphone', 'sendoperatorphone', 'sendername', 'sendercardid',
                                       'senderpostcode', 'senderphone', 'sendernation', 'sendercountry',
                                       'senderaddress', 'receivername', 'receivercardid', 'receiverpostcode',
                                       'receiverphone', 'receivernation', 'receivercountry', 'receiveraddress',
                                       'amountreceivable', 'currency', 'itemname', 'itemquantity', 'totalweight',
                                       'totalsize', 'insurancevalue', 'itemremark', 'senderlongitude', 'senderlatitude',
                                       'receiverlongitude', 'receiverlatitude', 'tablename']
vertex_table_info['vertex_qq'] = ['account', 'tablename']
vertex_table_info['vertex_wechat'] = ['account', 'tablename']
vertex_table_info['vertex_certificate'] = ['certificateid', 'zjhm', 'zjlx', 'zjname', 'dz']
vertex_table_info['vertex_school'] = ['school_id', 'name', 'history_name','school_address']
vertex_table_info['vertex_company'] = ['company', 'dwmc', 'dwdz', 'zzjgdm', 'tyshdm', 'clrq', 'jyzt', 'jyfw',
                                       'dwlx', 'zczb', 'cyrs', 'hylb', 'dwscdz', 'bgxq', 'bgrq']

# 边字段配置
edge_table_info = OrderedDict()
edge_table_info['edge_person_reserve_airline'] = ['sfzh', 'airlineid', 'hbh', 'hbrq', 'sfd', 'mdd']
edge_table_info['edge_person_spouse_person'] = ['sfzhnan', 'sfzhnv', 'start_time', 'end_time']
edge_table_info['edge_person_checkin_train'] = ['sfzh', 'tgkkdm', 'tgkkmc', 'tgsj', 'tablename']
edge_table_info['edge_certificate_open_phone'] = ['certificateid', 'phone', 'start_time', 'end_time', 'is_xiaohu',
                                                  'tablename']
edge_table_info['edge_person_has_ssaccount'] = ['sfzh', 'sbkh', 'link_time', 'tablename']
edge_table_info['edge_find_common_greatgrand'] = ['sfzh1', 'sfzh2']
edge_table_info['edge_find_spouse'] = ['sfzhnan', 'sfzhnnv']
edge_table_info['edge_person_work_school'] = ['start_person', 'end_school_id', 'start_time', 'end_time', 'jszc',
                                              'cjgzsj', 'gzxz', 'sjkm', 'tablename']
edge_table_info['edge_find_father'] = ['sfzh', 'fqsfzh']
edge_table_info['edge_find_grandparents'] = ['sfzh', 'grandsfzh']
edge_table_info['edge_find_greatgrands'] = ['sfzh', 'greatgrandsfzh']
edge_table_info['edge_find_mother'] = ['sfzh', 'mqsfzh']
edge_table_info['edge_person_attend_college'] = ['start_person', 'end_college_id', 'start_time', 'end_time', 'zymc',
                                                 'xuehao', 'xuezhi', 'tablename']
edge_table_info['edge_person_fatherIs_person'] = ['sfzh', 'fqsfzh', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_case_SuspectIs_person'] = ['asjbh', 'sfzh', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_certificate_link_phone'] = ['certificateid', 'phone', 'start_time', 'end_time', 'num',
                                                  'tablename']
edge_table_info['edge_vehicle_sameguarantee'] = ['start_person', 'end_person']
edge_table_info['edge_person_stay_hotel'] = ['sfzh', 'lgdm', 'start_time', 'end_time', 'num']
edge_table_info['edge_internetbar_use_ip'] = ['siteid', 'ip', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_reserve_trainline'] = ['sfzh', 'trianid', 'cc', 'spsj', 'fcrq', 'zwh', 'fz', 'dz',
                                                    'tablename']
edge_table_info['edge_wechat_msg_wechat'] = ['start_account', 'end_account', 'start_time', 'end_time', 'num',
                                             'tablename']
edge_table_info['edge_person_arrived_airline'] = ['sfzh', 'airlineid', 'hbh', 'cwdj', 'zwh', 'hbrq', 'yjddsj', 'sfd',
                                                  'mdd', 'tablename']
edge_table_info['edge_phone_send_package'] = ['phone', 'packageid', 'send_time', 'tablename']
edge_table_info['edge_person_guardianIs_person'] = ['sfzh', 'jhrsfzh', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_guaranteeis_person'] = ['start_person', 'end_person', 'start_time', 'end_time',
                                                     'tablename']
edge_table_info['edge_find_common_grand'] = ['sfzh1', 'sfzh2']
edge_table_info['edge_vehicle_has_autolpn'] = ['start_vehicle', 'end_autolpn', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_email_use_phone'] = ['email', 'phone', 'start_time', 'end_time', 'num', 'tablename']
edge_table_info['edge_person_hire_house'] = ['start_person', 'end_house_id', 'fwdz', 'start_time', 'end_time', 'htbh',
                                             'xm', 'hx', 'fwlx', 'phone', 'mph', 'tablename']
edge_table_info['edge_person_maintain_vehicle'] = ['start_person', 'end_vehicle', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_link_hukou'] = ['sfzh', 'hukouid', 'start_time', 'end_time', 'isfl', 'yhzgx', 'tablename']
edge_table_info['edge_person_surfing_internetbar'] = ['sfzh', 'siteid', 'start_time', 'end_time', 'num', 'tablename']
edge_table_info['edge_phone_list_phone'] = ['start_phone', 'end_phone', 'start_time', 'end_time', 'relationship_name',
                                            'relationship_email', 'relationship_address', 'tablename']
edge_table_info['edge_person_attend_school'] = ['start_person', 'end_school_id', 'start_time', 'end_time', 'fymc', 'bj',
                                                'zymc', 'xuezhi', 'tablename']
edge_table_info['edge_person_own_autolpn'] = ['start_person', 'end_autolpn', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_smz_phone'] = ['start_person', 'end_phone', 'start_time', 'end_time', 'xm', 'cnt',
                                            'matches', 'match_src', 'time_rn', 'matches_rn', 'cert_src', 'cert_weigth']
edge_table_info['edge_case_reportby_person'] = ['asjbh', 'sfzh', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_motherIs_person'] = ['sfzh', 'mqsfzh', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_case_suspicious_person'] = ['asjbh', 'sfzh', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_stay_jail_qjs'] = ['sfzh', 'jsbh', 'sylx', 'rsyy', 'start_time', 'end_time', 'fjh',
                                                'tablename']
edge_table_info['edge_person_stay_jail_jls'] = ['sfzh', 'jsbh', 'sylx', 'rsyy', 'start_time', 'end_time', 'fjh',
                                                'rsgjmjid', 'csgjmjid', 'tablename']
edge_table_info['edge_person_open_phone'] = ['sfzh', 'phone', 'start_time', 'end_time', 'is_xiaohu', 'tablename']
edge_table_info['edge_qq_msg_qq'] = ['start_account', 'end_account', 'start_time', 'end_time', 'num', 'tablename']
edge_table_info['edge_vehicle_sameintroducer'] = ['start_person', 'end_person']
edge_table_info['edge_qq_use_phone'] = ['account', 'phone', 'start_time', 'end_time', 'num', 'tablename']
edge_table_info['edge_person_work_com'] = ['start_person', 'end_company', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_visit_person'] = ['tfrsfzh', 'sfzh', 'start_time', 'end_time', 'num', 'tablename']
edge_table_info['edge_person_reserve_coachline'] = ['sfzh', 'coachid', 'spsj', 'zwh', 'gpczmc', 'tablename']
edge_table_info['edge_person_principal_com'] = ['start_person', 'end_company', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_introduceris_person'] = ['start_person', 'end_person', 'start_time', 'end_time',
                                                      'tablename']
edge_table_info['edge_wechat_use_phone'] = ['account', 'phone', 'start_time', 'end_time', 'num', 'tablename']
edge_table_info['edge_find_common_mother'] = ['sfzh1', 'sfzh2']
edge_table_info['edge_same_kss'] = ['sfzh1', 'sfzh2', 'data', 'num', 'start_time', 'end_time']
edge_table_info['edge_same_qjs'] = ['sfzh1', 'sfzh2', 'data', 'num', 'start_time', 'end_time']
edge_table_info['edge_person_buy_vehicle'] = ['start_person', 'end_vehicle', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_invest_com'] = ['start_person', 'end_company', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_link_phone'] = ['sfzh', 'phone', 'start_time', 'end_time', 'linknum', 'tablename']
edge_table_info['edge_phonenumber_use_imei'] = ['phone', 'imei', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_school_sub_college'] = ['start_school_id', 'end_college_id', 'tablename']
edge_table_info['edge_person_hire_autolpn'] = ['start_person', 'end_autolpn', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_phonenumber_use_imsi'] = ['phone', 'imsi', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_phone_sendpackage_phone'] = ['start_phone', 'end_phone', 'start_time', 'end_time', 'num',
                                                   'tablename']
edge_table_info['edge_person_checkin_airline'] = ['sfzh', 'airlineid', 'hbh', 'cwdj', 'zwh', 'hbrq', 'zjsj', 'djk',
                                                  'sfd', 'mdd', 'tablename']
edge_table_info['edge_person_stay_jail_kss'] = ['sfzh', 'jsbh', 'sylx', 'sxzm', 'start_time', 'end_time', 'fjh',
                                                'tablename']
edge_table_info['edge_certificate_link_person'] = ['certificateid', 'sfzh', 'start_time', 'end_time', 'num',
                                                   'tablename']
edge_table_info['edge_person_link_drivinglicense'] = ['sfzh', 'drivinglicense', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_has_nchca'] = ['sfzh', 'nchca', 'link_time', 'tablename']
edge_table_info['edge_find_common_father'] = ['sfzh1', 'sfzh2']
edge_table_info['edge_phone_receive_package'] = ['phone', 'packageid', 'receive_time', 'tablename']
edge_table_info['edge_find_common_guardian'] = ['sfzh1', 'sfzh2']
edge_table_info['edge_person_legal_com'] = ['start_person', 'end_company', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_work_com2'] = ['start_person', 'end_company', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_same_kss_house'] = ['sfzh1', 'sfzh2', 'data', 'num', 'start_time', 'end_time']
edge_table_info['edge_same_qjs_house'] = ['sfzh1', 'sfzh2', 'data', 'num', 'start_time', 'end_time']
edge_table_info['edge_email_send_email'] = ['start_email', 'end_email', 'start_time', 'end_time', 'num', 'tablename']
edge_table_info['edge_case_VictimIs_person'] = ['asjbh', 'sfzh', 'start_time', 'end_time', 'tablename']
edge_table_info['edge_person_open_vbn'] = ['sfzh', 'kdzh', 'start_time', 'end_time', 'is_xiaohu', 'khxm', 'azdz',
                                           'tablename']
edge_table_info['edge_house_hiremate'] = ['start_person', 'end_person', 'fwdz', 'start_time', 'end_time']
edge_table_info['edge_company_workmate'] = ['start_person', 'end_person', 'start_time', 'end_time', 'dwmc']
edge_table_info['edge_find_common_hukou'] = ['sfzh1', 'sfzh2', 'hukou']
edge_table_info['edge_groupmsg'] = ['start_phone', 'end_phone', 'start_time', 'end_time', 'message_number']
edge_table_info['edge_groupcall'] = ['start_phone', 'end_phone', 'start_time', 'end_time', 'call_total_duration',
                                     'call_total_times']
edge_table_info['edge_same_hotel_house'] = ['sfzh1', 'sfzh2', 'num']
edge_table_info['edge_same_hotel'] = ['sfzh1', 'sfzh2', 'num']
edge_table_info['edge_schoolmate'] = ['start_person', 'end_person', 'start_time', 'end_time', 'name']
edge_table_info['edge_phone_sameimei'] = ['start_phone', 'end_phone', 'start_time', 'end_time', 'times']
edge_table_info['edge_person_with_coach_travel'] = ['sfzh1', 'sfzh2', 'start_time', 'end_time', 'num']
edge_table_info['edge_person_with_airline_travel'] = ['sfzh1', 'sfzh2', 'start_time', 'end_time', 'num']
edge_table_info['edge_person_with_trainline_travel'] = ['sfzh1', 'sfzh2', 'start_time', 'end_time', 'num']
edge_table_info['edge_person_with_internetbar_surfing'] = ['sfzh1', 'sfzh2', 'start_time', 'end_time', 'num']
edge_table_info['edge_phone_wechat_phone'] = ['start_phone', 'end_phone', 'account', 'start_time', 'end_time']
edge_table_info['edge_phone_qq_phone'] = ['start_phone', 'end_phone', 'account', 'start_time', 'end_time']
edge_table_info['edge_phone_email_phone'] = ['start_phone', 'end_phone', 'account', 'start_time', 'end_time']
edge_table_info['edge_person_with_car'] = ['start_person', 'end_autolpn', 'count', 'fullcount']
edge_table_info['relation_score_res'] = ['sfzh1', 'sfzh2', 'data', 'last_tb_score']

# 所有边关系配置
# key为边表，值依次为，start_id 的依赖表，end_id 的依赖表，start_id 的依赖表，end_id 的依赖表字段
edge_info = {'edge_person_reserve_airline': ['vertex_person_jg', 'zjhm', 'vertex_airline_jg', 'airlineid'],
    'edge_person_spouse_person': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_checkin_train': ['vertex_person_jg', 'zjhm', 'vertex_traincheckin_jg', 'tgkkdm'],
    'edge_certificate_open_phone': ['vertex_certificate_jg', 'certificateid', 'vertex_phonenumber_jg', 'phone'],
    'edge_person_has_ssaccount': ['vertex_person_jg', 'zjhm', 'vertex_ssaccount_jg', 'sbkh'],
    'edge_find_common_greatgrand': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_groupcall': ['vertex_phonenumber_jg', 'phone', 'vertex_phonenumber_jg', 'phone'],
    'edge_find_spouse': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_work_school': ['vertex_person_jg', 'zjhm', 'vertex_school_jg', 'school_id'],
    'edge_find_father': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_find_mother': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_find_grandparents': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_find_greatgrands': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_attend_college': ['vertex_person_jg', 'zjhm', 'vertex_college_jg', 'college_id'],
    'edge_person_fatherIs_person': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_case_SuspectIs_person': ['vertex_case_jg', 'asjbh', 'vertex_person_jg', 'zjhm'],
    'edge_certificate_link_phone': ['vertex_certificate_jg', 'certificateid', 'vertex_phonenumber_jg', 'phone'],
    'edge_vehicle_sameguarantee': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_stay_hotel': ['vertex_person_jg', 'zjhm', 'vertex_hotel_jg', 'lgdm'],
    'edge_internetbar_use_ip': ['vertex_internetbar_jg', 'siteid', 'vertex_ip_jg', 'ip'],
    'edge_person_reserve_trainline': ['vertex_person_jg', 'zjhm', 'vertex_trainline_jg', 'trianid'],
    'edge_wechat_msg_wechat': ['vertex_wechat_jg', 'account', 'vertex_wechat_jg', 'account'],
    'edge_person_arrived_airline': ['vertex_person_jg', 'zjhm', 'vertex_airline_jg', 'airlineid'],
    'edge_phone_send_package': ['vertex_phonenumber_jg', 'phone', 'vertex_package_jg', 'packageid'],
    'edge_person_guardianIs_person': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_guaranteeis_person': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_find_common_grand': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_vehicle_has_autolpn': ['vertex_vehicle_jg', 'vehicle', 'vertex_autolpn_jg', 'autolpn'],
    'edge_email_use_phone': ['vertex_email_jg', 'email', 'vertex_phonenumber_jg', 'phone'],
    'edge_person_hire_house': ['vertex_person_jg', 'zjhm', 'vertex_house_jg', 'house_id'],
    'edge_person_maintain_vehicle': ['vertex_person_jg', 'zjhm', 'vertex_vehicle_jg', 'vehicle'],
    'edge_person_link_hukou': ['vertex_person_jg', 'zjhm', 'vertex_hukou_jg', 'hukouid'],
    'edge_person_surfing_internetbar': ['vertex_person_jg', 'zjhm', 'vertex_internetbar_jg', 'siteid'],
    'edge_phone_list_phone': ['vertex_phonenumber_jg', 'phone', 'vertex_phonenumber_jg', 'phone'],
    'edge_person_attend_school': ['vertex_person_jg', 'zjhm', 'vertex_school_jg', 'school_id'],
    'edge_person_own_autolpn': ['vertex_person_jg', 'zjhm', 'vertex_autolpn_jg', 'autolpn'],
    'edge_person_stay_jail_jls': ['vertex_person_jg', 'zjhm', 'vertex_jail_jg', 'unitid'],
    'edge_person_smz_phone': ['vertex_person_jg', 'zjhm', 'vertex_phonenumber_jg', 'phone'],
    'edge_case_reportby_person': ['vertex_case_jg', 'asjbh', 'vertex_person_jg', 'zjhm'],
    'edge_person_motherIs_person': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_case_suspicious_person': ['vertex_case_jg', 'asjbh', 'vertex_person_jg', 'zjhm'],
    'edge_person_stay_jail_qjs': ['vertex_person_jg', 'zjhm', 'vertex_jail_jg', 'unitid'],
    'edge_person_open_phone': ['vertex_person_jg', 'zjhm', 'vertex_phonenumber_jg', 'phone'],
    'edge_qq_msg_qq': ['vertex_qq_jg', 'account', 'vertex_qq_jg', 'account'],
    'edge_vehicle_sameintroducer': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_qq_use_phone': ['vertex_qq_jg', 'account', 'vertex_phonenumber_jg', 'phone'],
    'edge_person_work_com': ['vertex_person_jg', 'zjhm', 'vertex_company_jg', 'company'],
    'edge_person_visit_person': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_reserve_coachline': ['vertex_person_jg', 'zjhm', 'vertex_coachline_jg', 'coachid'],
    'edge_person_principal_com': ['vertex_person_jg', 'zjhm', 'vertex_company_jg', 'company'],
    'edge_person_introduceris_person': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_same_hotel_house': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_same_hotel': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_groupmsg': ['vertex_phonenumber_jg', 'phone', 'vertex_phonenumber_jg', 'phone'],
    'edge_wechat_use_phone': ['vertex_wechat_jg', 'account', 'vertex_phonenumber_jg', 'phone'],
    'edge_company_workmate': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_find_common_mother': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_same_kss': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_same_qjs': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_buy_vehicle': ['vertex_person_jg', 'zjhm', 'vertex_vehicle_jg', 'vehicle'],
    'edge_person_invest_com': ['vertex_person_jg', 'zjhm', 'vertex_company_jg', 'company'],
    'edge_house_hiremate': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_link_phone': ['vertex_person_jg', 'zjhm', 'vertex_phonenumber_jg', 'phone'],
    'edge_phonenumber_use_imei': ['vertex_phonenumber_jg', 'phone', 'vertex_imei_jg', 'imei'],
    'edge_school_sub_college': ['vertex_school_jg', 'school_id', 'vertex_college_jg', 'college_id'],
    'edge_person_hire_autolpn': ['vertex_person_jg', 'zjhm', 'vertex_autolpn_jg', 'autolpn'],
    'edge_phonenumber_use_imsi': ['vertex_phonenumber_jg', 'phone', 'vertex_imsi_jg', 'imsi'],
    'edge_phone_sendpackage_phone': ['vertex_phonenumber_jg', 'phone', 'vertex_phonenumber_jg', 'phone'],
    'edge_person_checkin_airline': ['vertex_person_jg', 'zjhm', 'vertex_airline_jg', 'airlineid'],
    'edge_person_stay_jail_kss': ['vertex_person_jg', 'zjhm', 'vertex_jail_jg', 'unitid'],
    'edge_certificate_link_person': ['vertex_certificate_jg', 'certificateid', 'vertex_person_jg', 'zjhm'],
    'edge_person_link_drivinglicense': ['vertex_person_jg', 'zjhm', 'vertex_drivinglicense_jg', 'drivinglicense'],
    'edge_person_has_nchca': ['vertex_person_jg', 'zjhm', 'vertex_nchcAccount_jg', 'nchca'],
    'edge_find_common_father': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_phone_receive_package': ['vertex_phonenumber_jg', 'phone', 'vertex_package_jg', 'packageid'],
    'edge_find_common_guardian': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_legal_com': ['vertex_person_jg', 'zjhm', 'vertex_company_jg', 'company'],
    'edge_person_work_com2': ['vertex_person_jg', 'zjhm', 'vertex_company_jg', 'company'],
    'edge_same_kss_house': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_same_qjs_house': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_email_send_email': ['vertex_email_jg', 'email', 'vertex_email_jg', 'email'],
    'edge_case_VictimIs_person': ['vertex_case_jg', 'asjbh', 'vertex_person_jg', 'zjhm'],
    'edge_person_open_vbn': ['vertex_person_jg', 'zjhm', 'vertex_vbn_jg', 'kdzh'],
    'edge_find_common_hukou': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_phone_sameimei': ['vertex_phonenumber_jg', 'phone', 'vertex_phonenumber_jg', 'phone'],
    'edge_schoolmate': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_with_coach_travel': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_with_airline_travel': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_with_trainline_travel': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_person_with_internetbar_surfing': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'],
    'edge_phone_wechat_phone': ['vertex_phonenumber_jg', 'phone', 'vertex_phonenumber_jg', 'phone'],
    'edge_phone_qq_phone': ['vertex_phonenumber_jg', 'phone', 'vertex_phonenumber_jg', 'phone'],
    'edge_phone_email_phone': ['vertex_phonenumber_jg', 'phone', 'vertex_phonenumber_jg', 'phone'],
    'edge_person_with_car': ['vertex_person_jg', 'zjhm', 'vertex_autolpn_jg', 'autolpn'],
    'relation_score_res': ['vertex_person_jg', 'zjhm', 'vertex_person_jg', 'zjhm'], }


from typing import List, Optional

from sqlalchemy.orm import Query

from app.lib import db
from app.models import Grade
from app.schemas import GradeCreate

def create(grade: GradeCreate) -> Grade:
    session = db.get_session()
    db_grade = Grade(**grade.dict())
    session.add(db_grade)
    session.commit()
    session.refresh(db_grade)
    return db_grade

def get(sort: str, student_id: Optional[int]) -> List[Grade]:
    session = db.get_session()
    query = session.query(Grade)
    query = _filter(query, student_id)
    return _sort(query, sort)

def get_one(id: int) -> Optional[Grade]:
    session = db.get_session()
    return session.query(Grade).filter(Grade.id == id).first()

def update(id: int, grade: GradeCreate) -> Grade:
    session = db.get_session()
    db_grade = session.query(Grade).filter(Grade.id == id).first()
    db_grade.student_id = grade.student_id
    db_grade.test_id = grade.test_id
    db_grade.value = grade.value
    session.commit()
    session.refresh(db_grade)
    return db_grade

def delete(id: int) -> None:
    session = db.get_session()
    db_grade = session.query(Grade).filter(Grade.id == id).first()
    session.delete(db_grade)
    session.commit()

def _filter(query: Query, student_id: Optional[int]) -> Query:
    if student_id:
        query = query.filter(Grade.student_id == student_id)
    return query

def _sort(query: Query, sort: str) -> List[Grade]:
    order_by = sort_params[sort]
    return query.order_by(*(order_by,)).all()

sort_params = {
    "student_id:asc": Grade.student_id.asc(),
    "student_id:desc": Grade.student_id.desc(),
}


from datetime import timedelta

from django.utils import timezone

from config import celery_app
from eaptekahack.constants import DRUG_IS_LOW_QUANTITY_CHOICES, EventChoices
from eaptekahack.controllers.treatment_course_controller import PlannedEventCreator
from eaptekahack.models import EventForReminder, MedicationAvailable, TreatmentCourse

COUNT_OF_EVENTS_FOR_DRUG_IS_RUNNING_OUT = 10
COUNT_OF_EVENTS_FOR_MAKE_ORDER = 3

def send_push(event, user):
    # Здесь должна быть логика по отправке пуша в пуш-сервис.
    # Т.к. фронт не успевает, то заглушка
    pass

@celery_app.task(ignore_result=True)
def generate_event_for_taking_medical_drugs():
    courses = TreatmentCourse.objects.filter(is_enabled_for_generation=True)

    for course in courses:
        event_creator = PlannedEventCreator(course)
        event_creator.create_events_taking_medical_drugs()

@celery_app.task(ignore_result=True)
def generate_event_for_drug_is_running_out():
    med_availables = MedicationAvailable.objects.filter(courses__is_active=True)

    for med_available in med_availables:
        event_creator = PlannedEventCreator(med_available.course)
        # если текущего кол-ва не хватает на 3 процедур приема, то генерим события для пуша на уведомление об окончании кол-ва препаратов
        if med_available.number_of_pills < med_available.course.quantity * COUNT_OF_EVENTS_FOR_MAKE_ORDER:
            event_creator.create_events_drug_is_running_out()
        # если текущего кол-ва не хватает на 10 процедур приема, то генерим события для пуша на дозаказ
        elif med_available.number_of_pills < med_available.course.quantity * COUNT_OF_EVENTS_FOR_DRUG_IS_RUNNING_OUT:
            event_creator.create_events_make_order()

@celery_app.task(ignore_result=True)
def send_push_reminder_taking_medical_drugs():
    # Осталось 15 минут до планового приема задачи
    time_start = timezone.now() + timedelta(minutes=14)
    time_finish = timezone.now() + timedelta(minutes=15)
    events = EventForReminder.objects.filter(
        event_code=EventChoices.TAKING_MEDICAL_DRUGS, planned_datetime__gte=time_start, planned_datetime__lt=time_finish
    )
    for event in events:
        send_push(event=event, user=event.course.user)

    # плановый прием лекарств!
    time_start = timezone.now() - timedelta(minutes=1)
    time_finish = timezone.now()
    events = EventForReminder.objects.filter(
        event_code=EventChoices.TAKING_MEDICAL_DRUGS,
        is_success=False,
        planned_datetime__gte=time_start,
        planned_datetime__lt=time_finish,
    )
    for event in events:
        send_push(event=event, user=event.course.user)

@celery_app.task(ignore_result=True)
def send_push_reminder_for_drug_is_running_out():
    # Осталось 15 минут до планового приема задачи
    time_start = timezone.now() + timedelta(minutes=14)
    time_finish = timezone.now() + timedelta(minutes=15)
    events = EventForReminder.objects.filter(
        event_code__in=DRUG_IS_LOW_QUANTITY_CHOICES, planned_datetime__gte=time_start, planned_datetime__lt=time_finish
    )
    for event in events:
        send_push(event=event, user=event.course.user)


#%%
# load package
import warnings
warnings.filterwarnings("ignore")
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import shap
#import shap.explainers.deep.deep_tf
import lime
from lime.lime_image import LimeImageExplainer
from PIL import Image
#import PIL.Image
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
import matplotlib.cm as c_map
#from IPython.display import Image, display
import keras.backend as K
from keras.models import Sequential
import ssl
from skimage.segmentation import mark_boundaries

import argparse
import numpy as np
import pandas as pd
from tensorflow.keras.models import load_model
from tensorflow.keras.metrics import AUC
from tensorflow.keras.applications.xception import Xception, preprocess_input, decode_predictions
from tensorflow.keras.applications.efficientnet import preprocess_input, decode_predictions
from tensorflow.keras.preprocessing import image

#print(np.__version__) 1.24.3

TARGET_COLUMNS = ['Atelectasis', 'Cardiomegaly', 'Consolidation', 'Edema', 'Pleural Effusion']
NUM_CLASSES = len(TARGET_COLUMNS)
AUTOTUNE = tf.data.AUTOTUNE
IMAGE_SIZE = [384, 384]
BATCH_SIZE = 1

# return : dataframe with columns [Path, label_1, label_2, ... label_n]
def csv_to_df(dataset_dir, dataset_name):
    # TODO : load data into pandas dataframe
    data = pd.read_csv(f"{dataset_dir}/valid.csv").fillna(0.0)

    # TODO : fix image path:
    # ex. CheXpert-v1.0/patient0000...
    # ex.       > /mnt/e/dataset/chexpert/CheXpert-v1.0/patient0000...
    data['Path'] = data['Path'].str.replace(dataset_name, dataset_dir, regex=False)

    # TODO : fill null values
    for col_idx in range(5, len(data.columns)):
        data[data.columns[col_idx]] = data[data.columns[col_idx]].astype(str)
        data[data.columns[col_idx]] = data[data.columns[col_idx]].str.replace("-1.0", "0.0").astype(float).fillna(0.0)

    # TODO : column reduction
    target_columns = ['Path'] + TARGET_COLUMNS
    dataframe = data[target_columns].reset_index(drop=True)

    return dataframe

def process_path_validation(image_path, label):
    # Read the image from the path
    image = tf.io.read_file(image_path)
    image = tf.image.decode_jpeg(image, channels=3)

    image = tf.cast(image, tf.float32) / 255.0

    return image, label

def transform(dataset_dir):
    # make train_images
    train_images = []
    for i in range(64541, 64741):
        img_path = f"{dataset_dir}/valid/patient{i}/study1/view1_frontal.jpg"
        img = image.load_img(img_path, target_size=(384, 384))
        x = image.img_to_array(img)
        x = np.expand_dims(x, axis=0)
        x = x / 255.0
        train_images.append(x)

    # Convert the list to a numpy array
    # (number_of_images, height, width, channels). (48, 384, 384, 3).
    train_images = np.concatenate(train_images, axis=0)
    #print(train_images.shape)

    return train_images

def auroc(y_true, y_pred):
    return AUC(curve='ROC')(y_true, y_pred)

# def predict_fn(images):
# 	return session.run(probabilities, feed_dict={processed_images: images})

def get_model_predictions(data):
    model_prediction = model.predict(data)
    print(f"The predicted class is : {decode_predictions(model_prediction, top=1)[0][0][1]}")
    return decode_predictions(model_prediction, top=1)[0][0][1]

def generate_prediction_sample(exp, exp_class, weight = 0.1, show_positive = False, hide_background = False):
    '''
    Method to display and highlight super-pixels used by the black-box model to make predictions
    '''
    image, mask = exp.get_image_and_mask(exp_class,
                                         positive_only=show_positive,
                                         num_features=6,
                                         hide_rest=hide_background,
                                         min_weight=weight
                                        )
    plt.imshow(mark_boundaries(image, mask))
    plt.savefig('lime_pred.png')
    #plt.axis('off')
    #plt.show()

def map2layer(x, layer):
    feed_dict = dict(zip([model.layers[0].input], [preprocess_input(x.copy())]))
    return K.get_session().run(model.layers[layer].input, feed_dict)

# lr scheduler
class CosineDecayWithWarmup(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, learning_rate, total_steps, warmup_steps=0):
        super(CosineDecayWithWarmup, self).__init__()
        assert warmup_steps < total_steps
        self.learning_rate = learning_rate
        self.total_steps = total_steps
        self.warmup_steps = warmup_steps
        self.cosine_decay = tf.keras.experimental.CosineDecay(
            learning_rate, total_steps - warmup_steps)

    @tf.function
    def __call__(self, step):
        if step < self.warmup_steps:
            return (self.learning_rate /
                    tf.cast(self.warmup_steps, tf.float32) *
                    tf.cast((step + 1), tf.float32))
        return self.cosine_decay(step - self.warmup_steps)

    def get_config(self):
        return {"learning_rate": np.array(self.learning_rate),
                "total_steps": np.array(self.total_steps),
                "warmup_steps": np.array(self.warmup_steps)}

class AUROC(AUC):
    def __init__(self, name='auroc', **kwargs):
        super().__init__(name=name, **kwargs)

class AUROC_convert(AUC):
    def __init__(self, name='auroc', **kwargs):
        super().__init__(name=name, **kwargs)

if __name__ == "__main__":
    # Input Parsing
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_dir', type=str, default='/home/n0/a2i006/xai/dataset', help='dataset directory')
    parser.add_argument('--dataset_name', type=str, default='CheXpert-v1.0-small', help='dataset name')
    parser.add_argument('--model_dir', type=str, default='/home/n0/a2i006/xai/models/EfficientNetV2B0-best_model.h5', help='model directory')
    parser.add_argument('--batch_size', type=int, default=1, help='batch size')
    args = parser.parse_args()

    # GPU Setting
    physical_devices = tf.config.list_physical_devices('GPU')
    print('Num_GPUs:{}, List:{}'.format(len(physical_devices), physical_devices))

    # load df from csv
    df = csv_to_df(args.dataset_dir, args.dataset_name)

    # Filter filenames ending with "_frontal.jpg"
    valid_df = df[df['Path'].str.endswith('_frontal.jpg')]
    #print(valid_df)

    # Dataset for validation
    list_ds_valid = tf.data.Dataset.from_tensor_slices((valid_df['Path'].values, valid_df.iloc[:, 1:].values))
    #print(list_ds_valid)
    ds_valid = list_ds_valid.map(process_path_validation, num_parallel_calls=AUTOTUNE)
    ds_valid = ds_valid.batch(args.batch_size)
    ds_valid = ds_valid.prefetch(AUTOTUNE)

    # load model
    #model = load_model(args.model_dir, custom_objects={'CosineDecayWithWarmup': CosineDecayWithWarmup})
    model = load_model(args.model_dir, custom_objects={'CosineDecayWithWarmup': CosineDecayWithWarmup, 'AUROC': AUROC})

    # Perform inference on the validation dataset
    predictions = model.predict(ds_valid)

    # Convert the predictions to labels
    predicted_labels = (predictions > 0.5).astype(int)

    # Get the true labels from the validation dataset
    true_labels = np.array([label.numpy() for _, label in ds_valid])
    true_labels = np.reshape(true_labels, (-1, NUM_CLASSES))  # Reshape to 2D array

    #image = Image.open('/home/n0/a2i006/xai/dataset/valid/patient64541/study1/view1_frontal.jpg')
    dataset_dir= '/home/n0/a2i006/xai/dataset'
    images = transform(dataset_dir)

    #plt.imshow(images[0])
    #pred_orig = get_model_predictions(image)

    explainer = LimeImageExplainer()
    explainer_l = explainer.explain_instance(images[0], model.predict, hide_color=0, top_labels=2, num_samples=1000)
    #temp, mask = explanation.get_image_and_mask(240, positive_only=True, num_features=5, hide_rest=True)
    #plt.imshow(mark_boundaries(temp / 2 + 0.5, mask))
    plt.imshow(explainer_l.segments)
    plt.savefig('lime.png')

    #generate_prediction_sample(exp, exp.top_labels[0], show_positive = True, hide_background = True)
    generate_prediction_sample(explainer_l, explainer_l.top_labels[0], show_positive = False, hide_background = False)

    # Initialize a list to store dataset elements
    elements = []

    # Iterate over the dataset and collect elements
    for element in ds_valid:
        resized_image = tf.image.resize(element[0].numpy(), IMAGE_SIZE)
        print(element[1])
        elements.append(resized_image.numpy())

    # Convert the elements list to a NumPy array
    elements = np.array(elements)
    # Reshape the elements array
    X = np.squeeze(elements, axis=1)

    # define a masker that is used to mask out partitions of the input image, this one uses a blurred background
    masker = shap.maskers.Image("inpaint_telea", X[0].shape)

    def f(X):
        tmp = X.copy()
        tmp = tmp / 255.0
        return model(tmp)

    #explainer_s = shap.DeepExplainer(model, map2layer(preprocess_input(X.copy()), 7))
    #shap_values,indexes = e.shap_values(map2layer(to_explain, 7), ranked_outputs=2)

    # select backgroud for shap
    #background = images[np.random.choice(images.shape[0], 1, replace=False)]
    # DeepExplainer to explain predictions of the model
    explainer_s = shap.Explainer(f, masker, output_names=TARGET_COLUMNS)
    #explainer_sd = shap.DeepExplainer(f, masker)
    # compute shap values
    #shap_values = explainer_s.shap_values(images[0])
    shap_values = explainer_s(X[1:3], max_evals = 100, batch_size=BATCH_SIZE, outputs=shap.Explanation.argsort.flip[:1])

    #print(shap_values)
    shap.image_plot(shap_values)

    # multi-label


import sys
import os
import pandas as pd
import warnings
import numpy as np
from datetime import date
from sklearn.linear_model import LassoCV
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error

sys.path.append(os.path.abspath("/projects/genomic-ml/da2343/ml_project_1/shared"))
from model_header import *
from constants import *

warnings.filterwarnings("ignore")
np.set_printoptions(threshold=np.inf)

params_df = pd.read_csv("params.csv")

if len(sys.argv) == 2:
    prog_name, task_str = sys.argv
    param_row = int(task_str)
else:
    print("len(sys.argv)=%d so trying first param" % len(sys.argv))
    param_row = 0

param_dict = dict(params_df.iloc[param_row, :])
data_set_name = param_dict["Dataset"]
n_sub_samples = param_dict["# of Total Samples"]
index_of_pred_col = param_dict["Index of Prediction Col"]

dataset_path = dataset_dict[data_set_name]
n_splits = 3
# Import the csv file of the dataset
dataset_pd = pd.read_csv(dataset_path, header=0)

threshold_param_dict = [
    {"threshold": [threshold]}
    for threshold in np.concatenate(
        (np.linspace(0, 0.4, 5), np.linspace(0.41, 0.6, 21), np.arange(0.7, 1.01, 0.1))
    )
]
learner_dict = {
    "Featureless": Featureless(),
    "Spearman": SpearmanRankRegressor(),
    "Pearson": MyPearsonRegressor(),
    "LASSO": LassoCV(random_state=1),
    "GGM": GaussianGraphicalModel(),
}

test_err_list = []

n_sections = int(np.floor(dataset_pd.shape[0] / n_sub_samples))
shuffled_df = dataset_pd.sample(frac=1, random_state=1)
total_samples = n_sub_samples * n_sections
shuffled_df_updated = shuffled_df.iloc[:total_samples, :]
shuffled_arr = np.split(shuffled_df_updated, n_sections)

for ss_index, sub_section in enumerate(shuffled_arr):
    # drop only one column per every iteration to form the input matrix
    # make the column you removed the output
    output_vec = sub_section.iloc[:, index_of_pred_col].to_numpy().ravel()
    input_mat = sub_section.drop(
        sub_section.columns[index_of_pred_col], axis=1
    ).to_numpy()

    k_fold = KFold(n_splits=n_splits, shuffle=True, random_state=1)
    for fold_id, indices in enumerate(k_fold.split(input_mat)):
        index_dict = dict(zip(["train", "test"], indices))
        set_data_dict = {}
        for set_name, index_vec in index_dict.items():
            set_data_dict[set_name] = {
                "X": input_mat[index_vec],
                "y": output_vec[index_vec],
            }
        # Loop through the learners
        # Fit the learner to the training data
        # Predict the test data
        # Calculate the test error
        for learner_name, learner in learner_dict.items():
            learner.fit(**set_data_dict["train"])
            pred_y = learner.predict(set_data_dict["test"]["X"])
            actual_y = set_data_dict["test"]["y"]
            mse = mean_squared_error(actual_y, pred_y)
            # r2_coef = r2_score(actual_y, pred_y)

            test_err_list.append(
                pd.DataFrame(
                    {
                        "Mean Squared Error": mse,
                        # "Root Mean Squared Error": np.sqrt(mse),
                        # "R Squared": pearsonr(actual_y, pred_y)[0] ** 2,
                        # "R2 Score": r2_coef,
                        "FoldID": fold_id,
                        "# of Total Samples": n_sub_samples,
                        "Index of Subsample": ss_index,
                        "Dataset": data_set_name,
                        "Index of Predicted Column": index_of_pred_col,
                        "Algorithm": learner_name,
                    },
                    index=[0],
                )
            )
main_test_err_df = pd.concat(test_err_list)

# Save dataframe as a csv to output directory
out_file = f"results/{param_row}.csv"
main_test_err_df.to_csv(out_file, encoding="utf-8", index=False)
print("Done!!")


# 核心思想：其实这题就是找多叉树的中心点，根据无环图的定义，中心点最多不会超过2（3个中心点就会形成环）
# 参考 https://leetcode.com/problems/minimum-height-trees/discuss/76132/Iterative-remove-leaves-Python-solution # 的代码
from collections import defaultdict
class Solution:
    def findMinHeightTrees(self, n: int, edges: List[List[int]]) -> List[int]:
        dic = defaultdict(set)
        for u, v in edges:
            dic[u].add(v)
            dic[v].add(u)
        nodes = set(range(n))
        while len(nodes) > 2:
            leaves = set(i for i in nodes if len(dic[i]) == 1)
            for node in leaves:
                for pair_node in dic[node]: # 因为dic的values是集合，集合无法直接取出元素，所以不能写成dic[dic[node]].discard(node)
                    dic[pair_node].discard(node)
            nodes -= leaves
        return list(nodes)

# 直接BFS会超时
# from collections import defaultdict
# class Solution:
#     def findMinHeightTrees(self, n: int, edges: List[List[int]]) -> List[int]:
#         ret, m_dep = [], float('inf')
#         dic = defaultdict(set)
#         for u, v in edges:
#             dic[u].add(v)
#             dic[v].add(u)
#         def bfs(root):
#             check = [1] * n
#             check[root] = 0
#             dep = 1
#             layer = [root]
#             next_layer = []
#             while layer:
#                 t = layer.pop(0)
#                 for it in dic[t]:
#                     if check[it]:
#                         next_layer.append(it)
#                         check[it] = 0
#                 if not layer and next_layer:
#                     layer = next_layer
#                     next_layer = []
#                     dep += 1
#             return dep
#         for i in range(n):
#             dep = bfs(i)
#             if dep < m_dep:
#                 m_dep = dep
#                 ret = [i]
#             elif dep == m_dep:
#                 ret.append(i)
#         return ret


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_

class RewardModel(nn.Module):
    def __init__(self, hidden_state_size, mlp_hidden_size, embed_size, vocab_size):
        super(RewardModel, self).__init__()
        # Size of hidden state of generator model.
        self.hidden_state_size = hidden_state_size
        self.mlp_hidden_size = mlp_hidden_size
        self.vocab_size = vocab_size

        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.fc_i = nn.Linear(hidden_state_size + embed_size, mlp_hidden_size)
        self.fc_h = nn.Linear(mlp_hidden_size, mlp_hidden_size)
        self.fc_o = nn.Linear(mlp_hidden_size, 1)

    def forward(self, x, a):
        """
        x : (batch_size, hidden_state_size) hidden state representing previously generated words.
        a : (batch_size,) words submitted at current timestep.
        """

        a_embed = self.embedding(a)
        z = torch.cat((x, a_embed), dim=len(x.shape) - 1)

        # TODO: Potentially change number of layers.
        z = F.relu(self.fc_i(z))
        z = F.relu(self.fc_h(z))
        output = self.fc_o(z)

        return output

class Rewarder:
    def __init__(
        self,
        seq_length,
        vocab_size,
        hidden_state_size,
        embed_dim,
        mlp_hidden_size,
        learning_rate,
        clip_max_norm,
        momentum,
        baseline=False,
    ):

        self.seq_length = seq_length
        self.vocab_size = vocab_size
        self.hidden_state_size = hidden_state_size  # hidden state of generator
        self.embed_dim = embed_dim  # action embedding
        self.mlp_hidden_size = mlp_hidden_size  # hidden layers of reward model
        self.learning_rate = learning_rate
        self.clip_max_norm = clip_max_norm
        self.momentum = momentum
        self.baseline = baseline

        self.model = RewardModel(
            hidden_state_size, mlp_hidden_size, embed_dim, vocab_size
        ).cuda()
        self.optimizer = torch.optim.SGD(
            self.model.parameters(),
            self.learning_rate,
            self.momentum,
            weight_decay=1e-5,
        )

    def restore_model(self, checkpoint_file):
        checkpoint = torch.load(checkpoint_file)
        self.model.load_state_dict(checkpoint)
        self.optimizer = torch.optim.SGD(
            self.model.parameters(), self.learning_rate, self.momentum
        )

    def compute_rewards_to_go(self, trajectories, generator, roll_num=4):
        """
        Compute reward for each partitial trajectory t:seq_length
        for all t 1:seq_length

        Parameters:
            trajectories : (batch_size, seq_len) type: int
            generator : Generator
            roll_num : int number of Monte Carlo estimates to compute roll-out estimate of value.

        Returns:
            rewards_to_go : (num_batches, batch_size, seq_length)
        """

        init_shape = trajectories.shape
        trajectories = trajectories.reshape((-1, self.seq_length))
        trajectories_real = generator.get_hidden_state(trajectories)
        batch_size = init_shape[0]

        rewards_to_go = torch.zeros(batch_size, self.seq_length).cuda()

        for t in range(self.seq_length):
            # Compute reward to go for each trajectory at s_t
            #   using MCMC sampling

            current_traj = trajectories[:, 0 : (t + 1)]
            expected_reward = 0
            for k in range(roll_num):
                rollouts, rollout_hidden_states = generator.generate(
                    batch_size,
                    1,
                    current_traj,
                    inc_hidden_state=True,
                    inc_probs=False,
                    decode=False,
                )

                rewards = self.model(
                    rollout_hidden_states.view(-1, self.hidden_state_size),
                    rollouts[0].view(-1),
                )
                result = rewards.view(batch_size, self.seq_length).sum(axis=1)

                expected_reward += result

            expected_reward /= roll_num

            # Compute the reward actually incurred by the trajectory.
            # (batch_size, remaining_seq_len)
            if self.baseline:
                remaining_trajectory = trajectories[:, t:].cuda()
                remaining_trajectory_hidden = trajectories_real[:, t:].cuda()
                realized_reward = self.model(
                    remaining_trajectory_hidden.view(-1, self.hidden_state_size),
                    remaining_trajectory.view(-1),
                )
                realized_reward = rewards.view(batch_size, self.seq_length).sum(axis=1)

                rewards_to_go[:, t] = realized_reward - expected_reward
            else:
                rewards_to_go[:, t] = expected_reward

        return rewards_to_go

    def train_step(self, real_batch, generator, generator_batch_size):
        """
        Perform one step of stochastic gradient descent for the Reward objective,
        as per equation (6) in https://arxiv.org/pdf/1804.11258.pdf.

        Parameters:
            real_batch : (batch_size, seq_len) Data from training set.
            generator : Generator
            generator_batch_size : int
        """

        # Compute reward for real sequences
        # Obtain batch of trajectories from real data. Each token is an embedding of the
        # state (context) at that index, embedded by GPT2 pre-trained layer.
        # Also store the actions taken at each timestep.
        real_batch_size = real_batch.shape[0]
        hidden_states_real = generator.get_hidden_state(real_batch)

        x_real = hidden_states_real.view(-1, self.hidden_state_size).cuda()
        a_real = real_batch.view(-1).cuda()

        # Compute reward for each state, action pair in the trajectories.
        reward_real = self.model(x_real, a_real).sum() / real_batch_size

        actions_gen, hidden_states_gen, log_probs = generator.generate(
            generator_batch_size,
            1,
            None,
            inc_hidden_state=True,
            inc_probs=True,
            decode=False,  # TODO: not sure about this
        )

        actions_gen = actions_gen[0]

        # Get rewards as a function of example, by summing over the sequence but not the batch size.
        reward_by_example = self.model(
            hidden_states_gen.view(-1, self.hidden_state_size), actions_gen.view(-1),
        )
        reward_by_example = reward_by_example.view(
            generator_batch_size, self.seq_length
        ).sum(axis=1)

        # (batch_size, seq_len, 1)
        indices = actions_gen.unsqueeze(-1)
        log_q = torch.gather(log_probs, 2, indices).sum(axis=1).cpu().data.numpy()
        log_w = reward_by_example.cpu().data.numpy() - log_q

        log_w -= log_w.max()
        w = torch.exp(torch.from_numpy(log_w)).cuda()
        reward_gen = torch.sum((w * reward_by_example) / w.sum())

        loss = -(reward_real - reward_gen)
        self.optimizer.zero_grad()
        loss.backward()
        clip_grad_norm_(self.model.parameters(), self.clip_max_norm)
        self.optimizer.step()

        return loss.cpu().data.numpy()


import datetime
import re

from ... import processors
from ... import sql, exc
from ... import types as sqltypes, schema as sa_schema
from ... import util
from ...engine import default, reflection
from ...sql import compiler

from ...types import (BLOB, BOOLEAN, CHAR, DATE, DECIMAL, FLOAT,
                      INTEGER, REAL, NUMERIC, SMALLINT, TEXT,
                      TIMESTAMP, VARCHAR)

class _DateTimeMixin(object):
    _reg = None
    _storage_format = None

    def __init__(self, storage_format=None, regexp=None, **kw):
        super(_DateTimeMixin, self).__init__(**kw)
        if regexp is not None:
            self._reg = re.compile(regexp)
        if storage_format is not None:
            self._storage_format = storage_format

    def adapt(self, cls, **kw):
        if issubclass(cls, _DateTimeMixin):
            if self._storage_format:
                kw["storage_format"] = self._storage_format
            if self._reg:
                kw["regexp"] = self._reg
        return super(_DateTimeMixin, self).adapt(cls, **kw)

    def literal_processor(self, dialect):
        bp = self.bind_processor(dialect)

        def process(value):
            return "'%s'" % bp(value)
        return process

class DATETIME(_DateTimeMixin, sqltypes.DateTime):
    """Represent a Python datetime object in SQLite using a string.

    The default string storage format is::

        "%(year)04d-%(month)02d-%(day)02d %(hour)02d:%(min)02d:\
%(second)02d.%(microsecond)06d"

    e.g.::

        2011-03-15 12:05:57.10558

    The storage format can be customized to some degree using the
    ``storage_format`` and ``regexp`` parameters, such as::

        import re
        from sqlalchemy.dialects.sqlite import DATETIME

        dt = DATETIME(
            storage_format="%(year)04d/%(month)02d/%(day)02d %(hour)02d:\
%(min)02d:%(second)02d",
            regexp=r"(\d+)/(\d+)/(\d+) (\d+)-(\d+)-(\d+)"
        )

    :param storage_format: format string which will be applied to the dict
     with keys year, month, day, hour, minute, second, and microsecond.

    :param regexp: regular expression which will be applied to incoming result
     rows. If the regexp contains named groups, the resulting match dict is
     applied to the Python datetime() constructor as keyword arguments.
     Otherwise, if positional groups are used, the datetime() constructor
     is called with positional arguments via
     ``*map(int, match_obj.groups(0))``.
    """

    _storage_format = (
        "%(year)04d-%(month)02d-%(day)02d "
        "%(hour)02d:%(minute)02d:%(second)02d.%(microsecond)06d"
    )

    def __init__(self, *args, **kwargs):
        truncate_microseconds = kwargs.pop('truncate_microseconds', False)
        super(DATETIME, self).__init__(*args, **kwargs)
        if truncate_microseconds:
            assert 'storage_format' not in kwargs, "You can specify only "\
                "one of truncate_microseconds or storage_format."
            assert 'regexp' not in kwargs, "You can specify only one of "\
                "truncate_microseconds or regexp."
            self._storage_format = (
                "%(year)04d-%(month)02d-%(day)02d "
                "%(hour)02d:%(minute)02d:%(second)02d"
            )

    def bind_processor(self, dialect):
        datetime_datetime = datetime.datetime
        datetime_date = datetime.date
        format = self._storage_format

        def process(value):
            if value is None:
                return None
            elif isinstance(value, datetime_datetime):
                return format % {
                    'year': value.year,
                    'month': value.month,
                    'day': value.day,
                    'hour': value.hour,
                    'minute': value.minute,
                    'second': value.second,
                    'microsecond': value.microsecond,
                }
            elif isinstance(value, datetime_date):
                return format % {
                    'year': value.year,
                    'month': value.month,
                    'day': value.day,
                    'hour': 0,
                    'minute': 0,
                    'second': 0,
                    'microsecond': 0,
                }
            else:
                raise TypeError("SQLite DateTime type only accepts Python "
                                "datetime and date objects as input.")
        return process

    def result_processor(self, dialect, coltype):
        if self._reg:
            return processors.str_to_datetime_processor_factory(
                self._reg, datetime.datetime)
        else:
            return processors.str_to_datetime

class DATE(_DateTimeMixin, sqltypes.Date):
    """Represent a Python date object in SQLite using a string.

    The default string storage format is::

        "%(year)04d-%(month)02d-%(day)02d"

    e.g.::

        2011-03-15

    The storage format can be customized to some degree using the
    ``storage_format`` and ``regexp`` parameters, such as::

        import re
        from sqlalchemy.dialects.sqlite import DATE

        d = DATE(
                storage_format="%(month)02d/%(day)02d/%(year)04d",
                regexp=re.compile("(?P<month>\d+)/(?P<day>\d+)/(?P<year>\d+)")
            )

    :param storage_format: format string which will be applied to the
     dict with keys year, month, and day.

    :param regexp: regular expression which will be applied to
     incoming result rows. If the regexp contains named groups, the
     resulting match dict is applied to the Python date() constructor
     as keyword arguments. Otherwise, if positional groups are used, the
     date() constructor is called with positional arguments via
     ``*map(int, match_obj.groups(0))``.
    """

    _storage_format = "%(year)04d-%(month)02d-%(day)02d"

    def bind_processor(self, dialect):
        datetime_date = datetime.date
        format = self._storage_format

        def process(value):
            if value is None:
                return None
            elif isinstance(value, datetime_date):
                return format % {
                    'year': value.year,
                    'month': value.month,
                    'day': value.day,
                }
            else:
                raise TypeError("SQLite Date type only accepts Python "
                                "date objects as input.")
        return process

    def result_processor(self, dialect, coltype):
        if self._reg:
            return processors.str_to_datetime_processor_factory(
                self._reg, datetime.date)
        else:
            return processors.str_to_date

class TIME(_DateTimeMixin, sqltypes.Time):
    """Represent a Python time object in SQLite using a string.

    The default string storage format is::

        "%(hour)02d:%(minute)02d:%(second)02d.%(microsecond)06d"

    e.g.::

        12:05:57.10558

    The storage format can be customized to some degree using the
    ``storage_format`` and ``regexp`` parameters, such as::

        import re
        from sqlalchemy.dialects.sqlite import TIME

        t = TIME(
            storage_format="%(hour)02d-%(minute)02d-%(second)02d-\
%(microsecond)06d",
            regexp=re.compile("(\d+)-(\d+)-(\d+)-(?:-(\d+))?")
        )

    :param storage_format: format string which will be applied to the dict
     with keys hour, minute, second, and microsecond.

    :param regexp: regular expression which will be applied to incoming result
     rows. If the regexp contains named groups, the resulting match dict is
     applied to the Python time() constructor as keyword arguments. Otherwise,
     if positional groups are used, the time() constructor is called with
     positional arguments via ``*map(int, match_obj.groups(0))``.
    """

    _storage_format = "%(hour)02d:%(minute)02d:%(second)02d.%(microsecond)06d"

    def __init__(self, *args, **kwargs):
        truncate_microseconds = kwargs.pop('truncate_microseconds', False)
        super(TIME, self).__init__(*args, **kwargs)
        if truncate_microseconds:
            assert 'storage_format' not in kwargs, "You can specify only "\
                "one of truncate_microseconds or storage_format."
            assert 'regexp' not in kwargs, "You can specify only one of "\
                "truncate_microseconds or regexp."
            self._storage_format = "%(hour)02d:%(minute)02d:%(second)02d"

    def bind_processor(self, dialect):
        datetime_time = datetime.time
        format = self._storage_format

        def process(value):
            if value is None:
                return None
            elif isinstance(value, datetime_time):
                return format % {
                    'hour': value.hour,
                    'minute': value.minute,
                    'second': value.second,
                    'microsecond': value.microsecond,
                }
            else:
                raise TypeError("SQLite Time type only accepts Python "
                                "time objects as input.")
        return process

    def result_processor(self, dialect, coltype):
        if self._reg:
            return processors.str_to_datetime_processor_factory(
                self._reg, datetime.time)
        else:
            return processors.str_to_time

colspecs = {
    sqltypes.Date: DATE,
    sqltypes.DateTime: DATETIME,
    sqltypes.Time: TIME,
}

ischema_names = {
    'BIGINT': sqltypes.BIGINT,
    'BLOB': sqltypes.BLOB,
    'BOOL': sqltypes.BOOLEAN,
    'BOOLEAN': sqltypes.BOOLEAN,
    'CHAR': sqltypes.CHAR,
    'DATE': sqltypes.DATE,
    'DATETIME': sqltypes.DATETIME,
    'DOUBLE': sqltypes.FLOAT,
    'DECIMAL': sqltypes.DECIMAL,
    'FLOAT': sqltypes.FLOAT,
    'INT': sqltypes.INTEGER,
    'INTEGER': sqltypes.INTEGER,
    'NUMERIC': sqltypes.NUMERIC,
    'REAL': sqltypes.REAL,
    'SMALLINT': sqltypes.SMALLINT,
    'TEXT': sqltypes.TEXT,
    'TIME': sqltypes.TIME,
    'TIMESTAMP': sqltypes.TIMESTAMP,
    'VARCHAR': sqltypes.VARCHAR,
    'NVARCHAR': sqltypes.NVARCHAR,
    'NCHAR': sqltypes.NCHAR,
}

class SQLiteCompiler(compiler.SQLCompiler):
    extract_map = util.update_copy(
        compiler.SQLCompiler.extract_map,
        {
            'month': '%m',
            'day': '%d',
            'year': '%Y',
            'second': '%S',
            'hour': '%H',
            'doy': '%j',
            'minute': '%M',
            'epoch': '%s',
            'dow': '%w',
            'week': '%W',
        })

    def visit_now_func(self, fn, **kw):
        return "CURRENT_TIMESTAMP"

    def visit_localtimestamp_func(self, func, **kw):
        return 'DATETIME(CURRENT_TIMESTAMP, "localtime")'

    def visit_true(self, expr, **kw):
        return '1'

    def visit_false(self, expr, **kw):
        return '0'

    def visit_char_length_func(self, fn, **kw):
        return "length%s" % self.function_argspec(fn)

    def visit_cast(self, cast, **kwargs):
        if self.dialect.supports_cast:
            return super(SQLiteCompiler, self).visit_cast(cast, **kwargs)
        else:
            return self.process(cast.clause, **kwargs)

    def visit_extract(self, extract, **kw):
        try:
            return "CAST(STRFTIME('%s', %s) AS INTEGER)" % (
                self.extract_map[extract.field],
                self.process(extract.expr, **kw)
            )
        except KeyError:
            raise exc.CompileError(
                "%s is not a valid extract argument." % extract.field)

    def limit_clause(self, select):
        text = ""
        if select._limit is not None:
            text += "\n LIMIT " + self.process(sql.literal(select._limit))
        if select._offset is not None:
            if select._limit is None:
                text += "\n LIMIT " + self.process(sql.literal(-1))
            text += " OFFSET " + self.process(sql.literal(select._offset))
        else:
            text += " OFFSET " + self.process(sql.literal(0))
        return text

    def for_update_clause(self, select):
        # sqlite has no "FOR UPDATE" AFAICT
        return ''

class SQLiteDDLCompiler(compiler.DDLCompiler):

    def get_column_specification(self, column, **kwargs):
        coltype = self.dialect.type_compiler.process(column.type)
        colspec = self.preparer.format_column(column) + " " + coltype
        default = self.get_column_default_string(column)
        if default is not None:
            colspec += " DEFAULT " + default

        if not column.nullable:
            colspec += " NOT NULL"

        if (column.primary_key and
                column.table.dialect_options['sqlite']['autoincrement'] and
                len(column.table.primary_key.columns) == 1 and
                issubclass(column.type._type_affinity, sqltypes.Integer) and
                not column.foreign_keys):
            colspec += " PRIMARY KEY AUTOINCREMENT"

        return colspec

    def visit_primary_key_constraint(self, constraint):
        # for columns with sqlite_autoincrement=True,
        # the PRIMARY KEY constraint can only be inline
        # with the column itself.
        if len(constraint.columns) == 1:
            c = list(constraint)[0]
            if (c.primary_key and
                    c.table.dialect_options['sqlite']['autoincrement'] and
                    issubclass(c.type._type_affinity, sqltypes.Integer) and
                    not c.foreign_keys):
                return None

        return super(SQLiteDDLCompiler, self).visit_primary_key_constraint(
            constraint)

    def visit_foreign_key_constraint(self, constraint):

        local_table = list(constraint._elements.values())[0].parent.table
        remote_table = list(constraint._elements.values())[0].column.table

        if local_table.schema != remote_table.schema:
            return None
        else:
            return super(
                SQLiteDDLCompiler,
                self).visit_foreign_key_constraint(constraint)

    def define_constraint_remote_table(self, constraint, table, preparer):
        """Format the remote table clause of a CREATE CONSTRAINT clause."""

        return preparer.format_table(table, use_schema=False)

    def visit_create_index(self, create):
        return super(SQLiteDDLCompiler, self).visit_create_index(
            create, include_table_schema=False)

class SQLiteTypeCompiler(compiler.GenericTypeCompiler):
    def visit_large_binary(self, type_):
        return self.visit_BLOB(type_)

class SQLiteIdentifierPreparer(compiler.IdentifierPreparer):
    reserved_words = set([
        'add', 'after', 'all', 'alter', 'analyze', 'and', 'as', 'asc',
        'attach', 'autoincrement', 'before', 'begin', 'between', 'by',
        'cascade', 'case', 'cast', 'check', 'collate', 'column', 'commit',
        'conflict', 'constraint', 'create', 'cross', 'current_date',
        'current_time', 'current_timestamp', 'database', 'default',
        'deferrable', 'deferred', 'delete', 'desc', 'detach', 'distinct',
        'drop', 'each', 'else', 'end', 'escape', 'except', 'exclusive',
        'explain', 'false', 'fail', 'for', 'foreign', 'from', 'full', 'glob',
        'group', 'having', 'if', 'ignore', 'immediate', 'in', 'index',
        'indexed', 'initially', 'inner', 'insert', 'instead', 'intersect',
        'into', 'is', 'isnull', 'join', 'key', 'left', 'like', 'limit',
        'match', 'natural', 'not', 'notnull', 'null', 'of', 'offset', 'on',
        'or', 'order', 'outer', 'plan', 'pragma', 'primary', 'query',
        'raise', 'references', 'reindex', 'rename', 'replace', 'restrict',
        'right', 'rollback', 'row', 'select', 'set', 'table', 'temp',
        'temporary', 'then', 'to', 'transaction', 'trigger', 'true', 'union',
        'unique', 'update', 'using', 'vacuum', 'values', 'view', 'virtual',
        'when', 'where',
    ])

    def format_index(self, index, use_schema=True, name=None):
        """Prepare a quoted index and schema name."""

        if name is None:
            name = index.name
        result = self.quote(name, index.quote)
        if (not self.omit_schema and
                use_schema and
                getattr(index.table, "schema", None)):
            result = self.quote_schema(
                index.table.schema, index.table.quote_schema) + "." + result
        return result

class SQLiteExecutionContext(default.DefaultExecutionContext):
    @util.memoized_property
    def _preserve_raw_colnames(self):
        return self.execution_options.get("sqlite_raw_colnames", False)

    def _translate_colname(self, colname):
        # adjust for dotted column names.  SQLite
        # in the case of UNION may store col names as
        # "tablename.colname", or if using an attached database,
        # "database.tablename.colname", in cursor.description
        if not self._preserve_raw_colnames and "." in colname:
            return colname.split(".")[-1], colname
        else:
            return colname, None

class SQLiteDialect(default.DefaultDialect):
    name = 'sqlite'
    supports_alter = False
    supports_unicode_statements = True
    supports_unicode_binds = True
    supports_default_values = True
    supports_empty_insert = False
    supports_cast = True
    supports_multivalues_insert = True
    supports_right_nested_joins = False

    default_paramstyle = 'qmark'
    execution_ctx_cls = SQLiteExecutionContext
    statement_compiler = SQLiteCompiler
    ddl_compiler = SQLiteDDLCompiler
    type_compiler = SQLiteTypeCompiler
    preparer = SQLiteIdentifierPreparer
    ischema_names = ischema_names
    colspecs = colspecs
    isolation_level = None

    supports_cast = True
    supports_default_values = True

    construct_arguments = [
        (sa_schema.Table, {
            "autoincrement": False
        })
    ]

    _broken_fk_pragma_quotes = False

    def __init__(self, isolation_level=None, native_datetime=False, **kwargs):
        default.DefaultDialect.__init__(self, **kwargs)
        self.isolation_level = isolation_level

        # this flag used by pysqlite dialect, and perhaps others in the
        # future, to indicate the driver is handling date/timestamp
        # conversions (and perhaps datetime/time as well on some hypothetical
        # driver ?)
        self.native_datetime = native_datetime

        if self.dbapi is not None:
            self.supports_default_values = (
                self.dbapi.sqlite_version_info >= (3, 3, 8))
            self.supports_cast = (
                self.dbapi.sqlite_version_info >= (3, 2, 3))
            self.supports_multivalues_insert = (
                # http://www.sqlite.org/releaselog/3_7_11.html
                self.dbapi.sqlite_version_info >= (3, 7, 11))
            # see http://www.sqlalchemy.org/trac/ticket/2568
            # as well as http://www.sqlite.org/src/info/600482d161
            self._broken_fk_pragma_quotes = (
                self.dbapi.sqlite_version_info < (3, 6, 14))

    _isolation_lookup = {
        'READ UNCOMMITTED': 1,
        'SERIALIZABLE': 0,
    }

    def set_isolation_level(self, connection, level):
        try:
            isolation_level = self._isolation_lookup[level.replace('_', ' ')]
        except KeyError:
            raise exc.ArgumentError(
                "Invalid value '%s' for isolation_level. "
                "Valid isolation levels for %s are %s" %
                (level, self.name, ", ".join(self._isolation_lookup))
            )
        cursor = connection.cursor()
        cursor.execute("PRAGMA read_uncommitted = %d" % isolation_level)
        cursor.close()

    def get_isolation_level(self, connection):
        cursor = connection.cursor()
        cursor.execute('PRAGMA read_uncommitted')
        res = cursor.fetchone()
        if res:
            value = res[0]
        else:
            # http://www.sqlite.org/changes.html#version_3_3_3
            # "Optional READ UNCOMMITTED isolation (instead of the
            # default isolation level of SERIALIZABLE) and
            # table level locking when database connections
            # share a common cache.""
            # pre-SQLite 3.3.0 default to 0
            value = 0
        cursor.close()
        if value == 0:
            return "SERIALIZABLE"
        elif value == 1:
            return "READ UNCOMMITTED"
        else:
            assert False, "Unknown isolation level %s" % value

    def on_connect(self):
        if self.isolation_level is not None:
            def connect(conn):
                self.set_isolation_level(conn, self.isolation_level)
            return connect
        else:
            return None

    @reflection.cache
    def get_table_names(self, connection, schema=None, **kw):
        if schema is not None:
            qschema = self.identifier_preparer.quote_identifier(schema)
            master = '%s.sqlite_master' % qschema
            s = ("SELECT name FROM %s "
                 "WHERE type='table' ORDER BY name") % (master,)
            rs = connection.execute(s)
        else:
            try:
                s = ("SELECT name FROM "
                     " (SELECT * FROM sqlite_master UNION ALL "
                     "  SELECT * FROM sqlite_temp_master) "
                     "WHERE type='table' ORDER BY name")
                rs = connection.execute(s)
            except exc.DBAPIError:
                s = ("SELECT name FROM sqlite_master "
                     "WHERE type='table' ORDER BY name")
                rs = connection.execute(s)

        return [row[0] for row in rs]

    def has_table(self, connection, table_name, schema=None):
        quote = self.identifier_preparer.quote_identifier
        if schema is not None:
            pragma = "PRAGMA %s." % quote(schema)
        else:
            pragma = "PRAGMA "
        qtable = quote(table_name)
        statement = "%stable_info(%s)" % (pragma, qtable)
        cursor = _pragma_cursor(connection.execute(statement))
        row = cursor.fetchone()

        # consume remaining rows, to work around
        # http://www.sqlite.org/cvstrac/tktview?tn=1884
        while not cursor.closed and cursor.fetchone() is not None:
            pass

        return row is not None

    @reflection.cache
    def get_view_names(self, connection, schema=None, **kw):
        if schema is not None:
            qschema = self.identifier_preparer.quote_identifier(schema)
            master = '%s.sqlite_master' % qschema
            s = ("SELECT name FROM %s "
                 "WHERE type='view' ORDER BY name") % (master,)
            rs = connection.execute(s)
        else:
            try:
                s = ("SELECT name FROM "
                     " (SELECT * FROM sqlite_master UNION ALL "
                     "  SELECT * FROM sqlite_temp_master) "
                     "WHERE type='view' ORDER BY name")
                rs = connection.execute(s)
            except exc.DBAPIError:
                s = ("SELECT name FROM sqlite_master "
                     "WHERE type='view' ORDER BY name")
                rs = connection.execute(s)

        return [row[0] for row in rs]

    @reflection.cache
    def get_view_definition(self, connection, view_name, schema=None, **kw):
        if schema is not None:
            qschema = self.identifier_preparer.quote_identifier(schema)
            master = '%s.sqlite_master' % qschema
            s = ("SELECT sql FROM %s WHERE name = '%s'"
                 "AND type='view'") % (master, view_name)
            rs = connection.execute(s)
        else:
            try:
                s = ("SELECT sql FROM "
                     " (SELECT * FROM sqlite_master UNION ALL "
                     "  SELECT * FROM sqlite_temp_master) "
                     "WHERE name = '%s' "
                     "AND type='view'") % view_name
                rs = connection.execute(s)
            except exc.DBAPIError:
                s = ("SELECT sql FROM sqlite_master WHERE name = '%s' "
                     "AND type='view'") % view_name
                rs = connection.execute(s)

        result = rs.fetchall()
        if result:
            return result[0].sql

    @reflection.cache
    def get_columns(self, connection, table_name, schema=None, **kw):
        quote = self.identifier_preparer.quote_identifier
        if schema is not None:
            pragma = "PRAGMA %s." % quote(schema)
        else:
            pragma = "PRAGMA "
        qtable = quote(table_name)
        statement = "%stable_info(%s)" % (pragma, qtable)
        c = _pragma_cursor(connection.execute(statement))

        rows = c.fetchall()
        columns = []
        for row in rows:
            (name, type_, nullable, default, primary_key) = (
                row[1], row[2].upper(), not row[3], row[4], row[5])

            columns.append(self._get_column_info(name, type_, nullable,
                                                 default, primary_key))
        return columns

    def _get_column_info(self, name, type_, nullable, default, primary_key):
        coltype = self._resolve_type_affinity(type_)

        if default is not None:
            default = util.text_type(default)

        return {
            'name': name,
            'type': coltype,
            'nullable': nullable,
            'default': default,
            'autoincrement': default is None,
            'primary_key': primary_key,
        }

    def _resolve_type_affinity(self, type_):
        """Return a data type from a reflected column, using affinity tules.

        SQLite's goal for universal compatibility introduces some complexity
        during reflection, as a column's defined type might not actually be a
        type that SQLite understands - or indeed, my not be defined *at all*.
        Internally, SQLite handles this with a 'data type affinity' for each
        column definition, mapping to one of 'TEXT', 'NUMERIC', 'INTEGER',
        'REAL', or 'NONE' (raw bits). The algorithm that determines this is
        listed in http://www.sqlite.org/datatype3.html section 2.1.

        This method allows SQLAlchemy to support that algorithm, while still
        providing access to smarter reflection utilities by regcognizing
        column definitions that SQLite only supports through affinity (like
        DATE and DOUBLE).

        """
        match = re.match(r'([\w ]+)(\(.*?\))?', type_)
        if match:
            coltype = match.group(1)
            args = match.group(2)
        else:
            coltype = ''
            args = ''

        if coltype in self.ischema_names:
            coltype = self.ischema_names[coltype]
        elif 'INT' in coltype:
            coltype = sqltypes.INTEGER
        elif 'CHAR' in coltype or 'CLOB' in coltype or 'TEXT' in coltype:
            coltype = sqltypes.TEXT
        elif 'BLOB' in coltype or not coltype:
            coltype = sqltypes.NullType
        elif 'REAL' in coltype or 'FLOA' in coltype or 'DOUB' in coltype:
            coltype = sqltypes.REAL
        else:
            coltype = sqltypes.NUMERIC

        if args is not None:
            args = re.findall(r'(\d+)', args)
            try:
                coltype = coltype(*[int(a) for a in args])
            except TypeError:
                util.warn(
                    "Could not instantiate type %s with "
                    "reflected arguments %s; using no arguments." %
                    (coltype, args))
                coltype = coltype()
        else:
            coltype = coltype()

        return coltype

    @reflection.cache
    def get_pk_constraint(self, connection, table_name, schema=None, **kw):
        cols = self.get_columns(connection, table_name, schema, **kw)
        pkeys = []
        for col in cols:
            if col['primary_key']:
                pkeys.append(col['name'])
        return {'constrained_columns': pkeys, 'name': None}

    @reflection.cache
    def get_foreign_keys(self, connection, table_name, schema=None, **kw):
        quote = self.identifier_preparer.quote_identifier
        if schema is not None:
            pragma = "PRAGMA %s." % quote(schema)
        else:
            pragma = "PRAGMA "
        qtable = quote(table_name)
        statement = "%sforeign_key_list(%s)" % (pragma, qtable)
        c = _pragma_cursor(connection.execute(statement))
        fkeys = []
        fks = {}
        while True:
            row = c.fetchone()
            if row is None:
                break
            (numerical_id, rtbl, lcol, rcol) = (
                row[0], row[2], row[3], row[4])

            self._parse_fk(fks, fkeys, numerical_id, rtbl, lcol, rcol)
        return fkeys

    def _parse_fk(self, fks, fkeys, numerical_id, rtbl, lcol, rcol):
        # sqlite won't return rcol if the table was created with REFERENCES
        # <tablename>, no col
        if rcol is None:
            rcol = lcol

        if self._broken_fk_pragma_quotes:
            rtbl = re.sub(r'^[\"\[`\']|[\"\]`\']$', '', rtbl)

        try:
            fk = fks[numerical_id]
        except KeyError:
            fk = {
                'name': None,
                'constrained_columns': [],
                'referred_schema': None,
                'referred_table': rtbl,
                'referred_columns': [],
            }
            fkeys.append(fk)
            fks[numerical_id] = fk

        if lcol not in fk['constrained_columns']:
            fk['constrained_columns'].append(lcol)
        if rcol not in fk['referred_columns']:
            fk['referred_columns'].append(rcol)
        return fk

    @reflection.cache
    def get_indexes(self, connection, table_name, schema=None, **kw):
        quote = self.identifier_preparer.quote_identifier
        if schema is not None:
            pragma = "PRAGMA %s." % quote(schema)
        else:
            pragma = "PRAGMA "
        include_auto_indexes = kw.pop('include_auto_indexes', False)
        qtable = quote(table_name)
        statement = "%sindex_list(%s)" % (pragma, qtable)
        c = _pragma_cursor(connection.execute(statement))
        indexes = []
        while True:
            row = c.fetchone()
            if row is None:
                break
            # ignore implicit primary key index.
            # http://www.mail-archive.com/sqlite-users@sqlite.org/msg30517.html
            elif (not include_auto_indexes and
                  row[1].startswith('sqlite_autoindex')):
                continue

            indexes.append(dict(name=row[1], column_names=[], unique=row[2]))
        # loop thru unique indexes to get the column names.
        for idx in indexes:
            statement = "%sindex_info(%s)" % (pragma, quote(idx['name']))
            c = connection.execute(statement)
            cols = idx['column_names']
            while True:
                row = c.fetchone()
                if row is None:
                    break
                cols.append(row[2])
        return indexes

    @reflection.cache
    def get_unique_constraints(self, connection, table_name,
                               schema=None, **kw):
        UNIQUE_SQL = """
            SELECT sql
            FROM
                sqlite_master
            WHERE
                type='table' AND
                name=:table_name
        """
        c = connection.execute(UNIQUE_SQL, table_name=table_name)
        table_data = c.fetchone()[0]

        UNIQUE_PATTERN = 'CONSTRAINT (\w+) UNIQUE \(([^\)]+)\)'
        return [
            {'name': name,
             'column_names': [col.strip(' "') for col in cols.split(',')]}
            for name, cols in re.findall(UNIQUE_PATTERN, table_data)
        ]

def _pragma_cursor(cursor):
    """work around SQLite issue whereby cursor.description
    is blank when PRAGMA returns no rows."""

    if cursor.closed:
        cursor.fetchone = lambda: None
        cursor.fetchall = lambda: []
    return cursor


import os
from notion_client import Client
from langchain.document_loaders import PyPDFLoader, BSHTMLLoader, NotionDBLoader # for loading the pdf
from langchain.embeddings import OpenAIEmbeddings # for creating embeddings
from langchain.vectorstores import Chroma # for the vectorization part
from langchain.chains import ChatVectorDBChain # for chatting with the pdf
from langchain.chains import ConversationalRetrievalChain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.llms import OpenAI # the LLM model we'll use (CHatGPT)
from langchain.chat_models import ChatOpenAI
import textwrap
from os.path import join
import pickle as pkl

# os.environ["NOTION_TOKEN"] =
# os.environ["OPENAI_API_KEY"] =
#%%
database_id = "d3e3be7fc96a45de8e7d3a78298f9ccd"
notion = Client(auth=os.environ["NOTION_TOKEN"])
loader = NotionDBLoader(os.environ["NOTION_TOKEN"], database_id)
#%%
entries_return = notion.databases.query(database_id=database_id, filter={
      "property": "Name", "title": { "contains": "Progressive" }})
#%%
# print (entries_return)
def print_entries(entries_return):
    # formating the output, so Name starts at the same column
    # pad the string to be 36 character
    print("id".ljust(36), "\t", "Name",)
    for entry in entries_return["results"]:
        print(entry["id"], "\t", entry["properties"]["Name"]["title"][0]["plain_text"], )

print_entries(entries_return)
#%%
# clean up metadata
def clean_metadata(metadata):
    metadata_new = {}
    for k, v in metadata.items():
        if v is None or v == []:
            continue
        metadata_new[k] = metadata[k]
    return metadata_new
#%%
page_entry = entries_return["results"][0]
page_title = page_entry["properties"]["Name"]["title"][0]["plain_text"]
doc = loader.load_page(page_entry["id"])
#%%
doc.metadata = clean_metadata(doc.metadata)
page_save_name = page_title.replace("|", " ").replace(":", " ").replace(" ", "_")
pages = RecursiveCharacterTextSplitter().split_documents([doc])
#%%
# wrap the text to 80 characters and print it
print(textwrap.fill(pages[0].page_content, 80))
#%%
embeddings = OpenAIEmbeddings()
vectordb = Chroma.from_documents(pages, embedding=embeddings,
                                 persist_directory=page_save_name, )
vectordb.persist()
#%%
notion_qa = ConversationalRetrievalChain.from_llm(ChatOpenAI(temperature=0.9, model_name="gpt-3.5-turbo"),
                                    vectordb.as_retriever(), return_source_documents=True, max_tokens_limit=4097)
#%%
qa_path = page_save_name + "_qa_history"
os.makedirs(qa_path, exist_ok=True)
#%%
def save_qa_history(query, result, qa_path,):
    uid = 0
    while os.path.exists(join(qa_path, f"QA{uid:05d}.pkl")):
        uid += 1
    pkl.dump((query, result), open(join(qa_path, f"QA{uid:05d}.pkl"), "wb"))

    pkl_path = join(qa_path, "chat_history.pkl")
    if os.path.exists(pkl_path):
        chat_history = pkl.load(open(pkl_path, "rb"))
    else:
        chat_history = []
    chat_history.append((query, result))
    pkl.dump(chat_history, open(pkl_path, "wb"))

    with open(os.path.join(qa_path, "QA.md"), "a", encoding="utf-8") as f:
        f.write("\n**Question:**\n\n")
        f.write(query)
        f.write("\n\n**Answer:**\n\n")
        f.write(result["answer"])
        f.write("\n\nReferences:\n\n")
        for doc in result["source_documents"]:
            f.write("> ")
            f.write(doc.page_content[:250])
            f.write("\n\n")
        f.write("-------------------------\n\n")

#%%
# query = "What's the difference of novel stimuli and familiar stimuli in their response?"
query = "How does the paper define novel and faimilar stimuli in their method?"
query = "Why do they decide to use the intermittent stimuli set? What's conclusion about them? "
query = "How is this paper related to sleep and memory consolidation?"
query = "What is difference between the neurons with different time constants?"
query = "How is the neural response adaptation related to sleep?"
result = notion_qa({"question": query, "chat_history": ""})

print("Answer:")
print(textwrap.fill(result["answer"], 80))
save_qa_history(query, result, qa_path)
# print("\nReferences")
# for doc in result["source_documents"]:
#     print(doc.page_content[:100])
#     print("\n")


from modules.PickUpUrl import PickUpUrlClass
from modules.SaveAsPDF import SaveAsPDFClass
import sys
sys.dont_write_bytecode = True
try:
    pick_up_url = PickUpUrlClass()
    pick_up_url.read_cards()
    pick_up_url.serch_cards()
    pick_up_url.show_nohit_cards()
except:
    print("URLを読み込む過程でエラーが発生し、すべてのカードが検索できませんでした。")
    print("最後に検索したカード：" + pick_up_url.url_list[-1])

try:
    save_as_PDF = SaveAsPDFClass(pick_up_url.url_list)
    save_as_PDF.make_images()
except:
    print("ブラウザを動かせませんでした")
    print("取得したurl:")
    print(pick_up_url.url_list)