# 项目文件：training_project/src/readingtorch/cli.py
import argparse
import json

import numpy as np

from .engine import evaluate, train
from .predict import Predictor


def main():
    # 各子命令登记自己的输入，公共计算留给对应模块。
    parser = argparse.ArgumentParser(description="离线训练、评估和预测")
    actions = parser.add_subparsers(dest="action", required=True)
    fit = actions.add_parser("train")
    fit.add_argument("--task", choices=["regression", "images"], required=True)
    fit.add_argument("--data", required=True)
    fit.add_argument("--output", required=True)
    fit.add_argument("--epochs", type=int, default=30)
    fit.add_argument("--seed", type=int, default=7)
    fit.add_argument("--batch-size", type=int, default=24)
    fit.add_argument("--lr", type=float, default=0.03)
    fit.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
    fit.add_argument("--resume")
    check = actions.add_parser("evaluate")
    check.add_argument("--checkpoint", required=True)
    check.add_argument("--data", required=True)
    predict = actions.add_parser("predict")
    predict.add_argument("--checkpoint", required=True)
    predict.add_argument("--minutes", type=float)
    predict.add_argument("--sessions", type=float)
    predict.add_argument("--image")

# 续写同一文件：training_project/src/readingtorch/cli.py
    args = vars(parser.parse_args())
    action = args.pop("action")
    # 子命令只负责组织参数，训练与推理仍可由普通 Python 函数调用。
    try:
        if action == "train":
            report = train(**args)
            result = {key: value for key, value in report.items() if key != "history"}
        elif action == "evaluate":
            result = evaluate(**args)
        elif args["image"]:
            result = Predictor(args["checkpoint"]).symbol(
                np.load(args["image"], allow_pickle=False)
            )
        elif args["minutes"] is not None and args["sessions"] is not None:
            result = {
                "pages": Predictor(args["checkpoint"]).pages(
                    args["minutes"], args["sessions"]
                )
            }
        else:
            parser.error("预测时请填写 --minutes 与 --sessions，或 --image")
    except (ValueError, OSError) as error:
        parser.error(str(error))
    print(json.dumps(result, ensure_ascii=False, indent=2))

# 项目文件：training_project/src/readingtorch/__main__.py
from .cli import main

if __name__ == "__main__":
    main()  # 模块入口和安装后的 reading-train 使用同一组命令。