"""Training script generated by BilgeAI v{{ schema_version }}
Model: {{ model_name }}
Training: DPO (Direct Preference Optimization)
Generated for: {{ hardware_summary }}
"""

import random
import numpy as np
import torch
from datasets import load_dataset
from peft import LoraConfig, TaskType
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from trl import DPOTrainer

# ============================================================
# Reproducibility
# ============================================================

SEED = {{ seed }}
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
    torch.mps.manual_seed(SEED)

# ============================================================
# Model & Tokenizer
# ============================================================

model_name = "{{ model_name }}"
{% if quantization == "4bit" %}
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.{{ "bfloat16" if bf16 else "float16" }},
    bnb_4bit_use_double_quant=True,
)
{% elif quantization == "8bit" %}
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
)
{% else %}
bnb_config = None
{% endif %}

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    {% if quantization != "none" %}quantization_config=bnb_config,{% endif %}
    device_map="auto",
    torch_dtype=torch.{{ "bfloat16" if bf16 else "float16" }},
    trust_remote_code=True,
)

# Reference model (for DPO loss computation)
ref_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    {% if quantization != "none" %}quantization_config=bnb_config,{% endif %}
    device_map="auto",
    torch_dtype=torch.{{ "bfloat16" if bf16 else "float16" }},
    trust_remote_code=True,
)

# ============================================================
# LoRA Configuration
# ============================================================
{% if lora_enabled %}

lora_config = LoraConfig(
    r={{ lora_rank }},
    lora_alpha={{ lora_alpha }},
    lora_dropout=0.05,
    target_modules={{ lora_target_modules }},
    task_type=TaskType.CAUSAL_LM,
    bias="none",
)
{% endif %}

# ============================================================
# Dataset
# ============================================================

dataset = load_dataset("json", data_files="{{ dataset_path }}", split="train")

# DPO expects columns: prompt, chosen, rejected
# If your data uses different column names, map them here:
# dataset = dataset.rename_columns({
#     "question": "prompt",
#     "preferred": "chosen",
#     "dispreferred": "rejected",
# })

# ============================================================
# Training Arguments
# ============================================================

training_args = TrainingArguments(
    output_dir="./output/{{ model_name | replace('/', '_') }}_dpo",
    num_train_epochs={{ num_epochs }},
    per_device_train_batch_size={{ batch_size }},
    gradient_accumulation_steps={{ gradient_accumulation_steps }},
    learning_rate={{ learning_rate }},
    lr_scheduler_type="{{ lr_scheduler }}",
    warmup_ratio={{ warmup_ratio }},
    weight_decay=0.01,
    optim="{{ optimizer_name }}",
    fp16={{ fp16 }},
    bf16={{ bf16 }},
    max_grad_norm=1.0,
    logging_steps=10,
    save_strategy="epoch",
    save_total_limit=2,
    {% if use_gradient_checkpointing %}gradient_checkpointing=True,{% endif %}
    {% if distributed_strategy == "ddp" %}ddp_find_unused_parameters=False,{% endif %}
    {% if distributed_strategy == "fsdp" %}fsdp="full_shard auto_wrap",
    fsdp_config={"transformer_layer_cls_to_wrap": "LlamaDecoderLayer"},{% endif %}
    {% if report_to %}report_to="{{ report_to }}",{% else %}report_to="none",{% endif %}
    seed={{ seed }},
)

# ============================================================
# DPO Trainer
# ============================================================

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
    {% if lora_enabled %}peft_config=lora_config,{% endif %}
    max_length={{ max_seq_length }},
    beta=0.1,
)

# ============================================================
# Train
# ============================================================

print("Starting DPO training...")
trainer.train()

# Save
trainer.save_model("./output/{{ model_name | replace('/', '_') }}_dpo/final")
tokenizer.save_pretrained("./output/{{ model_name | replace('/', '_') }}_dpo/final")
print("DPO training complete! Model saved to ./output/{{ model_name | replace('/', '_') }}_dpo/final")
