from sklearn import datasets
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_extraction.text import CountVectorizer

import joblib
import csv

'''
Gets the full feature list to be used as a training dataset
    Parameters:
            path(string): path of the file to be used
    Returns:
        f: the list of features that would be used for training
        labels: the labels of what the paper was classified as
'''
def train(path):
    features = pd.read_csv(path)
    labels = np.array(features['label'])

    # drops column keyword and label
    features = features.drop('keyword', axis = 1)
    features = features.drop('label', axis = 1)
    # use the title vector as a feature by converting it to a counter
    #title_counter = convert_to_counter(np.array(features['vector_t']))
    features = features.drop('vector_t', axis = 1)
    features = features.drop('vector_s', axis = 1)
    features = features.drop('position_k', axis = 1)
    features = features.drop('position_d', axis = 1)
    features = np.array(features)
    f = []
    # for i in range(len(features)):
    #     f.append(np.concatenate([features[i], title_counter[i]]))
    return features, labels

'''
Gets the full feature list to be used as a training dataset
    Parameters:
        feature(list): The dataset of features
        label(list): List of labels
    Returns:
        The accuarcy of the classifier
'''
def test_scholars(feature, label):
    train_features, test_features, train_labels, test_labels = train_test_split(feature, label, test_size = 0.20, random_state = 15)
    rfc = RandomForestClassifier(n_estimators = 500, random_state = 42)
    rfc.fit(train_features, train_labels)
    joblib.dump(rfc, "./random_forest.joblib")
    predicitons = rfc.predict(test_features)
    correct = 0
    for i in range(len(test_labels)):
        if test_labels[i] == predicitons[i]:
            correct += 1
    print(correct)
    return correct * 100.0 / len(test_labels)

'''
Store the random forrest tree classifier.
    Parameters:
        feature(list): The dataset of features
        label(list): List of labels
'''
def train_scholars(feature, label):
    rfc = RandomForestClassifier(n_estimators = 500, random_state = 42)
    rfc.fit(feature, label)
    joblib.dump(rfc, "./random_forest.joblib")

'''
Fill in the missing data by dummy values that could help them get ignored
    Parameters:
        path(string): path of the file to be modified
'''
def fill_missing_data(path):
    df = pd.read_csv(path)
    citation_num = [i for i in df['citation#'] if i != -999]
    year = [i for i in df['year'] if i != -999]
    k_list = [i for i in df['position_k'] if i != -999]
    d_list = [i for i in df['position_d'] if i != -999]
    avgs = []
    avgs.append(sum(citation_num) / len(citation_num))
    avgs.append(sum(k_list) / len(k_list))
    avgs.append(sum(d_list) / len(d_list))
    avgs.append(sum(year) / len(year))
    print(avgs)
    df.loc[df['year'] == -999, 'year'] = avgs[3]
    df.loc[df['position_d'] == -999, 'position_d'] = avgs[2]
    df.loc[df['position_k'] == -999, 'position_k'] = avgs[1]
    df.loc[df['citation#'] == -999, 'citation#'] = avgs[0]

    df.to_csv(path, index = False)

'''
Convert a list of strings to a vector of strings that could be utilized as features
    Parameters:
        vec(list): Vector that would be converted
'''
def convert_to_counter(vec):
    strs_list = []
    i = 0
    for v in vec:
        v = v.replace('[', '')
        v= v.replace(']', '')
        v = v.replace("',", "")
        v = v.replace("'", "")
        strs_list.append(v)
    vectorizer = CountVectorizer()
    return vectorizer.fit_transform(strs_list).toarray()

# f, l = train('scholars_updated_3 copy.csv')
# train_scholars(f, l)
# print(test_scholars(f, l))
# fill_missing_data('scholars_updated_3 copy.csv')


from pathlib import Path

from bug_changing.types.placeholder import Placeholder
from bug_changing.utils.file_util import FileUtil
from bug_changing.utils.path_util import PathUtil
import pandas as pd
import numpy as np

from config import OUTPUT_DIR, MOZILLA_PROJ, ECLIPSE_PROJ

def convert_summary_pairs_into_dataframe(summary_pairs, folder_name='mozilla'):
    columns = ['bug id', 'bug link', 'summary id', 'summary',
               Placeholder.SUGGESTED_SOLUTION, Placeholder.ACTUAL_BEHAVIOR, Placeholder.TRIGGER_ACTION,
               Placeholder.IS_ACTUAL_BEHAVIOR_VAGUE, Placeholder.IS_TRIGGER_ACTION_VAGUE,
               Placeholder.PLATFORM, Placeholder.COMPONENT,
               Placeholder.SPELLING_ERROR, Placeholder.SYNTAX_ERROR,
               Placeholder.DECLARATIVE_STATEMENT, Placeholder.PRESENT_TENSE,
               Placeholder.PERSONAL_PRONOUN, Placeholder.EMOTIONAL_REACTION,
               Placeholder.OFFENCE, Placeholder.WISHY_WASHY_WORD]

    # for index in range(5):
    #     columns.extend(['change from', 'change what', 'change for'])
    rows = []
    for summary_pair in summary_pairs:
        if folder_name == MOZILLA_PROJ:
            bug_link = f"https://bugzilla.mozilla.org/show_bug.cgi?id={summary_pair.bug.id}"
        else:
            bug_link = f"https://bugs.eclipse.org/bugs/show_bug.cgi?id={summary_pair.bug.id}"

        summaries = [summary_pair.rm_summary, summary_pair.add_summary]
        for summary in summaries:
            row = [summary_pair.bug.id, bug_link, summary.id, summary.text, ]
            for add_column in range(15):
                row.extend([''])
            # row.append('')
            rows.append(row)
    df = pd.DataFrame(rows, columns=columns)
    return df

def get_labels_from_row(row):
    labels = row.values[5:20].tolist()  # [5, 20)
    return labels

if __name__ == "__main__":
    folder_name = MOZILLA_PROJ
    # folder_name = ECLIPSE_PROJ
    summary_pairs_filepath = PathUtil.get_instance_summary_pairs_filepath(folder_name)
    summary_pairs = FileUtil.load_pickle(summary_pairs_filepath)
    # print(summary_pairs)
    summary_pairs.remove_tags()

    df = pd.read_excel(Path(OUTPUT_DIR, folder_name, f'instance_summary_pairs_labels.xlsx'),
                       sheet_name=f'{folder_name}')
    df = df.replace(np.nan, None)

    for index, two_rows in df.groupby(np.arange(len(df)) // 2):
        # print(index)
        first_row = two_rows.iloc[0]
        second_row = two_rows.iloc[1]
        bug_id = int(first_row.values[1])
        rm_summary_id = int(first_row.values[3])
        add_summary_id = int(second_row.values[3])
        summary_pair = summary_pairs.get_summary_pair_by_bug_id_summary_ids(bug_id, rm_summary_id, add_summary_id)
        labels_1 = get_labels_from_row(first_row)
        labels_2 = get_labels_from_row(second_row)
        print(labels_1)
        print(labels_2)
        summary_pair.rm_summary.initiate_from_labels(labels_1)
        summary_pair.add_summary.initiate_from_labels(labels_2)
        print(summary_pair)
        # print(labels_1)
        # print(labels_2)
        # print(type(first_row.values[0:5].tolist()))
        # print(first_row.values[1])
        # print(second_row)
        # print(type(g))
        # input()
    if folder_name == MOZILLA_PROJ:
        bug_id_summary_ids_tuple_list = [(551591, 3, 4),
                                         (1388414, 0, 1), (697762, 0, 1), (1393563, 0, 1), (831452, 0, 1),
                                         (1611864, 0, 1), (641322, 1, 2), (1302671, 0, 1), (1698580, 1, 2)]
    else:
        bug_id_summary_ids_tuple_list = [(478512, 0, 1), (560060, 0, 1), (357540, 0, 1)]
    print(len(summary_pairs))
    summary_pairs = summary_pairs.filter_for_goal_model_instances(bug_id_summary_ids_tuple_list)
    summary_pairs.complete_vague_labels()
    print(len(summary_pairs))
    # print(summary_pairs.get_summary_pair_by_bug_id_summary_ids(1403460, 0, 1))
    FileUtil.dump_pickle(summary_pairs_filepath, summary_pairs)


import numpy as np
from primeqa.ir.dense.dpr_top.dpr.simple_mmap_dataset import Corpus
from primeqa.ir.dense.dpr_top.util.args_help import fill_from_args
import faiss
from primeqa.ir.dense.dpr_top.util.reporting import Reporting
import os
import time
import math
import logging

logger = logging.getLogger(__name__)

def l2_convert_indexed_vectors(vectors: np.ndarray, max_norm_sqrd: float):
    norms = np.linalg.norm(vectors, axis=1)
    aux_dims_sqrd = max_norm_sqrd - norms * norms
    if np.min(aux_dims_sqrd) < 0:
        print(f'WARNING: max_norm_sqrd = {max_norm_sqrd} but it was less '
              f'({np.min(aux_dims_sqrd)}) than a vectors norm_sqrd')
        aux_dims_sqrd = np.maximum(aux_dims_sqrd, 0)
    aux_dims = np.sqrt(aux_dims_sqrd)
    converted_vectors = np.hstack((vectors, aux_dims.reshape(-1, 1)))
    return converted_vectors

def l2_convert_query_vectors(vectors: np.ndarray):
    aux_dim = np.zeros(vectors.shape[0], dtype=np.float32)
    converted_vectors = np.hstack((vectors, aux_dim.reshape(-1, 1)))
    return converted_vectors

class ANNIndex:
    def __init__(self, index_file):
        self.index = faiss.read_index(index_file)
        self.is_l2 = type(self.index) == faiss.IndexHNSWSQ

    def search(self, query_vectors, k):
        if self.is_l2:
            query_vectors = l2_convert_query_vectors(query_vectors)
        scores, indexes = self.index.search(query_vectors, k)
        if self.is_l2:
            scores = -1 * scores  # make higher scores better always
        return scores, indexes

    def dim(self):
        if self.is_l2:
            return self.index.d - 1
        else:
            return self.index.d

class IndexOptions():
    def __init__(self):
        self.d = 768
        self.m = 128
        self.ef_search = 128
        self.ef_construction = 200
        self.index_batch_size = 100000
        self.scalar_quantizer = -1
        self.product_quantizer_m = -1  # suggested values of 8, 16, 32 - maybe we need much higher though
        self.product_quantizer_sv_bits = 8  # probably don't change this
        self.is_l2 = False
        self.num_vectors = -1
        self.max_norm = -1

    def _post_argparse(self):
        self.is_l2 = self.scalar_quantizer > 0

def build_index(corpus_dir, output_file, opts: IndexOptions):
    logger.info(f'building index, reading data from {corpus_dir}, writing to {output_file}')
    corpus = Corpus(corpus_dir)
    if opts.is_l2:
        print(f'Using L2 distance conversion')

    if (opts.num_vectors <= 0 and opts.product_quantizer_m > 0) or (opts.max_norm <= 0 and opts.is_l2):
        max_norm = 0
        num_vectors = 0
        start_time = time.time()
        for psg in corpus:
            vector = psg['vector']
            max_norm = max(max_norm, np.linalg.norm(vector))
            num_vectors += 1
        print(f'found max norm = {max_norm} over {num_vectors} vectors in {(time.time()-start_time)/60} min.')
        opts.max_norm = max_norm
        opts.num_vectors = num_vectors
    max_norm_sqrd = opts.max_norm * opts.max_norm

    opts.is_trained = False
    if opts.scalar_quantizer > 0:
        if opts.scalar_quantizer == 16:
            sq = faiss.ScalarQuantizer.QT_fp16
        elif opts.scalar_quantizer == 8:
            sq = faiss.ScalarQuantizer.QT_8bit
        elif opts.scalar_quantizer == 4:
            sq = faiss.ScalarQuantizer.QT_4bit
        elif opts.scalar_quantizer == 6:
            sq = faiss.ScalarQuantizer.QT_6bit
        else:
            raise ValueError(f'unknown --scalar_quantizer {opts.scalar_quantizer}')
        # see https://github.com/facebookresearch/faiss/blob/master/benchs/bench_hnsw.py
        index = faiss.IndexHNSWSQ(opts.d+1, sq, opts.m)
        index.hnsw.efSearch = opts.ef_search
        index.hnsw.efConstruction = opts.ef_construction
    elif opts.product_quantizer_m > 0:
        # product quant: https://github.com/matsui528/faiss_tips
        # seems awful - maybe product_quantizer_m should be much higher?
        nlist = int(math.sqrt(opts.num_vectors))
        quantizer = faiss.IndexHNSWFlat(opts.d, opts.m, faiss.METRIC_INNER_PRODUCT)
        index = faiss.IndexIVFPQ(quantizer, opts.d, nlist, opts.product_quantizer_m,
                                 opts.product_quantizer_sv_bits, faiss.METRIC_INNER_PRODUCT)
        index.nprobe = opts.ef_search  # matsui528 recommended 8
        opts.index_batch_size = max(opts.index_batch_size, 50 * nlist)
    else:
        index = faiss.IndexHNSWFlat(opts.d, opts.m, faiss.METRIC_INNER_PRODUCT)
        # defaults are 16 and 40
        index.hnsw.efSearch = opts.ef_search
        index.hnsw.efConstruction = opts.ef_construction
        opts.is_trained = True  # doesn't need training

    vectors = np.zeros((opts.index_batch_size, opts.d), dtype=np.float32)
    vector_ndx = 0

    def add_to_index(vectors):
        if opts.is_l2:
            to_index = l2_convert_indexed_vectors(vectors, max_norm_sqrd)
        else:
            to_index = vectors
        if not opts.is_trained:
            index.train(to_index)
            opts.is_trained = True
        logger.info(f'calling index.add with {len(to_index)} vectors')
        index.add(to_index)

    report = Reporting()

    for pndx, psg in enumerate(corpus):
        if (pndx % 100000) == 0:
            logger.info(f'processed {pndx} passages')
        if report.is_time():
            print(report.progress_str(instance_name='vector'))
        vector = psg['vector']
        vectors[vector_ndx] = vector
        vector_ndx += 1
        if vector_ndx == opts.index_batch_size:
            logger.info(f'processed {pndx} passages')
            add_to_index(vectors)
            vector_ndx = 0
    if vector_ndx > 0:
        add_to_index(vectors[:vector_ndx])
    logger.info(f'processed {len(corpus)} passages')
    logger.info(f'finished building index, writing index file to {output_file}')
    #print(f'finished building index, writing index file to {output_file}')
    faiss.write_index(index, output_file)
    logger.info(f'took {report.elapsed_time_str()}')
    #print(f'took {report.elapsed_time_str()}')

if __name__ == "__main__":
    class CmdOptions(IndexOptions):
        def __init__(self):
            super().__init__()
            self.collection = ''  # can be a directory with passages*.json.gz.records or a single such file
            self.__required_args__ = ['collection']

    opts = CmdOptions()
    fill_from_args(opts)

    if os.path.isdir(opts.collection):
        output_file = os.path.join(opts.collection, 'index.faiss')
    else:
        base_dir, filename = os.path.split(opts.collection)
        assert filename.startswith('passages') and filename.endswith('.json.gz.records')
        index_fname = f'index{filename[len("passages"):-len(".json.gz.records")]}.faiss'
        output_file = os.path.join(base_dir, index_fname)

    build_index(opts.collection, output_file, opts)


import matplotlib.pyplot as plt
import pynmea2
import cgi
import math
import numpy as np
from math import radians, cos, sin, asin, sqrt
import pyproj
import re
import pymap3d as pm

EARTH_REDIUS = 6378.137

def dms2dd(degrees, minutes, seconds, direction):
    dd = float(degrees) + float(minutes)/60 + float(seconds)/(60*60);
    if direction == 'S' or direction == 'W':
        dd *= -1
    return dd;

def dd2dms(deg):
    d = int(deg)
    md = abs(deg - d) * 60
    m = int(md)
    sd = (md - m) * 60
    return [d, m, sd]

def parse_dms(dms):
    parts = re.split('[^\d\w]+', dms)
    lat = dms2dd(parts[0], parts[1], parts[2], parts[3])
    lng = dms2dd(parts[4], parts[5], parts[6], parts[7])

    return (lat, lng)

## ==========================================================================================================================
def np_getDistance(A , B ):# 先緯度後經度
    ra = 6378140  # radius of equator: meter
    rb = 6356755  # radius of polar: meter
    flatten = 0.003353 # Partial rate of the earth
    # change angle to radians

    radLatA = np.radians(A[:,0])
    radLonA = np.radians(A[:,1])
    radLatB = np.radians(B[:,0])
    radLonB = np.radians(B[:,1])

    pA = np.arctan(rb / ra * np.tan(radLatA))
    pB = np.arctan(rb / ra * np.tan(radLatB))

    x = np.arccos( np.multiply(np.sin(pA),np.sin(pB)) + np.multiply(np.multiply(np.cos(pA),np.cos(pB)),np.cos(radLonA - radLonB)))
    c1 = np.multiply((np.sin(x) - x) , np.power((np.sin(pA) + np.sin(pB)),2)) / np.power(np.cos(x / 2),2)
    c2 = np.multiply((np.sin(x) + x) , np.power((np.sin(pA) - np.sin(pB)),2)) / np.power(np.sin(x / 2),2)
    dr = flatten / 8 * (c1 - c2)
    distance = 0.001 * ra * (x + dr)
    return distance

def rad(d):
    return d * np.pi / 180.0

def getDistance(lat1, lng1, lat2, lng2):
    radLat1 = rad(lat1)
    radLat2 = rad(lat2)
    a = radLat1 - radLat2
    b = rad(lng1) - rad(lng2)
    s = 2 * math.asin(math.sqrt(math.pow(sin(a/2), 2) + cos(radLat1) * cos(radLat2) * math.pow(sin(b/2), 2)))
    s = s * EARTH_REDIUS
    return s

def deg2num(lat_deg=0.0, lon_deg=0.0, zoom=12):
  lat_rad = math.radians(lat_deg)
  n = 2.0 ** zoom
  xtile = int((lon_deg + 180.0) / 360.0 * n)
  ytile = int((1.0 - math.log(math.tan(lat_rad) + (1 / math.cos(lat_rad))) / math.pi) / 2.0 * n)
  return (xtile, ytile)

def Write_RTKlibFormat(file_name,data):
    file_name.write(data)

rangeUpdateTime = [0.0]

if __name__=='__main__':
    input = open('serverlog.txt', 'r')
    file_name='tesst.nmea'
    input_nmea = open(file_name, 'r')

    RTKlog_name = open(file_name+'_log', 'w')
    line = '$GPGGA,092750.000,5321.6802,N,00630.3372,W,1,8,1.03,61.7,M,55.2,M,,*76'
    #record = pynmea2.parse(line)
    for line in input_nmea:

        if 'GGA' in line:
            record = pynmea2.parse(line)
            Write_RTKlibFormat(RTKlog_name,line)

            '''
            print('Latitude: {:02d}°{:02d}′{:07.4f}″'.format(int(record.latitude), int(record.latitude_minutes), record.latitude_seconds))
            print('Longitude: {:02d}°{:02d}′{:07.4f}″'.format(int(record.longitude), int(record.longitude_minutes), record.longitude_seconds))
            print('Altitude: {:.3f}'.format(record.altitude))
            '''

        #if 'GPGSA' in line:
            #Write_RTKlibFormat(RTKlog_name,line)
        if 'GPRMC' in line:
            Write_RTKlibFormat(RTKlog_name,line)
    RTKlog_name.close()

    A=np.matrix([[25.060835, 121.544242]])
    B=np.matrix([[25.062247, 121.544236]])

    Distance_2=np_getDistance(A,B)
    Distance_3=getDistance(25.060835, 121.544242,25.062247, 121.544236)
    x = 0
    y = 45
    z = 1000

    ecef = pyproj.Proj(proj='geocent', ellps='WGS84', datum='WGS84')
    lla = pyproj.Proj(proj='latlong', ellps='WGS84', datum='WGS84')
    lon, lat, alt = pyproj.transform(ecef, lla, x, y, z, radians=True)

    print ( Distance_2, Distance_3)
    dd = parse_dms("36°57'9' N 110°4'21' W")

    x,y,z=pm.geodetic2ecef(22.99665875 ,120.222584889,98.211)
    x,y,z=pm.ecef2geodetic(x,y,z)
    print(x,y,z)
    print(dd)
    print(dd2dms(dd[0]))

    for line in input:
        line = line.split()
        if 'update' in line:
            rangeUpdateTime.append(float(line[-1]))

    plt.figure('frame time')
    plt.subplot(211)
    plt.plot(rangeUpdateTime, '.r',)
    plt.grid(True)
    plt.subplot(212)
    plt.plot(rangeUpdateTime)
    plt.grid(True)
    plt.show()


import time
import json
import csv
import pandas as pd
from selenium import webdriver
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.chrome.service import Service

from webdriver_manager.chrome import ChromeDriverManager

driver = webdriver.Chrome(ChromeDriverManager().install())

#website
driver.get('https://opensea.io/collection/art')

driver.maximize_window()
#get title
driver.implicitly_wait(10)

print("--->title",driver.title)

#/html/body/div/div[1]/main/div/div/div[2]/div/div/div[1]/aside/div/div[2]/div/div/div/div/div[2]/div[1]/section/div/div/input
driver.execute_script("window.scrollBy(0,300)","")
obj=driver.find_element_by_xpath("//input[@placeholder='Min']")
s=obj.send_keys("1.25")
obj=driver.find_element_by_xpath("//input[@placeholder='Max']")
s=obj.send_keys("3.25")

#button
#//button[@class='Blockreact__Block-sc-1xf18x6-0 Buttonreact__StyledButton-sc-glfma3-0 cWgiBW kaikhX ActionButtonreact__StyledButton-sc-7jpoey-0 bwNKOH']
apply=driver.find_element_by_xpath("//button[@class='Blockreact__Block-sc-1xf18x6-0 Buttonreact__StyledButton-sc-glfma3-0 fAtwh kaikhX ActionButtonreact__StyledButton-sc-7jpoey-0 bwNKOH']")
s=apply.click()

time.sleep(3)

#aricle
# //article[@class="Assetreact__DivContainer-sc-bnjqwy-0 gwTKfX Asset--loaded AssetSearchList--asset"]

#image url data
image_url=[]
while True:
    driver.execute_script("window.scrollBy(0,300)","")

    obj = driver.find_elements_by_xpath("//img[@class='Image--image']")
    for i in obj:
        img_src=i.get_attribute("src")
        image_url.append(img_src)
        print(img_src)

    if len(img_src)<5000:
        break

"""
<img class="Image--image" src="https://lh3.googleusercontent.com/xUgm8yFV61OGAtdBz73D5TZpNXsZJTYXjfjHKYZuW9poum0aChIxu6W1LquWXsj5JqCSuLuqfI1ypVau5byDWduhEdOZveDn7PKx=w333" style="object-fit: contain; width: auto; height: auto; border-radius: 5px; max-width: 100%; max-height: 100%;">
 """
"""
lisst=[]
condition= True

while condition :
    obj = driver.find_elements_by_xpath("//*[@id='__next']/div[1]/main/div/div/div/div[2]/div/div[5]/div/div/div/div/div/div/div[2]/div/div[2]/div[3]/div/a/self::a")

    for ai in obj:
        #print(obj)
        out = obj.text
        lisst.append(out)
        print(lisst)  """
#1.collection
#//div[@class='AssetCardFooter--collection']
#//a[@class='Linkreact__StyledA-sc-18se2b0-0 iWVmEt AssetCardFooter--collection-name']
li=[]
txt=i.text
while True:
    driver.execute_script("window.scrollBy(0,200)","")

    col = driver.find_elements_by_xpath("//a[@class='Linkreact__StyledA-sc-18se2b0-0 iWVmEt AssetCardFooter--collection-name']")
    for i in col:
        txt=i.text
        li.append(txt)
        print(txt)

    if len(txt)<5000:
        #print(len(text1))
        break

#2.name
#//div[@class='AssetCardFooter--name']
name=[]
name1=i.text
while True:
    driver.execute_script("window.scrollBy(0,200)","")

    nam = driver.find_elements_by_xpath("//div[@class='AssetCardFooter--name']")
    for i in nam:
        name1=i.text
        name.append(name1)
        print(name1)

    if len(name1)<5000:
        #print(len(name1))
        break

#3.bid value
# //div[@class='Overflowreact__OverflowContainer-sc-10mm0lu-0 fqMVjm Price--amount']
bid_value=[]
while True:
    driver.execute_script("window.scrollBy(0,200)","")

    bid_ = driver.find_elements_by_xpath("//div[@class='Overflowreact__OverflowContainer-sc-10mm0lu-0 fqMVjm Price--amount']")
    for i in bid_:
        sad=i.text
        bid_value.append(sad)
        print(sad)

    if len(sad)<5000:
        #print(len(sad))
        break

#4.type of bid
# //div[@class='AssetCardFooter--price-header']
type_bid=[]
while True:
    driver.execute_script("window.scrollBy(0,200)","")

    type_bi = driver.find_elements_by_xpath("//div[@class='AssetCardFooter--price-header']")
    for i in type_bi:
        na=i.text
        type_bid.append(na)
        print(na)

    if len(na)<5000:
        #print(len(na))
        break

#5.favorite
#//div[@class="Blockreact__Block-sc-1xf18x6-0 Flexreact__Flex-sc-1twd32i-0 FlexVerticalreact__FlexVertical-sc-x35rw8-0 VerticalAlignedreact__VerticalAligned-sc-1v4gkkl-0 beUSAq jYqxGr gXnSUK gJMdLt"]
#//span[@class='Blockreact__Block-sc-1xf18x6-0 Textreact__Text-sc-1w94ul3-0 jEeFPO gwsEKa']

#//div[@class='Blockreact__Block-sc-1xf18x6-0 Flexreact__Flex-sc-1twd32i-0 FlexVerticalreact__FlexVertical-sc-x35rw8-0 VerticalAlignedreact__VerticalAligned-sc-1v4gkkl-0 beUSAq jYqxGr gXnSUK gJMdLt']
#//div[@class='Blockreact__Block-sc-1xf18x6-0 Flexreact__Flex-sc-1twd32i-0 FlexVerticalreact__FlexVertical-sc-x35rw8-0 VerticalAlignedreact__VerticalAligned-sc-1v4gkkl-0 fAzvQd jYqxGr gXnSUK gJMdLt']
fav=[]
while True:
    driver.execute_script("window.scrollBy(0,200)","")

    fA = driver.find_elements_by_xpath("//div[@class='Blockreact__Block-sc-1xf18x6-0 Flexreact__Flex-sc-1twd32i-0 FlexVerticalreact__FlexVertical-sc-x35rw8-0 VerticalAlignedreact__VerticalAligned-sc-1v4gkkl-0 fAzvQd jYqxGr gXnSUK gJMdLt']")
    for i in fA:
        s=i.text
        fav.append(s)
        print(s)

    if len(s)<1000:
        #print(len(s))
        break

#token

# //a[@class='Linkreact__StyledA-sc-18se2b0-0 iWVmEt Asset--anchor']
tok=[]
for j in range(100):
    #driver.execute_script("window.scrollBy(0,200)","")

    Tk = driver.find_elements_by_xpath("//a[@class='Blockreact__Block-sc-1xf18x6-0 fAtwh AssetsSearchView--assets']")
    #Blockreact__Block-sc-1xf18x6-0 fAtwh AssetsSearchView--assets
    for i in Tk: #Linkreact__StyledA-sc-18se2b0-0 iWVmEt Asset--anchor
        da=i.get_attribute("href")
        tok.append(da)
        print(da)

""" dic={
    'image':image_url[0:4],
    'collection':li[0:4],
    'name':name[0:4],
    'type_of_bid':type_bid[0:4],
    'favourite':fav[0:4],
    'eth_value':bid_value[0:4],
    'token_id':tok[0:4]

}

print(dic)

 with open('openseaee.csv','w')as f:
    w=csv.DictWriter(f,dic.keys())
    w.writeheader()

df =pd.DataFrame(dic)
print(df)
df.to_csv('sasea.csv',encoding='utf-8', index=False)      """


import argparse
import os

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--file', default=None, type=str, help="path to txt file that we want to split")
parser.add_argument('--num', default=2, type=int, help="number of target files")
parser.add_argument('--output-dir', default='output', type=str, help='dir to store output files')

args = parser.parse_args()
assert args.file, "file path is nesserary."

if not os.path.exists(args.output_dir):
    os.makedirs(args.output_dir)

input_filename = args.file.split('/')[-1].split('.')[0]
output_filenames = [input_filename + '_' + str(i) for i in range(args.num)]

file_counter = 0
with open(args.file, 'r') as f_in:
    full_lines = f_in.readlines()
    assert len(full_lines) % args.num == 0, 'cannot split into %d files with %d lines input.' % (args.num, len(full_lines))
    each_num = len(full_lines) // args.num
    line_counter = 0
    for line in full_lines:
        line = line.strip()
        if line_counter >= each_num:
            file_counter += 1
            line_counter = 0
        output_file = os.path.join(args.output_dir, output_filenames[file_counter] + '.txt')
        with open(output_file, 'a') as f_out:
            f_out.writelines(line + '\n')
        line_counter += 1


import os
import json
import subprocess

git_dir = subprocess.check_output(["git","rev-parse","--show-toplevel"]).decode().strip()
text_files = [
    "SingleMuon_2018A", "SingleMuon_2018B", "TTJets_UL2018", "SingleMuon_2018C",
    "SingleMuon_2018D", "WJets_UL2018"
]
mc_dict = {
    "TTJets_UL2018": {
        "nevts_total": 296_610_126
    },
    "WJets_UL2018": {
        "nevts_total": 81_051_269
    }
}
d = {}

for dataset_name in text_files:
    file_path = os.path.join(git_dir,"analysis/samples",dataset_name + ".txt")
    fin = open(file_path, mode='r')
    is_mc = dataset_name in mc_dict

    if dataset_name.find("2016") >= 0:
        year = 2016
    elif dataset_name.find("2017") >= 0:
        year = 2017
    elif dataset_name.find("2018") >= 0:
        year = 2018
    else:
        print("No valid year found")

    # d[dataset_name] = {
        # "variations": {},
        # "is_data": True,
    # }

    d[dataset_name] = {}
    d[dataset_name]["variations"] = {}
    d[dataset_name]["is_data"] = not is_mc
    d[dataset_name]["year"] = year

    # d[dataset_name]["variations"] = {"nominal": {"nevts_total": -1, "files": []} }
    # d[dataset_name]["variations"]["nominal"] = {"nevts_total": -1, "files": []}

    if is_mc:
        numb_evts = mc_dict[dataset_name]["nevts_total"]
    else:
        numb_evts = -1

    d[dataset_name]["variations"]["nominal"] = {"nevts_total": numb_evts, "files": []}

    for x in fin.readlines():
        x = x.strip()

        y = {"path" : "root://xcache/" + x, "nevts": -1}
        print(y)
        d[dataset_name]["variations"]["nominal"]["files"].append(y)

    print(d[dataset_name]["variations"]["nominal"]["files"])

fout = open("Run2_data.json", mode = "w")
json.dump(d,fout,indent = 4)


import numpy as np
import pandas as pd
import pytest
import os
import sys

current_path = os.path.dirname(os.path.abspath(__file__))
new_path = os.path.join(current_path, "..")
sys.path.append(new_path)

from sbc import sbc_class

data = pd.read_csv("data/sales_train_clean.csv")

def test_raise_empty_data_df():
    ## ndim == 2
    input_data = pd.DataFrame()
    with pytest.raises(ValueError):
        sbc_class.sbc_class(input_data)

def test_raise_empty_data_list():
    ## ndim == 1
    input_data = []
    with pytest.raises(ValueError):
        sbc_class.sbc_class(input_data)

def test_null_object():
    input_data = None
    with pytest.raises(ValueError):
        sbc_class.sbc_class(input_data)

def test_plot_value():
    df = data.iloc[:,1:3]
    type = ['line_plot']
    with pytest.raises(ValueError):
        sbc_class.sbc_class(df, plot_type=type)

def test_reproducible():
    df = data.iloc[:, 1]
    df1 = data.iloc[:,1]
    out1 = sbc_class.sbc_class(df)
    out2 = sbc_class.sbc_class(df1)

    assert out1.equals(out2)


import sys
import os
from dataclasses import dataclass

import numpy as np
import pandas as pd

from src.exception import CustomException
from src.logger import logging
from src.utils import save_object, evaluate_model

from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor, AdaBoostRegressor
from sklearn.linear_model import LinearRegression
from sklearn.neighbors import KNeighborsRegressor
from catboost import CatBoostRegressor
from xgboost import XGBRegressor

from sklearn.metrics import r2_score
from src.config.hyper_parameter_tuning import params

@dataclass
class ModelTrainingConfig:
    trained_model_file_path = os.path.join("artifacts","model.pkl")

class ModelTrainer:
    def __init__(self):
        self.model_trainer_config=ModelTrainingConfig()

    def initiate_model_trainer(self,train_array,test_array):
        try:
            logging.info("Splitting train and test data")
            X_train,y_train,X_test,y_test = (
                train_array[:,:-1],
                train_array[:,-1],
                test_array[:,:-1],
                test_array[:,-1]
            )
            logging.info("Splitting train and test data Completed")

            models = {
                "Decision Tree": DecisionTreeRegressor(),
                "Random Forest":RandomForestRegressor(),
                "Linear Regression":LinearRegression(),
                "Gradient Boosting":GradientBoostingRegressor(),
                "AdaBoost Regressor": AdaBoostRegressor(),
                "CatBoosting Regressor": CatBoostRegressor(verbose=False),
                "XGBoost Regressor": XGBRegressor()
            }

            model_report:dict = evaluate_model(X_train=X_train, y_train=y_train, X_test=X_test, y_test=y_test,
                                               models=models, params=params)

            best_model_score = max(sorted(model_report.values()))

            best_model_name = list(model_report.keys())[
                list(model_report.values()).index(best_model_score)
            ]

            best_model = models[best_model_name]

            if best_model_score<0.6:
                raise CustomException("No good model found")
            logging.info(f"Best found model on both training and testing dataset")
            logging.info("Best model is [{0}] with r2_score [{1}]".format(best_model_name,best_model_score))
            save_object(
                file_path=self.model_trainer_config.trained_model_file_path,
                obj=best_model
            )

            predicted=best_model.predict(X_test)

            r2_square = r2_score(y_test, predicted)
            return r2_square

        except Exception as e:
            CustomException(e,sys)


import arcpy
from tempfile import mkdtemp
from os import remove, rmdir
from urllib2 import Request, urlopen, HTTPHandler, build_opener
from base64 import b64encode
from json import dumps
previewPage = """
<!doctype html>

		<meta charset='utf-8'/>
		<title>
			ESRI To Couch
		</title>
		<link rel="stylesheet" href="http://cdn.leafletjs.com/leaflet-0.6.4/leaflet.css" />

"""
design = dumps({
   "_id": "_design/geojson",
   "spatial": {
       "geo": "function(doc){\n       if(doc.geometry){ emit(doc.geometry,doc)};\n            }"
   },
   "views": {
       "geometry": {
           "map": "function(doc) {\n  if(doc.geometry){\nemit(doc.geometry,doc);\n}\n}"
       }
   },
   "rewrites": [
       {
           "to": "/_list/geojson/geometry",
           "from": "/all"
       },
       {
           "to": "/index.html",
           "from": "/"
       }
   ],
   "_attachments": {
        "index.html" : {
            "content_type" : "text/html",
            "data" : b64encode(previewPage)
        }
    },
   "lists": {
       "geojson": "function(head, req) {\n    var row, out, sep = '\\n';\n    if (req.headers.Accept.indexOf('application/json')!=-1) {\n        start({\"headers\":{\"Content-Type\" : \"application/json\"}});\n    }    else {\n        start({\"headers\":{\"Content-Type\" : \"text/plain\"}});\n    }\n    if ('callback' in req.query) {\n        send(req.query['callback'] + \"(\");\n    }    send('{\"type\": \"FeatureCollection\", \"features\":[');\n    while (row = getRow()) {\n        out = {type: \"Feature\", geometry: row.value.geometry};\n\t\t\t\tdelete row.value.geometry;\n               out.properties= row.value\n        send(sep + JSON.stringify(out));        sep = ',\\n';\n    }\n    send(\"]}\");\n    if ('callback' in req.query) {\n        send(\")\");}};\n"
   },
   "validate_doc_update": "function(newDoc, oldDoc, userCtx) {  if (userCtx.roles.indexOf('_admin') !== -1) {    return;  } else {    throw({forbidden: 'Only admins may edit the database'});  }}"
})

def doStuff(inFile,dbUrl,tbx,login,createDB):
    arcpy.ImportToolbox(tbx)
    base = mkdtemp()
    outFile = base+'//output.json'
    arcpy.esri2open(inFile,outFile)
    postToCouch(dbUrl,outFile,login,createDB)
    cleanUp(base,outFile)

def cleanUp(p,f):
    remove(f)
    rmdir(p)

def dealLogin(req,login):
    if login == 'NONE':
        return
    req.add_header('Authorization','Basic '+b64encode(login))

def postToCouch(url,outFile,login,createDB):
    if createDB=='true' and login != 'NONE':
        opener = build_opener(HTTPHandler)
        res3 = Request(url)
        dealLogin(res3,login)
        res3.get_method = lambda: 'PUT'
        opener.open(res3)
        arcpy.AddMessage("Created DB")
    newUrl = url + '//_bulk_docs'
    baseFile = open(outFile,'r')
    data = baseFile.read()
    baseFile.close()
    req = Request(newUrl)
    dealLogin(req,login)
    req.add_header('Content-Type', 'application/json')
    urlopen(req,data)
    arcpy.AddMessage("Added Rows")
    if login != 'NONE':
        req2 = Request(url)
        dealLogin(req2,login)
        req2.add_header('Content-Type', 'application/json')
        urlopen(req2,design)
        arcpy.AddMessage("Added Design Doc")
        arcpy.AddMessage("View it at "+url+"/_design/geojson/_rewrite/")

doStuff(arcpy.GetParameterAsText(0),arcpy.GetParameterAsText(1),arcpy.GetParameterAsText(2),arcpy.GetParameterAsText(3),arcpy.GetParameterAsText(4))


# http://stackoverflow.com/questions/10657383/stitching-photos-together
from PIL import Image

def merge_images(file1, file2, file3):
    """Merge two images into one, displayed side by side
    :param file1: path to first input image file
    :param file2: path to second input image file
    :param file3: path to output image file
    """
    image1 = Image.open(file1)
    image2 = Image.open(file2)

    (width1, height1) = image1.size
    (width2, height2) = image2.size

    result_width = width1 + width2
    result_height = max(height1, height2)

    result = Image.new('RGB', (result_width, result_height), color=(255,255,255,255))
    result.paste(im=image1, box=(0, 0))
    result.paste(im=image2, box=(width1, 0))

    result.save (file3)

def merge_horizontal(input_paths, outpath):
    """Horizontally merge several images into one, displayed side by side
    :param input_paths: list of path to input image files
    :param outpath: path to output image file
    """
    result_width = 0
    result_height = 0

    for path in input_paths:
      image = Image.open(path)
      (width, height) = image.size
      result_width += width
      result_height = max(result_height, height)

    result = Image.new('RGB', (result_width, result_height), color=(255,255,255,255))

    xpos = 0
    for path in input_paths:
      image = Image.open(path)
      (width, height) = image.size
      result.paste(im=image, box=(xpos, 0))
      xpos += width

    result.save (outpath)


#!/usr/bin/python
# coding=utf-8
import xml.etree.ElementTree

def get_words(file_path):
    with open(file_path, 'r') as xml_file:
        e = xml.etree.ElementTree.parse(xml_file).getroot()

        inner_dictionary = {}
        for chunk in e.findall('chunk'):
            for sentence in chunk.findall('sentence'):
                for tok in sentence.findall('tok'):
                    is_proper_name = False
                    for ann in tok.findall('ann'):
                        if ann.attrib['chan'] == 'nam' and ann.text != '0':
                            is_proper_name = True
                    for orth in tok.find('orth'):
                        if orth.text not in inner_dictionary:
                            inner_dictionary[orth.text] = is_proper_name
                    for lex in tok.findall('lex'):
                        if lex.find('base').text not in inner_dictionary:
                            inner_dictionary[lex.find('base').text] = is_proper_name
        return inner_dictionary

with open('../data/kpwr-1.2.7-names-disamb-nam-flatten/index_names.txt', 'rU') as f:
    file_ = open('../data/prepared/AllFlattenNamWords.csv', 'w')
    file_.truncate()
    book = {}
    for line in f:
        words = get_words("../data/kpwr-1.2.7-names-disamb-nam-flatten/" + line.decode("utf-8").rstrip('\n'))
        for key in words:
            if key not in book:
                book[key] = words[key]
        print ("Finished crawling " + line.decode("utf-8").rstrip('\n'))
    for word in book:
        # if word != "," and word != "." and word != "(" and word != ")" and word != "?" and word != "-" and word != ":" and word != "′" and word != ">" and word != "<" and word != "_":
        if len(word) > 1 and "." not in word and "-" not in word and "/" not in word:
            if book[word]:
                if word.istitle():
                    file_.write((word + ",1,1\n").encode("utf-8"))
                else:
                    file_.write((word + ",0,1\n").encode("utf-8"))
            else:
                if word.istitle():
                    file_.write((word + ",1,0\n").encode("utf-8"))
                else:
                    file_.write((word + ",0,0\n").encode("utf-8"))
    file_.close()


# -*- coding: utf-8 -*-
"""
Created on Wed Feb 22 17:14:10 2023

@author: Micha
"""

###############################################################################
#Libraries
###############################################################################
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines

###############################################################################
#Data process
###############################################################################
#Read excel spreadsheet
sheet = ("Gibbs (Water) (SMD)", "Gibbs (Ethanol) (SMD)", "Gibbs (Acetonitrile) (SMD)")
for i in sheet:

    df = pd.read_excel("Vanillin Data.xlsx", sheet_name=i, index_col="Functional")

    dfPBE0 = df.iloc[9:16, 1:15]

    #Equation for element-wise multiplication (coefficient * energy of molecule)
    def Gibbs_step(x, y):
        G = x*y
        return sum(G)

    ###########################################################################
    #Databases
    ###########################################################################
    ######
    #PBE0#
    ######
    #Collect each step as a list of coefficients and the energies as a seperate list
    step_1_PBE0 = pd.to_numeric(dfPBE0.iloc[0, :])
    step_2_PBE0 = pd.to_numeric(dfPBE0.iloc[1, :])
    step_3_PBE0 = pd.to_numeric(dfPBE0.iloc[2, :])
    step_4_PBE0 = pd.to_numeric(dfPBE0.iloc[3, :])
    step_5_PBE0 = pd.to_numeric(dfPBE0.iloc[4, :])
    # PBE0energy = pd.to_numeric(dfPBE0.iloc[5, :])
    PBE0energy = pd.to_numeric(dfPBE0.iloc[6, :])

    #Find relative energy level for each reaction step
    PBE0G = Gibbs_step(step_1_PBE0, PBE0energy)
    PBE0G1 = (Gibbs_step(step_1_PBE0, PBE0energy)) - PBE0G
    PBE0G2 = Gibbs_step(step_2_PBE0, PBE0energy) - PBE0G
    PBE0G3 = Gibbs_step(step_3_PBE0, PBE0energy) - PBE0G
    PBE0G4 = Gibbs_step(step_4_PBE0, PBE0energy) - PBE0G
    PBE0G5 = Gibbs_step(step_5_PBE0, PBE0energy) - PBE0G

    #Find energy difference for each reaction step
    PBE0DG1 = round((PBE0G2 - PBE0G1),2)
    PBE0DG2 = round((PBE0G3 - PBE0G2),2)
    PBE0DG3 = round((PBE0G4 - PBE0G3),2)
    PBE0DG4 = round((PBE0G5 - PBE0G4),2)

    if i == "Gibbs (Water) (SMD)":
        PBE0WaterG = PBE0G
        PBE0WaterG1 = PBE0G1
        PBE0WaterG2 = PBE0G2
        PBE0WaterG3 = PBE0G3
        PBE0WaterG4 = PBE0G4
        PBE0WaterG5 = PBE0G5

        PBE0WaterDG1 = PBE0DG1
        PBE0WaterDG2 = PBE0DG2
        PBE0WaterDG3 = PBE0DG3
        PBE0WaterDG4 = PBE0DG4

    elif i == "Gibbs (Ethanol) (SMD)":
        PBE0EthanolG = PBE0G
        PBE0EthanolG1 = PBE0G1
        PBE0EthanolG2 = PBE0G2
        PBE0EthanolG3 = PBE0G3
        PBE0EthanolG4 = PBE0G4
        PBE0EthanolG5 = PBE0G5

        PBE0EthanolDG1 = PBE0DG1
        PBE0EthanolDG2 = PBE0DG2
        PBE0EthanolDG3 = PBE0DG3
        PBE0EthanolDG4 = PBE0DG4

    elif i == "Gibbs (Acetonitrile) (SMD)":
        PBE0AcetonitrileG = PBE0G
        PBE0AcetonitrileG1 = PBE0G1
        PBE0AcetonitrileG2 = PBE0G2
        PBE0AcetonitrileG3 = PBE0G3
        PBE0AcetonitrileG4 = PBE0G4
        PBE0AcetonitrileG5 = PBE0G5

        PBE0AcetonitrileDG1 = PBE0DG1
        PBE0AcetonitrileDG2 = PBE0DG2
        PBE0AcetonitrileDG3 = PBE0DG3
        PBE0AcetonitrileDG4 = PBE0DG4

sheet = ("Gibbs (Water)", "Gibbs (Ethanol)", "Gibbs (Acetonitrile)")
for i in sheet:

    df = pd.read_excel("Z:/Michael Nicolaou/Projects/Vanillin-veratric acid/Workbench.xlsx", sheet_name=i, index_col="Functional")

    dfPBE0 = df.iloc[9:16, 1:15]

    #Equation for element-wise multiplication (coefficient * energy of molecule)
    def Gibbs_step(x, y):
        G = x*y
        return sum(G)

    ###########################################################################
    #Databases
    ###########################################################################
    ######
    #PBE0#
    ######
    #Collect each step as a list of coefficients and the energies as a seperate list
    step_1_PBE0 = pd.to_numeric(dfPBE0.iloc[0, :])
    step_2_PBE0 = pd.to_numeric(dfPBE0.iloc[1, :])
    step_3_PBE0 = pd.to_numeric(dfPBE0.iloc[2, :])
    step_4_PBE0 = pd.to_numeric(dfPBE0.iloc[3, :])
    step_5_PBE0 = pd.to_numeric(dfPBE0.iloc[4, :])
    # PBE0energy = pd.to_numeric(dfPBE0.iloc[5, :])
    PBE0energy = pd.to_numeric(dfPBE0.iloc[6, :])

    #Find relative energy level for each reaction step
    if i == "Gibbs (Water)":
        PBE0G = Gibbs_step(step_1_PBE0, PBE0energy)
        PBE0G1 = (Gibbs_step(step_1_PBE0, PBE0energy)) - PBE0WaterG
        PBE0G2 = Gibbs_step(step_2_PBE0, PBE0energy) - PBE0WaterG
        PBE0G3 = Gibbs_step(step_3_PBE0, PBE0energy) - PBE0WaterG
        PBE0G4 = Gibbs_step(step_4_PBE0, PBE0energy) - PBE0WaterG
        PBE0G5 = Gibbs_step(step_5_PBE0, PBE0energy) - PBE0WaterG
    elif i == "Gibbs (Ethanol)":
        PBE0G = Gibbs_step(step_1_PBE0, PBE0energy)
        PBE0G1 = (Gibbs_step(step_1_PBE0, PBE0energy)) - PBE0EthanolG
        PBE0G2 = Gibbs_step(step_2_PBE0, PBE0energy) - PBE0EthanolG
        PBE0G3 = Gibbs_step(step_3_PBE0, PBE0energy) - PBE0EthanolG
        PBE0G4 = Gibbs_step(step_4_PBE0, PBE0energy) - PBE0EthanolG
        PBE0G5 = Gibbs_step(step_5_PBE0, PBE0energy) - PBE0EthanolG
    elif i == "Gibbs (Acetonitrile)":
        PBE0G = Gibbs_step(step_1_PBE0, PBE0energy)
        PBE0G1 = (Gibbs_step(step_1_PBE0, PBE0energy)) - PBE0AcetonitrileG
        PBE0G2 = Gibbs_step(step_2_PBE0, PBE0energy) - PBE0AcetonitrileG
        PBE0G3 = Gibbs_step(step_3_PBE0, PBE0energy) - PBE0AcetonitrileG
        PBE0G4 = Gibbs_step(step_4_PBE0, PBE0energy) - PBE0AcetonitrileG
        PBE0G5 = Gibbs_step(step_5_PBE0, PBE0energy) - PBE0AcetonitrileG

    #Find energy difference for each reaction step
    PBE0DG1 = round((PBE0G2 - PBE0G1),2)
    PBE0DG2 = round((PBE0G3 - PBE0G2),2)
    PBE0DG3 = round((PBE0G4 - PBE0G3),2)
    PBE0DG4 = round((PBE0G5 - PBE0G4),2)

    if i == "Gibbs (Water)":
        PBE0WaterPCMG1 = PBE0G1
        PBE0WaterPCMG2 = PBE0G2
        PBE0WaterPCMG3 = PBE0G3
        PBE0WaterPCMG4 = PBE0G4
        PBE0WaterPCMG5 = PBE0G5

        PBE0WaterPCMDG1 = PBE0DG1
        PBE0WaterPCMDG2 = PBE0DG2
        PBE0WaterPCMDG3 = PBE0DG3
        PBE0WaterPCMDG4 = PBE0DG4

    elif i == "Gibbs (Ethanol)":
        PBE0EthanolPCMG1 = PBE0G1
        PBE0EthanolPCMG2 = PBE0G2
        PBE0EthanolPCMG3 = PBE0G3
        PBE0EthanolPCMG4 = PBE0G4
        PBE0EthanolPCMG5 = PBE0G5

        PBE0EthanolPCMDG1 = PBE0DG1
        PBE0EthanolPCMDG2 = PBE0DG2
        PBE0EthanolPCMDG3 = PBE0DG3
        PBE0EthanolPCMDG4 = PBE0DG4

    elif i == "Gibbs (Acetonitrile)":
        PBE0AcetonitrilePCMG1 = PBE0G1
        PBE0AcetonitrilePCMG2 = PBE0G2
        PBE0AcetonitrilePCMG3 = PBE0G3
        PBE0AcetonitrilePCMG4 = PBE0G4
        PBE0AcetonitrilePCMG5 = PBE0G5

        PBE0AcetonitrilePCMDG1 = PBE0DG1
        PBE0AcetonitrilePCMDG2 = PBE0DG2
        PBE0AcetonitrilePCMDG3 = PBE0DG3
        PBE0AcetonitrilePCMDG4 = PBE0DG4

###############################################################################
#Plots
###############################################################################
#Define axes
xaxis = (1, 2, 3, 4, 5)
PBE0Wateryaxis = (PBE0WaterG1, PBE0WaterG2, PBE0WaterG3, PBE0WaterG4, PBE0WaterG5)
PBE0Ethanolyaxis = (PBE0EthanolG1, PBE0EthanolG2, PBE0EthanolG3, PBE0EthanolG4, PBE0EthanolG5)
PBE0Acetonitrileyaxis = (PBE0AcetonitrileG1, PBE0AcetonitrileG2, PBE0AcetonitrileG3, PBE0AcetonitrileG4, PBE0AcetonitrileG5)
PBE0WaterPCMyaxis = (PBE0WaterPCMG1, PBE0WaterPCMG2, PBE0WaterPCMG3, PBE0WaterPCMG4, PBE0WaterPCMG5)
PBE0EthanolPCMyaxis = (PBE0EthanolPCMG1, PBE0EthanolPCMG2, PBE0EthanolPCMG3, PBE0EthanolPCMG4, PBE0EthanolPCMG5)
PBE0AcetonitrilePCMyaxis = (PBE0AcetonitrilePCMG1, PBE0AcetonitrilePCMG2, PBE0AcetonitrilePCMG3, PBE0AcetonitrilePCMG4, PBE0AcetonitrilePCMG5)

#Plot figure
#Add plot
fig, ax = plt.subplots(figsize = (20,10))

#Plot font
plt.rcParams['font.family'] = "serif"
plt.rcParams.update({'font.size': 18})

#Remove main plot default axis ticks
plt.yticks([])
plt.xticks([])

#Remove main plot frame
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.spines['left'].set_visible(False)

#Adjust subplot positions
# plt.subplots_adjust(bottom = -0.2, left = 0, top = 1., right = 1)
sub1 = fig.add_subplot(3, 3, (1, 3))
sub2 = fig.add_subplot(3, 3, (4, 6))
sub3 = fig.add_subplot(3, 3, (7, 9))

#Plot for each step's energy level
######
#BP86#
######
#Scatterplot
sub1.scatter(xaxis, PBE0Wateryaxis, s = 2000, c = "blue", marker = "_", linewidths = 3)
sub1.scatter(xaxis, PBE0WaterPCMyaxis, s = 2000, c = "cyan", marker = "_", linewidths = 3)

sub2.scatter(xaxis, PBE0Ethanolyaxis, s = 2000, c = "blue", marker = "_", linewidths = 3)
sub2.scatter(xaxis, PBE0EthanolPCMyaxis, s = 2000, c = "cyan", marker = "_", linewidths = 3)

sub3.scatter(xaxis, PBE0Acetonitrileyaxis, s = 2000, c = "blue", marker = "_", linewidths = 3)
sub3.scatter(xaxis, PBE0AcetonitrilePCMyaxis, s = 2000, c = "cyan", marker = "_", linewidths = 3)

#Energy change line
sub1.plot((1.08, 1.92), (PBE0WaterG1, PBE0WaterG2), c = "blue", linestyle = "dashed")
sub1.plot((2.08, 2.92), (PBE0WaterG2, PBE0WaterG3), c = "blue", linestyle = "dashed")
sub1.plot((3.08, 3.92), (PBE0WaterG3, PBE0WaterG4), c = "blue", linestyle = "dashed")
sub1.plot((4.08, 4.92), (PBE0WaterG4, PBE0WaterG5), c = "blue", linestyle = "dashed")

sub1.plot((1.08, 1.92), (PBE0WaterPCMG1, PBE0WaterPCMG2), c = "cyan", linestyle = "dashed")
sub1.plot((2.08, 2.92), (PBE0WaterPCMG2, PBE0WaterPCMG3), c = "cyan", linestyle = "dashed")
sub1.plot((3.08, 3.92), (PBE0WaterPCMG3, PBE0WaterPCMG4), c = "cyan", linestyle = "dashed")
sub1.plot((4.08, 4.92), (PBE0WaterPCMG4, PBE0WaterPCMG5), c = "cyan", linestyle = "dashed")

sub2.plot((1.08, 1.92), (PBE0EthanolG1, PBE0EthanolG2), c = "blue", linestyle = "dashed")
sub2.plot((2.08, 2.92), (PBE0EthanolG2, PBE0EthanolG3), c = "blue", linestyle = "dashed")
sub2.plot((3.08, 3.92), (PBE0EthanolG3, PBE0EthanolG4), c = "blue", linestyle = "dashed")
sub2.plot((4.08, 4.92), (PBE0EthanolG4, PBE0EthanolG5), c = "blue", linestyle = "dashed")

sub2.plot((1.08, 1.92), (PBE0EthanolPCMG1, PBE0EthanolPCMG2), c = "cyan", linestyle = "dashed")
sub2.plot((2.08, 2.92), (PBE0EthanolPCMG2, PBE0EthanolPCMG3), c = "cyan", linestyle = "dashed")
sub2.plot((3.08, 3.92), (PBE0EthanolPCMG3, PBE0EthanolPCMG4), c = "cyan", linestyle = "dashed")
sub2.plot((4.08, 4.92), (PBE0EthanolPCMG4, PBE0EthanolPCMG5), c = "cyan", linestyle = "dashed")

sub3.plot((1.08, 1.92), (PBE0AcetonitrileG1, PBE0AcetonitrileG2), c = "blue", linestyle = "dashed")
sub3.plot((2.08, 2.92), (PBE0AcetonitrileG2, PBE0AcetonitrileG3), c = "blue", linestyle = "dashed")
sub3.plot((3.08, 3.92), (PBE0AcetonitrileG3, PBE0AcetonitrileG4), c = "blue", linestyle = "dashed")
sub3.plot((4.08, 4.92), (PBE0AcetonitrileG4, PBE0AcetonitrileG5), c = "blue", linestyle = "dashed")

sub3.plot((1.08, 1.92), (PBE0AcetonitrilePCMG1, PBE0AcetonitrilePCMG2), c = "cyan", linestyle = "dashed")
sub3.plot((2.08, 2.92), (PBE0AcetonitrilePCMG2, PBE0AcetonitrilePCMG3), c = "cyan", linestyle = "dashed")
sub3.plot((3.08, 3.92), (PBE0AcetonitrilePCMG3, PBE0AcetonitrilePCMG4), c = "cyan", linestyle = "dashed")
sub3.plot((4.08, 4.92), (PBE0AcetonitrilePCMG4, PBE0AcetonitrilePCMG5), c = "cyan", linestyle = "dashed")

#Plot font
plt.rcParams['font.family'] = "serif"
plt.rcParams.update({'font.size': 18})

#Axis labels
sub1.set_ylabel("Gibbs free energy (kcal mol$^{-1}$)", fontsize = 12)
sub2.set_ylabel("Gibbs free energy (kcal mol$^{-1}$)", fontsize = 12)
sub3.set_ylabel("Gibbs free energy (kcal mol$^{-1}$)", fontsize = 12)

#Scatter labels
sub1.text(0.95, (float(PBE0WaterPCMG1)+3), round(float(PBE0WaterPCMG1),1), color = "cyan", fontsize = 12)
sub1.text(1.945, (float(PBE0WaterPCMG2)+3), round(float(PBE0WaterPCMG2),1), color = "cyan", fontsize = 12)
sub1.text(2.965, (float(PBE0WaterPCMG3)+3), round(float(PBE0WaterPCMG3),1), color = "cyan", fontsize = 12)
sub1.text(3.94, (float(PBE0WaterPCMG4)+3), round(float(PBE0WaterPCMG4),1), color = "cyan", fontsize = 12)
sub1.text(4.94, (float(PBE0WaterPCMG5)+3), round(float(PBE0WaterPCMG5),1), color = "cyan", fontsize = 12)

sub1.text(0.96, (float(PBE0WaterG1)-9.5), round(float(PBE0WaterG1),1), color = "blue", fontsize = 12)
sub1.text(1.96, (float(PBE0WaterG2)-9.5), round(float(PBE0WaterG2),1), color = "blue", fontsize = 12)
sub1.text(2.95, (float(PBE0WaterG3)-9.5), round(float(PBE0WaterG3),1), color = "blue", fontsize = 12)
sub1.text(3.94, (float(PBE0WaterG4)-9.5), round(float(PBE0WaterG4),1), color = "blue", fontsize = 12)
sub1.text(4.94, (float(PBE0WaterG5)-9.5), round(float(PBE0WaterG5),1), color = "blue", fontsize = 12)

sub2.text(0.95, (float(PBE0EthanolPCMG1)+3), round(float(PBE0EthanolPCMG1),1), color = "cyan", fontsize = 12)
sub2.text(1.945, (float(PBE0EthanolPCMG2)+3), round(float(PBE0EthanolPCMG2),1), color = "cyan", fontsize = 12)
sub2.text(2.945, (float(PBE0EthanolPCMG3)+3), round(float(PBE0EthanolPCMG3),1), color = "cyan", fontsize = 12)
sub2.text(3.94, (float(PBE0EthanolPCMG4)+3), round(float(PBE0EthanolPCMG4),1), color = "cyan", fontsize = 12)
sub2.text(4.94, (float(PBE0EthanolPCMG5)+3), round(float(PBE0EthanolPCMG5),1), color = "cyan", fontsize = 12)

sub2.text(0.96, (float(PBE0EthanolG1)-9.5), round(float(PBE0EthanolG1),1), color = "blue", fontsize = 12)
sub2.text(1.96, (float(PBE0EthanolG2)-9.5), round(float(PBE0EthanolG2),1), color = "blue", fontsize = 12)
sub2.text(2.95, (float(PBE0EthanolG3)-9.5), round(float(PBE0EthanolG3),1), color = "blue", fontsize = 12)
sub2.text(3.94, (float(PBE0EthanolG4)-9.5), round(float(PBE0EthanolG4),1), color = "blue", fontsize = 12)
sub2.text(4.94, (float(PBE0EthanolG5)-9.5), round(float(PBE0EthanolG5),1), color = "blue", fontsize = 12)

sub3.text(0.95, (float(PBE0AcetonitrilePCMG1)+3), round(float(PBE0AcetonitrilePCMG1),1), color = "cyan", fontsize = 12)
sub3.text(1.945, (float(PBE0AcetonitrilePCMG2)+3), round(float(PBE0AcetonitrilePCMG2),1), color = "cyan", fontsize = 12)
sub3.text(2.967, (float(PBE0AcetonitrilePCMG3)+3), round(float(PBE0AcetonitrilePCMG3),1), color = "cyan", fontsize = 12)
sub3.text(3.94, (float(PBE0AcetonitrilePCMG4)+3), round(float(PBE0AcetonitrilePCMG4),1), color = "cyan", fontsize = 12)
sub3.text(4.94, (float(PBE0AcetonitrilePCMG5)+3), round(float(PBE0AcetonitrilePCMG5),1), color = "cyan", fontsize = 12)

sub3.text(0.96, (float(PBE0AcetonitrileG1)-9.5), round(float(PBE0AcetonitrileG1),1), color = "blue", fontsize = 12)
sub3.text(1.96, (float(PBE0AcetonitrileG2)-9.5), round(float(PBE0AcetonitrileG2),1), color = "blue", fontsize = 12)
sub3.text(2.95, (float(PBE0AcetonitrileG3)-9.5), round(float(PBE0AcetonitrileG3),1), color = "blue", fontsize = 12)
sub3.text(3.94, (float(PBE0AcetonitrileG4)-9.5), round(float(PBE0AcetonitrileG4),1), color = "blue", fontsize = 12)
sub3.text(4.94, (float(PBE0AcetonitrileG5)-9.5), round(float(PBE0AcetonitrileG5),1), color = "blue", fontsize = 12)

#Ticks
sub1.set_xticks([])
sub2.set_xticks([])
sub3.set_xticks([])
#Limits
sub1.set_ylim([-90, 30])
sub2.set_ylim([-90, 30])
sub3.set_ylim([-90, 30])

#Ticks
sub1.set_yticks([20, 10, 0, -10, -20, -30, -40, -50, -60, -70 , -80])
sub2.set_yticks([20, 10, 0, -10, -20, -30, -40, -50, -60, -70 , -80])
sub3.set_yticks([20, 10, 0, -10, -20, -30, -40, -50, -60, -70 , -80])

#Plot titles
sub1.text(0.84, -82.5, "Water", bbox = {'facecolor': "white", "pad": 4}, fontsize = 20)
sub2.text(0.84, -82.5, "Ethanol", bbox = {'facecolor': "white", "pad": 4}, fontsize = 20)
sub3.text(0.84, -82.5, "Acetonitrile", bbox = {'facecolor': "white", "pad": 4}, fontsize = 20)

#Colour parts of plots
sub1.fill_between((1, 2), -100, 30, facecolor = "green", alpha = 0.1)
sub1.fill_between((3, 4), -100, 30, facecolor = "blue", alpha = 0.1)
sub1.fill_between((4, 5), -100, 30, facecolor = "orange", alpha = 0.1)

sub2.fill_between((1, 2), -100, 30, facecolor = "green", alpha = 0.1)
sub2.fill_between((3, 4), -100, 30, facecolor = "blue", alpha = 0.1)
sub2.fill_between((4, 5), -100, 30, facecolor = "orange", alpha = 0.1)

sub3.fill_between((1, 2), -100, 30, facecolor = "green", alpha = 0.1)
sub3.fill_between((3, 4), -100, 30, facecolor = "blue", alpha = 0.1)
sub3.fill_between((4, 5), -100, 30, facecolor = "orange", alpha = 0.1)

#Legend
PBE0Waterleg = mlines.Line2D([], [], color = "blue", linestyle = "dashed", label = "SMD")
PBE0Ethanolleg = mlines.Line2D([], [], color = "blue", linestyle = "dashed", label = "SMD")
PBE0Acetonitrileleg = mlines.Line2D([], [], color = "blue", linestyle = "dashed", label = "SMD")
PBE0WaterPCMleg = mlines.Line2D([], [], color = "cyan", linestyle = "dashed", label = "PCM")
PBE0EthanolPCMleg = mlines.Line2D([], [], color = "cyan", linestyle = "dashed", label = "PCM")
PBE0AcetonitrilePCMleg = mlines.Line2D([], [], color = "cyan", linestyle = "dashed", label = "PCM")
Waterlegends = (PBE0Waterleg, PBE0WaterPCMleg)
Ethanollegends = (PBE0Ethanolleg, PBE0EthanolPCMleg)
Acetonitrilelegends = (PBE0Acetonitrileleg, PBE0AcetonitrilePCMleg)
sub1.legend(handles=Waterlegends, loc="upper right")
sub2.legend(handles=Ethanollegends, loc="upper right")
sub3.legend(handles=Acetonitrilelegends, loc="upper right")

plt.savefig("Solvation models PBE0.png", format = "png", bbox_inches='tight', dpi = 300)
plt.show()

Vanillin analysis/Figure (Energy profile - Solvation model - PBE0).py


# -*- coding: utf-8 -*-
import argparse
from datetime import datetime

import pandas as pd

from zvt import init_log, zvt_config
from zvt.api.data_type import Region, Provider, EntityType
from zvt.api.quote import get_kdata_schema, get_kdata
from zvt.domain import Stock, StockKdataCommon, Stock1dHfqKdata
from zvt.contract import IntervalLevel, AdjustType
from zvt.contract.recorder import FixedCycleDataRecorder
from zvt.recorders.joinquant.common import to_jq_trading_level, to_jq_entity_id
from zvt.database.api import get_db_session
from zvt.networking.request import jq_get_token, jq_get_bars
from zvt.utils.pd_utils import pd_is_not_null
from zvt.utils.time_utils import to_time_str, now_pd_timestamp, PD_TIME_FORMAT_DAY, PD_TIME_FORMAT_ISO8601

class JqChinaStockKdataRecorder(FixedCycleDataRecorder):
    # 数据来自jq
    region = Region.CHN
    provider = Provider.JoinQuant
    entity_schema = Stock
    # 只是为了把recorder注册到data_schema
    data_schema = StockKdataCommon

    def __init__(self,
                 exchanges=['sh', 'sz'],
                 entity_ids=None,
                 codes=None,
                 batch_size=10,
                 force_update=True,
                 sleeping_time=0,
                 default_size=zvt_config['batch_size'],
                 real_time=False,
                 fix_duplicate_way='ignore',
                 start_timestamp=None,
                 end_timestamp=None,
                 level=IntervalLevel.LEVEL_1WEEK,
                 kdata_use_begin_time=False,
                 close_hour=15,
                 close_minute=0,
                 one_day_trading_minutes=4 * 60,
                 adjust_type=AdjustType.qfq,
                 share_para=None) -> None:
        level = IntervalLevel(level)
        adjust_type = AdjustType(adjust_type)
        self.data_schema = get_kdata_schema(entity_type=EntityType.Stock, level=level, adjust_type=adjust_type)
        self.jq_trading_level = to_jq_trading_level(level)

        super().__init__(EntityType.Stock, exchanges, entity_ids, codes, batch_size, force_update, sleeping_time,
                         default_size, real_time, fix_duplicate_way, start_timestamp, end_timestamp, close_hour,
                         close_minute, level, kdata_use_begin_time, one_day_trading_minutes, share_para=share_para)
        self.adjust_type = adjust_type

        jq_get_token(zvt_config['jq_username'], zvt_config['jq_password'], force=True)

    def recompute_qfq(self, entity, qfq_factor, last_timestamp):
        # 重新计算前复权数据
        if qfq_factor != 0:
            kdatas = get_kdata(region=self.region,
                               provider=self.provider,
                               entity_id=entity.id,
                               level=self.level.value,
                               order=self.data_schema.timestamp.asc(),
                               return_type='domain',
                               filters=[self.data_schema.timestamp < last_timestamp])
            if kdatas:
                self.logger.info('recomputing {} qfq kdata,factor is:{}'.format(entity.code, qfq_factor))
                for kdata in kdatas:
                    kdata.open = round(kdata.open * qfq_factor, 2)
                    kdata.close = round(kdata.close * qfq_factor, 2)
                    kdata.high = round(kdata.high * qfq_factor, 2)
                    kdata.low = round(kdata.low * qfq_factor, 2)
                session = get_db_session(region=self.region,
                                         provider=self.provider,
                                         data_schema=self.data_schema)
                session.bulk_save_objects(kdatas)
                session.commit()

    def generate_domain_id(self, entity, df, time_fmt=PD_TIME_FORMAT_DAY):
        format = PD_TIME_FORMAT_DAY if self.level >= IntervalLevel.LEVEL_1DAY else PD_TIME_FORMAT_ISO8601
        return entity.id + '_' + df[self.get_evaluated_time_field()].dt.strftime(format)

    def record(self, entity, start, end, size, timestamps, http_session):
        if self.adjust_type == AdjustType.hfq:
            fq_ref_date = '2000-01-01'
        else:
            fq_ref_date = to_time_str(now_pd_timestamp(Region.CHN))

        if not self.end_timestamp:
            df = jq_get_bars(to_jq_entity_id(entity),
                             count=size,
                             unit=self.jq_trading_level,
                             # fields=['date', 'open', 'close', 'low', 'high', 'volume', 'money'],
                             fq_ref_date=fq_ref_date)
        else:
            end_timestamp = to_time_str(self.end_timestamp)
            df = jq_get_bars(to_jq_entity_id(entity),
                             count=size,
                             unit=self.jq_trading_level,
                             # fields=['date', 'open', 'close', 'low', 'high', 'volume', 'money'],
                             end_date=end_timestamp,
                             fq_ref_date=fq_ref_date)
        # self.logger.info("record {} for {}, size:{}".format(self.data_schema.__name__, entity.id, len(df)))

        if pd_is_not_null(df):
            # start_timestamp = to_time_str(df.iloc[1]['timestamp'])
            # end_timestamp = to_time_str(df.iloc[-1]['timestamp'])

            # 判断是否需要重新计算之前保存的前复权数据
            if self.adjust_type == AdjustType.qfq:
                check_df = df.head(1)
                check_date = check_df['timestamp'][0]
                current_df = get_kdata(region=self.region,
                                       entity_id=entity.id,
                                       provider=self.provider,
                                       start_timestamp=check_date,
                                       end_timestamp=check_date,
                                       limit=1,
                                       level=self.level,
                                       adjust_type=self.adjust_type)
                if pd_is_not_null(current_df):
                    old = current_df.iloc[0, :]['close']
                    new = check_df['close'][0]
                    # 相同时间的close不同，表明前复权需要重新计算
                    if round(old, 2) != round(new, 2):
                        qfq_factor = new / old
                        last_timestamp = pd.Timestamp(check_date)
                        self.recompute_qfq(entity, qfq_factor=qfq_factor, last_timestamp=last_timestamp)
            return df

        return None

    def format(self, entity, df):
        df.rename(columns={'money': 'turnover', 'date': 'timestamp'}, inplace=True)

        if 'timestamp' not in df.columns:
            df['timestamp'] = pd.to_datetime(df[self.get_original_time_field()])
        elif not isinstance(df['timestamp'].dtypes, datetime):
            df['timestamp'] = pd.to_datetime(df['timestamp'])

        # df.update(df.select_dtypes(include=[np.number]).fillna(0))

        df['entity_id'] = entity.id
        df['provider'] = self.provider.value
        df['code'] = entity.code
        df['name'] = entity.name
        df['level'] = self.level.value

        df['id'] = self.generate_domain_id(entity, df)
        return df

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--level', help='trading level', default='1d', choices=[item.value for item in IntervalLevel])
    parser.add_argument('--codes', help='codes', default=['000001'], nargs='+')

    args = parser.parse_args()

    level = IntervalLevel(args.level)
    codes = args.codes

    init_log('jq_china_stock_{}_kdata.log'.format(args.level))
    JqChinaStockKdataRecorder(level=level, sleeping_time=0, codes=codes, real_time=False,
                              adjust_type=AdjustType.hfq).run()

    print(get_kdata(region=Region.CHN, entity_id='stock_sz_000001', limit=10, order=Stock1dHfqKdata.timestamp.desc(),
                    adjust_type=AdjustType.hfq))

# the __all__ is generated
__all__ = ['JqChinaStockKdataRecorder']


import os, sys

sys.path.insert(0, os.getcwd())
from src.client.comet import add_metrics
from src.visualizations.visualization import plots
from src.client import start_experiment
from src.features import load_df_shots
from sklearn.model_selection import train_test_split
import pandas as pd
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score, roc_auc_score, precision_score, recall_score

df_2016 = load_df_shots(2016)
df_2017 = load_df_shots(2017)
df_2018 = load_df_shots(2018)
df_2019 = load_df_shots(2019)

# len(pd.concat([df_2016, df_2017, df_2018]))

df = pd.concat([df_2016, df_2017, df_2018]).reset_index(drop=True)
df.Shot_angle = df.Shot_angle.abs()
df_2019.Shot_angle = df_2019.Shot_angle.abs()

# print(len(pd.concat([df_2016, df_2017, df_2018])) / len(df))

# X_train, X_val, y_train, y_val = train_test_split(df.Shot_distance, df.Goal, test_size=0.23, random_state=42)

X_train = df[['Shot_distance','Shot_angle']]
# X_train = df[['Shot_angle']]
y_train = df.Goal
X_val = df_2019[['Shot_distance','Shot_angle']]
# X_val = df_2019[['Shot_angle']]
y_val = df_2019.Goal
val_labels = y_val.values.reshape(-1,1)

bst = XGBClassifier()
bst.fit(X_train, y_train)

preds = bst.predict(X_val)
probs = bst.predict_proba(X_val)[:,1]

# accuracy_score(y_val, preds)
# roc_auc_score(y_val, probs)
# precision_score(y_val, preds)
# recall_score(y_val, preds)

df = pd.concat([df_2016, df_2017, df_2018, df_2019]).reset_index(drop=True)
X_train = df[['Shot_distance','Shot_angle']].abs()
y_train = df.Goal

bst = XGBClassifier()
bst.fit(X_train, y_train)
bst.save_model('model/xgb.model.json')

# print(y_val.values)

exp = start_experiment()
exp.set_name(f'xgd-distance-angle-{exp.get_name()}')
add_metrics(
        exp,
        accuracy_score(y_val, preds),
        roc_auc_score(y_val, probs),
        precision_score(y_val, preds),
        recall_score(y_val, preds))
exp.log_model("xgb", "model/xgb.model.json")
exp.add_tags(['XGB', 'Distance', 'Angle'])
plots(val_labels, probs, f'xgd-distance-angle', exp)

# plots(val_labels, probs, f'xgd-distance')


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# This script is used for testing on an arm machine.
# Please change the constant values according to your environment.
import os
import subprocess
import sys
import time
from glob import glob

from parse import parse

OK = '\033[92m'
FAIL = '\033[91m'
END = '\033[0m'

TIMEOUT = 100
ROUND = 1

scenes = ['function_test2020', 'function_test2021',
          'performance_test2021-public', 'performance_test2021-private']
scene = scenes[int(sys.argv[1])]

testcase = os.path.join('/home/pi/compiler2021/公开用例与运行时库', scene)
compiler = '/home/pi/kobayashi-compiler/build/compiler'
runtime = '/home/pi/kobayashi-compiler/runtime/sylib.c /home/pi/kobayashi-compiler/runtime/armv7/thread.S'
cc = 'gcc -march=armv7ve'

try:
    testcases = sorted([os.path.splitext(os.path.basename(file))[0]
                        for file in glob(os.path.join(testcase, '*.sy'))], key=lambda x: int(x.split('_')[0]))
except:
    testcases = sorted(os.path.splitext(os.path.basename(file))[
                       0] for file in glob(os.path.join(testcase, '*.sy')))

for base in testcases:
    fails = list()

    # names
    try:
        case = base.split('_')
        case[0] = str(int(case[0])).zfill(3)
        case = '_'.join(case)
    except:
        case = base
    print(f'{case}  ', end='', flush=True)

    sy_file = os.path.join(testcase, f'{base}.sy')
    in_file = os.path.join(testcase, f'{base}.in')
    ans_file = os.path.join(testcase, f'{base}.out')

    # compile
    os.system(
        ' '.join([compiler, sy_file, '-o', 'asm.s'] + sys.argv[2:]))
    os.system(f'{cc} asm.s {runtime} -o main')
    os.system(f'touch {in_file}')

    # time
    avg = 0
    for i in range(ROUND):
        print(f'\b{i}', end='', flush=True)
        p = subprocess.Popen(f'./main', stdin=open(in_file),
                             stdout=open('out.txt', 'w'), stderr=open('time.txt', 'w'))
        try:
            p.wait(TIMEOUT)
        except subprocess.TimeoutExpired:
            p.kill()
        ret = p.returncode
        with open('time.txt') as f:
            timing = list(filter(None, f.read().strip().split('\n')))
        if timing:
            h, m, s, us = map(int, parse(
                'TOTAL: {}H-{}M-{}S-{}us', timing[-1]))
            timing = ((h * 60 + m) * 60 + s) * 1000000 + us
            timing /= 1000000
        else:
            timing = 0.0
        avg += timing
    avg /= ROUND

    # judge
    with open('out.txt') as f:
        out = f.read().strip()
    with open(ans_file) as f:
        ans = f.read().strip().split('\n')
    ans_out, ans_ret = '\n'.join(ans[:-1]).strip(), int(ans[-1])
    if ans_out == out and ans_ret == ret:
        print(f'\r[{OK}OK{END}]      {round(avg, 4):7.3f} | {case}')
    else:
        print(f'\r[{FAIL}FAILED{END}]  {round(avg, 4):7.3f} | {case}')
        fails.append(case)
    os.remove('main')

if not fails:
    print(f'All {scene} passed!')
else:
    for fail in fails:
        print(f'{FAIL}{fail}{END} failed!')


import uuid
from enum import Enum
from http import HTTPStatus
from typing import (
    TYPE_CHECKING,
    Annotated,
    Any,
    Dict,
    Iterable,
    List,
    Optional,
    Type,
    Union,
)

import pytest
from fastapi import APIRouter, FastAPI, Query
from fastapi.params import Depends
from pydantic import BaseModel, Field
from pydantic.main import ModelMetaclass
from starlette.testclient import TestClient

from furiousapi.core.api.controllers import CBV, ModelController, action
from furiousapi.core.db.fields import SortableFieldEnum
from furiousapi.core.db.models import FuriousPydanticConfig
from furiousapi.core.db.repository import BaseRepository, RepositoryConfig
from furiousapi.core.pagination import AllPaginationStrategies, PaginatedResponse
from furiousapi.core.responses import BulkResponseModel
from furiousapi.core.types import TEntity, TModelFields

if TYPE_CHECKING:
    from fastapi.dependencies.models import Dependant
    from fastapi.routing import APIRoute

class MyCBV(CBV):
    @action("/endpoint1")
    def endpoint1(self, query_param: str = Query(...)):
        pass

    @action("/endpoint2")
    def endpoint2(self, query_param: str = Query(...)):
        pass

def test_cbv__when_no_api_router__then_initialize():
    assert isinstance(MyCBV.api_router, APIRouter)

static_api_router = APIRouter()

class MyCBV2(MyCBV):
    my_dep = Depends(lambda x: x)
    api_router = static_api_router

def test_cbv__when_api_router__then_use_it():
    assert id(MyCBV2.api_router) == id(static_api_router)

def test_cbv__actions_are_defined():
    my = MyCBV()
    assert len(MyCBV.api_router.routes) == 2  # noqa: PLR2004
    route1: APIRoute = my.api_router.routes[0]
    route2: APIRoute = my.api_router.routes[1]
    assert route1.path == "/endpoint1"
    assert route2.path == "/endpoint2"
    dependant: Dependant = route1.dependant

    assert dependant.query_params[0].type_ == str
    assert dependant.query_params[0].name == "query_param"
    assert dependant.query_params[0].required is True

def test_cbv__when_class_dependency_defined_and_not_passed__then_raise_type_error():
    with pytest.raises(TypeError):
        MyCBV2()

def test_cbv__when_class_dependency_defined__then_init_changed():
    try:
        MyCBV2(my_dep=Depends(repository_dependency))  # type: ignore[call-arg]
    except Exception:  # noqa: BLE001
        pytest.fail("something happened")

class MyModel(BaseModel):
    id: Optional[str] = Field(alias="_id")
    my_param: str

    class Config(FuriousPydanticConfig):
        pass

# TODO:
# mypy issue with Generic[TEntity]
class InMemoryDBRepository(BaseRepository[TEntity]):  # type: ignore[misc]
    def __init__(self) -> None:
        self._store: Dict[Union[str, int, Dict[str, Any], tuple], TEntity] = {}

    async def get(
        self,
        identifiers: Union[int, str, Dict[str, Any], tuple],
        fields: Optional[Iterable[Enum]] = None,  # noqa: ARG002
        *,
        should_error: bool = True,  # noqa: ARG002
    ) -> Optional[TEntity]:
        return self._store.get(identifiers)

    async def list(
        self,
        pagination: AllPaginationStrategies,  # noqa: ARG002
        fields: Optional[Iterable[TModelFields]] = None,  # noqa: ARG002
        sorting: Optional[List[SortableFieldEnum]] = None,  # noqa: ARG002
        filtering: Optional[TEntity] = None,  # noqa: ARG002
    ) -> PaginatedResponse[TEntity]:
        return PaginatedResponse[TEntity](total=len(self._store), items=list(self._store.values()), index=0, next=None)

    async def add(self, entity: TEntity) -> TEntity:

        if entity.id in self._store:  # type: ignore[attr-defined]
            raise ValueError(f"Key {entity.id} already exists")  # type: ignore[attr-defined]
        entity.id = str(uuid.uuid4())  # type: ignore[attr-defined]
        self._store[entity.id] = entity  # type: ignore[attr-defined]
        return entity

    async def update(self, entity: TEntity, **kwargs) -> Optional[TEntity]:
        if entity.id not in self._store:  # type: ignore[attr-defined]
            raise KeyError(f"Key {entity.id} does not exist")  # type: ignore[attr-defined]
        self._store[entity.id] = entity  # type: ignore[attr-defined]
        return entity

    async def delete(self, entity: Union[TEntity, str, int], **kwargs) -> None:

        if entity.id not in self._store:  # type: ignore[union-attr]
            raise KeyError(f"Key {entity.id} does not exist")  # type: ignore[union-attr]
        del self._store[entity.id]  # type: ignore[union-attr]

    async def bulk_create(self, bulk: List[TEntity]) -> BulkResponseModel:  # type: ignore[empty-body]
        pass

    async def bulk_delete(self, bulk: List[Union[TEntity, Any]]) -> List:  # type: ignore[empty-body]
        pass

    async def bulk_update(self, bulk: List[TEntity]) -> List:  # type: ignore[empty-body]
        pass

class MyRepository(InMemoryDBRepository[MyModel]):  # type: ignore[type-arg]
    class Config(RepositoryConfig):
        @staticmethod
        def model_to_query(x: Type[MyModel]) -> Type[MyModel]:
            return x

        filter_model = ModelMetaclass

repository = MyRepository()

def repository_dependency() -> MyRepository:
    return repository

class MyController(ModelController):
    repository: Depends = Depends(repository_dependency, use_cache=True)

class MyController2(ModelController):
    repository: Annotated[MyRepository, Depends] = Depends(repository_dependency)  # type: ignore[assignment]
    __enabled_routes__ = ("get",)

@pytest.fixture()
def controller() -> MyController:
    return MyController(repository=MyRepository())

def test_controller_repository__is_initialized(controller: MyController):
    assert isinstance(controller.repository, MyRepository)

def test_create(controller: MyController):
    create: APIRoute = next(filter(lambda x: x.name == "create", controller.api_router.routes))
    assert create.dependant.body_params[0].type_ == MyModel
    assert create.response_model == MyModel

def test_list(controller: MyController):
    route: APIRoute = next(filter(lambda x: x.name == "list", controller.api_router.routes))
    pagination = route.dependant.dependencies[1]
    assert pagination.name == "pagination"
    for param, dep in zip(["limit", "type", "next"], pagination.query_params):
        assert param == dep.name

    filtering = route.dependant.query_params[2]
    assert filtering.name == "filtering"

    fields = route.dependant.query_params[0]
    assert fields.name == "fields"
    assert fields.type_ is controller.repository.__fields__
    sorting = route.dependant.query_params[1]
    assert sorting.name == "sorting"
    assert sorting.type_ is controller.repository.__sort__
    assert route.response_model == PaginatedResponse[MyModel]

def test_get(controller: MyController):
    get: APIRoute = next(filter(lambda x: x.name == "get", controller.api_router.routes))
    fields = get.dependant.query_params[0]
    assert get.response_model is MyModel
    assert fields.type_ is controller.repository.__fields__

@pytest.mark.asyncio()
async def test_app():
    app = FastAPI()
    app.include_router(MyController.api_router)
    model = MyModel(my_param="1")

    client = TestClient(app)

    create_response = client.post("/", json=model.dict())
    assert create_response.status_code == HTTPStatus.OK
    doc = create_response.json()
    model.id = doc["_id"]
    get_response = client.get(f"/{doc['_id']}")
    assert get_response.status_code == HTTPStatus.OK
    assert get_response.json() == model.dict(by_alias=True)

    list_response = client.get("/")
    assert list_response.status_code == HTTPStatus.OK
    assert list_response.json()["items"][0] == model.dict(by_alias=True)

def test_api_router():
    assert id(MyController.api_router) != id(MyController2.api_router)

def test_model_controller__when__routes_defined__then_only_routes_defined():
    routes = MyController2.api_router.routes
    assert len(routes) == 1
    assert routes[0].name == "get"

def test_controller_no_overlapped_mixin():
    routes = [i.name for i in MyController.api_router.routes]
    duplicates = {x for x in routes if routes.count(x) > 1}
    assert not duplicates


import os
import os.path as op
import shutil
import yaml
from pybombs.commands import CommandBase
from pybombs.utils import dict_merge
from pybombs.utils import subproc
from pybombs.utils import confirm
from pybombs import fetcher

### Parser Helpers
def setup_subsubparser_initsdk(parser):
    parser.add_argument(
        'sdkname',
        help="SDK Package Name",
        nargs=1,
    )

def setup_subsubparser_init(parser):
    parser.add_argument(
        'path',
        help="Path to the new prefix (defaults to CWD).",
        nargs="?",
        default=os.getcwd(),
    )
    parser.add_argument(
        '--sdk',
        dest='sdkname',
        help="Initialize prefix with SDK Package",
    )
    parser.add_argument(
        '-a', '--alias',
        help="If specified, store an alias to this new prefix in the local config file.",
    )
    parser.add_argument(
        '--virtualenv', action='store_true',
        help="Use this to make the new prefix also a virtualenv. Args for virtualenv may be passed in here",
    )

### Class definition
class Prefix(CommandBase):
    """
    Prefix operations
    """
    cmds = {
        'prefix': 'Prefix commands',
    }

    @staticmethod
    def setup_subparser(parser, cmd=None):
        """
        Set up a subparser for a specific command
        """
        ###### Start of setup_subparser()
        subparsers = parser.add_subparsers(
                help="Prefix Commands:",
                dest='prefix_command',
        )
        for cmd_name, cmd_info in Prefix.prefix_cmd_name_list.iteritems():
            subparser = subparsers.add_parser(
                    cmd_name,
                    help=cmd_info['help']
            )
            cmd_info['parser'](subparser)
        return parser

    def __init__(self, cmd, args):
        CommandBase.__init__(self,
                cmd, args,
                load_recipes=True,
                require_prefix=(args.prefix_command != 'init'),
        )

    def run(self):
        """ Go, go, go! """
        self.prefix_cmd_name_list[self.args.prefix_command]['run'](self)

    def _print_prefix_info(self):
        """
        pybombs prefix info
        """
        self.log.info('Prefix dir: {}'.format(self.prefix.prefix_dir))

    def _print_prefix_env(self):
        """
        pybombs prefix env
        """
        print('Prefix env:')
        for k, v in self.prefix.env.iteritems():
            print("{}={}".format(k, v))

    def _run_installsdk(self):
        """
        pybombs prefix install-sdk
        """
        self._install_sdk_to_prefix(
            self.args.sdkname[0],
        )

    def _run_init(self):
        """
        pybombs prefix init
        """
        # Make sure the directory is writable
        path = op.abspath(op.normpath(self.args.path))
        if not op.isdir(path):
            self.log.info("Creating directory `{0}'".format(path))
            os.mkdir(path)
        assert op.isdir(path)
        if not os.access(path, os.W_OK|os.X_OK):
            self.log.error("Cannot write to prefix path `{0}'.".format(path))
            exit(1)

        # Make sure that a pybombs directory doesn't already exist
        test_path = op.join(path, ".pybombs")
        if op.exists(test_path):
            self.log.error("Ignoring. A prefix already exists in `{0}'".format(path))
            return

        # Copy template
        # TODO: I'm not too happy about this, all the hard coded stuff. Needs
        # cleaning up, for sure. Especially that setup_env.sh stuff at the end.
        # Ideally, we could switch prefix templates.
        self.log.info("Initializing PyBOMBS prefix in `{0}'...".format(path))
        skel_dir = op.join(self.cfg.module_dir, 'skel')
        for p in os.listdir(skel_dir):
            if op.exists(op.join(path, p)):
                self.log.obnoxious("Skipping {0}".format(p))
                continue
            self.log.obnoxious("Copying {0}".format(p))
            p_full = op.join(skel_dir, p)
            if op.isdir(p_full):
                shutil.copytree(
                    p_full, op.join(path, p),
                    ignore=shutil.ignore_patterns('.ignore')
                )
            else:
                shutil.copy(p_full, path)
        open(op.join(path, 'setup_env.sh'), 'w').write(
            open(op.join(skel_dir, 'setup_env.sh')).read().format(
                prefix_dir=path,
            )
        )
        # Register alias
        if self.args.alias is not None:
            if self.prefix is not None and \
                self.prefix.prefix_aliases.get(self.args.alias) is not None \
                and not confirm("Alias `{0}' already exists, overwrite?".format(self.args.alias)):
                    self.log.warn('Aborting.')
                    return 1
            self.cfg.update_cfg_file({'prefix_aliases': {self.args.alias: path}})
        # Create virtualenv if so desired
        if self.args.virtualenv:
            self.log.info("Creating Python virtualenv in prefix...")
            venv_args = ['virtualenv']
            venv_args.append(path)
            subproc.monitor_process(args=venv_args)
        # Install SDK if so desired
        if self.args.sdkname is not None:
            self.log.info("Reloading configuration...")
            self.cfg.load(select_prefix=path)
            self.prefix = self.cfg.get_active_prefix()
            self.inventory = self.prefix.inventory
            self._install_sdk_to_prefix(self.args.sdkname)

    def _install_sdk_to_prefix(self, sdkname):
        """
        Read recipe for sdkname, and install the SDK to the prefix.
        """
        from pybombs import recipe
        src_dir = self.prefix.src_dir
        cfg_file = self.prefix.cfg_file
        ### Get the recipe
        r = recipe.get_recipe(sdkname, target='sdk')
        try:
            self.log.obnoxious("Switching CWD to {0}".format(src_dir))
            if not op.isdir(src_dir):
                os.mkdir(src_dir)
            os.chdir(src_dir)
        except:
            self.log.error("Source dir required to install SDK.")
            return -1
        ### Install the actual SDK file
        self.log.debug("Fetching SDK `{sdk}'".format(sdk=sdkname))
        fetcher.Fetcher().fetch(r)
        self.log.info("Installing SDK `{sdk}'".format(sdk=sdkname))
        # Install command
        cmd = r.var_replace_all(r.get_command('install'))
        if subproc.monitor_process(cmd, shell=True, env=os.environ) == 0:
            self.log.debug("Installation successful")
        else:
            self.log.error("Error installing SDK. Aborting.")
            return -1
        # Clean up
        files_to_delete = [op.normpath(op.join(src_dir, r.var_replace_all(x))) for x in r.clean]
        if len(files_to_delete):
            self.log.info("Cleaning up files...")
        for ftd in files_to_delete:
            if op.commonprefix((src_dir, ftd)) != src_dir:
                self.log.warn("Not removing {ftd} -- outside source dir!".format(ftd=ftd))
                continue
            self.log.debug("Removing {ftd}...".format(ftd=ftd))
            if op.isdir(ftd):
                shutil.rmtree(ftd)
            elif op.isfile(ftd):
                os.remove(ftd)
            else:
                self.log.error("Not sure what this is: {ftd}".format(ftd=ftd))
                return -1
        ### Update the prefix-local config file
        self.log.debug("Updating config file with SDK recipe info.")
        try:
            old_cfg_data = yaml.safe_load(open(cfg_file).read()) or {}
        except IOError:
            self.log.debug("There doesn't seem to be a config file yet for this prefix.")
            old_cfg_data = {}
        # Filter out keys we don't care about:
        sdk_recipe_keys_for_config = ('config', 'packages', 'categories', 'env')
        sdk_cfg_data = {k: v for k, v in r.get_dict().iteritems() if k in sdk_recipe_keys_for_config}
        self.log.obnoxious("New data: {new}".format(new=sdk_cfg_data))
        cfg_data = dict_merge(old_cfg_data, sdk_cfg_data)
        self.log.debug("Writing updated prefix config to `{0}'".format(cfg_file))
        open(cfg_file, 'wb').write(yaml.dump(cfg_data, default_flow_style=False))

    #########################################################################
    # Sub-commands:
    prefix_cmd_name_list = {
        'info': {
            'help': 'Display information on the currently used prefix.',
            'parser': lambda p: None,
            'run': _print_prefix_info,
        },
        'init': {
            'help': 'Create and initialize a prefix.',
            'parser': setup_subsubparser_init,
            'run': _run_init,
        },
        'env': {
            'help': 'Print the environment variables used in this prefix.',
            'parser': lambda p: None,
            'run': _print_prefix_info,
        },
        'install-sdk': {
            'help': 'Install an SDK into the prefix.',
            'parser': setup_subsubparser_initsdk,
            'run': _run_installsdk,
        },
    }
    #########################################################################


import argparse as ap
import re
import sys
import pyperclip
from datetime import datetime

CONVERTER = {"2": 2, "1": 1, "0": 0, "-": -1, "=": -2}
NEGATE = {"2": "=", "1": "-", "0": "0", "-": "1", "=": "2"}
DECONVERTER = {v: k for k, v in CONVERTER.items()}

def snafu_to_decimal(snafu: str) -> int:
    total = 0
    for c in snafu:
        total += CONVERTER[c]
        total *= 5

    return total // 5

def decimal_to_snafu(decimal: int, pos: int = 0) -> str:
    if -2 <= decimal <= 2 and pos <= 1:
        return DECONVERTER[decimal]

    mul = 1
    if decimal < 0:
        mul = -1
        decimal *= -1

    # First find the largest power of 5 that fits in decimal
    power = pos
    while 5**power <= decimal:
        power += 1

    power -= 1
    # Does it start with a 2 or 1?
    result = ""
    if decimal <= sum(2 * 5**i for i in range(power)):
        result = DECONVERTER[0] + decimal_to_snafu(decimal, power)
    elif decimal <= 5**power + sum(2 * 5**i for i in range(power)):
        result = DECONVERTER[1] + decimal_to_snafu(decimal - 5**power, power)
    elif decimal <= sum(2 * 5**i for i in range(power + 1)):
        result = DECONVERTER[2] + decimal_to_snafu(decimal - 2 * 5**power, power)
    elif decimal <= 5 ** (power + 1):
        result = DECONVERTER[1] + decimal_to_snafu(
            decimal - 5 ** (power + 1), power + 1
        )
    else:
        result = DECONVERTER[2] + decimal_to_snafu(
            decimal - 2 * 5 ** (power + 1), power + 1
        )

    if mul == -1:
        result = "".join(NEGATE[c] for c in result)

    return result

class Solution:
    def __init__(self, filename: str):
        with open(filename, "r") as f:
            self.input = f.read()
            self.input_lines = self.input.splitlines()

    def solve_part1(self) -> str:
        total = 0
        for line in self.input_lines:
            total += snafu_to_decimal(line)

        return decimal_to_snafu(total)

    def solve_part2(self) -> int:
        pass

if __name__ == "__main__":
    parser = ap.ArgumentParser()
    parser.add_argument(
        "-e",
        "--example",
        help="Use the example file for input instead of main",
        action="store_true",
    )

    args = parser.parse_args()
    day = re.search(r"\d+", sys.argv[0])
    if day:
        day = day.group(0)
    else:
        day = datetime.now().day + 1

    filename = f"../data/day{day}.txt"
    if args.example:
        filename = f"../data/day{day}ex.txt"

    for i in range(500):
        try:
            assert snafu_to_decimal(decimal_to_snafu(i)) == i
        except AssertionError:
            print(f"Failed for {i}")
            raise

    sol = Solution(filename)
    x = sol.solve_part1()
    print(f"Part 1: {x}")
    if x is not None:
        pyperclip.copy(x)

    sol = Solution(filename)
    x = sol.solve_part2()
    print(f"Part 2: {x}")
    if x is not None:
        pyperclip.copy(x)


import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config.database import get_database_dict_info
from config.university_list import get_university_list
import requests
from bs4 import BeautifulSoup
from pymongo import MongoClient
# pip install readability/readability-lxml/readability-api
from readability import Document
import html2text
import json
import threading

MAX_PAGE_NUMBERS = 39

# 过滤一些无效的新闻数据
def filter(doc):
    if doc["body"] == "error":
        return "false"

    if len(doc["body"]) <= 50:
        return "false"

    if len(doc["date"]) <= 10:
        return "false"

    if doc["body"].find('ä') != -1:
        return "false"

    if doc["title"].find('ä') != -1:
        return "false"

    return "true"

# 提取各种网页中的主体正文
def get_website_body_from_html(url):

    user_agent = 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/56.0.2924.87 Safari/537.36'
    request_headers = {'User-Agent': user_agent}

    try:
        response = requests.get(url, headers=request_headers)

        # print("url", url)
        # print(response.headers['content-type'])
        # print(response.encoding)
        # print(response.apparent_encoding)
        # print(requests.utils.get_encodings_from_content(response.text))

        html = response.text

        # 判定编码，解决爬取中文网页乱码问题
        # 乱码解决方法：http://blog.chinaunix.net/uid-13869856-id-5747417.html

        if response.encoding == 'ISO-8859-1':
            encodings = requests.utils.get_encodings_from_content(response.text)
            if encodings:
                encoding = encodings[0]
            else:
                encoding = response.apparent_encoding
            content = response.content      # bytes流

            if encoding == 'gb2312' or encoding == 'GB2312' or ((encoding == 'gbk' or encoding == 'GBK') and (response.apparent_encoding == 'GB2312' or response.apparent_encoding == 'gb2312')):
                html = content.decode("gb2312", 'replace')
            else:
                html = content.decode(encoding, 'replace').encode('utf-8', 'replace')

        readable_article = Document(html).summary()
        body = html2text.html2text(readable_article)

    except BaseException as e:
        print('Download error', e)
        body = "false"

    return body

def request_baidu_news(university_name,start_page,end_page,university_abbr):

    # 请求地址模板
    base_url = "http://news.baidu.com/ns?word={university_name}&pn={pn}&rn=20&cl=2"

    user_agent = 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/56.0.2924.87 Safari/537.36'

    request_headers = {'User-Agent': user_agent}

    news_documents_list = []

    threads = []

    for i in range(start_page,end_page):

        pn = i * 20 - 20

        url = base_url.format(university_name=university_name, pn=pn)

        response = requests.get(url, headers=request_headers)

        # 拿到 BS4 的解析结果
        soup = BeautifulSoup(response.text,'lxml')

        for j in range(1,21):

            id_of_news = str(j + pn)

            result_div = soup.find(id=id_of_news)

            if not result_div:
                break

            # 为了实现多线程而增加的函数get_detail_news
            thread = threading.Thread(
                target=get_detail_news,
                args=(
                    result_div, id_of_news, university_name,
                    university_abbr, news_documents_list
                )
            )
            thread.start()
            threads.append(thread)

    for thread in threads:
        print("剩余:" + str(threading.active_count()))
        thread.join()

    print("完成了",university_name,"的新闻信息爬取,共 ",str(len(news_documents_list))," 条")

    return news_documents_list

# 为了多线程而增加的函数
def get_detail_news(result_div, id_of_news, university_name,
                    university_abbr, news_documents_list):
    media_date_text = result_div.p.text
    index_of_year = media_date_text.find("20")

    title = result_div.a.text
    url = result_div.a['href']
    media = media_date_text[0:index_of_year - 2]
    date = media_date_text[index_of_year:]

    body = get_website_body_from_html(url)

    current_news = {
        "title":title.replace(" \r\t\n", "").strip(),
        "url":url,
        "media":media.replace(" \r\t\n", "").strip(),
        "Uname":university_name.replace(" \r\t\n", "").strip(),
        "date":date.replace("\r\t\n", "").strip(),
        "abbr":university_abbr,
        "body":body,
        "ranking":id_of_news,
        "classification":None,
        "sentiment":None
    }
    if(filter(current_news) == "true"):
        news_documents_list.append(current_news)
    print("thread" + id_of_news + " of " + university_name + " is ended")

def save_newslist_to_db():

    # 获取学校列表，数据库配置信息
    university_list = get_university_list()
    db_config = get_database_dict_info()

    #建立数据库连接
    conn = MongoClient(db_config["host"],db_config["port"])

    NewsPOA = conn.NewsPOA

    # NewsPOA["newslist"].drop()

    for i in range(0,len(university_list)):
    # for i in range(0,1):

        uni = university_list[i]
        if NewsPOA['newslist'].find({"Uname": uni["zh_name"]}).count() != 0:
            continue
        news_documents_list = request_baidu_news(uni["zh_name"],1,MAX_PAGE_NUMBERS,uni["en_name"])
        NewsPOA["newslist"].insert(news_documents_list)
        print(uni["zh_name"],"的新闻列表保存成功")

    print("新闻全部爬取完毕")

# save_newslist_to_db()

def save_newslist_into_file():

    university_list = get_university_list()
    for i in range(0,len(university_list)):

        uni = university_list[i]
        news_documents_list = request_baidu_news(uni["zh_name"],1,MAX_PAGE_NUMBERS,uni["en_name"])

        news_path = "./news_result/" + uni["zh_name"] + ".json"
        with open(news_path, 'w', encoding='utf-8') as json_file:
            json.dump(news_documents_list, json_file, ensure_ascii=False)

            print(uni["zh_name"],"的新闻列表保存成功")

# save_newslist_into_file()