"""
Embedded sample benchmark datasets for quick testing of the experiment pipeline.

Contains both English and Chinese subsets:
- GSM8K    (10 problems, English grade-school math)
- MATH     (10 problems, English competition math)
- AIME     (5 problems, English high-difficulty math)
- C_GSM8K  (10 problems, Chinese primary/middle-school math)
- CMATH    (10 problems, Chinese high-school / competition math)
- C_AIME   (5 problems, Chinese olympiad / Gaokao hard problems)

Chinese datasets are designed to test reasoning in Chinese,
which is the native training language of Qwen-family models.
"""

GSM8K_SAMPLES = [
    {
        "question": "Janet's ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins with four every day. She sells the remaining eggs at the farmers' market daily for $2 per dozen. How much money does she make per day at the farmers' market?",
        "answer": "1.50",
    },
    {
        "question": "A robe takes 2 1/2 yards of fabric and 1/8 yard of trim. If you have 10 yards of fabric and 1 yard of trim, how many robes can you make?",
        "answer": "4",
    },
    {
        "question": "Kyle is a collector of rare coins. One day he found a rare coin online and bought it for $50. The next month, he sold it at a 20% profit. Six months later, he bought the same coin back, but had to pay 10% more than he sold it for. How much did Kyle pay to buy the coin back?",
        "answer": "66",
    },
    {
        "question": "A pet shelter has 9 kittens available for adoption. The shelter normally has 2 cats available. If each cat treats 2 kittens as their own, how many animals are there in all?",
        "answer": "7",
    },
    {
        "question": "Faye is a flower farmer. She plants 12 rows of daisies with 100 daisies in each row. She also plants 7 rows of roses with 12 roses in each row. How many more daisies than roses does she plant?",
        "answer": "1116",
    },
    {
        "question": "A farmer had 27 watermelons. After selling some watermelons, he had 13 left. Then he bought 28 more watermelons. Later, he sold 15 watermelons. How many watermelons does he have now?",
        "answer": "26",
    },
    {
        "question": "James buys a bag of 40 cookies. He eats 5 cookies and then divides the rest evenly among his 7 friends. How many cookies does each friend get?",
        "answer": "5",
    },
    {
        "question": "Tom's mother had 15 apples. She gave Tom 3 apples and gave Sarah twice as many as Tom. How many apples does she have left?",
        "answer": "6",
    },
    {
        "question": "A school bus can hold 48 students. If 3 buses are used for a field trip and each bus is 75% full, how many students are on the field trip?",
        "answer": "108",
    },
    {
        "question": "Maria bought 6 notebooks for $2 each and 4 pens for $1.50 each. She paid with a $20 bill. How much change did she receive?",
        "answer": "2",
    },
]

MATH_SAMPLES = [
    {
        "question": "Find the value of $x$ such that $\\sqrt{4x + 5} = 7$.",
        "answer": "11",
    },
    {
        "question": "A box contains 5 red balls and 7 blue balls. If two balls are drawn at random without replacement, what is the probability that both balls are red? Express your answer as a fraction in simplest form.",
        "answer": "5/33",
    },
    {
        "question": "What is the sum of the first 20 positive integers?",
        "answer": "210",
    },
    {
        "question": "If $f(x) = 2x^2 + 3x - 5$, find $f(-2)$.",
        "answer": "-3",
    },
    {
        "question": "The sum of two numbers is 20 and their difference is 4. What is the product of the two numbers?",
        "answer": "96",
    },
    {
        "question": "How many distinguishable permutations can be made from the letters in the word 'MISSISSIPPI'?",
        "answer": "34650",
    },
    {
        "question": "Find the area of a circle with circumference $8\\pi$.",
        "answer": "16\\pi",
    },
    {
        "question": "If $\\log_2(x) + \\log_2(x-2) = 3$, find $x$.",
        "answer": "4",
    },
    {
        "question": "A geometric sequence has first term 3 and common ratio 2. What is the 8th term?",
        "answer": "384",
    },
    {
        "question": "How many ways can 5 people be arranged in a line?",
        "answer": "120",
    },
]

AIME_SAMPLES = [
    {
        "question": "Find the number of positive integers $n \\le 1000$ such that $n$ is divisible by neither 2 nor 5.",
        "answer": "400",
    },
    {
        "question": "Let $a$, $b$, $c$ be real numbers such that $a + b + c = 6$, $ab + bc + ca = 11$, and $abc = 6$. Find $a^2 + b^2 + c^2$.",
        "answer": "14",
    },
    {
        "question": "How many ordered pairs $(x, y)$ of positive integers satisfy $\\frac{1}{x} + \\frac{1}{y} = \\frac{1}{6}$?",
        "answer": "5",
    },
    {
        "question": "A fair coin is tossed 10 times. What is the probability that exactly 5 heads appear? Express your answer as a fraction in simplest form.",
        "answer": "63/256",
    },
    {
        "question": "Find the smallest positive integer $n$ such that $n!$ ends with exactly 100 zeros in base 10.",
        "answer": "405",
    },
]

# ── Chinese Datasets (中文数学测试) ────────────────────────────────

C_GSM8K_SAMPLES = [
    {
        "question": "小明有 12 个苹果，他给了小红 5 个，又给了小刚 3 个。请问小明还剩下几个苹果？",
        "answer": "4",
    },
    {
        "question": "一个长方形的长是 8 厘米，宽是 5 厘米。如果长增加 2 厘米，宽不变，新的长方形面积比原来的大多少平方厘米？",
        "answer": "10",
    },
    {
        "question": "超市里一瓶牛奶卖 6 元，一袋面包卖 4 元。小华买了 3 瓶牛奶和 2 袋面包，付了 50 元，请问应找回多少元？",
        "answer": "24",
    },
    {
        "question": "一本书有 200 页，小明第一天看了全书的 1/4，第二天看了余下的 1/3。请问小明两天一共看了多少页？",
        "answer": "100",
    },
    {
        "question": "一个水池有甲、乙两个进水管。单开甲管 6 小时可以注满水池，单开乙管 8 小时可以注满水池。如果两管同时开放，多少小时可以注满水池？",
        "answer": "24/7",
    },
    {
        "question": "某商品原价 200 元，先涨价 10%，再打九折出售。请问现在的售价是多少元？",
        "answer": "198",
    },
    {
        "question": "学校图书馆有科技书和故事书共 360 本，其中科技书的本数是故事书的 4/5。请问科技书有多少本？",
        "answer": "160",
    },
    {
        "question": "一个等腰三角形的顶角是 40 度，请问它的一个底角是多少度？",
        "answer": "70",
    },
    {
        "question": "甲乙两地相距 240 千米。一辆汽车从甲地开往乙地，前 2 小时行驶了 96 千米。照这样的速度，还要几小时才能到达乙地？",
        "answer": "3",
    },
    {
        "question": "一个圆柱的底面半径是 3 厘米，高是 10 厘米。请问这个圆柱的体积是多少立方厘米？（π 取 3.14）",
        "answer": "282.6",
    },
]

CMATH_SAMPLES = [
    {
        "question": "已知函数 $f(x) = 2x^2 - 4x + 1$，求 $f(x)$ 的最小值。",
        "answer": "-1",
    },
    {
        "question": "在等差数列 $\\{a_n\\}$ 中，$a_1 = 3$，$a_3 + a_5 = 22$，求 $a_7$ 的值。",
        "answer": "25",
    },
    {
        "question": "从 3 名男生和 2 名女生中选出 2 人参加志愿者活动，求选出的两人恰好是一男一女的概率。",
        "answer": "3/5",
    },
    {
        "question": "已知双曲线 $\\frac{x^2}{a^2} - \\frac{y^2}{b^2} = 1$ 的离心率为 $\\sqrt{3}$，且过点 $(2, \\sqrt{2})$，求双曲线的标准方程。",
        "answer": "\\frac{x^2}{2} - \\frac{y^2}{4} = 1",
    },
    {
        "question": "已知 $\\sin\\alpha + \\cos\\alpha = \\frac{1}{5}$，且 $\\alpha \\in (0, \\pi)$，求 $\\tan\\alpha$ 的值。",
        "answer": "-4/3",
    },
    {
        "question": "不等式 $|x-1| + |x+2| < 5$ 的解集是什么？",
        "answer": "(-3, 2)",
    },
    {
        "question": "已知向量 $\\vec{a} = (2, -1)$，$\\vec{b} = (3, 4)$，求 $\\vec{a}$ 在 $\\vec{b}$ 方向上的投影。",
        "answer": "2/5",
    },
    {
        "question": "从 1, 2, 3, 4, 5 中任取两个不同的数字组成两位数，求这个两位数大于 30 的概率。",
        "answer": "3/5",
    },
    {
        "question": "已知 $\\log_2(x-1) + \\log_2(x+1) = 3$，求 $x$ 的值。",
        "answer": "3",
    },
    {
        "question": "一个袋子中有 4 个红球和 6 个白球，每次从中随机取出一个球，取出后不放回。求第二次取到红球的概率。",
        "answer": "2/5",
    },
]

C_AIME_SAMPLES = [
    {
        "question": "已知函数 $f(x) = x^3 - 3x + 1$，过点 $P(2, 3)$ 作曲线 $y = f(x)$ 的切线，求切线的方程。提示：P点可能在曲线上也可能不在。",
        "answer": "y = 9x - 15",
    },
    {
        "question": "设 $a, b, c$ 是正实数，满足 $a + b + c = 3$，求 $\\frac{a}{b+c} + \\frac{b}{c+a} + \\frac{c}{a+b}$ 的最小值。",
        "answer": "3/2",
    },
    {
        "question": "已知复数 $z$ 满足 $|z| = 1$，求 $|z^2 - z + 1|$ 的最大值。",
        "answer": "3",
    },
    {
        "question": "将 5 个相同的小球放入 3 个不同的盒子中，每个盒子至少放一个球，有多少种不同的放法？",
        "answer": "6",
    },
    {
        "question": "已知 $F_1, F_2$ 是椭圆 $\\frac{x^2}{25} + \\frac{y^2}{9} = 1$ 的两个焦点，$P$ 是椭圆上一点，且 $\\angle F_1PF_2 = 60^\\circ$，求 $\\triangle F_1PF_2$ 的面积。",
        "answer": "3\\sqrt{3}",
    },
]


def get_dataset(name: str, size: int = -1) -> list:
    """Get a sample dataset by name.  size=-1 returns all samples.

    Available datasets:
      gsm8k, math, aime        — English
      c_gsm8k, cmath, c_aime   — Chinese (中文)
    """
    datasets = {
        "gsm8k": GSM8K_SAMPLES,
        "math": MATH_SAMPLES,
        "aime": AIME_SAMPLES,
        "c_gsm8k": C_GSM8K_SAMPLES,
        "cmath": CMATH_SAMPLES,
        "c_aime": C_AIME_SAMPLES,
    }
    data = datasets.get(name.lower())
    if data is None:
        raise ValueError(f"Unknown dataset '{name}'. Choose from: {list(datasets.keys())}")
    if size > 0:
        return data[:size]
    return data


def get_all_small() -> dict:
    """Return a dict of dataset_name -> samples for all datasets."""
    return {
        "gsm8k": GSM8K_SAMPLES,
        "math": MATH_SAMPLES,
        "aime": AIME_SAMPLES,
        "c_gsm8k": C_GSM8K_SAMPLES,
        "cmath": CMATH_SAMPLES,
        "c_aime": C_AIME_SAMPLES,
    }


if __name__ == "__main__":
    print(f"GSM8K samples: {len(GSM8K_SAMPLES)}")
    print(f"MATH samples:  {len(MATH_SAMPLES)}")
    print(f"AIME samples:  {len(AIME_SAMPLES)}")
    print(f"C_GSM8K samples: {len(C_GSM8K_SAMPLES)}")
    print(f"CMATH samples:  {len(CMATH_SAMPLES)}")
    print(f"C_AIME samples: {len(C_AIME_SAMPLES)}")
    print("\nExample Chinese question:")
    print(C_GSM8K_SAMPLES[0]["question"][:80] + "...")
