from __future__ import annotations

import argparse


def train_main() -> None:
    from train import _build_argparser, train

    args = _build_argparser().parse_args()
    train(args)


def train_fast_main() -> None:
    from train_fast import train

    parser = argparse.ArgumentParser(description="Chimera 5.2 Fast CPU training")
    parser.add_argument("--config", default="config.json")
    parser.add_argument("--scale", default="tiny", choices=["tiny", "small", "medium", "full"])
    parser.add_argument("--seq_len", type=int, default=32)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--warmup", type=int, default=100)
    parser.add_argument("--max_steps", type=int, default=1000)
    parser.add_argument("--max_samples", type=int, default=5000)
    parser.add_argument("--bf16", action="store_true", default=False)
    parser.add_argument("--compile", action="store_true", default=False)
    parser.add_argument("--cache_dir", default="./cache")
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--save_every", type=int, default=500)
    parser.add_argument("--output_dir", default="./chimera_output")
    train(parser.parse_args())


def train_hyper_main() -> None:
    from train_hyper import benchmark, cli, train_hyper

    args = cli().parse_args()
    if args.max_samples and not args.max_tokens:
        args.max_tokens = args.max_samples * (args.seq_len + 1)
    if args.all:
        args.growlength = True
        args.reservoir = True
        args.progressive_unfreeze = True
    if args.benchmark:
        args.growlength = True
        args.reservoir = True
        args.progressive_unfreeze = True
        benchmark(args)
        return
    train_hyper(args)


def infer_main() -> None:
    from inference import main

    main()


def import_gguf_main() -> None:
    from gguf_import import main

    main()