#!/usr/bin/env python3
"""
Train a Claude Code-level coding model via QLoRA SFT on Qwen2.5-Coder-7B-Instruct.

Recipe (based on published SOTA results):
- Base: Qwen2.5-Coder-7B-Instruct (88.4% HumanEval baseline)
- Data: KodCode-V1-SFT-R1 (verified competitive programming with R1-style CoT)
       + Code-Feedback (multi-turn code dialogue)
       + Magicoder-OSS-Instruct (diverse code generation)
       + Magicoder-Evol-Instruct (evolved code instructions)
- Method: QLoRA (4-bit NF4 + LoRA r=64, all-linear)
- Target: Push past 90%+ HumanEval, maximize LiveCodeBench with CoT reasoning

References:
- rStar-Coder (arxiv:2505.21297): Qwen2.5-Coder-7B → 57.3% LiveCodeBench
- KodCode (arxiv:2503.02951): Verified coding dataset with R1-style reasoning
- Qwen2.5-Coder (arxiv:2409.12186): Base model technical report
"""

import os
import torch
from datasets import load_dataset, concatenate_datasets
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, prepare_model_for_kbit_training
from trl import SFTConfig, SFTTrainer

# ─── Configuration ────────────────────────────────────────────────────────────
MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"
OUTPUT_DIR = "./qwen25-coder-7b-mythos"
HUB_MODEL_ID = "ashhhhhh26/qwen25-coder-32b-mythos"

SYSTEM_PROMPT = """You are an elite software engineer with deep expertise across all programming languages, frameworks, and paradigms. You write clean, efficient, well-documented code. You think step-by-step through complex problems, consider edge cases, and provide production-quality solutions. When debugging, you methodically trace through the code to identify root causes. You explain your reasoning clearly and concisely."""

# ─── 1. Load datasets ────────────────────────────────────────────────────────
print("=" * 60)
print("Loading datasets...")
print("=" * 60)

# Dataset 1: KodCode-V1-SFT-R1 (verified competitive programming + R1-style reasoning)
ds_kodcode = load_dataset("KodCode/KodCode-V1-SFT-R1", split="train")
print(f"KodCode-V1-SFT-R1 (raw): {len(ds_kodcode)} samples")

# Filter for R1-verified correct solutions only
ds_kodcode = ds_kodcode.filter(lambda x: x["r1_correctness"] is True, num_proc=4)
print(f"KodCode-V1-SFT-R1 (r1_correctness=True): {len(ds_kodcode)} samples")

# Dataset 2: Code-Feedback (~66K) - already in messages format
ds_feedback = load_dataset("m-a-p/Code-Feedback", split="train")
print(f"Code-Feedback: {len(ds_feedback)} samples")

# Dataset 3: Magicoder-OSS-Instruct-75K
ds_magicoder_oss = load_dataset("ise-uiuc/Magicoder-OSS-Instruct-75K", split="train")
print(f"Magicoder-OSS-Instruct: {len(ds_magicoder_oss)} samples")

# Dataset 4: Magicoder-Evol-Instruct-110K
ds_magicoder_evol = load_dataset("ise-uiuc/Magicoder-Evol-Instruct-110K", split="train")
print(f"Magicoder-Evol-Instruct: {len(ds_magicoder_evol)} samples")


# ─── 2. Convert all datasets to ChatML messages format ───────────────────────
print("\\nConverting datasets to ChatML format...")

def convert_kodcode(example):
    """Convert KodCode conversations (human/gpt) to standard ChatML messages."""
    role_map = {"human": "user", "gpt": "assistant"}
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in example["conversations"]:
        role = role_map.get(msg["from"], msg["from"])
        messages.append({"role": role, "content": msg["value"]})
    return {"messages": messages}

def convert_feedback(example):
    """Code-Feedback already has messages, just add system prompt."""
    messages = example["messages"]
    if messages and messages[0]["role"] != "system":
        messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
    return {"messages": messages}

def convert_magicoder_oss(example):
    """Convert problem/solution to messages format."""
    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": example["problem"]},
            {"role": "assistant", "content": example["solution"]},
        ]
    }

def convert_magicoder_evol(example):
    """Convert instruction/response to messages format."""
    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": example["instruction"]},
            {"role": "assistant", "content": example["response"]},
        ]
    }

# Apply conversions
ds_kodcode = ds_kodcode.map(convert_kodcode, num_proc=4, remove_columns=ds_kodcode.column_names)
ds_feedback = ds_feedback.map(convert_feedback, num_proc=4, remove_columns=[c for c in ds_feedback.column_names if c != "messages"])
ds_magicoder_oss = ds_magicoder_oss.map(convert_magicoder_oss, num_proc=4, remove_columns=ds_magicoder_oss.column_names)
ds_magicoder_evol = ds_magicoder_evol.map(convert_magicoder_evol, num_proc=4, remove_columns=ds_magicoder_evol.column_names)

# Combine all datasets
combined_dataset = concatenate_datasets([ds_kodcode, ds_feedback, ds_magicoder_oss, ds_magicoder_evol])
combined_dataset = combined_dataset.shuffle(seed=42)
print(f"\\nTotal combined dataset: {len(combined_dataset)} samples")

# Quality filter
def filter_quality(example):
    """Remove examples with very short responses."""
    msgs = example["messages"]
    assistant_msgs = [m for m in msgs if m["role"] == "assistant"]
    if not assistant_msgs:
        return False
    total_assistant_len = sum(len(m["content"]) for m in assistant_msgs)
    return total_assistant_len >= 50

combined_dataset = combined_dataset.filter(filter_quality, num_proc=4)
print(f"After quality filter: {len(combined_dataset)} samples")


# ─── 3. Load model with QLoRA (4-bit quantization) ───────────────────────────
print("\\n" + "=" * 60)
print(f"Loading {MODEL_ID} with 4-bit quantization...")
print("=" * 60)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model = prepare_model_for_kbit_training(model)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

print(f"Model loaded. Parameters: {model.num_parameters():,}")


# ─── 4. LoRA config ──────────────────────────────────────────────────────────
peft_config = LoraConfig(
    r=64,
    lora_alpha=128,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules="all-linear",
)


# ─── 5. Training config ──────────────────────────────────────────────────────
training_args = SFTConfig(
    output_dir=OUTPUT_DIR,

    # Data
    max_length=4096,
    packing=True,
    dataset_num_proc=8,

    # Training hyperparams
    num_train_epochs=2,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.05,
    weight_decay=0.01,
    max_grad_norm=1.0,
    optim="paged_adamw_8bit",

    # Memory optimization
    gradient_checkpointing=True,
    bf16=True,
    tf32=True,

    # Logging
    logging_steps=5,
    logging_first_step=True,
    disable_tqdm=True,

    # Saving & Hub
    save_strategy="steps",
    save_steps=1000,
    save_total_limit=3,
    push_to_hub=True,
    hub_model_id=HUB_MODEL_ID,
    hub_strategy="every_save",

    # Monitoring
    report_to="trackio",
    run_name="qwen25-coder-7b-mythos-sft",
    project="code-mythos",

    # Misc
    seed=42,
    dataloader_num_workers=4,
    remove_unused_columns=False,
)


# ─── 6. Create trainer and train ─────────────────────────────────────────────
print("\\n" + "=" * 60)
print("Initializing SFTTrainer...")
print("=" * 60)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=combined_dataset,
    processing_class=tokenizer,
    peft_config=peft_config,
)

trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = model.num_parameters()
print(f"Trainable: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)")

print("\\n" + "=" * 60)
print("Starting training...")
print("=" * 60)

trainer.train()

# ─── 7. Save and push final model ────────────────────────────────────────────
print("\\n" + "=" * 60)
print("Saving final model...")
print("=" * 60)

trainer.save_model(OUTPUT_DIR)
trainer.push_to_hub()

print("\\n" + "=" * 60)
print(f"✅ Training complete! Model pushed to: https://huggingface.co/{HUB_MODEL_ID}")
print("=" * 60)