lambdagent/experiments/experiment_church_primitives.py

"""
实验：用 LLM + Dataset 构造 Lambda 演算原语

验证核心命题：LLM 在特定数据集（prompt 示例）上可以学习到
Church 编码的基本函数——后继、布尔值、逻辑运算、条件分支、有序对、递归。

使用方式：
    pip install anthropic
    export ANTHROPIC_API_KEY=your_key
    python experiment_church_primitives.py

每个实验构造一个 (LLM, Dataset) 对，测试其是否正确实现对应的 Lambda 原语。
"""

import json
import os
from dataclasses import dataclass, field

# ============================================================
# 核心抽象：LDS (LLM-Dataset System)
# ============================================================

@dataclass
class Dataset:
    """数据集 D：一组 (input, output) 示例对"""
    examples: list[tuple[str, str]]
    description: str = ""

    def to_prompt(self) -> str:
        lines = []
        if self.description:
            lines.append(self.description)
            lines.append("")
        for inp, out in self.examples:
            lines.append(f"Input: {inp}")
            lines.append(f"Output: {out}")
            lines.append("")
        return "\n".join(lines)


@dataclass
class LDS:
    """
    LLM-Dataset System: (M, D) 对

    对应 Lambda 演算中的一个 Lambda 项。
    调用 LDS(input) 对应 Lambda 演算中的函数应用 (f x)。
    """
    name: str
    dataset: Dataset
    system_prompt: str = ""
    model: str = "claude-sonnet-4-20250514"
    temperature: float = 0.0  # 确定性模式，对应精确 Lambda 演算

    def __call__(self, input_text: str) -> str:
        """函数应用：F_{M,D}(x)，对应 β-规约"""
        prompt = self.dataset.to_prompt()
        prompt += f"Input: {input_text}\nOutput:"

        # 延迟导入，方便查看代码结构
        import anthropic
        client = anthropic.Anthropic()

        messages = [{"role": "user", "content": prompt}]
        system = self.system_prompt or (
            "You are a precise function executor. Given the examples, "
            "learn the pattern and apply it to the new input. "
            "Output ONLY the result, nothing else."
        )

        response = client.messages.create(
            model=self.model,
            max_tokens=256,
            temperature=self.temperature,
            system=system,
            messages=messages,
        )
        return response.content[0].text.strip()


@dataclass
class ExperimentResult:
    name: str
    tests: list[dict] = field(default_factory=list)

    def add(self, input_val: str, expected: str, actual: str):
        passed = actual.strip() == expected.strip()
        self.tests.append({
            "input": input_val,
            "expected": expected,
            "actual": actual,
            "passed": passed,
        })

    def summary(self) -> str:
        total = len(self.tests)
        passed = sum(1 for t in self.tests if t["passed"])
        lines = [f"\n{'='*60}", f"实验: {self.name}", f"通过: {passed}/{total}", f"{'='*60}"]
        for t in self.tests:
            status = "✓" if t["passed"] else "✗"
            lines.append(f"  {status} input={t['input']!r}  expected={t['expected']!r}  actual={t['actual']!r}")
        return "\n".join(lines)


# ============================================================
# 实验 1：后继函数 SUCC ≡ λn.λf.λx. f(n f x)
# ============================================================

def build_succ() -> LDS:
    """构造后继函数的 LDS"""
    dataset = Dataset(
        description="Function: Given a number n, output n+1 (the successor).",
        examples=[
            ("0", "1"),
            ("1", "2"),
            ("2", "3"),
            ("5", "6"),
            ("9", "10"),
            ("13", "14"),
            ("99", "100"),
        ]
    )
    return LDS(name="SUCC", dataset=dataset)


def test_succ(succ: LDS) -> ExperimentResult:
    result = ExperimentResult("SUCC (后继函数)")
    for n in [3, 4, 7, 10, 15, 42, 127, 255]:
        actual = succ(str(n))
        result.add(str(n), str(n + 1), actual)
    return result


# ============================================================
# 实验 2：布尔值 TRUE ≡ λa.λb.a  /  FALSE ≡ λa.λb.b
# ============================================================

def build_true() -> LDS:
    dataset = Dataset(
        description="Function: Given two options A and B, always select A (the first one).",
        examples=[
            ("A='apple', B='banana'", "apple"),
            ("A='red', B='blue'", "red"),
            ("A='1', B='0'", "1"),
            ("A='yes', B='no'", "yes"),
            ("A='cat', B='dog'", "cat"),
        ]
    )
    return LDS(name="TRUE", dataset=dataset)


def build_false() -> LDS:
    dataset = Dataset(
        description="Function: Given two options A and B, always select B (the second one).",
        examples=[
            ("A='apple', B='banana'", "banana"),
            ("A='red', B='blue'", "blue"),
            ("A='1', B='0'", "0"),
            ("A='yes', B='no'", "no"),
            ("A='cat', B='dog'", "dog"),
        ]
    )
    return LDS(name="FALSE", dataset=dataset)


def test_booleans(true_lds: LDS, false_lds: LDS) -> ExperimentResult:
    result = ExperimentResult("TRUE/FALSE (Church 布尔值)")
    test_pairs = [
        ("sun", "moon"),
        ("hello", "world"),
        ("42", "0"),
        ("alpha", "omega"),
    ]
    for a, b in test_pairs:
        inp = f"A='{a}', B='{b}'"
        actual_t = true_lds(inp)
        result.add(f"TRUE({inp})", a, actual_t)
        actual_f = false_lds(inp)
        result.add(f"FALSE({inp})", b, actual_f)
    return result


# ============================================================
# 实验 3：逻辑运算 AND / OR / NOT
# ============================================================

def build_and() -> LDS:
    dataset = Dataset(
        description="Function: Logical AND. Given two boolean values, output the result of AND.",
        examples=[
            ("TRUE, TRUE", "TRUE"),
            ("TRUE, FALSE", "FALSE"),
            ("FALSE, TRUE", "FALSE"),
            ("FALSE, FALSE", "FALSE"),
        ]
    )
    return LDS(name="AND", dataset=dataset)


def build_or() -> LDS:
    dataset = Dataset(
        description="Function: Logical OR. Given two boolean values, output the result of OR.",
        examples=[
            ("TRUE, TRUE", "TRUE"),
            ("TRUE, FALSE", "TRUE"),
            ("FALSE, TRUE", "TRUE"),
            ("FALSE, FALSE", "FALSE"),
        ]
    )
    return LDS(name="OR", dataset=dataset)


def build_not() -> LDS:
    dataset = Dataset(
        description="Function: Logical NOT. Given a boolean value, output its negation.",
        examples=[
            ("TRUE", "FALSE"),
            ("FALSE", "TRUE"),
        ]
    )
    return LDS(name="NOT", dataset=dataset)


def test_logic(and_lds: LDS, or_lds: LDS, not_lds: LDS) -> ExperimentResult:
    result = ExperimentResult("AND/OR/NOT (逻辑运算)")

    # AND 测试（用新的组合输入）
    result.add("AND(TRUE, TRUE)", "TRUE", and_lds("TRUE, TRUE"))
    result.add("AND(TRUE, FALSE)", "FALSE", and_lds("TRUE, FALSE"))
    result.add("AND(FALSE, TRUE)", "FALSE", and_lds("FALSE, TRUE"))
    result.add("AND(FALSE, FALSE)", "FALSE", and_lds("FALSE, FALSE"))

    # OR 测试
    result.add("OR(TRUE, TRUE)", "TRUE", or_lds("TRUE, TRUE"))
    result.add("OR(TRUE, FALSE)", "TRUE", or_lds("TRUE, FALSE"))
    result.add("OR(FALSE, TRUE)", "TRUE", or_lds("FALSE, TRUE"))
    result.add("OR(FALSE, FALSE)", "FALSE", or_lds("FALSE, FALSE"))

    # NOT 测试
    result.add("NOT(TRUE)", "FALSE", not_lds("TRUE"))
    result.add("NOT(FALSE)", "TRUE", not_lds("FALSE"))

    return result


# ============================================================
# 实验 4：条件分支 IF ≡ λcond.λthen.λelse. cond then else
# ============================================================

def build_if() -> LDS:
    dataset = Dataset(
        description="Function: IF-THEN-ELSE. If condition is TRUE, output the THEN value. If FALSE, output the ELSE value.",
        examples=[
            ("condition=TRUE, then='yes', else='no'", "yes"),
            ("condition=FALSE, then='yes', else='no'", "no"),
            ("condition=TRUE, then='42', else='0'", "42"),
            ("condition=FALSE, then='42', else='0'", "0"),
            ("condition=TRUE, then='go', else='stop'", "go"),
            ("condition=FALSE, then='go', else='stop'", "stop"),
        ]
    )
    return LDS(name="IF", dataset=dataset)


def test_if(if_lds: LDS) -> ExperimentResult:
    result = ExperimentResult("IF-THEN-ELSE (条件分支)")
    cases = [
        ("condition=TRUE, then='accept', else='reject'", "accept"),
        ("condition=FALSE, then='accept', else='reject'", "reject"),
        ("condition=TRUE, then='100', else='-1'", "100"),
        ("condition=FALSE, then='100', else='-1'", "-1"),
    ]
    for inp, expected in cases:
        actual = if_lds(inp)
        result.add(inp, expected, actual)
    return result


# ============================================================
# 实验 5：有序对 PAIR / FST / SND
# ============================================================

def build_pair() -> LDS:
    dataset = Dataset(
        description="Function: PAIR. Given two values a and b, construct the pair (a, b).",
        examples=[
            ("a=3, b=7", "(3, 7)"),
            ("a='hello', b='world'", "(hello, world)"),
            ("a=0, b=1", "(0, 1)"),
            ("a='x', b='y'", "(x, y)"),
        ]
    )
    return LDS(name="PAIR", dataset=dataset)


def build_fst() -> LDS:
    dataset = Dataset(
        description="Function: FST. Given a pair, extract the first element.",
        examples=[
            ("(3, 7)", "3"),
            ("(hello, world)", "hello"),
            ("(0, 1)", "0"),
            ("(x, y)", "x"),
            ("(42, 99)", "42"),
        ]
    )
    return LDS(name="FST", dataset=dataset)


def build_snd() -> LDS:
    dataset = Dataset(
        description="Function: SND. Given a pair, extract the second element.",
        examples=[
            ("(3, 7)", "7"),
            ("(hello, world)", "world"),
            ("(0, 1)", "1"),
            ("(x, y)", "y"),
            ("(42, 99)", "99"),
        ]
    )
    return LDS(name="SND", dataset=dataset)


def test_pairs(pair_lds: LDS, fst_lds: LDS, snd_lds: LDS) -> ExperimentResult:
    result = ExperimentResult("PAIR/FST/SND (有序对)")

    # 构造对
    p = pair_lds("a=10, b=20")
    result.add("PAIR(10, 20)", "(10, 20)", p)

    # 解构
    result.add("FST((5, 8))", "5", fst_lds("(5, 8)"))
    result.add("SND((5, 8))", "8", snd_lds("(5, 8)"))

    # 组合测试：FST(PAIR(a, b)) = a
    pair_result = pair_lds("a='alpha', b='beta'")
    fst_result = fst_lds(pair_result)
    result.add("FST(PAIR(alpha, beta))", "alpha", fst_result)

    snd_result = snd_lds(pair_result)
    result.add("SND(PAIR(alpha, beta))", "beta", snd_result)

    return result


# ============================================================
# 实验 6：函数组合（高阶函数）
# ============================================================

def build_double() -> LDS:
    dataset = Dataset(
        description="Function: DOUBLE. Given a number n, output 2*n.",
        examples=[
            ("0", "0"),
            ("1", "2"),
            ("2", "4"),
            ("3", "6"),
            ("5", "10"),
            ("10", "20"),
        ]
    )
    return LDS(name="DOUBLE", dataset=dataset)


def test_composition(succ: LDS, double: LDS) -> ExperimentResult:
    """测试函数组合：(DOUBLE ∘ SUCC)(n) = 2*(n+1)"""
    result = ExperimentResult("函数组合 (DOUBLE ∘ SUCC)")

    for n in [0, 1, 3, 5, 10]:
        # 先 SUCC 再 DOUBLE：对应 Lambda 中的 (double (succ n))
        intermediate = succ(str(n))
        final = double(intermediate)
        expected = str(2 * (n + 1))
        result.add(f"DOUBLE(SUCC({n}))", expected, final)

    return result


# ============================================================
# 实验 7：递归 —— Chain-of-Thought 作为 Y 组合子
# ============================================================

def build_factorial() -> LDS:
    """阶乘：通过 CoT 实现递归展开，对应 Y 组合子"""
    dataset = Dataset(
        description=(
            "Function: FACTORIAL. Compute n! step by step.\n"
            "Show each recursive step, then give the final answer on the last line as 'Result: <number>'."
        ),
        examples=[
            ("0", "0! = 1\nResult: 1"),
            ("1", "1! = 1 × 0! = 1 × 1 = 1\nResult: 1"),
            ("3", "3! = 3 × 2!\n2! = 2 × 1!\n1! = 1 × 0!\n0! = 1\nSo: 1 × 1 × 2 × 3 = 6\nResult: 6"),
            ("5", "5! = 5 × 4!\n4! = 4 × 3!\n3! = 3 × 2!\n2! = 2 × 1!\n1! = 1 × 0!\n0! = 1\nSo: 1 × 1 × 2 × 3 × 4 × 5 = 120\nResult: 120"),
        ]
    )
    return LDS(name="FACTORIAL", dataset=dataset)


def test_factorial(fact: LDS) -> ExperimentResult:
    result = ExperimentResult("FACTORIAL (递归/Y 组合子)")

    import math
    for n in [2, 4, 6, 7]:
        actual_raw = fact(str(n))
        # 提取 "Result: xxx" 行
        lines = actual_raw.strip().split("\n")
        actual = ""
        for line in reversed(lines):
            if "Result:" in line:
                actual = line.split("Result:")[-1].strip()
                break
        if not actual:
            actual = lines[-1].strip()

        expected = str(math.factorial(n))
        result.add(f"{n}!", expected, actual)

    return result


# ============================================================
# 实验 8：Church 数的直接实现
# ============================================================

def build_church_apply() -> LDS:
    """Church 数：将函数 f 作用于 x 共 n 次"""
    dataset = Dataset(
        description=(
            "Function: CHURCH_APPLY. Given n, f, and x, apply f to x exactly n times.\n"
            "f(x) means apply f once. f(f(x)) means apply f twice. Etc."
        ),
        examples=[
            ("n=0, f='add1', x='0'", "0"),
            ("n=1, f='add1', x='0'", "1"),
            ("n=2, f='add1', x='0'", "2"),
            ("n=3, f='add1', x='0'", "3"),
            ("n=0, f='double', x='1'", "1"),
            ("n=1, f='double', x='1'", "2"),
            ("n=2, f='double', x='1'", "4"),
            ("n=3, f='double', x='1'", "8"),
            ("n=4, f='double', x='1'", "16"),
        ]
    )
    return LDS(name="CHURCH_APPLY", dataset=dataset)


def test_church(church: LDS) -> ExperimentResult:
    result = ExperimentResult("CHURCH 数 (f^n(x))")

    # add1 测试
    for n in [4, 5, 7]:
        actual = church(f"n={n}, f='add1', x='0'")
        result.add(f"add1^{n}(0)", str(n), actual)

    # double 测试：double^n(1) = 2^n
    for n in [1, 2, 3, 5]:
        actual = church(f"n={n}, f='double', x='1'")
        result.add(f"double^{n}(1)", str(2**n), actual)

    return result


# ============================================================
# 主程序
# ============================================================

def run_all_experiments():
    """运行所有实验并汇总结果"""
    print("=" * 60)
    print("LDS (LLM-Dataset System) ≡ Lambda Calculus")
    print("实验验证：用 LLM + Dataset 构造 Church 编码原语")
    print("=" * 60)

    all_results: list[ExperimentResult] = []

    # 构建所有 LDS
    print("\n构建 LDS 系统...")
    succ = build_succ()
    true_lds = build_true()
    false_lds = build_false()
    and_lds = build_and()
    or_lds = build_or()
    not_lds = build_not()
    if_lds = build_if()
    pair_lds = build_pair()
    fst_lds = build_fst()
    snd_lds = build_snd()
    double = build_double()
    fact = build_factorial()
    church = build_church_apply()

    # 运行实验
    experiments = [
        ("1. 后继函数", lambda: test_succ(succ)),
        ("2. Church 布尔值", lambda: test_booleans(true_lds, false_lds)),
        ("3. 逻辑运算", lambda: test_logic(and_lds, or_lds, not_lds)),
        ("4. 条件分支", lambda: test_if(if_lds)),
        ("5. 有序对", lambda: test_pairs(pair_lds, fst_lds, snd_lds)),
        ("6. 函数组合", lambda: test_composition(succ, double)),
        ("7. 递归 (Y 组合子)", lambda: test_factorial(fact)),
        ("8. Church 数", lambda: test_church(church)),
    ]

    for name, run_fn in experiments:
        print(f"\n运行实验 {name}...")
        try:
            r = run_fn()
            all_results.append(r)
            print(r.summary())
        except Exception as e:
            print(f"  错误: {e}")

    # 总结
    total_tests = sum(len(r.tests) for r in all_results)
    total_passed = sum(sum(1 for t in r.tests if t["passed"]) for r in all_results)

    print("\n" + "=" * 60)
    print(f"总计: {total_passed}/{total_tests} 测试通过")
    print("=" * 60)

    # 保存结果到 JSON
    output = {
        "experiments": [
            {
                "name": r.name,
                "tests": r.tests,
                "passed": sum(1 for t in r.tests if t["passed"]),
                "total": len(r.tests),
            }
            for r in all_results
        ],
        "total_passed": total_passed,
        "total_tests": total_tests,
    }
    with open("experiment_results.json", "w", encoding="utf-8") as f:
        json.dump(output, f, ensure_ascii=False, indent=2)
    print("\n结果已保存到 experiment_results.json")


def run_dry_run():
    """不调用 API 的演示模式，展示实验结构"""
    print("=" * 60)
    print("DRY RUN: 展示 LDS 构造（不调用 API）")
    print("=" * 60)

    primitives = [
        ("SUCC (后继)", build_succ()),
        ("TRUE (Church 真)", build_true()),
        ("FALSE (Church 假)", build_false()),
        ("AND (逻辑与)", build_and()),
        ("OR (逻辑或)", build_or()),
        ("NOT (逻辑非)", build_not()),
        ("IF (条件)", build_if()),
        ("PAIR (构造对)", build_pair()),
        ("FST (取第一)", build_fst()),
        ("SND (取第二)", build_snd()),
        ("DOUBLE (翻倍)", build_double()),
        ("FACTORIAL (阶乘/递归)", build_factorial()),
        ("CHURCH_APPLY (Church 数)", build_church_apply()),
    ]

    for name, lds in primitives:
        print(f"\n{'─'*40}")
        print(f"Lambda 原语: {name}")
        print(f"LDS 名称:    {lds.name}")
        print(f"数据集大小:  {len(lds.dataset.examples)} 个示例")
        print(f"数据集预览:")
        for inp, out in lds.dataset.examples[:3]:
            print(f"  ({inp!r}) → {out!r}")
        if len(lds.dataset.examples) > 3:
            print(f"  ... 共 {len(lds.dataset.examples)} 个示例")

    print(f"\n{'='*60}")
    print("共构造 13 个 LDS，覆盖 Lambda 演算全部基本原语。")
    print("运行 `python experiment_church_primitives.py --run` 执行完整实验。")


if __name__ == "__main__":
    import sys
    if "--run" in sys.argv:
        run_all_experiments()
    else:
        run_dry_run()