# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %load_ext nb_black

# source：
# 1. Reddit, Analysing 2650 Artifacts and their Main/Sub stat distributions. [Data compiled from user submissions led by /u/Acheron-X]， https://www.reddit.com/r/Genshin_Impact/comments/kn0r0s/analysing_2650_artifacts_and_their_mainsub_stat/
# 2. NGA, [数据讨论]《圣遗物数值学导论》 https://nga.178.com/read.php?tid=24270728

# +
import numpy as np
from enum import Enum
from collections import namedtuple
from tqdm import tqdm


class Stat(int, Enum):
    DEF = 1
    ATK = 2
    HP = 3
    ATK_PCT = 4
    HP_PCT = 5
    DEF_PCT = 6
    EM = 7
    ER = 8
    #  correct elem damage pcts
    ED = 9
    # unwanted elem damage pcts
    ED_OTHER = 10
    #  Phys damage pcts
    PD = 11
    # heal bonus pcts
    HB = 12
    CR = 13
    CD = 14


class ArtifactType(int, Enum):
    FLOWER = 101
    FEATHER = 102
    SANDS = 103
    GOBLET = 104
    CIRCLET = 105


main_stat_max = {
    Stat.ATK: 322,
    Stat.HP: 4780,
    Stat.ATK_PCT: 46.6,
    Stat.HP_PCT: 46.6,
    Stat.DEF_PCT: 58.3,
    Stat.EM: 187,
    Stat.ER: 51.8,
    Stat.ED: 46.6,
    Stat.ED_OTHER: 46.6,
    Stat.PD: 58.3,
    Stat.HB: 35.9,
    Stat.CR: 31.1,
    Stat.CD: 62.2,
}

sub_stat_ranges = {
    Stat.DEF: [16, 19, 21, 23],
    Stat.ATK: [14, 16, 18, 19],
    Stat.HP: [209, 239, 269, 299],
    Stat.ATK_PCT: [4.1, 4.7, 5.3, 5.8],
    Stat.HP_PCT: [4.1, 4.7, 5.3, 5.8],
    Stat.DEF_PCT: [5.1, 5.8, 6.6, 7.3],
    Stat.EM: [16, 19, 21, 23],
    Stat.ER: [4.5, 5.2, 5.8, 6.5],
    Stat.CR: [2.7, 3.1, 3.5, 3.9],
    Stat.CD: [5.4, 6.2, 7.0, 7.9],
}


stats_name = {
    "CN": {
        Stat.ATK: "攻击力",
        Stat.HP: "生命值",
        Stat.DEF: "防御力",
        Stat.ATK_PCT: "攻击力%",
        Stat.HP_PCT: "生命值%",
        Stat.DEF_PCT: "防御力%",
        Stat.EM: "元素精通",
        Stat.ER: "元素充能%",
        Stat.ED: "元素伤害%",
        Stat.ED_OTHER: "元素伤害%（歪）",
        Stat.PD: "物理伤害%",
        Stat.HB: "治疗加成%",
        Stat.CR: "暴击率%",
        Stat.CD: "暴击伤害%",
        ArtifactType.FLOWER: "生之花",
        ArtifactType.FEATHER: "死之羽",
        ArtifactType.SANDS: "时之沙",
        ArtifactType.GOBLET: "空之杯",
        ArtifactType.CIRCLET: "理之冠",
    }
}
# -

# ## Replace the default numbers with your estimation

# +
# initialize parameters on drop rate, main/sub-stats distribution etc
# most numbers here are based on estimation/guess


# the pdf for the number of level 5 artifact per run
default_num_lvl5_pdf = {
    1: 0.9,
    2: 0.1,
}

# the probability of an artifact with 3 initial substats
default_p_substat_cnt_3 = 0.8

# the probability of dropping the correct set
p_correct_set = 0.5

# the pdf for the artifact type
default_artifact_type_pdf = {
    ArtifactType.FLOWER: 0.2,
    ArtifactType.FEATHER: 0.2,
    ArtifactType.SANDS: 0.2,
    ArtifactType.GOBLET: 0.2,
    ArtifactType.CIRCLET: 0.2,
}

# ratio between different substats
default_substat_pdf = {
    Stat.DEF: 0.15,
    Stat.ATK: 0.15,
    Stat.HP: 0.15,
    Stat.ATK_PCT: 0.1,
    Stat.HP_PCT: 0.1,
    Stat.DEF_PCT: 0.1,
    Stat.EM: 0.09,
    Stat.ER: 0.09,
    Stat.CR: 0.08,
    Stat.CD: 0.08,
}

# main stat distribution for each type of artifact
main_stats_pdfs = {
    ArtifactType.FLOWER: {Stat.HP: 1.0},
    ArtifactType.FEATHER: {Stat.ATK: 1.0},
    ArtifactType.SANDS: {
        Stat.HP_PCT: 0.28,
        Stat.DEF_PCT: 0.28,
        Stat.ATK_PCT: 0.25,
        Stat.EM: 0.10,
        Stat.ER: 0.09,
    },
    ArtifactType.GOBLET: {
        Stat.HP_PCT: 0.19,
        Stat.DEF_PCT: 0.21,
        Stat.ATK_PCT: 0.21,
        Stat.PD: 0.04,
        Stat.EM: 0.02,
        # 0.33 shared evenly between 5 elements
        Stat.ED: 0.33 / 5,
        Stat.ED_OTHER: 0.33 * 4 / 5,
    },
    ArtifactType.CIRCLET: {
        Stat.HP_PCT: 0.22,
        Stat.DEF_PCT: 0.24,
        Stat.ATK_PCT: 0.22,
        Stat.HB: 0.07,
        Stat.EM: 0.05,
        Stat.CR: 0.10,
        Stat.CD: 0.10,
    },
}

# +
Artifact = namedtuple("Artifact", ["ttype", "main_stat", "substats"])


def normalize(pdf_p):
    return pdf_p / np.sum(pdf_p)


def prepare_pdf(pdf_dict):
    items = []
    pdf = []
    for k, v in pdf_dict.items():
        items.append(k)
        pdf.append(v)
    pdf = np.array(pdf)
    pdf = normalize(pdf)
    return items, pdf


class ArtifactGenerator(object):
    def __init__(self, ttype, main_stats_pdf, substat_pdf=None, p_substat_cnt_3=None):
        self.ttype = ttype
        self.main_stats, self.main_stats_pdf = prepare_pdf(main_stats_pdf)
        if substat_pdf is None:
            substat_pdf = default_substat_pdf
        self.sub_stats, self.sub_stats_pdf = prepare_pdf(substat_pdf)
        if p_substat_cnt_3 is None:
            self.p_substat_cnt_3 = default_p_substat_cnt_3

    def gen(self):
        # gen main stat
        main_stat = np.random.choice(self.main_stats, p=self.main_stats_pdf)
        num_init_substats = 3 if np.random.rand() < self.p_substat_cnt_3 else 4
        substats = self.gen_substats(
            main_stat=main_stat, num_substats=5 + num_init_substats,
        )
        return Artifact(ttype=self.ttype, main_stat=main_stat, substats=substats,)

    def gen_substats(self, main_stat, num_substats):
        # first draw the four types with rejection sampling
        substat_set = set()
        while len(substat_set) < min(4, num_substats):
            substat = np.random.choice(self.sub_stats, p=self.sub_stats_pdf)
            if substat == main_stat or substat in substat_set:
                continue
            substat_set.add(substat)
        # fill the numbers for each substats
        num_substat_types = len(substat_set)
        remaining_cnt = num_substats - num_substat_types
        upgrade_cnt = (
            np.random.multinomial(
                remaining_cnt, [1 / num_substat_types] * num_substat_types,
            )
            + 1
        )
        substats = {
            ttype: np.sum(
                np.random.choice(sub_stat_ranges[ttype], size=cnt, replace=True)
            )
            for cnt, ttype in zip(upgrade_cnt, substat_set)
        }
        return substats


class Simulator(object):
    def __init__(self, generators, artifact_type_pdf=None, num_lvl5_pdf=None):
        self.generators = generators
        if artifact_type_pdf is None:
            artifact_type_pdf = default_artifact_type_pdf
        self.artifact_types, self.artifact_type_pdf = prepare_pdf(artifact_type_pdf)

        if num_lvl5_pdf is None:
            num_lvl5_pdf = default_num_lvl5_pdf
        self.num_lvl5s, self.num_lvl5_pdf = prepare_pdf(num_lvl5_pdf)

    def run(self):
        cnt = np.random.choice(self.num_lvl5s, p=self.num_lvl5_pdf)
        artifacts = []
        for _ in range(cnt):
            if np.random.rand() >= p_correct_set:
                continue
            ttype = np.random.choice(self.artifact_types, p=self.artifact_type_pdf)
            artifacts.append(self.generators[ttype].gen())
        return artifacts

    def run_batch(self, num_runs):
        artifacts = {}
        for _ in tqdm(range(num_runs)):
            for v in self.run():
                if v.ttype not in artifacts:
                    artifacts[v.ttype] = []
                artifacts[v.ttype].append(v)
        return artifacts


def display(artifact, lang="CN"):
    print("===================")
    print(
        "{}:\n{:15}:\t{:3.1f}".format(
            stats_name[lang][artifact.ttype],
            stats_name[lang][artifact.main_stat],
            main_stat_max[artifact.main_stat],
        )
    )
    print("-------------------")
    for t, v in artifact.substats.items():
        print("{:15}:\t{:3.1f}".format(stats_name[lang][t], v,))
    print("===================")


# -

# ### how to compare between different artifacts

generators = {
    ttype: ArtifactGenerator(ttype=ttype, main_stats_pdf=main_stats_pdf,)
    for ttype, main_stats_pdf in main_stats_pdfs.items()
}
simulator = Simulator(generators=generators)

# # simulate the grind!

num_days = 365
num_runs = 180 * num_days // 20
print("Total number of trails: {}".format(num_runs))

artifacts = simulator.run_batch(num_runs=num_runs)


class ArtifactScore(object):
    # using the main stat max to normalize the scores
    stat_weights = {
        Stat.ATK: 1.0 / 322,
        Stat.HP: 0.0 / 4780,
        Stat.DEF: 0.0 / 187,  # no sure here
        Stat.ATK_PCT: 3.0 / 46.6,
        Stat.HP_PCT: 0.0 / 46.6,
        Stat.DEF_PCT: 0.0 / 58.3,
        Stat.EM: 0.0 / 187,
        Stat.ER: 0.0 / 51.8,
        Stat.ED: 10.0 / 46.6,
        Stat.ED_OTHER: 0.0 / 46.6,
        Stat.PD: 0.0 / 58.3,
        Stat.HB: 0.0 / 35.9,
        Stat.CR: 10.0 / 31.1,
        Stat.CD: 10.0 / 62.2,
    }

    def __init__(self, stat_weights=None):
        if stat_weights is not None:
            self.stat_weights = stat_weights

    def __call__(self, artifact):
        score = 0
        score += main_stat_max[artifact.main_stat] * self.stat_weights.get(
            artifact.main_stat, 0.0
        )
        for t, v in artifact.substats.items():
            score += v * self.stat_weights.get(t, 0.0)
        return score


# sort the artifact within each type
comparator = ArtifactScore()
for ttype in artifacts:
    artifacts[ttype].sort(key=comparator, reverse=True)

# +
lang = "CN"
topk = 3

for ttype in {
    ArtifactType.FLOWER,
    ArtifactType.FEATHER,
    ArtifactType.SANDS,
    ArtifactType.GOBLET,
    ArtifactType.CIRCLET,
}:
    print("Got {} {}.".format(len(artifacts[ttype]), stats_name[lang][ttype]))
    for i in range(topk):
        display(artifacts[ttype][i], lang=lang)
    print("\n\n")
# -

# ## how many runs do I need to get a circlet with CR/CD as main and CD/CR as good substat

# +
num_runs_needed = []
num_trials = 500

scorer1 = ArtifactScore(
    stat_weights = {       
        Stat.CD: 1.0 
    }
)
# # cd in substats increased more than 3 times
substat_thd = 20
for _ in tqdm(range(num_trials)):
    cnt = 0
    found = False
    while not found:
        cnt += 1
        for v in simulator.run():
            if v.ttype != ArtifactType.CIRCLET or v.main_stat != Stat.CR:
                continue
            if scorer1(v) > substat_thd:
                num_runs_needed.append(cnt)
                found = True
                break           


# +
num_runs_needed.sort()

print("10-percentile: {}".format(num_runs_needed[num_trials // 10]))
print("50-percentile: {}".format(num_runs_needed[num_trials // 2]))
print("90-percentile: {}".format(num_runs_needed[num_trials - num_trials // 10]))
# -

# ## how many runs do I need to get a SANDS with ATK_PCT as main and CD&CR as good substat

# +
num_runs_needed = []
num_trials = 500

scorer1 = ArtifactScore(
    stat_weights = {       
        Stat.CD: 1.0,
        Stat.CR: 2.0,
    }
)
# # cd & cr in substats increased more than 3 times
substat_thd = 28
for _ in tqdm(range(num_trials)):
    cnt = 0
    found = False
    while not found:
        cnt += 1
        for v in simulator.run():
            if v.ttype != ArtifactType.SANDS or v.main_stat != Stat.ATK_PCT:
                continue
            if scorer1(v) > substat_thd:
                num_runs_needed.append(cnt)
                found = True
                break       


# +
display(v)

num_runs_needed.sort()

print("10-percentile: {}".format(num_runs_needed[num_trials // 10]))
print("50-percentile: {}".format(num_runs_needed[num_trials // 2]))
print("90-percentile: {}".format(num_runs_needed[num_trials - num_trials // 10]))
# -

# ## how many runs do I need to get a GOBLET with Elem Damage as main and CD&CR as good substat

# +
num_runs_needed = []
num_trials = 500

scorer1 = ArtifactScore(
    stat_weights = {       
        Stat.CD: 1.0,
        Stat.CR: 2.0,
    }
)
# # cd & cr in substats increased more than 3 times
substat_thd = 28
for _ in tqdm(range(num_trials)):
    cnt = 0
    found = False
    while not found:
        cnt += 1
        for v in simulator.run():
            if v.ttype != ArtifactType.GOBLET or v.main_stat != Stat.ED:
                continue
            if scorer1(v) > substat_thd:
                num_runs_needed.append(cnt)
                found = True
                break       


# +
display(v)

num_runs_needed.sort()

print("10-percentile: {}".format(num_runs_needed[num_trials // 10]))
print("50-percentile: {}".format(num_runs_needed[num_trials // 2]))
print("90-percentile: {}".format(num_runs_needed[num_trials - num_trials // 10]))
# -

# ## how many runs do I need to get a FLOWER with CD&CR as good substat

# +
num_runs_needed = []
num_trials = 500

scorer1 = ArtifactScore(
    stat_weights = {       
        Stat.CD: 1.0,
        Stat.CR: 2.0,
    }
)
# # cd & cr in substats increased more than 3 times
substat_thd = 28
for _ in tqdm(range(num_trials)):
    cnt = 0
    found = False
    while not found:
        cnt += 1
        for v in simulator.run():
            if v.ttype != ArtifactType.FLOWER:
                continue
            if scorer1(v) > substat_thd:
                num_runs_needed.append(cnt)
                found = True
                break       


# +
display(v)

num_runs_needed.sort()

print("10-percentile: {}".format(num_runs_needed[num_trials // 10]))
print("50-percentile: {}".format(num_runs_needed[num_trials // 2]))
print("90-percentile: {}".format(num_runs_needed[num_trials - num_trials // 10]))
# -



StarcoderdataJupyterScriptsDedupFiltered