"""Curriculum helpers."""

from __future__ import annotations

from app.common.enums import Difficulty, SubEnvironment


def pick_difficulty(episode_index: int) -> Difficulty:
    if episode_index < 20:
        return Difficulty.EASY
    if episode_index < 40:
        return Difficulty.MEDIUM
    return Difficulty.HARD


def pick_sub_environment(episode_index: int, difficulty: Difficulty) -> SubEnvironment:
    # Curriculum starts in DDI/risk-heavy cases then introduces bandits/dosing.
    if difficulty == Difficulty.EASY:
        return SubEnvironment.DDI if episode_index % 2 == 0 else SubEnvironment.REGIMEN_RISK
    if difficulty == Difficulty.MEDIUM:
        if episode_index % 4 == 0:
            return SubEnvironment.BANDIT_MINING
        if episode_index % 4 == 1:
            return SubEnvironment.ALTERNATIVE_SUGGESTION
        return SubEnvironment.REGIMEN_RISK
    if episode_index % 5 == 0:
        return SubEnvironment.PRECISION_DOSING
    if episode_index % 5 == 1:
        return SubEnvironment.LONGITUDINAL_DEPRESCRIBING
    if episode_index % 5 == 2:
        return SubEnvironment.WEB_SEARCH_MISSING_DATA
    if episode_index % 5 == 3:
        return SubEnvironment.NEW_DRUG_DECOMPOSITION
    return SubEnvironment.REGIMEN_RISK