from datetime import datetime

import requests
import pandas as pd
import json
from os.path import expanduser
from requests.auth import HTTPBasicAuth
from time import sleep
import logging
logging.basicConfig(filename='simulation.log',level = logging.INFO,
                    format ='%(asctime)s-%(levelname)s-%(message)s')


with open(expanduser("brain_credentials.txt")) as f:
    credentials = json.load(f)

username, password = credentials
sess = requests.Session()
sess.auth = HTTPBasicAuth(username, password)

resp = sess.post('https://api.worldquantbrain.com/authentication')

print(resp.status_code)
print(resp.json())


def judge_session_status():
    status_resp = sess.get('https://api.worldquantbrain.com/authentication')
    if str(status_resp.status_code)[0] != '2':
        sess.auth = HTTPBasicAuth(username, password)
        sess.post('https://api.worldquantbrain.com/authentication')


def get_datafields(
    sess,
    searchscope,
    dataset_id: str = "",
    search: str = '',
):
    instrument_type = searchscope['instrumentType']
    region = searchscope['region']
    delay = searchscope['delay']
    universe = searchscope['universe']
    if len(search) == 0:
        url_template = "https://api.worldquantbrain.com/data-fields?" +\
            f"&instrumentType={instrument_type}" +\
            f"&region={region}&delay={str(delay)}&universe={universe}&dataset.id={dataset_id}&limit=50" +\
            "&offset={x}"
        count = sess.get(url_template.format(x=0)).json()['count']
    else:
        url_template = "https://api.worldquantbrain.com/data-fields?" +\
            f"&instrumentType={instrument_type}" +\
            f"&region={region}&delay={str(delay)}&universe={universe}&limit=50" +\
            f"&search={search}" +\
            "&offset={x}"
        count = 100
    datafields_list = []
    for x in range(0, count, 50):
        datafields = sess.get(url_template.format(x=x))
        datafields_list.append(datafields.json()['results'])
        datafields_list_flat = [item for sublist in datafields_list for item in sublist]
    datafields_df = pd.DataFrame(datafields_list_flat)

    return datafields_df


searchscope = {'region': 'USA', "delay": '1', 'universe': 'TOP3000', 'instrumentType': 'EQUITY'}
fundamental6 = get_datafields(sess=sess, searchscope=searchscope, dataset_id='fundamental6')
fundamental6 = fundamental6[fundamental6['type'] == "MATRIX"]

print(fundamental6.head())


def get_datafields(
    sess,
    searchscope,
    dataset_id: str = "",
    search: str = '',
):
    instrument_type = searchscope['instrumentType']
    region = searchscope['region']
    delay = searchscope['delay']
    universe = searchscope['universe']
    if len(search) == 0:
        url_template = "https://api.worldquantbrain.com/data-fields?" +\
            f"&instrumentType={instrument_type}" +\
            f"&region={region}&delay={str(delay)}&universe={universe}&dataset.id={dataset_id}&limit=50" +\
            "&offset={x}"
        count = sess.get(url_template.format(x=0)).json()['count']
    else:
        url_template = "https://api.worldquantbrain.com/data-fields?" +\
            f"&instrumentType={instrument_type}" +\
            f"&region={region}&delay={str(delay)}&universe={universe}&limit=50" +\
            f"&search={search}" +\
            "&offset={x}"
        count = 100
    datafields_list = []
    for x in range(0, count, 50):
        datafields = sess.get(url_template.format(x=x))
        datafields_list.append(datafields.json()['results'])
        datafields_list_flat = [item for sublist in datafields_list for item in sublist]
    datafields_df = pd.DataFrame(datafields_list_flat)

    return datafields_df


group_compare_op = ['group_rank','group_zscore','group_neutralize']
ts_compare_op = ['rank','ts_rank','ts_zscore','ts_av_diff']
datafields_list_fundamental6 = fundamental6['id'].values
company_fundamentals = datafields_list_fundamental6
list(company_fundamentals).extend(['cap', 'revenue'])
days =[66,126,252,504]
alpha_expressions = []

for gco in group_compare_op:
    for tco in ts_compare_op:
        for cf in company_fundamentals:
            for d in days:
                for grp in group_compare_op:
                    alpha_expressions.append(f"Turn20_ = ts_mean(volume/sharesout, {d});"
                                             f"Turn20 = {gco}(Turn20_, bucket({tco}({cf}), range='0.1,1,0.1'));"
                                             f"STR_ = ts_std_dev(volume/sharesout, {d});"
                                             f"STR = {gco}(STR_, bucket({tco}({cf}), range='0.1,1,0.1'));"
                                             f"score2 = {tco}(- nan_mask(Turn20, if_else({tco}(STR) < 0.5, 1, -1))) * 0.5;"
                                             f"score3 = {tco}(nan_mask(Turn20, if_else({tco}(STR) >= 0.5, 1, -1))) * 0.5;"
                                             f"signal_ = add({tco}(STR), score2, score3, filter = true);"
                                             f"signal = left_tail({tco}(signal_), maximum=0.98);"
                                             f"- {grp}(signal, bucket({tco}({cf}), range='0.1,1,0.1'))"
                                             )

print(f'There are total {len(alpha_expressions)} alpha expressions')


alpha_list = []

for alpha_expression in alpha_expressions:
    simulation_data = {
        'type': 'REGULAR',
        'settings': {
            'instrumentType': 'EQUITY',
            'region': 'USA',
            'universe': 'TOP3000',
            'delay': 1,
            'decay': 0,
            'neutralization': 'SUBINDUSTRY',
            'truncation': 0.01,
            'pasteurization': 'ON',
            'unitHandling': 'VERIFY',
            'nanHandling': 'OFF',
            'language': 'FASTEXPR',
            'visualization': False,
        },
        'regular': alpha_expression
    }
    alpha_list.append(simulation_data)

alpha_fail_attempt_tolerance =15

for alpha in alpha_list:
    keep_trying = True
    failure_count = 0

    while keep_trying:
        try:
            judge_session_status()
            #尝试发送模拟回测请求
            sim_resp = sess.post(
                'https://api.worldquantbrain.com/simulations',
                json=alpha   #将当前一个alpha发送到服务器！
            )
            sim_progress_url =sim_resp.headers['Location']  #从响应header中获取位置
            logging.info(f'Alpha location is: {sim_progress_url}')
            print(f'Alpha location is: {sim_progress_url}')
            keep_trying =False

        except Exception as e:
            #对于异常情况：记录错误，休眠并增加失败次数计数
            logging.error(f"No location, sleep 15 and retry,error message:{str(e)}")
            print("No location, sleep 15 and retry")
            sleep(15)
            failure_count += 1

            # #检查失败次数是否达到忍耐限度，哈哈哈
            # if failure_count>= alpha_fail_attempt_tolerance:
            #     sess = sign_in()
            #     failure_count = 0  #重置失败次数为0
            #
            #     #记录并打印信息，指示将移动到下一个alpha
            #     logging.error(f"No location for too many times, move to next alpha {alpha['regular']}")
            #     print(f"No location for too many times, move to next alpha{alpha['regular']}")
            #     break #退出while循环，移动到for循环中的下一个alpha