"""
Qwen3-8B Coding & Agentic Reasoning Expert — Multi-Dataset SFT Training
========================================================================
Base: Qwen/Qwen3-8B (Apache 2.0, 8.2B params, 32K context)
Method: QLoRA SFT with assistant-only loss masking
Datasets:
  - TIGER-Lab/VisCode-200K (visualization/chart generation) — ChatML ready
  - m-a-p/CodeFeedback-Filtered-Instruction (code instruction tuning)
  - nvidia/OpenCodeReasoning (reasoning with <think> blocks)
  - glaiveai/glaive-function-calling-v2 (tool calling)
  - ise-uiuc/Magicoder-OSS-Instruct-75K (code generation)

Recipe: Based on Qwen3-Coder-Next + LoRA Without Regret papers
Target: Coding + agentic reasoning + visualization + tool-use expert

Usage:
  pip install transformers>=4.51.0 trl>=1.3.0 peft>=0.15.0 datasets accelerate bitsandbytes torch trackio
  HUB_MODEL_ID=your-username/model-name python train_coding_agent.py
"""

import os
import re
import json
import torch
import trackio
from datasets import load_dataset, concatenate_datasets, Dataset
from transformers import AutoTokenizer, BitsAndBytesConfig, TrainerCallback
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig, TaskType

# ============================================================
# Configuration
# ============================================================
MODEL_ID = "Qwen/Qwen3-8B"
OUTPUT_DIR = "./qwen3-8b-coding-agent"
HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "sukritvemula/Qwen3-8B-CodeAgent")

# Training hyperparameters (from Qwen3 + LoRA Without Regret papers)
LEARNING_RATE = 2e-4
NUM_EPOCHS = 2
BATCH_SIZE = 2
GRAD_ACCUM = 8
MAX_LENGTH = 4096
LORA_R = 64
LORA_ALPHA = 16
WARMUP_RATIO = 0.05

# Dataset proportions (~50K samples)
MAX_VISCODE = 12000
MAX_CODEFEEDBACK = 10000
MAX_OPENCODE = 10000
MAX_GLAIVE = 8000
MAX_MAGICODER = 10000

SYSTEM_PROMPT = """You are an expert AI assistant specialized in coding, agentic reasoning, data visualization, and tool use. You can:
1. Write, debug, and explain code in any programming language
2. Reason step-by-step through complex problems using <think>...</think> blocks
3. Generate charts, graphs, and data visualizations using matplotlib, plotly, seaborn
4. Call functions and tools when needed, returning structured JSON for tool invocations
5. Search the web and read research papers to provide accurate, up-to-date information
6. Replicate images and diagrams programmatically

Always think carefully before responding. Be precise, avoid hallucination, and cite sources when possible."""


class AlertCallback(TrainerCallback):
    def __init__(self):
        self.best_loss = float('inf')
        self.initial_loss = None
        self.steps_since_improvement = 0

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None:
            return
        loss = logs.get("loss")
        if loss is None:
            return
        step = state.global_step
        if self.initial_loss is None:
            self.initial_loss = loss
            trackio.alert(title="Training Started", text=f"Initial loss={loss:.4f} at step {step}. Model: {MODEL_ID}, lr={LEARNING_RATE}, batch={BATCH_SIZE}x{GRAD_ACCUM}={BATCH_SIZE*GRAD_ACCUM}", level="INFO")
        if loss != loss or loss > 20.0:
            trackio.alert(title="DIVERGENCE DETECTED", text=f"loss={loss} at step {step} — training has diverged. lr likely too high, try lr={LEARNING_RATE*0.1:.1e}", level="ERROR")
            return
        if loss < self.best_loss:
            self.best_loss = loss
            self.steps_since_improvement = 0
        else:
            self.steps_since_improvement += 1
        if step > 100 and loss > self.initial_loss * 0.9:
            trackio.alert(title="Slow Convergence", text=f"loss={loss:.4f} at step {step}, only {((self.initial_loss - loss) / self.initial_loss * 100):.1f}% reduction from initial {self.initial_loss:.4f}. Consider lr={LEARNING_RATE*2:.1e}", level="WARN")
        if self.steps_since_improvement > 200:
            trackio.alert(title="Loss Stagnation", text=f"No improvement for {self.steps_since_improvement} steps. Best loss={self.best_loss:.4f}, current={loss:.4f}.", level="WARN")
        if step > 0 and step % 500 == 0:
            trackio.alert(title="Training Milestone", text=f"Step {step}: loss={loss:.4f}, best_loss={self.best_loss:.4f}, lr={logs.get('learning_rate', 'N/A')}", level="INFO")


def process_viscode(max_samples):
    print(f"Loading VisCode-200K (max {max_samples})...")
    ds = load_dataset("TIGER-Lab/VisCode-200K", split=f"train[:{max_samples}]")
    def add_system(example):
        messages = example["messages"]
        if messages and messages[0]["role"] != "system":
            messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
        return {"messages": messages}
    ds = ds.map(add_system, num_proc=4)
    print(f"  VisCode: {len(ds)} samples loaded")
    return ds


def process_codefeedback(max_samples):
    print(f"Loading CodeFeedback (max {max_samples})...")
    ds = load_dataset("m-a-p/CodeFeedback-Filtered-Instruction", split=f"train[:{max_samples}]")
    def to_messages(example):
        return {"messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": example["query"]},
            {"role": "assistant", "content": example["answer"]}
        ]}
    ds = ds.map(to_messages, remove_columns=ds.column_names, num_proc=4)
    print(f"  CodeFeedback: {len(ds)} samples loaded")
    return ds


def process_opencode_reasoning(max_samples):
    print(f"Loading OpenCodeReasoning (max {max_samples})...")
    ds = load_dataset("nvidia/OpenCodeReasoning", "split_0", split=f"split_0[:{max_samples}]")
    def to_messages(example):
        return {"messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": example["input"]},
            {"role": "assistant", "content": example["output"]}
        ]}
    ds = ds.map(to_messages, remove_columns=ds.column_names, num_proc=4)
    print(f"  OpenCodeReasoning: {len(ds)} samples loaded")
    return ds


def process_glaive_function_calling(max_samples):
    print(f"Loading Glaive Function Calling (max {max_samples})...")
    ds = load_dataset("glaiveai/glaive-function-calling-v2", split=f"train[:{max_samples}]")
    def to_messages(example):
        system_content = re.sub(r'^SYSTEM:\s*', '', example["system"])
        chat = example["chat"]
        messages = [{"role": "system", "content": system_content}]
        parts = re.split(r'\n*(USER:|ASSISTANT:|FUNCTION RESPONSE:)', chat)
        current_role, current_content = None, ""
        for part in parts:
            part = part.strip()
            if not part:
                continue
            if part == "USER:":
                if current_role and current_content.strip():
                    messages.append({"role": current_role, "content": current_content.strip()})
                current_role, current_content = "user", ""
            elif part == "ASSISTANT:":
                if current_role and current_content.strip():
                    messages.append({"role": current_role, "content": current_content.strip()})
                current_role, current_content = "assistant", ""
            elif part == "FUNCTION RESPONSE:":
                if current_role and current_content.strip():
                    messages.append({"role": current_role, "content": current_content.strip()})
                current_role, current_content = "user", "[Function Response] "
            else:
                current_content += part
        if current_role and current_content.strip():
            messages.append({"role": current_role, "content": current_content.strip()})
        merged = [messages[0]]
        for msg in messages[1:]:
            if merged and msg["role"] == merged[-1]["role"]:
                merged[-1]["content"] += "\n" + msg["content"]
            else:
                merged.append(msg)
        messages = merged
        if len(messages) < 3 or messages[-1]["role"] != "assistant":
            return {"messages": []}
        return {"messages": messages}
    ds = ds.map(to_messages, remove_columns=ds.column_names, num_proc=4)
    ds = ds.filter(lambda x: len(x["messages"]) >= 3 and any(m["role"] == "assistant" for m in x["messages"]))
    print(f"  Glaive Function Calling: {len(ds)} samples loaded")
    return ds


def process_magicoder(max_samples):
    print(f"Loading Magicoder (max {max_samples})...")
    ds = load_dataset("ise-uiuc/Magicoder-OSS-Instruct-75K", split=f"train[:{max_samples}]")
    def to_messages(example):
        return {"messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": example["problem"]},
            {"role": "assistant", "content": example["solution"]}
        ]}
    ds = ds.map(to_messages, remove_columns=ds.column_names, num_proc=4)
    print(f"  Magicoder: {len(ds)} samples loaded")
    return ds


def main():
    print("=" * 60)
    print("Qwen3-8B Coding & Agentic Reasoning Expert Training")
    print("=" * 60)

    datasets_list = []
    for loader in [
        lambda: process_viscode(MAX_VISCODE),
        lambda: process_codefeedback(MAX_CODEFEEDBACK),
        lambda: process_opencode_reasoning(MAX_OPENCODE),
        lambda: process_glaive_function_calling(MAX_GLAIVE),
        lambda: process_magicoder(MAX_MAGICODER),
    ]:
        try:
            datasets_list.append(loader())
        except Exception as e:
            print(f"  ⚠️ Failed: {e}")

    if not datasets_list:
        raise ValueError("No datasets loaded!")

    combined = concatenate_datasets(datasets_list).shuffle(seed=42)
    print(f"✅ Total training samples: {len(combined)}")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True,
    )
    peft_config = LoraConfig(
        r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=0.05, bias="none",
        task_type=TaskType.CAUSAL_LM, target_modules="all-linear", use_rslora=True,
    )
    training_args = SFTConfig(
        output_dir=OUTPUT_DIR, push_to_hub=True, hub_model_id=HUB_MODEL_ID,
        hub_strategy="every_save", max_length=MAX_LENGTH, packing=False,
        assistant_only_loss=True, num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE, gradient_accumulation_steps=GRAD_ACCUM,
        learning_rate=LEARNING_RATE, lr_scheduler_type="cosine", warmup_ratio=WARMUP_RATIO,
        weight_decay=0.01, max_grad_norm=1.0, bf16=True, tf32=True,
        gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False},
        logging_steps=10, logging_first_step=True, disable_tqdm=True,
        save_strategy="steps", save_steps=500, save_total_limit=3, eval_strategy="no",
        report_to="trackio", run_name="sft-qwen3-8b-coding-agent-v1",
        model_init_kwargs={
            "quantization_config": bnb_config, "device_map": "auto",
            "use_cache": False, "torch_dtype": torch.bfloat16,
        },
        seed=42, dataloader_num_workers=4, dataloader_pin_memory=True,
    )

    trainer = SFTTrainer(
        model=MODEL_ID, args=training_args, train_dataset=combined,
        peft_config=peft_config, callbacks=[AlertCallback()],
    )

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

    train_result = trainer.train()
    trainer.save_model(OUTPUT_DIR)
    trainer.push_to_hub(commit_message="Training complete: Qwen3-8B Coding Agent v1")

    metrics = train_result.metrics
    trackio.alert(title="Training Complete", text=f"Final loss={metrics.get('train_loss', 'N/A')}, hub_model={HUB_MODEL_ID}", level="INFO")
    print(f"✅ DONE! Model: https://huggingface.co/{HUB_MODEL_ID}")


if __name__ == "__main__":
    main()