"""Score functions."""

from __future__ import annotations

from collections.abc import Callable

ScoreFn = Callable[[bool, int, int], int]
SCORE_FNS: dict[str, ScoreFn] = {}


def register(name: str) -> Callable[[ScoreFn], ScoreFn]:
    def decorator(func: ScoreFn) -> ScoreFn:
        SCORE_FNS[name] = func
        return func

    return decorator


@register("linear_decay")
def linear_decay(correct: bool, elapsed_ms: int, time_limit_ms: int) -> int:
    if not correct:
        return 0
    elapsed_ms = max(0, min(elapsed_ms, time_limit_ms))
    return round(1000 * (1 - 0.5 * elapsed_ms / time_limit_ms))


@register("flat")
def flat(correct: bool, elapsed_ms: int, time_limit_ms: int) -> int:
    return 1000 if correct else 0


@register("exponential_decay")
def exponential_decay(correct: bool, elapsed_ms: int, time_limit_ms: int) -> int:
    if not correct:
        return 0
    import math

    elapsed_ms = max(0, min(elapsed_ms, time_limit_ms))
    decay = math.exp(-2 * elapsed_ms / time_limit_ms)
    return round(1000 * (0.5 + 0.5 * decay))