import random
import copy
import threading
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np


def single_pull(no_six_pull_count=0):
    x = random.randint(1, 100)
    six_star_threshold = (no_six_pull_count - 50) * 2 if no_six_pull_count > 50 else 0
    factor = (100 - (2 + six_star_threshold)) / 98
    prob_list = [2 + six_star_threshold, 8 * factor, 50 * factor, 40 * factor]
    sum_prob_list = [prob_list[0], prob_list[0] + prob_list[1], prob_list[0] + prob_list[1] + prob_list[2], prob_list[0] + prob_list[1] + prob_list[2] + prob_list[3]]
    if 0 < x <= sum_prob_list[0]:
        result = 6
    elif sum_prob_list[0] < x <= sum_prob_list[1]:
        result = 5
    elif sum_prob_list[1] < x <= sum_prob_list[2]:
        result = 4
    else:
        result = 3
    if result == 6:
        no_six_pull_count = 0
    else:
        no_six_pull_count += 1
    return result, no_six_pull_count


def ten_pull(no_six_pull_count=0):
    result_list = []
    is_african = True
    for i in range(10):
        res, no_six_pull_count = single_pull(no_six_pull_count)
        result_list.append(res)
        if res > 4:
            is_african = False
    return result_list, no_six_pull_count, is_african


def run_simulation(total_pull=200):
    ten_pull_count = int(total_pull / 10)
    single_pull_count = ten_pull_count - ten_pull_count
    no_six_pull_count = 0
    african_count = 0
    full_res_list = []
    for i in range(ten_pull_count):
        ten_pull_res, no_six_pull_count, is_african = ten_pull(no_six_pull_count)
        full_res_list += copy.deepcopy(ten_pull_res)
        if is_african:
            african_count += 1

    for i in range(single_pull_count):
        single_res, no_six_pull_count = single_pull(no_six_pull_count)
    six_star_list = [i for i in full_res_list if i == 6]

    # pull_result = {'Result': full_res_list, 'african': african_count, 'six star count': len(six_star_list)}
    pull_result = {'african': african_count, 'six star count': len(six_star_list)}
    return pull_result


if __name__ == '__main__':

    hit = 0
    hit1 = 0
    hit2 = 0

    total_pull_count = 180
    african_threshold = 10
    six_threshold = 4
    simulation_count = 50000

    status_dict = {}

    for i in range(simulation_count):
        k = run_simulation(total_pull_count)
        key = f"{k['african']}-{k['six star count']}"
        if key in status_dict:
            status_dict[key][0] += 1
        else:
            status_dict[key] = [0, k['african'], k['six star count']]
        if k['african'] >= african_threshold and k['six star count'] <= six_threshold:
            hit += 1
            if k['african'] == african_threshold and k['six star count'] == six_threshold:
                hit1 += 1
        if k['african'] >= african_threshold:
            hit2 += 1
    print(hit / simulation_count * 1000, hit1 / simulation_count * 1000, hit2 / simulation_count * 1000)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    max_color = max([item[0] for item in status_dict.values()])
    for key, item in status_dict.items():
        color = np.array([1, 1 - item[0] / max_color, 0])  # 颜色 其中每个元素在0~1之间
        ax.bar3d(item[1], item[2], 0, 1, 1, item[0], color=color)
        ax.text(x=item[1], y=item[2], z=item[0] + 30, s=f"{item[1]}, {item[2]}, {item[0]}",
                 ha='center', fontsize=8, family='Calibri', va='center')
    ax.set_xlabel('African ten pull')
    ax.set_ylabel('Six star count')
    ax.set_zlabel('Count')
    ax.set_title('Pull Simulation Result')
    plt.show()

    print('done')