"""
SFT Training: math | lora | vision=freeze
"""
import os
import json
import torch
from PIL import Image
from torch.utils.data import Dataset
from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    AutoProcessor,
    TrainingArguments,
    Trainer,
)
from peft import LoraConfig, get_peft_model, TaskType


MODEL_NAME = "/workspace/rl4phyx/models/Qwen2.5-VL-3B-Instruct"
DATA_PATH = "/workspace/rl4phyx/RL4Phyx/SFT/sft_train/coldstart_formatted.jsonl"
OUTPUT_DIR = "/workspace/rl4phyx/RL4Phyx/SFT/checkpoints/lora_math_f"

NUM_EPOCHS = 3
LEARNING_RATE = 1e-5
PER_DEVICE_BATCH_SIZE = 1
GRAD_ACCUM_STEPS = 16
MAX_LENGTH = 4096
FREEZE_VISION = True
LORA_R = 64
LORA_ALPHA = 128
LORA_DROPOUT = 0.05


# === Dataset and Collator classes imported from template ===
exec(open('/workspace/rl4phyx/RL4Phyx/SFT/_sft_classes.py').read())


def main():
    print("Config: math | lora | vision=freeze")
    print(f"Data: {DATA_PATH}")
    print(f"Output: {OUTPUT_DIR}")

    processor = AutoProcessor.from_pretrained(
        MODEL_NAME, min_pixels=3136, max_pixels=1204224,
    )

    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_NAME, torch_dtype=torch.bfloat16, attn_implementation="sdpa",
    )

    if FREEZE_VISION:
        for name, param in model.named_parameters():
            if 'visual' in name:
                param.requires_grad = False
        print("Froze vision encoder")
    else:
        print("Vision encoder trainable")

    lora_config = LoraConfig(
        r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT,
        target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
        task_type=TaskType.CAUSAL_LM,
    )
    model = get_peft_model(model, lora_config)
    model.enable_input_require_grads()
    model.print_trainable_parameters()

    dataset = PhysicsCoTDataset(data_path=DATA_PATH, processor=processor, max_length=MAX_LENGTH)

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=PER_DEVICE_BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM_STEPS,
        learning_rate=LEARNING_RATE,
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        weight_decay=0.01,
        bf16=True,
        logging_steps=10,
        save_strategy="steps",
        save_steps=20,
        save_total_limit=2,
        dataloader_num_workers=4,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={'use_reentrant': False},
        remove_unused_columns=False,
        report_to="none",
        deepspeed="ds_zero2.json",
        save_only_model=True,
    )

    collator = VLMDataCollator(processor)
    trainer = Trainer(model=model, args=training_args, train_dataset=dataset, data_collator=collator)

    print("\n===== Starting SFT Training =====")
    trainer.train()

    print("\n===== Saving final model =====")
    trainer.save_model(os.path.join(OUTPUT_DIR, "final"))
    processor.save_pretrained(os.path.join(OUTPUT_DIR, "final"))
    print("Merging LoRA weights...")
    merged_model = model.merge_and_unload()
    merged_output = os.path.join(OUTPUT_DIR, "merged")
    merged_model.save_pretrained(merged_output)
    processor.save_pretrained(merged_output)
    import shutil
    for fn in ['preprocessor_config.json', 'chat_template.json']:
        src = os.path.join(MODEL_NAME, fn)
        if os.path.exists(src): shutil.copy2(src, os.path.join(merged_output, fn))
    print(f"Merged model saved to: {merged_output}")

    print("\n===== SFT Training Complete =====")


if __name__ == "__main__":
    main()