Generation.py

import random
from typing import NamedTuple, Dict, Any, Callable, Optional
import pandas as pd
import numpy as np
from enum import Enum
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt


class StepKind(Enum):
    COOL_DOWN = "COOL_DOWN"
    RUN = "RUN"
    ENEMY_1 = "ENEMY_1"
    ENEMY_1_COMBO = "ENEMY_1_COMBO"
    ENEMY_2 = "ENEMY_2"
    STAR = "STAR"
    STAR_COMBO = "STAR_COMBO"


MAPPERS: Dict[StepKind, Dict[StepKind, Callable[[int, int], int]]] = {
    StepKind.COOL_DOWN: {
        StepKind.RUN: lambda x, _t: 1 if x >= 2 else 0,
    },
    StepKind.RUN: {
        StepKind.ENEMY_1: lambda x, t: x * 0.2 + (t / 1000),
        StepKind.ENEMY_2: lambda x, t: x * 0.08 + (t / 1000),
        StepKind.STAR: lambda x, t: x * 0.08 + (t / 1000),
    },
    StepKind.ENEMY_1: {
        StepKind.COOL_DOWN: lambda _x, _t: 1,
        StepKind.ENEMY_1_COMBO: lambda _x, t: 0.2 + (t / 1000)
    },
    StepKind.ENEMY_1_COMBO: {
        StepKind.COOL_DOWN: lambda _x, _t: 1,
    },
    StepKind.ENEMY_2: {
        StepKind.COOL_DOWN: lambda _x, _t: 1,
    },
    StepKind.STAR: {
        StepKind.COOL_DOWN: lambda _x, _t: 1,
        StepKind.STAR_COMBO: lambda x, t: x * 0.3 + (t / 1000),
    },
    StepKind.STAR_COMBO: {
        StepKind.COOL_DOWN: lambda _x, _t: 1,
    }
}


def _get_next_kind(kinds: Dict[StepKind, float]) -> Optional[StepKind]:
    norm = max(sum(kinds.values()), 1)
    rand = random.random()

    cumulative = 0
    for key, value in kinds.items():
        cumulative += value / norm
        if rand < cumulative:
            return key

    return None


def step():
    x = 0
    t = 0
    kind = StepKind.COOL_DOWN

    while True:
        agg = {}
        for key, fun in MAPPERS[kind].items():
            agg[key] = fun(x, t)

        next_kind = _get_next_kind(agg)
        if next_kind is not None:
            x = 0
            kind = next_kind
        else:
            x += 1

        t += 1
        yield kind, agg


def main():
    gen = step()

    all_agg = {}

    for index, (kind, agg) in zip(range(100), gen):
        print(index, kind, agg)
        all_agg[index] = agg

    df = pd.DataFrame(all_agg).fillna(0).T

    df.plot()
    plt.show()
    print(df)


if __name__ == '__main__':
    main()