from datetime import datetime
from tqdm import tqdm
import requests
import json
from os.path import expanduser
from requests.auth import HTTPBasicAuth
from time import sleep
from multiprocessing import Pool

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

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

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

print(resp.status_code)
print(resp.json())
alpha_list = []

def simulate_alpha(alpha):
    # global alpha_list
    # index = alpha_list.index(alpha)
    sim_resp = session.post('https://api.worldquantbrain.com/simulations', json=alpha)

    try:
        sim_progress_url = sim_resp.headers['Location']
        while True:
            sim_progress_resp = session.get(sim_progress_url)
            retry_after_sec = float(sim_progress_resp.headers.get("Retry-After", 0))
            if retry_after_sec == 0:  # simulation done!
                break
            sleep(retry_after_sec)
        alpha_id = sim_progress_resp.json()["alpha"]  # the final simulation result
        print(f"{str(datetime.now())} alpha_id_:{alpha_id}")
    except:
        print(f"no location, sleep for 10 seconds and try next alpha, {alpha}")
        sleep(10)


if __name__ == '__main__':
    # 读取dataset.txt文件，获取alpha列表
    dataset = []
    # 打开文件并读取
    with open('./dataset.txt', 'r', encoding='utf-8') as file:
        for line in file:
            # 去除每行末尾的换行符并添加到列表中
            dataset.append(line.strip())

    # 此时，lines列表中包含了文件的每一行作为一个字符串
    print(len(dataset))
    for datafield in dataset:
        alpha_expression = f'{datafield}'
        print(alpha_expression)
        simulation_data = {
            'type': 'REGULAR',
            'settings': {
                'instrumentType': 'EQUITY',
                'region': 'USA',
                'universe': 'TOP3000',
                'delay': 1,
                'decay': 0,
                'neutralization': 'SUBINDUSTRY',
                'truncation': 0.08,
                'pasteurization': 'ON',
                'unitHandling': 'VERIFY',
                'nanHandling': 'ON',
                'language': 'FASTEXPR',
                'visualization': False,
            },
            'regular': alpha_expression
        }
        alpha_list.append(simulation_data)
    pool_size = 3
    with Pool(pool_size) as pool:
        results = list(tqdm(pool.map(simulate_alpha, alpha_list), total=len(alpha_list)))
        print(results)