Equipment Expansion Supply

import numpy as np


class Gacha:
    def __init__(self, rate_r, rate_ssr, rate_target, guarantee, pity):
        self.rate_r = rate_r
        self.rate_ssr = rate_ssr
        self.rate_target = rate_target
        self.guarantee = guarantee
        self.pity = pity
        self.limit = 200
        self.ops_count = 0

    def __init__(self):
        self.hi3rd_expansion()

    def set_limit(self,limit):
        self.limit = limit

    def prob(self,table,state,pulls):

##        self.ops_count += 1
##        if self.ops_count == 1000000:
##            print(self.ops_count)

        wp_w8 = (state.targets_goal[0] - state.targets_counter[0])*2
        stg_w8 = (state.targets_goal[1] - state.targets_counter[1])
        max_w8 = wp_w8 + stg_w8

        if state.are_all_targets_acquired():
            return [1,0]

        if pulls == self.limit:
            return [0,0]

        if table.get(state.get_tuple()):
            return table[state.get_tuple()]

        if self.pity == state.pity_counter+1:
            val = [0,0]

            if state.targets_counter[0] < state.targets_goal[0]:
                pity_state = state.copy()
                pity_state.pulls += 1
                pity_state.got_target()
                pity_state.targets_counter = (pity_state.targets_counter[0]+1,
                                              pity_state.targets_counter[1])
                temp_val = self.prob(table,pity_state,pulls+1)
                val[0] += temp_val[0]*wp_w8/max_w8
                val[1] += (temp_val[1]+1)*wp_w8/max_w8

            if state.targets_counter[1] < state.targets_goal[1]:
                pity_state = state.copy()
                pity_state.pulls += 1
                pity_state.got_target()
                pity_state.targets_counter = (pity_state.targets_counter[0],
                                              pity_state.targets_counter[1]+1)
                temp_val = self.prob(table,pity_state,pulls+1)
                val[0] += temp_val[0]*stg_w8/max_w8
                val[1] += (temp_val[1]+1)*stg_w8/max_w8

            table[state.get_tuple()] = val
            return val

        if self.guarantee == state.guarantee_counter+1:
            val = [0,0]

            if state.targets_counter[0] < state.targets_goal[0]:
                brag_state = state.copy()
                brag_state.pulls += 1
                brag_state.got_target()
                brag_state.targets_counter = (brag_state.targets_counter[0]+1,
                                              brag_state.targets_counter[1])
                temp_val = self.prob(table,brag_state,pulls+1)
                brag_mult = self.rate_target*wp_w8/self.rate_ssr
                val[0] += temp_val[0]*brag_mult
                val[1] += (temp_val[1]+1)*brag_mult

            if state.targets_counter[1] < state.targets_goal[1]:
                brag_state = state.copy()
                brag_state.pulls += 1
                brag_state.got_target()
                brag_state.targets_counter = (brag_state.targets_counter[0],
                                              brag_state.targets_counter[1]+1)
                temp_val = self.prob(table,brag_state,pulls+1)
                brag_mult = self.rate_target*stg_w8/self.rate_ssr
                val[0] += temp_val[0]*brag_mult
                val[1] += (temp_val[1]+1)*brag_mult


            salt_state = state.copy()
            salt_state.pulls += 1
            salt_state.got_salt()
            temp_val = self.prob(table,salt_state,pulls+1)
            salt_mult = 1-(self.rate_target*max_w8/self.rate_ssr)
            val[0] += temp_val[0]*salt_mult
            val[1] += (temp_val[1]+1)*salt_mult

            table[state.get_tuple()] = val
            return val

        val = [0,0]

        if state.targets_counter[0] < state.targets_goal[0]:
                brag_state = state.copy()
                brag_state.pulls += 1
                brag_state.got_target()
                brag_state.targets_counter = (brag_state.targets_counter[0]+1,
                                              brag_state.targets_counter[1])
                temp_val = self.prob(table,brag_state,pulls+1)
                brag_mult = self.rate_target*wp_w8
                val[0] += temp_val[0]*brag_mult
                val[1] += (temp_val[1]+1)*brag_mult

        if state.targets_counter[1] < state.targets_goal[1]:
                brag_state = state.copy()
                brag_state.pulls += 1
                brag_state.got_target()
                brag_state.targets_counter = (brag_state.targets_counter[0],
                                              brag_state.targets_counter[1]+1)
                temp_val = self.prob(table,brag_state,pulls+1)
                brag_mult = self.rate_target*stg_w8
                val[0] += temp_val[0]*brag_mult
                val[1] += (temp_val[1]+1)*brag_mult

        salt_state = state.copy()
        salt_state.pulls += 1
        salt_state.got_salt()
        temp_val = self.prob(table,salt_state,pulls+1)
        salt_mult = self.rate_ssr - self.rate_target*max_w8
        val[0] += temp_val[0]*salt_mult
        val[1] += (temp_val[1]+1)*salt_mult

        uber_salt_state = state.copy()
        uber_salt_state.pulls += 1
        uber_salt_state.got_uber_salt()
        temp_val = self.prob(table,uber_salt_state,pulls+1)
        uber_salt_mult = self.rate_r
        val[0] += temp_val[0]*uber_salt_mult
        val[1] += (temp_val[1]+1)*uber_salt_mult

        table[state.get_tuple()] = val
        return val

    def hi3rd_expansion(self):
        self.rate_r = 0.876
        self.rate_ssr = 0.124
        self.rate_target = 0.0092
        self.guarantee = 10
        self.pity = 50
        self.limit = -1
        self.ops_count = 0


class State:
    def __init__(self, guarantee_counter, pity_counter, pulls):
        self.guarantee_counter = guarantee_counter
        self.pity_counter = pity_counter
        self.pulls = pulls
        self.targets_counter = (0,0)
        self.targets_goal = (1,3)

    def are_all_targets_acquired(self):
        return self.targets_counter == self.targets_goal

    def reset_counter(self):
        self.guarantee_counter = 0
        self.pity_counter = 0

    def copy(self):
        copied_state = State(self.guarantee_counter, self.pity_counter, self.pulls)
        copied_state.targets_counter = self.targets_counter
        copied_state.targets_goal = self.targets_goal
        return copied_state

    def got_target(self):
##      self.targets_counter += 1

        self.reset_counter()

    def got_salt(self):
        self.guarantee_counter = 0
        self.pity_counter += 1

    def got_uber_salt(self):
        self.guarantee_counter += 1
        self.pity_counter += 1

    def get_tuple(self):
        return (self.guarantee_counter, self.pity_counter, self.pulls, self.targets_counter)

def prob_table():
    gacha = Gacha()
    for i in range(40):
        for j in range(5):
            gacha.set_limit(i*5+j+1)
            table = {}
            init_state = State(0,0,0)
            val = gacha.prob(table,init_state,0)
            print("%3i : %6.5f " % (i*5+j+1,val[0]), end=" ")
        print(" ")