#!/usr/bin/env python3.11 from __future__ import annotations import argparse import asyncio import math import os import sys from collections.abc import Awaitable from collections.abc import Iterator from collections.abc import Sequence from dataclasses import dataclass from dataclasses import field from pathlib import Path from typing import Any from typing import TypeVar import databases from akatsuki_pp_py import Beatmap from akatsuki_pp_py import Calculator from redis import asyncio as aioredis sys.path.insert(0, os.path.abspath(os.pardir)) os.chdir(os.path.abspath(os.pardir)) try: import app.settings import app.state.services from app.constants.gamemodes import GameMode from app.constants.mods import Mods from app.constants.privileges import Privileges from app.objects.beatmap import ensure_osu_file_is_available except ModuleNotFoundError: print("\x1b[;91mMust run from tools/ directory\x1b[m") raise T = TypeVar("T") DEBUG = False BEATMAPS_PATH = Path.cwd() / ".data/osu" @dataclass class Context: database: databases.Database redis: aioredis.Redis beatmaps: dict[int, Beatmap] = field(default_factory=dict) def divide_chunks(values: list[T], n: int) -> Iterator[list[T]]: for i in range(0, len(values), n): yield values[i : i + n] async def recalculate_score( score: dict[str, Any], beatmap_path: Path, ctx: Context, ) -> None: beatmap = ctx.beatmaps.get(score["map_id"]) if beatmap is None: beatmap = Beatmap(path=str(beatmap_path)) ctx.beatmaps[score["map_id"]] = beatmap calculator = Calculator( mode=GameMode(score["mode"]).as_vanilla, mods=score["mods"], combo=score["max_combo"], n_geki=score["ngeki"], # Mania 320s n300=score["n300"], n_katu=score["nkatu"], # Mania 200s, Catch tiny droplets n100=score["n100"], n50=score["n50"], n_misses=score["nmiss"], ) attrs = calculator.performance(beatmap) new_pp: float = attrs.pp if math.isnan(new_pp) or math.isinf(new_pp): new_pp = 0.0 new_pp = min(new_pp, 9999.999) await ctx.database.execute( "UPDATE scores SET pp = :new_pp WHERE id = :id", {"new_pp": new_pp, "id": score["id"]}, ) if DEBUG: print( f"Recalculated score ID {score['id']} ({score['pp']:.3f}pp -> {new_pp:.3f}pp)", ) async def process_score_chunk( chunk: list[dict[str, Any]], ctx: Context, ) -> None: tasks: list[Awaitable[None]] = [] for score in chunk: osu_file_available = await ensure_osu_file_is_available( score["map_id"], expected_md5=score["map_md5"], ) if osu_file_available: tasks.append( recalculate_score( score, BEATMAPS_PATH / f"{score['map_id']}.osu", ctx, ), ) await asyncio.gather(*tasks) async def recalculate_user( id: int, game_mode: GameMode, ctx: Context, ) -> None: best_scores = await ctx.database.fetch_all( "SELECT s.pp, s.acc FROM scores s " "INNER JOIN maps m ON s.map_md5 = m.md5 " "WHERE s.userid = :user_id AND s.mode = :mode " "AND s.status = 2 AND m.status IN (2, 3) " # ranked, approved "ORDER BY s.pp DESC", {"user_id": id, "mode": game_mode}, ) total_scores = len(best_scores) if not total_scores: return # calculate new total weighted accuracy weighted_acc = sum(row["acc"] * 0.95**i for i, row in enumerate(best_scores)) bonus_acc = 100.0 / (20 * (1 - 0.95**total_scores)) acc = (weighted_acc * bonus_acc) / 100 # calculate new total weighted pp weighted_pp = sum(row["pp"] * 0.95**i for i, row in enumerate(best_scores)) bonus_pp = 416.6667 * (1 - 0.9994**total_scores) pp = round(weighted_pp + bonus_pp) await ctx.database.execute( "UPDATE stats SET pp = :pp, acc = :acc WHERE id = :id AND mode = :mode", {"pp": pp, "acc": acc, "id": id, "mode": game_mode}, ) user_info = await ctx.database.fetch_one( "SELECT country, priv FROM users WHERE id = :id", {"id": id}, ) if user_info is None: raise Exception(f"Unknown user ID {id}?") if user_info["priv"] & Privileges.UNRESTRICTED: await ctx.redis.zadd( f"bancho:leaderboard:{game_mode.value}", {str(id): pp}, ) await ctx.redis.zadd( f"bancho:leaderboard:{game_mode.value}:{user_info['country']}", {str(id): pp}, ) if DEBUG: print(f"Recalculated user ID {id} ({pp:.3f}pp, {acc:.3f}%)") async def process_user_chunk( chunk: list[int], game_mode: GameMode, ctx: Context, ) -> None: tasks: list[Awaitable[None]] = [] for id in chunk: tasks.append(recalculate_user(id, game_mode, ctx)) await asyncio.gather(*tasks) async def recalculate_mode_users(mode: GameMode, ctx: Context) -> None: user_ids = [ row["id"] for row in await ctx.database.fetch_all("SELECT id FROM users") ] for id_chunk in divide_chunks(user_ids, 100): await process_user_chunk(id_chunk, mode, ctx) async def recalculate_mode_scores(mode: GameMode, ctx: Context) -> None: scores = [ dict(row) for row in await ctx.database.fetch_all( """\ SELECT scores.id, scores.mode, scores.mods, scores.map_md5, scores.pp, scores.acc, scores.max_combo, scores.ngeki, scores.n300, scores.nkatu, scores.n100, scores.n50, scores.nmiss, maps.id as `map_id` FROM scores INNER JOIN maps ON scores.map_md5 = maps.md5 WHERE scores.status = 2 AND scores.mode = :mode ORDER BY scores.pp DESC """, {"mode": mode}, ) ] for score_chunk in divide_chunks(scores, 100): await process_score_chunk(score_chunk, ctx) async def main(argv: Sequence[str] | None = None) -> int: argv = argv if argv is not None else sys.argv[1:] parser = argparse.ArgumentParser( description="Recalculate performance for scores and/or stats", ) parser.add_argument( "-d", "--debug", help="Enable debug logging", action="store_true", ) parser.add_argument( "--no-scores", help="Disable recalculating scores", action="store_true", ) parser.add_argument( "--no-stats", help="Disable recalculating user stats", action="store_true", ) parser.add_argument( "-m", "--mode", nargs=argparse.ONE_OR_MORE, required=False, default=["0", "1", "2", "3", "4", "5", "6", "8"], # would love to do things like "vn!std", but "!" will break interpretation choices=["0", "1", "2", "3", "4", "5", "6", "8"], ) args = parser.parse_args(argv) global DEBUG DEBUG = args.debug db = databases.Database(app.settings.DB_DSN) await db.connect() redis = await aioredis.from_url(app.settings.REDIS_DSN) ctx = Context(db, redis) for mode in args.mode: mode = GameMode(int(mode)) if not args.no_scores: await recalculate_mode_scores(mode, ctx) if not args.no_stats: await recalculate_mode_users(mode, ctx) await app.state.services.http_client.aclose() await db.disconnect() await redis.aclose() return 0 if __name__ == "__main__": raise SystemExit(asyncio.run(main()))