import json
import random
import gevent
from db_helper import RedisHelper
from config import set_config
from gevent import monkey, pool
monkey.patch_all()

config_ = set_config()
redis_helper = RedisHelper()


def thompson_process(ad_idea_id):
    # 获取ad_idea_id对应的Thompson参数
    thompson_param = redis_helper.get_data_from_redis(key_name=f"{config_.THOMPSON_PARAM_KEY_PREFIX}{ad_idea_id}")
    if thompson_param is None or thompson_param == '':
        # 参数不存在，随机生成[0, 1)之间的浮点数
        score = random.random()
        random_flag = 'random'
    else:
        # 参数存在
        param_alpha, param_beta = json.loads(thompson_param.strip())
        param_alpha, param_beta = int(param_alpha), int(param_beta)
        if param_alpha + param_beta >= 100:
            # ad_idea_id 曝光数 >= 100，生成参数为(param_alpha+1, param_beta+1)的beta分布随机数
            score = random.betavariate(alpha=param_alpha+1, beta=param_beta+1)
            random_flag = 'beta'
        else:
            # ad_idea_id 曝光数 < 100，随机生成[0, 1)之间的浮点数
            score = random.random()
            random_flag = 'random'
    thompson_res = [ad_idea_id, score, thompson_param, random_flag]
    return thompson_res


def get_ad_idea_id_with_thompson(mid, ad_idea_id_list):
    """利用Thompson采样获取此次要展示的广告创意ID"""
    # 限制协程最大并发数：20
    gevent_pool = pool.Pool(20)
    tasks = [gevent_pool.spawn(thompson_process, ad_idea_id) for ad_idea_id in ad_idea_id_list]
    gevent.joinall(tasks)
    thompson_res_list = [t.get() for t in tasks]
    # 按照score排序
    thompson_res_rank = sorted(thompson_res_list, key=lambda x: x[1], reverse=True)
    rank_res = {
        'mid': mid,
        'ad_idea_id': thompson_res_rank[0][0],
        'score': thompson_res_rank[0][1],
        'thompson_param': thompson_res_rank[0][2],
        'random_flag': thompson_res_rank[0][3],
        'thompson_res_rank': thompson_res_rank
    }
    return rank_res