"""
Training script: Fine-tune Qwen2.5-0.5B-Instruct for game command recognition from noisy STT.

Based on:
- arxiv:2502.12923 approach (structured JSON output, synthetic command data)
- TRL v1.3.0 SFTTrainer with LoRA (PEFT)
- Target: llama.cpp GGUF compatible output (Qwen2 architecture)

Dataset: Marco333/game-commands-noisy-stt (conversational messages format)
Model: Qwen/Qwen2.5-0.5B-Instruct -> fine-tuned -> merge LoRA -> push to Hub

Usage:
    pip install trl peft transformers datasets accelerate torch
    python train.py
"""

import torch
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig

# =============================================================================
# Configuration - edit these for your use case
# =============================================================================

MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
DATASET_NAME = "Marco333/game-commands-noisy-stt"
OUTPUT_DIR = "./qwen2.5-0.5b-game-commands"
HUB_MODEL_ID = "Marco333/qwen2.5-0.5b-game-commands-stt"

# Training hyperparameters
LEARNING_RATE = 2e-4       # Higher LR for LoRA (10x base)
NUM_EPOCHS = 5
BATCH_SIZE = 8
GRAD_ACCUM_STEPS = 2
MAX_LENGTH = 512           # Short sequences for command recognition
WARMUP_RATIO = 0.05

# LoRA configuration
LORA_R = 32
LORA_ALPHA = 16
LORA_DROPOUT = 0.05


def main():
    print(f"Training {MODEL_NAME} on {DATASET_NAME}")
    dataset = load_dataset(DATASET_NAME)

    peft_config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    )

    training_args = SFTConfig(
        output_dir=OUTPUT_DIR,
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM_STEPS,
        learning_rate=LEARNING_RATE,
        lr_scheduler_type="cosine",
        warmup_ratio=WARMUP_RATIO,
        bf16=True,
        max_length=MAX_LENGTH,
        packing=False,
        logging_steps=10,
        logging_first_step=True,
        disable_tqdm=True,
        logging_strategy="steps",
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=2,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        push_to_hub=True,
        hub_model_id=HUB_MODEL_ID,
        hub_strategy="end",
        report_to="none",
        assistant_only_loss=True,
        seed=42,
        dataloader_num_workers=2,
    )

    trainer = SFTTrainer(
        model=MODEL_NAME,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        peft_config=peft_config,
    )

    print("Starting training...")
    train_result = trainer.train()
    print(f"Done! Loss: {train_result.training_loss:.4f}")
    trainer.save_model()

    # Merge LoRA into base model for GGUF compatibility
    print("Merging LoRA...")
    from peft import AutoPeftModelForCausalLM
    from transformers import AutoTokenizer

    merged = AutoPeftModelForCausalLM.from_pretrained(OUTPUT_DIR, dtype=torch.bfloat16, device_map="auto")
    merged = merged.merge_and_unload()
    merged_dir = f"{OUTPUT_DIR}-merged"
    merged.save_pretrained(merged_dir)
    tok = AutoTokenizer.from_pretrained(MODEL_NAME)
    tok.save_pretrained(merged_dir)

    #print("Pushing to Hub...")
    #merged.push_to_hub(HUB_MODEL_ID, commit_message="Merged model - LoRA into Qwen2.5-0.5B for game commands")
    #tok.push_to_hub(HUB_MODEL_ID, commit_message="Add tokenizer")
    #print(f"Model: https://huggingface.co/{HUB_MODEL_ID}")

    # Quick inference test
    print("\nInference test:")
    from transformers import pipeline
    pipe = pipeline("text-generation", model=merged_dir, dtype=torch.bfloat16, device_map="auto")
    system = "You are a game command interpreter. The user's input is noisy speech-to-text from a microphone in a gaming environment. Your job is to identify the intended game command from the noisy input.\n\nValid commands: Attack, CastSpell, CloseDoor, Crouch, Defend, DropItem, Heal, Interact, Jump, OpenDoor, OpenInventory, OpenMap, PauseGame, PickUp, Reload, SaveGame, Sprint, Stop, SwitchWeapon, UseItem\n\nRespond with ONLY a JSON object: {\"command\": \"<CommandName>\"}\nIf the input doesn't match any command, respond: {\"command\": \"Unknown\"}"

    for inp in ["opun da dur", "attak enami", "hel me", "pic it op", "hello how are you"]:
        out = pipe([{"role": "system", "content": system}, {"role": "user", "content": inp}], max_new_tokens=50, do_sample=False)
        print(f"  '{inp}' -> {out[0]['generated_text'][-1]['content']}")

    print("\n\nGGUF conversion (for llama.cpp):")
    print(f"  python llama.cpp/convert_hf_to_gguf.py {HUB_MODEL_ID} --outtype f16 --outfile model.gguf")
    print(f"  ./llama-quantize model.gguf model_q4.gguf Q4_K_M")


if __name__ == "__main__":
    main()