#!/usr/bin/env python3
"""
Frox AI — Migrate a Morph 1.0 checkpoint to Morph 1.1

Handles the breaking changes between versions:
  - Vocab: 32,000 → 64,000 (embedding/lm_head rows are zero-padded;
    the new rows are trained from scratch during your next SFT pass)
  - QK-norm: added fresh (was absent in 1.0), initialized to identity
  - Context: 8K → 16K max_position_embeddings (YaRN scaling recomputed
    automatically — no weight changes needed, it's a RoPE parameter)
  - Everything else (attention/MLP weights, layer norms) transfers
    1:1 since the core GQA + SwiGLU block is unchanged

Usage:
    python scripts/convert_from_v1.py --input ./morph-1.0-checkpoint \
        --output ./frox-morph-1-1-output/migrated
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from model.architecture.morph_model import MorphForCausalLM
from utils.common import print_banner


def main():
    parser = argparse.ArgumentParser(description="Migrate Morph 1.0 → 1.1")
    parser.add_argument("--input", type=str, required=True, help="Morph 1.0 checkpoint dir")
    parser.add_argument("--output", type=str, required=True, help="Where to save the 1.1 model")
    args = parser.parse_args()

    print_banner()
    print(f"🔄 Migrating {args.input} → Morph 1.1\n")

    model = MorphForCausalLM.from_morph_1_checkpoint(args.input)

    print(f"\n⚠  Post-migration checklist:")
    print(f"   1. The 32,000 new vocab rows (32000-63999) are randomly initialized.")
    print(f"      Run a short SFT pass before serving, or those tokens will be garbage.")
    print(f"   2. QK-norm layers are newly initialized (identity-like RMSNorm weights).")
    print(f"      A brief SFT warmup (~500 steps) lets the model adapt to them.")
    print(f"   3. Context length is now 16K via YaRN — no retraining needed for this part,")
    print(f"      but quality past ~4K tokens will be better after some long-context SFT data.")

    model.save(args.output)
    print(f"\n✅ Migrated model saved to {args.output}")
    print(f"   Recommended: python scripts/train.py --phase sft --from-checkpoint {args.output} --steps 2000")


if __name__ == "__main__":
    main()