from enum import Enum


class Treatment(Enum):
    BaselineOne = "BaselineOne"
    BaselineTwo = "BaselineTwo"
    AaPolicy = "AaPolicy"
    Lucky = "Lucky"

    @classmethod
    def choices(cls):
        return [_.value for _ in cls]


class TreatmentRange(Enum):
    # treatment applies to the whole session
    # participants in the same session must have same treatment
    InterSession = "inter_session"
    # treatment applies to individual participants
    # they can have different treatments
    InterParticipant = "inter_participant"