# Pricing per 1M tokens:
# -----------------------------------------------------------------------------
# Model               | Input (up to 128K tokens) | Output (up to 128K tokens) |
#                     | Input (over 128K tokens)  | Output (over 128K tokens)  |
# -----------------------------------------------------------------------------
# GPT-4o              | $5                        | $15                        |
# GPT-4o Mini         | $0.15                     | $0.6                       |
# Claude 3.5 Sonnet   | $3                        | $15                        |
# Gemini 1.5 Pro      | $3.5 / $7.0               | $10.5 / $21.0              |
# Gemini 1.5 Flash    | $0.35 / $0.7              | $1.05 / $2.10              |
# -----------------------------------------------------------------------------
#
# Request and Token Limits:
# https://docs.anthropic.com/en/api/rate-limits
# https://platform.openai.com/docs/guides/rate-limits/usage-tiers?context=tier-two
# -------------------------------------------------------------------------------------
# Model                       | Token Limits  | Request Limits            | Batch Queue Limits |
# -------------------------------------------------------------------------------------
# GPT-4o(Tier 1)              | 30,000 TPM    | 500 RPM                   | 90,000 TPD         |
# GPT-4o(Tier 2)              | 450,000 TPM   | 5000 RPM                  | 1,350,000 TPD      |
# GPT-4o(Tier 3)              | 800,000 TPM   | 5000 RPM                  | 50,000,000 TPD     |
# GPT-4o(Tier 4)              | 2,000,000 TPM | 10000 RPM                 | 200,000,000 TPD    |
# GPT-4o(Tier 5)              | 30,000,000 TPM| 10000 RPM                 | 5,000,000,000 TPD  |
# GPT-4o Mini                 | 200,000 TPM   | 500 RPM (10,000 RPD)      | 2,000,000 TPD      |
# Claude 3.5 Sonnet(Tier 2)   | 80,000 TPM    | 1,000 RPM (2,500,000 TPD) | Not specified      |
# Claude 3.5 Sonnet(Tier 3)   | 160,000 TPM   | 2,000 RPM (5,000,000 TPD) | Not specified      |
# Claude 3.5 Sonnet(Tier 4)   | 400,000 TPM   | 4,000 RPM (50,000,000 TPD)| Not specified      |
# Gemini 1.5 Pro              | 2,000,000 TPM | 360 RPM (10,000 RPD)      | Not specified      |
# Gemini 1.5 Flash            | 2,000,000 TPM | 1,000 RPM                 | Not specified      |
# -------------------------------------------------------------------------------------
#
# Vision
# https://ai.google.dev/gemini-api/docs/vision?lang=python
# https://platform.openai.com/docs/guides/vision
# https://docs.anthropic.com/en/docs/build-with-claude/vision#example-multiple-images
#
# Key:
# TPM: Tokens Per Minute
# RPM: Requests Per Minute
# RPD: Requests Per Day
# TPD: Tokens Per Day

# API Keys
OPENAI_API_KEY = "your_openai_api_key"
ANTHROPIC_API_KEY = "your_anthropic_api_key"
GOOGLE_API_KEY = "your_google_api_key"

MAX_OUTPUT_TOKENS = 4096

# Token limits Per Minute (TPM)
TPM_LIMITS = {
    # "GPT-4o": 30000,    # Tier 1
    # "GPT-4o": 450000,   # Tier 2
    "GPT-4o": 800000,   # Tier 3
    # "GPT-4o": 2000000,  # Tier 4
    # "GPT-4o": 3000000,  # Tier 5
    # "Claude-3.5 Sonnet": 80000,   # Tier 2
    # "Claude-3.5 Sonnet": 160000,  # Tier 3
    "Claude-3.5 Sonnet": 400000,  # Tier 4
    "Gemini-1.5 Pro": 2000000
}
# Request limits Per Minute (RPM)
RPM_LIMITS = {
    # "GPT-4o": 480,   # Tier 1
    # "GPT-4o": 4900,  # Tier 2
    "GPT-4o": 4900,  # Tier 3
    # "GPT-4o": 9900,  # Tier 4
    # "GPT-4o": 9900,  # Tier 5
    # "Claude-3.5 Sonnet": 980,   # Tier 2
    # "Claude-3.5 Sonnet": 1980,  # Tier 3
    "Claude-3.5 Sonnet": 3980,  # Tier 4
    "Gemini-1.5 Pro": 340
}
# Token limits Per Day (TPD)
TPD_LIMITS = {
    "GPT-4o": None,
    # "Claude-3.5 Sonnet": 2500000,   # Tier 2
    # "Claude-3.5 Sonnet": 5000000,   # Tier 3
    "Claude-3.5 Sonnet": 50000000,  # Tier 4
    "Gemini-1.5 Pro": None
}
# Request limits Per Day (RPD)
RPD_LIMITS = {
    "GPT-4o": None,
    "Claude-3.5 Sonnet": None,
    "Gemini-1.5 Pro": 10000
}

PROMPT_PRICE_MODELS = {
    "GPT-4o": 5,
    "Claude-3.5 Sonnet": 3,
    "Gemini-1.5 Pro": 3.5
}

COMPLETION_PRICE_MODELS = {
    "GPT-4o": 15,
    "Claude-3.5 Sonnet": 15,
    "Gemini-1.5 Pro": 10.5
}

# Task descriptions
TASK_DESCRIPTIONS = {
    "WCST": f"In this task, you will be presented with an image containing five cards. Each card has figures that differ in number, color, and shape. Four cards will be displayed at the top of the image, and one card will appear in the bottom left corner of the image. Your task will be to try to determine which of the four cards at the top the bottom left card belongs to. The correct answer depends upon a rule, which will be based solely on either the number of shapes, the color of the shapes, or the shape type itself. There will be no combination of these characteristics to define the rule. For example: If the rule is to match according to color, you should choose the card that matches the color of the bottom left card. If the rule is to match according to shape, you should choose the card that has the same shape as the bottom left card. If the rule is to match according to the number of shapes, you should choose the card that has the same number of shapes as the bottom left card. We will not tell you what the rule is. However, we will provide feedback after each selection, indicating only whether your selection was correct or incorrect, without any feedback into your reasoning process. The rule may change during the task. When the rule changes, you should figure out what the rule is as quickly as possible. When you make your selection, for each trial, respond ONLY in this format: Selection: [Card number, 1-4]. Example response: Selection: 2. You can just choose from the top four cards, not from the bottom left card. Feedback will be given after each selection, and you will be presented a different image. This task consists of 64 trials involving different images. Please proceed with the first image.",
    "WCST_without_restriction": f"In this task, you will be presented with an image containing five cards. Each card has figures that differ in number, color, and shape. Four cards will be displayed at the top of the image, and one card will appear in the bottom left corner of the image. Your task will be to try to determine which of the four cards at the top the bottom left card belongs to. The correct answer depends upon a rule, which will be based solely on either the number of shapes, the color of the shapes, or the shape type itself. We will not tell you what the rule is. However, we will provide feedback after each selection, indicating only whether your selection was correct or incorrect, without any feedback into your reasoning process. The rule may change during the task. When the rule changes, you should figure out what the rule is as quickly as possible. When you make your selection, for each trial, respond ONLY in this format: Selection: [Card number, 1-4]. Example response: Selection: 2. You can just choose from the top four cards, not from the bottom left card. Feedback will be given after each selection, and you will be presented a different image. This task consists of 64 trials involving different images. Please proceed with the first image.",
}
TASK_DESCRIPTIONS_CN = {
    "WCST": "在这个任务中，图片顶部有四张卡片，左下角有一张卡片。你的任务是确定左下角的卡片属于顶部的哪一张卡片。正确答案取决于一个规则，图形的数量、颜色或形状。例如：如果规则是颜色，你应该选择与左下角卡片颜色相同的卡片。如果规则是形状，你应该选择具有相同形状的卡片。如果规则是图形数量，你应该选择具有相同数量图形的卡片。任务不会告诉你规则是什么，但是，每次选择后会在屏幕上提示您的选择是正确还是错误，如果您的匹配不正确，您需要尝试其他规则。一旦您确定了要使用的规则，就可以放松一会儿了。但是，这还不是全部。匹配规则会不时更改，您需要仔细监控反馈，如果给您错误信息，则需要更改规则，就是这样。下方的四个按钮代表图片顶部的四张卡片，每次选择时，点击一个对应的按钮。这个任务包含64次试验，请开始第一张图片。",
    "WCST_without_restriction": "在这个任务中，你将会看到一张包含五张卡片的图片。每张卡片上的图形在数量、颜色和形状上各不相同。图片顶部将展示四张卡片，左下角会出现一张卡片。你的任务是确定左下角的卡片属于顶部的哪一张卡片。正确答案取决于一个规则，该规则仅基于图形的数量、颜色或形状类型之一。我们不会告诉你规则是什么。但是，我们将在每次选择后提供反馈，仅表明您的选择是正确的还是错误的，而不会反馈您的推理过程。规则可能在任务中改变。当规则改变时，你应该尽快找出新的规则是什么。每次选择时，请仅以以下格式回答：选择：[卡片编号, 1-4]。例如回答：选择：2。你只能从顶部四张卡片中选择，不能选择左下角的卡片。每次选择后将给出反馈，并且你将看到不同的图片。这个任务包含64次试验，涉及不同的图片。请开始第一张图片。",
}

VISION_QUESTIONS = {
    "Overall" : "First, state how many cards you see in the image. Then, describe each card in detail. The cards contain shapes that vary in color, type, and number. For each card, provide the following information: number of shapes, color of shapes, and shape type. Format your response as: Number of cards: [number], Card Label: [number, color, shape], ...",
}

ROLE_DESCRIPTIONS = {
    "Goal_Maint": "While performing this task, you have significant difficulty maintaining goals in your working memory. You often lose track of the current sorting rule you're supposed to be following. You struggle to consistently apply a rule across multiple trials. You may need to refer back to the instructions more frequently than usual. Despite these challenges, you still attempt to complete the task to the best of your ability. Your responses should reflect a genuine effort to perform the task, but with clear signs of difficulty in maintaining focus on the current goal or rule.",
    "Inhib_Ctrl": "As you work on this task, you have significant difficulty inhibiting irrelevant information and controlling impulsive responses. You're are easily distracted by unimportant details and may occasionally struggle to focus solely on the relevant aspects of the cards. You often make impulsive choices without fully considering all options. Despite these challenges, you still attempt to complete the task to the best of your ability. Your responses should reflect a genuine effort to perform the task, but with clear signs of difficulty in inhibiting irrelevant information and controlling impulsive reactions.",
    "Adapt_Upd": "During this task, you have significant difficulty in adapting to changes. You strongly tend to stick with a previously correct sorting rule, even when it's no longer effective. When faced with negative feedback, you struggle to recognize that a rule change has occurred and instead continue applying the old rule. Despite these challenges, you still attempt to complete the task to the best of your ability and you can eventually adapt after more attempts. Your responses should reflect a genuine effort to perform the task, but with clear signs of difficulty in adapting to changes and updating your behavior based on feedback.",
}

TASK_DIR = {
    "WCST": "./task_datasets/WSCT/",
    "WCST_without_restriction": "./task_datasets/WSCT/",
}
REQUEST_LOG_FILE = './request_log.json'

WCST_RULES = ["color_rule", "shape_rule", "number_rule"]
WCST_TOTAL_TRIALS = 64
WCST_CORRECT_TO_CHANGE = 10

# Prompt types
PROMPT_TYPES = {
    "STA": "Provide your final answer immediately without explaining your reasoning. Your response should be concise and direct, containing only the essential information required by the task. Do not include any additional explanations or thought processes.",
    "CoT": "Let's think step by step, first explain your thought process in detail, then clearly state your final response to the task.",
    }

PROMPT_TYPES_CN = {
    "STA": "直接提供最终答案，不解释推理过程。回应应简洁直接，只包含任务所需的关键信息，不包含任何额外的解释或思考过程",
    "CoT": "让我们逐步思考，首先详细解释你的思考过程，然后清楚地给出你对这个任务的最终回应。",
    }

# Presentation modes
PRESENTATION_MODES = ["OI", "OT"]

# Number of independent sessions per model per condition per task
NUM_SESSIONS = 20
NUM_HUMAN_SESSIONS = 5

# Models
MODELS = ["GPT-4o", "Gemini-1.5 Pro", "Claude-3.5 Sonnet"]