Genshin Impact Wishing model

# -*- coding: utf-8 -*-
"""
Created on Thu Sep  9 23:09:15 2021

@author: Bobrokrot @ Genshin Academy https://vk.com/genshinacademy
Statistical model: Sengalev https://www.hoyolab.com/genshin/article/497840
"""

from random import random, choice
import matplotlib.pyplot as plt
import numpy as np
import csv

class ItemGiver:
    def __init__(self, star = 5, banner = "character"):
        #Items: Upper_case = 5 star, lower case = 4 star
        #A,B are features, Z,Y,X and so on are standard
        self.star = star
        self.banner = banner
        if banner == "character" and star == 5:
            self.rate = 0.5
            self.featured_items = ["A"]
            self.standard_items = ["Q", "R", "S", "T", "U"]
            self.path = False
        elif banner == "weapon" and star == 5:
            self.rate = 0.75
            self.featured_items = ["A", "B"]
            self.standard_items = ["Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"]
            self.path = True
        elif banner == "character" and star == 4:
            self.rate = 0.5
            self.featured_items = ["a", "b", "c"]
            self.standard_items = ["m"] #too many
            self.path = False
        elif banner == "weapon" and star == 4:
            self.rate = 0.75
            self.featured_items = ["a", "b", "c", "d", "e"]
            self.standard_items = ["m"] #too many
            self.path = False
        elif banner == "standard" and star == 5:
            #untested
            self.rate = 0.5
            self.featured_items = ["A", "B", "C", "D", "E"]
            self.standard_items = ["V", "W", "X", "Y", "Z"]
            self.path = False
        elif banner == "standard" and star == 4:
            self.rate = 0.5
            self.featured_items = ["m"]
            self.standard_items = ["m"]
            self.path = False
        else:
            raise "star or banner is wrong"
        self.last_featured = True
        self.path_counter = 0
        self.history = []

    def drop(self):
        if self.path and self.path_counter == 2:
            item = "A"
            self.path_counter = 0
            self.last_featured = True
        else:
            if not self.last_featured or random() < self.rate: #give featured
                item = choice(self.featured_items)
                self.last_featured = True
            else:
                item = choice(self.standard_items)
                if self.banner != "standard":
                    self.last_featured = False
                else:
                    self.last_featured = True
            if self.path:
                if item == "A":
                    self.path_counter = 0
                else:
                    self.path_counter += 1

        self.history.append(item)
        return item


class Banner:
    def __init__(self, banner = "character"):
        self.banner = banner
        if banner == "character":
            self._rate5 = .006
            self._rate4 = 0.051
            self._pity5 = 73 #pity starts after 73
            self._pity4 = 8
        elif banner == "weapon":
            self._rate5 = .007
            self._rate4 = 0.06
            self._pity5 = 62 #pity starts after 62
            self._pity4 = 7
        else:
            raise "banner is wrong"
        self.giver5 = ItemGiver(5, banner)
        self.giver4 = ItemGiver(4, banner)
        self.roll_counter = 0
        self.pity_counter4 = 1
        self.pity_counter5 = 1
        self.history = ""

    def roll(self):
        self.roll_counter += 1

        prob5 = self._rate5 + max(0, (self.pity_counter5 - self._pity5)*10*self._rate5)
        prob4 = self._rate4 + max(0, (self.pity_counter4 - self._pity4)*10*self._rate4)

        x = random()
        if x < prob5: #5 star drop
            item = self.giver5.drop()
            self.pity_counter5 = 1
            self.pity_counter4 += 1
        elif x < prob4 + prob5: #4 star drop
            item = self.giver4.drop()
            self.pity_counter5 += 1
            self.pity_counter4 = 1

        else: #3 star drop
            item = "."
            self.pity_counter5 += 1
            self.pity_counter4 += 1

        self.history += item
        return item

    def rollN(self, X):
        for i in range(X):
            self.roll()
        return self.history

    def rollUntil(self, items = ["A"]):
        #input a list of items
        #stop rolling until an item from the list drops
        voc = self.giver4.featured_items + self.giver4.standard_items \
            + self.giver5.featured_items + self.giver5.standard_items
        for desired in items:
            if desired not in voc:
                raise "this item is not in the banner!"

        while(self.roll() not in items):
            pass
        return self.history

    def reset(self):
        self.roll_counter = 0
        self.pity_counter4 = 1
        self.pity_counter5 = 1
        self.history = ""

def pullUntil(banner, items = ["A"], constellation = 0, N = 1000000):
    vals = np.zeros(10000) #distribution
    for i in range(N):
        banner.reset()
        for j in range(0, constellation + 1):
            banner.rollUntil(items)
        vals[banner.roll_counter] += 1
    return np.trim_zeros(vals, "b")

if __name__ == '__main__':

    banner = Banner("weapon") #weapon for weapon banner, character for character banner
    N = 1000000 #size of sample
    desired_items = ["A"] #list of items you are aiming to pull for
    constellation = 0 #number of items+1 until stop
    #vals = pullUntil(banner, ["A"], 4, N) # — pull until R5 banner weapon if C is weapon banner
    #vals = pullUntil(banner, ["A, B"], 0, N) # — pull until either of 5 star weapons drop
    #C = Banner("character")
    #vals = pullUntil(banner, ["A"], 6, N) # — pull until C6 banner charater
    #vals = pullUntil(banner, ["A", "Q", "R", "S", "T", "U"], 0, N) # — pull until either
    # banner character (A) or any of the standard characters (QRSTU)
    #vals = pullUntil(banner, ["a"], 4, N) # — pull until a particular rate-up C4 4 star

    vals = pullUntil(banner, desired_items, constellation, N)
    n = vals.shape[0]
    sums = np.zeros(n) #cumulative
    for i in range(n):
        sums[i] = sum(vals[:i+1])

    """
    #export to file
    with open('desired5starweapon.csv', 'w', newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["N", "Распределение", "Кумулятивное"])
        for i in range(1,n):
            writer.writerow([i, "{:.7f}".format(vals[i] / N), sums[i] / N])
    """

    plt.figure()
    plt.bar(range(0,n), vals / N, width = 1)

    xran = np.rint(np.linspace(1, n, 37))
    plt.xticks(xran, list(map(lambda x: '{:.0f}'.format(x), xran)))
    #plt.xticks([1] + list(range(0,181,5))[1:])
    S = np.linspace(0,0.1, 11)
    plt.yticks(S, list(map(lambda x: '{:.0f}%'.format(x*100), S)))
    plt.xlim(0,n)
    plt.ylim(0,S[-1])
    plt.grid()
    plt.title("Распределение вероятности")


    plt.figure()
    plt.bar(range(0,n), sums / N, width = 1)
    plt.xticks(xran, list(map(lambda x: '{:.0f}'.format(x), xran)))
    S = np.linspace(0,1,21)
    plt.yticks(S, list(map(lambda x: '{:.0f}%'.format(x*100), S)))
    plt.xlim(0,n)
    plt.ylim(0,S[-1])
    plt.grid()
    plt.title("Кумулятивная вероятность")
    plt.show()


    """
    xran = list(range(1, 81))
    vals = np.zeros(90)
    _rate = .007
    _pity = 62
    sum = 0
    for x in xran:
        prob5 = min(1, _rate + max(0, (x - _pity)*10*_rate))
        vals[x-1] = prob5
    with open('5starweapon_rate.csv', 'w', newline="") as f:
        writer = csv.writer(f, quoting=csv.QUOTE_NONE)
        writer.writerow(["N", "Вероятность"])
        for x in xran:
            writer.writerow([x, vals[x-1]])
    """