from __future__ import annotations
import argparse
import logging
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from transformers import AutoTokenizer
from riemann.growth_ladder import get_stage, next_stage
from riemann.growth import load_checkpoint_and_grow
from riemann.paths import CHECKPOINT_ROOT

logger = logging.getLogger("riemann.grow_and_continue")
TOKENIZER_NAME = "gpt2"


def main() -> None:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--from-stage", required=True, help="the stage you just finished training"
    )
    parser.add_argument(
        "--to-stage",
        default=None,
        help="target stage name (default: the next one in the ladder)",
    )
    parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
    args = parser.parse_args()
    from_stage = get_stage(args.from_stage)
    to_stage = (
        get_stage(args.to_stage) if args.to_stage else next_stage(from_stage.name)
    )
    from_dir = Path(args.checkpoint_root) / from_stage.name
    to_dir = Path(args.checkpoint_root) / to_stage.name
    if not (from_dir / "config.json").exists():
        raise FileNotFoundError(
            f"No checkpoint found at {from_dir} -- train this stage first"
        )
    logger.info("Growing %s -> %s", from_stage.name, to_stage.name)
    grown_model = load_checkpoint_and_grow(str(from_dir), to_stage)
    to_dir.mkdir(parents=True, exist_ok=True)
    grown_model.save_pretrained(str(to_dir))
    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
    tokenizer.save_pretrained(str(to_dir))
    n_params = sum((p.numel() for p in grown_model.parameters()))
    logger.info(
        "Saved grown checkpoint to %s (%d params, ~%.2fB)",
        to_dir,
        n_params,
        n_params / 1000000000.0,
    )
    if not to_stage.feasible_on_this_machine:
        logger.warning(
            "%s is marked not feasible on this machine; expect this to be slow or run out of memory",
            to_stage.name,
        )
    logger.info("Next: python scripts/pretrain.py --stage %s --resume", to_stage.name)


if __name__ == "__main__":
    main()