Honkai Impact 3 gacha simulator. Implements expansion equipment, focus, and schicksal arsenal supply

sim.py

import random
import statistics

from collections import Counter
from dataclasses import dataclass
from enum import Enum
from multiprocessing import Pool
from typing import Dict, List
from tqdm import tqdm


class Gacha(Enum):
    FOCUS = 1
    EXPANSION_EQUIP = 2
    ARSENAL = 3


class PullType(Enum):
    PITY = 1
    EXPANSION_EQUIP_GUARANTEE = 2
    ARSENAL_GUARANTEE = 3
    NON_PITY = 4


class PullResult(Enum):
    STIG_T = 1
    STIG_M = 2
    STIG_B = 3
    WEAPON = 4
    FOUR_STAR_TRASH = 5
    TRASH = 6


@dataclass
class SimResult:
    ending: PullType
    wishing_well: bool
    pull_count: int
    four_star_count: int


# Rates derived with help from: https://walkthrough.honkaiimpact3.com/article/show/50/4672
RATES = {
    Gacha.FOCUS: {
        PullType.NON_PITY: {
            PullResult.WEAPON: .976,
            PullResult.STIG_T: .488,
            PullResult.STIG_M: .488,
            PullResult.STIG_B: .488,
            PullResult.FOUR_STAR_TRASH: 2.439,
            PullResult.TRASH: 0,
        },
        PullType.PITY: {
            PullResult.WEAPON: 20,
            PullResult.STIG_T: 10,
            PullResult.STIG_M: 10,
            PullResult.STIG_B: 10,

            # Pity is always 4*
            PullResult.FOUR_STAR_TRASH: 50,
            PullResult.TRASH: 0,
        },
    },
    Gacha.ARSENAL: {
        PullType.NON_PITY: {
            PullResult.WEAPON: 1.96339,
            PullResult.STIG_T: .50324,
            PullResult.STIG_M: .50324,
            PullResult.STIG_B: .50324,
            PullResult.FOUR_STAR_TRASH: 2.56815,
            PullResult.TRASH: 0,
        },
        PullType.PITY: {
            PullResult.WEAPON: 32.5,
            PullResult.STIG_T: 8.33,
            PullResult.STIG_M: 8.33,
            PullResult.STIG_B: 8.33,

            # Pity is always 4*
            PullResult.FOUR_STAR_TRASH: 42.51,
            PullResult.TRASH: 0,
        },
        PullType.ARSENAL_GUARANTEE: {
            PullResult.WEAPON: 100,
        },
    },
    Gacha.EXPANSION_EQUIP: {
        PullType.NON_PITY: {
            PullResult.WEAPON: .72004,
            PullResult.STIG_T: .36199,
            PullResult.STIG_M: .36199,
            PullResult.STIG_B: .36199,
            PullResult.FOUR_STAR_TRASH: 3.06903,
            PullResult.TRASH: 0,
        },
        PullType.PITY: {
            PullResult.WEAPON: 14.764,
            PullResult.STIG_T: 7.422,
            PullResult.STIG_M: 7.422,
            PullResult.STIG_B: 7.422,

            # Pity is always 4*
            PullResult.FOUR_STAR_TRASH: 62.97,
            PullResult.TRASH: 0,
        },
        PullType.EXPANSION_EQUIP_GUARANTEE: {
            PullResult.WEAPON: 40,
            PullResult.STIG_T: 20,
            PullResult.STIG_M: 20,
            PullResult.STIG_B: 20,
        },
    },
}


def simulate_full_set(gacha_type: Gacha):
    pull_counter = 0
    last_four_star = 0
    last_new_uprate = 0
    pull_result_count = {}
    sim_result = SimResult(
        pull_count=1000,
        wishing_well=False,
        ending=PullType.NON_PITY,
        four_star_count=0
    )
    while pull_counter < 1000:
        pull_counter += 1
        if pull_counter - last_four_star == 10:
            pull_type = PullType.PITY
        else:
            pull_type = PullType.NON_PITY

        if gacha_type == Gacha.EXPANSION_EQUIP and pull_counter - last_new_uprate == 50:
            pull_type = PullType.EXPANSION_EQUIP_GUARANTEE

        if gacha_type == Gacha.ARSENAL and pull_counter == 50 and PullResult.WEAPON not in pull_result_count:
            pull_type = PullType.ARSENAL_GUARANTEE

        result = simulate_single(gacha_type, pull_type, pull_result_count)

        if result == PullResult.FOUR_STAR_TRASH:
            last_four_star = pull_counter
            sim_result.four_star_count += 1
        elif result == PullResult.TRASH:
            continue
        else:
            sim_result.four_star_count += 1
            last_four_star = pull_counter
            if result not in pull_result_count:
                last_new_uprate = pull_counter
            pull_result_count[result] = pull_result_count.get(result, 0) + 1

            end_result = is_finished(pull_result_count)
            if end_result > 0:
                sim_result.pull_count = pull_counter
                sim_result.ending = pull_type
                if end_result == 2:
                    sim_result.wishing_well = True
                break

    return sim_result


def is_finished(pull_result_count: Dict[PullResult, int]):
    """Determines if a series of pulls has full gear or can wishing well for it.

        Returns 0 is note done, 1 if done with full set, or 2 if done with wishing well.
    """
    if PullResult.WEAPON not in pull_result_count:
        return 0

    stig_results = [
        pull_result_count.get(PullResult.STIG_T, 0),
        pull_result_count.get(PullResult.STIG_M, 0),
        pull_result_count.get(PullResult.STIG_B, 0)
    ]

    unique_count = len([x for x in stig_results if x > 0])

    if unique_count == 3:
        # We have all 3 stigs.
        return 1
    elif unique_count == 2 and sum(stig_results) >= 4:
        # We have two stigs and enough dupes to wishing well the third.
        return 2
    elif sum(stig_results) >= 5:
        # We have five of one stig and we can wishing well twice (takes two supplies)
        return 2

    return 0


def simulate_single(gacha_type: Gacha, pull_type: PullType, pull_result_count: Dict[PullResult, int]):
    rates = RATES[gacha_type][pull_type]
    if pull_type == PullType.EXPANSION_EQUIP_GUARANTEE:
        missing_pieces = []

        sum_remaining_rates = 0
        for result_type, rate in rates.items():
            if result_type not in pull_result_count:
                missing_pieces.append(result_type)
                sum_remaining_rates += rate

        scaled_rates = {}
        for result_type in missing_pieces:
            # If we think the rates on the 50 pity are proportional.
            scaled_rates[result_type] = rates[result_type] / \
                sum_remaining_rates * 100

            # If we think the rates on the 50 pity are equal.
            # scaled_rates[result_type] = 100/len(missing_pieces)

        return roll_single(missing_pieces, scaled_rates)

    if pull_type == PullType.ARSENAL_GUARANTEE:
        return PullResult.WEAPON


    return roll_single(PullResult, rates)


def roll_single(pull_result_choices: List[PullResult], rates: Dict[PullResult, float]):
    # Multiply by 100 since we specified rates in percentages.
    rand = random.random() * 100
    cumulative_rand = 0

    for result in pull_result_choices:
        result_rate = rates[result]
        cumulative_rand += result_rate
        if rand < cumulative_rand:
            return result
    return PullResult.TRASH


def simulate_with_stats(gacha_type: Gacha, sim_count: int):
    pool = Pool()
    print(f'=={gacha_type} supply, {sim_count} trials==')
    sim_results = []
    for result in tqdm(pool.imap(simulate_full_set, [gacha_type] * sim_count)):
        sim_results.append(result)
    sim_results.sort(key=lambda r: r.pull_count)
    pulls_needed = [r.pull_count for r in sim_results]
    four_stars_pulled = [r.four_star_count for r in sim_results]
    four_star_rate = sum([r.four_star_count for r in sim_results]) / sum([r.pull_count for r in sim_results])
    wishing_well_percent = sum(
        [True if r.wishing_well else False for r in sim_results]) / len(sim_results)
    max_pulls_percent = sum(
        [True if r.pull_count >= 200 else False for r in sim_results]) / len(sim_results)
    end_pull_types = Counter([r.ending for r in sim_results])
    print(f'Average: {statistics.mean(pulls_needed)}')
    print(f'Median: {statistics.median(pulls_needed)}')
    print(f'Stdev: {statistics.stdev(pulls_needed)}')
    print(f'Average 4* count: {statistics.mean(four_stars_pulled)}')
    print(f'Overall 4* rate: {four_star_rate:.2%}')
    # Covered by the cumulative distribution already
    # print(f'Min: {min(pulls_needed)}')
    # print(f'Max: {max(pulls_needed)}')
    print(f'200 or more pulls needed: {max_pulls_percent:.2%}')
    print(f'Wishing well %: {wishing_well_percent:.2%}')
    print(f'End pull types: {end_pull_types}')

    print('Outputting percentiles to file...')
    with open(f'{gacha_type}.txt', 'w') as f:
        for i in range(0, 101, 1):
            print(
                f'{i/100},{sim_results[round((sim_count - 1) * i / 100)].pull_count}', file=f)

    print('\n')
    return sim_results


def list_rindex(in_list, value):
    return len(in_list) - in_list[-1::-1].index(value) - 1


if __name__ == "__main__":
    while (pull_count_input := input("How many pulls to simulate?:")):
        try:
            pull_count = int(pull_count_input)
            break
        except:
            print('Is that a number???')

    focus_results = simulate_with_stats(Gacha.FOCUS, pull_count)
    expansion_equip_results = simulate_with_stats(
        Gacha.EXPANSION_EQUIP, pull_count)
    arsenal_results = simulate_with_stats(
        Gacha.ARSENAL, pull_count)

    print("Simulation done, choose pull count to calculate chance of finishing full set.")

    while True:
        while (pull_count_input := input("Enter the number of pulls:")):
            try:
                pull_count = int(pull_count_input)
                break
            except:
                print('Is that a number???')

        def success_stats(results, gacha_type: Gacha, pulls: int):
            print(f'=={pulls} pulls on {gacha_type}==')
            pulls_needed = [r.pull_count for r in results]

            # All pulls before this one are successes, all pulls after and including this one are failures.
            boundary_index = next(
                (i for i, p in enumerate(pulls_needed) if p > pulls),
                (len(pulls_needed) - 1)
            )
            successful_runs = results[:boundary_index]
            success_pulls_needed = pulls_needed[:boundary_index]
            percent_chance = len(successful_runs) / len(pulls_needed)
            print(f'Chance of full set: {percent_chance:.2%}')

            print(f'Stats for the successful runs:')
            wishing_well_percent = sum(
                [True if r.wishing_well else False for r in successful_runs]) / len(successful_runs)
            end_pull_types = Counter([r.ending for r in successful_runs])
            print(f'Average: {statistics.mean(success_pulls_needed)}')
            print(f'Median: {statistics.median(success_pulls_needed)}')
            print(f'Stdev: {statistics.stdev(success_pulls_needed)}')
            print(f'Wishing well %: {wishing_well_percent:.2%}')
            print(f'End pull types: {end_pull_types}')
            print('\n')

        success_stats(focus_results, Gacha.FOCUS, pull_count)
        success_stats(expansion_equip_results,
                      Gacha.EXPANSION_EQUIP, pull_count)
        success_stats(arsenal_results,
                      Gacha.ARSENAL, pull_count)

    # for i in range(1000):
    # print(simulate_full_set(Gacha.FOCUS))
    # print(simulate_full_set(Gacha.EXPANSION_EQUIP))