"""
实验：用 LLM + Dataset 构造 Lambda 演算原语（本地模拟版）

通过 subprocess 调用 claude CLI 逐个执行 LDS 测试。
每个测试将 dataset (few-shot examples) + 新输入组成 prompt，
然后让 LLM 输出结果，对比期望值。

使用方式：
    python run_experiment.py
"""

import subprocess
import json
import sys
import time
from dataclasses import dataclass, field


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

    def add(self, input_val: str, expected: str, actual: str):
        # 宽松匹配：去空格、去引号、大小写
        def normalize(s):
            return s.strip().strip("'\"").strip().lower()
        passed = normalize(actual) == normalize(expected)
        self.tests.append({
            "input": input_val,
            "expected": expected,
            "actual": actual,
            "passed": passed,
        })

    @property
    def passed_count(self):
        return sum(1 for t in self.tests if t["passed"])

    @property
    def total_count(self):
        return len(self.tests)

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


def call_llm(prompt: str, max_retries: int = 2) -> str:
    """调用 claude CLI 执行 LDS prompt"""
    for attempt in range(max_retries + 1):
        try:
            result = subprocess.run(
                ["claude", "-p", prompt, "--model", "sonnet"],
                capture_output=True,
                text=True,
                timeout=30,
            )
            if result.returncode == 0:
                return result.stdout.strip()
            else:
                if attempt < max_retries:
                    time.sleep(1)
                    continue
                return f"ERROR: {result.stderr.strip()}"
        except subprocess.TimeoutExpired:
            if attempt < max_retries:
                continue
            return "ERROR: timeout"
        except FileNotFoundError:
            return "ERROR: claude CLI not found"
    return "ERROR: max retries exceeded"


def build_prompt(description: str, examples: list[tuple[str, str]], test_input: str) -> str:
    """构造 LDS prompt: system instruction + dataset + test input"""
    lines = [
        "You are a precise function executor. Given the examples below, "
        "learn the pattern and apply it to the new input. "
        "Output ONLY the result, nothing else. No explanation, no quotes, no extra text.",
        "",
        f"Function: {description}",
        "",
        "Examples:",
    ]
    for inp, out in examples:
        lines.append(f"  Input: {inp}")
        lines.append(f"  Output: {out}")
        lines.append("")
    lines.append(f"Input: {test_input}")
    lines.append("Output:")
    return "\n".join(lines)


# ============================================================
# 实验定义
# ============================================================

def experiment_succ() -> ExperimentResult:
    """实验 1: SUCC 后继函数 ≡ λn.λf.λx. f(n f x)"""
    result = ExperimentResult("SUCC (后继函数)")
    dataset = [
        ("0", "1"), ("1", "2"), ("2", "3"),
        ("5", "6"), ("9", "10"), ("13", "14"), ("99", "100"),
    ]
    tests = [3, 4, 7, 10, 15, 42, 127, 255]
    for n in tests:
        prompt = build_prompt(
            "Given a number n, output n+1 (the successor).",
            dataset, str(n)
        )
        actual = call_llm(prompt)
        result.add(str(n), str(n + 1), actual)
    return result


def experiment_booleans() -> ExperimentResult:
    """实验 2: TRUE/FALSE ≡ λa.λb.a / λa.λb.b"""
    result = ExperimentResult("TRUE/FALSE (Church 布尔值)")

    true_dataset = [
        ("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"),
    ]
    false_dataset = [
        ("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"),
    ]

    test_pairs = [("sun", "moon"), ("hello", "world"), ("42", "0"), ("alpha", "omega")]
    for a, b in test_pairs:
        inp = f"A='{a}', B='{b}'"
        # TRUE
        prompt_t = build_prompt("Given two options A and B, always select A (the first one).", true_dataset, inp)
        actual_t = call_llm(prompt_t)
        result.add(f"TRUE({a},{b})", a, actual_t)
        # FALSE
        prompt_f = build_prompt("Given two options A and B, always select B (the second one).", false_dataset, inp)
        actual_f = call_llm(prompt_f)
        result.add(f"FALSE({a},{b})", b, actual_f)
    return result


def experiment_logic() -> ExperimentResult:
    """实验 3: AND/OR/NOT"""
    result = ExperimentResult("AND/OR/NOT (逻辑运算)")

    and_ds = [("TRUE, TRUE", "TRUE"), ("TRUE, FALSE", "FALSE"),
              ("FALSE, TRUE", "FALSE"), ("FALSE, FALSE", "FALSE")]
    or_ds = [("TRUE, TRUE", "TRUE"), ("TRUE, FALSE", "TRUE"),
             ("FALSE, TRUE", "TRUE"), ("FALSE, FALSE", "FALSE")]
    not_ds = [("TRUE", "FALSE"), ("FALSE", "TRUE")]

    # AND
    for inp, exp in and_ds:
        prompt = build_prompt("Logical AND of two boolean values.", and_ds, inp)
        actual = call_llm(prompt)
        result.add(f"AND({inp})", exp, actual)
    # OR
    for inp, exp in or_ds:
        prompt = build_prompt("Logical OR of two boolean values.", or_ds, inp)
        actual = call_llm(prompt)
        result.add(f"OR({inp})", exp, actual)
    # NOT
    for inp, exp in not_ds:
        prompt = build_prompt("Logical NOT of a boolean value.", not_ds, inp)
        actual = call_llm(prompt)
        result.add(f"NOT({inp})", exp, actual)
    return result


def experiment_if() -> ExperimentResult:
    """实验 4: IF-THEN-ELSE ≡ λcond.λthen.λelse. cond then else"""
    result = ExperimentResult("IF-THEN-ELSE (条件分支)")

    dataset = [
        ("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"),
    ]
    tests = [
        ("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"),
        ("condition=TRUE, then='open', else='close'", "open"),
        ("condition=FALSE, then='open', else='close'", "close"),
    ]
    for inp, exp in tests:
        prompt = build_prompt(
            "IF-THEN-ELSE: If condition is TRUE, output the THEN value. If FALSE, output the ELSE value.",
            dataset, inp
        )
        actual = call_llm(prompt)
        result.add(inp, exp, actual)
    return result


def experiment_pairs() -> ExperimentResult:
    """实验 5: PAIR/FST/SND"""
    result = ExperimentResult("PAIR/FST/SND (有序对)")

    pair_ds = [
        ("a=3, b=7", "(3, 7)"),
        ("a=hello, b=world", "(hello, world)"),
        ("a=0, b=1", "(0, 1)"),
    ]
    fst_ds = [
        ("(3, 7)", "3"), ("(hello, world)", "hello"),
        ("(0, 1)", "0"), ("(42, 99)", "42"),
    ]
    snd_ds = [
        ("(3, 7)", "7"), ("(hello, world)", "world"),
        ("(0, 1)", "1"), ("(42, 99)", "99"),
    ]

    # PAIR
    for inp, exp in [("a=10, b=20", "(10, 20)"), ("a=foo, b=bar", "(foo, bar)")]:
        prompt = build_prompt("Construct a pair from a and b.", pair_ds, inp)
        actual = call_llm(prompt)
        result.add(f"PAIR({inp})", exp, actual)

    # FST
    for inp, exp in [("(5, 8)", "5"), ("(alpha, beta)", "alpha"), ("(100, 200)", "100")]:
        prompt = build_prompt("Extract the first element of a pair.", fst_ds, inp)
        actual = call_llm(prompt)
        result.add(f"FST({inp})", exp, actual)

    # SND
    for inp, exp in [("(5, 8)", "8"), ("(alpha, beta)", "beta"), ("(100, 200)", "200")]:
        prompt = build_prompt("Extract the second element of a pair.", snd_ds, inp)
        actual = call_llm(prompt)
        result.add(f"SND({inp})", exp, actual)

    return result


def experiment_composition() -> ExperimentResult:
    """实验 6: 函数组合 DOUBLE ∘ SUCC ≡ (g ∘ f)(x) = g(f(x))"""
    result = ExperimentResult("DOUBLE ∘ SUCC (函数组合)")

    succ_ds = [("0", "1"), ("1", "2"), ("2", "3"), ("5", "6"), ("9", "10")]
    double_ds = [("0", "0"), ("1", "2"), ("2", "4"), ("3", "6"), ("5", "10"), ("10", "20")]

    for n in [0, 1, 3, 5, 7, 10]:
        # Step 1: SUCC(n)
        p1 = build_prompt("Given a number n, output n+1.", succ_ds, str(n))
        succ_result = call_llm(p1)

        # Step 2: DOUBLE(SUCC(n))
        p2 = build_prompt("Given a number n, output 2*n.", double_ds, succ_result)
        double_result = call_llm(p2)

        expected = str(2 * (n + 1))
        result.add(f"DOUBLE(SUCC({n}))", expected, double_result)
    return result


def experiment_factorial() -> ExperimentResult:
    """实验 7: FACTORIAL —— 递归通过 CoT 展开，等价于 Y 组合子"""
    result = ExperimentResult("FACTORIAL (递归/Y 组合子)")

    dataset = [
        ("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"),
    ]

    import math
    for n in [2, 4, 6, 7, 8, 10]:
        prompt = build_prompt(
            "Compute n! (factorial) step by step. Show recursive expansion. End with 'Result: <number>'.",
            dataset, str(n)
        )
        actual_raw = call_llm(prompt)
        # 提取 Result 行
        actual = actual_raw
        for line in reversed(actual_raw.split("\n")):
            if "Result:" in line:
                actual = line.split("Result:")[-1].strip()
                break
        expected = str(math.factorial(n))
        result.add(f"{n}!", expected, actual)
    return result


def experiment_church_numerals() -> ExperimentResult:
    """实验 8: Church 数 c_n ≡ λf.λx. f^n(x)"""
    result = ExperimentResult("CHURCH 数 (f^n(x))")

    dataset = [
        ("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"),
    ]

    # add1 tests
    for n in [4, 5, 7, 10]:
        prompt = build_prompt(
            "Apply function f to value x exactly n times. "
            "add1 adds 1 each time. double multiplies by 2 each time.",
            dataset, f"n={n}, f='add1', x='0'"
        )
        actual = call_llm(prompt)
        result.add(f"add1^{n}(0)", str(n), actual)

    # double tests: double^n(1) = 2^n
    for n in [5, 6, 7, 8]:
        prompt = build_prompt(
            "Apply function f to value x exactly n times. "
            "add1 adds 1 each time. double multiplies by 2 each time.",
            dataset, f"n={n}, f='double', x='1'"
        )
        actual = call_llm(prompt)
        result.add(f"double^{n}(1)", str(2**n), actual)

    return result


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

def main():
    print("=" * 60)
    print("LDS ≡ Lambda Calculus 实验验证")
    print("模式: 本地 Claude CLI 模拟")
    print("=" * 60)

    # 检查 claude CLI
    try:
        r = subprocess.run(["claude", "--version"], capture_output=True, text=True, timeout=5)
        print(f"Claude CLI: {r.stdout.strip()}")
    except Exception as e:
        print(f"错误: 无法找到 claude CLI: {e}")
        sys.exit(1)

    experiments = [
        ("1. 后继函数 SUCC", experiment_succ),
        ("2. Church 布尔值 TRUE/FALSE", experiment_booleans),
        ("3. 逻辑运算 AND/OR/NOT", experiment_logic),
        ("4. 条件分支 IF-THEN-ELSE", experiment_if),
        ("5. 有序对 PAIR/FST/SND", experiment_pairs),
        ("6. 函数组合 DOUBLE∘SUCC", experiment_composition),
        ("7. 递归/阶乘 (Y 组合子)", experiment_factorial),
        ("8. Church 数 f^n(x)", experiment_church_numerals),
    ]

    all_results = []
    for name, fn in experiments:
        print(f"\n>>> 运行实验 {name}...")
        r = fn()
        all_results.append(r)
        print(r.summary())

    # 汇总
    total = sum(r.total_count for r in all_results)
    passed = sum(r.passed_count for r in all_results)

    print("\n" + "=" * 60)
    print("汇总结果")
    print("=" * 60)
    for r in all_results:
        rate = r.passed_count / r.total_count * 100 if r.total_count else 0
        bar = "█" * int(rate / 5) + "░" * (20 - int(rate / 5))
        print(f"  {r.name:40s} {r.passed_count:2d}/{r.total_count:2d}  {bar} {rate:.0f}%")
    print(f"\n  {'总计':40s} {passed:2d}/{total:2d}  {passed/total*100:.1f}%")
    print("=" * 60)

    # 保存 JSON
    output = {
        "model": "claude-sonnet (via CLI)",
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "experiments": [
            {
                "name": r.name,
                "passed": r.passed_count,
                "total": r.total_count,
                "tests": r.tests,
            }
            for r in all_results
        ],
        "total_passed": passed,
        "total_tests": total,
    }
    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")


if __name__ == "__main__":
    main()