import argparse
import json
import torch
from .data import make_corpus
from .engine import train, evaluate_checkpoint, load_checkpoint
from .model import generate


def main():
    parser = argparse.ArgumentParser(description="Train and explain a small diffusion text model")
    sub = parser.add_subparsers(dest="command", required=True)
    corpus = sub.add_parser("prepare", help="Generate the original controlled English grammar corpus")
    corpus.add_argument("--output", default="data/synthetic.txt")
    training = sub.add_parser("train")
    training.add_argument("--data", default="data/synthetic.txt")
    training.add_argument("--output", default="runs/demo")
    for name, default in [("steps", 600), ("batch-size", 32), ("eval-every", 100), ("seed", 42), ("width", 64), ("heads", 4), ("layers", 2), ("length", 64), ("threads", 4)]:
        training.add_argument("--" + name, type=int, default=default)
    training.add_argument("--lr", type=float, default=0.001)
    training.add_argument("--device", choices=["cpu", "cuda", "auto"], default="cpu")
    evaluation = sub.add_parser("evaluate")
    evaluation.add_argument("--data", default="data/synthetic.txt")
    sampling = sub.add_parser("sample")
    sampling.add_argument("--template", default="The cat ~~~~~ in the garden today.")
    sampling.add_argument("--steps", type=int, default=12)
    sampling.add_argument("--temperature", type=float, default=0.0)
    sampling.add_argument("--seed", type=int, default=42)
    serving = sub.add_parser("serve")
    serving.add_argument("--port", type=int, default=7860)
    for cmd in (evaluation, sampling, serving):
        cmd.add_argument("--checkpoint", default="runs/demo/best.pt")
        cmd.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
        cmd.add_argument("--threads", type=int, default=4)
    args = parser.parse_args()
    if args.command == "prepare":
        print(f"Wrote {make_corpus(args.output)} unique sentences to {args.output}")
    elif args.command == "train":
        train(args)
    elif args.command == "evaluate":
        evaluate_checkpoint(args)
    else:
        torch.set_num_threads(args.threads)
        model, checkpoint = load_checkpoint(args.checkpoint, args.device)
        if args.command == "sample":
            print(json.dumps(generate(model, args.template, args.steps, args.temperature, args.seed), indent=2, ensure_ascii=False))
        else:
            from .demo import serve
            serve(model, checkpoint, args.port)


if __name__ == "__main__":
    main()