"""Training script generated by BilgeAI v{{ schema_version }}
Model: {{ model_name }}
Training: RLHF (PPO - Proximal Policy 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,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
)
from trl import PPOConfig, PPOTrainer, AutoModelForCausalLMWithValueHead

# ============================================================
# 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 }}"
reward_model_name = "{{ reward_model_name | default(model_name, true) }}"
{% 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
tokenizer.padding_side = "left"

# Policy model (with value head for PPO)
model = AutoModelForCausalLMWithValueHead.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,
    {% if lora_enabled %}peft_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 %}
)

# Reward model
reward_model = AutoModelForSequenceClassification.from_pretrained(
    reward_model_name,
    num_labels=1,
    device_map="auto",
    torch_dtype=torch.{{ "bfloat16" if bf16 else "float16" }},
    trust_remote_code=True,
)

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

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

# RLHF expects a "query" column with prompts
# Map your data format:
{% if dataset_format == "alpaca" %}

def format_query(example):
    """Extract prompts for RLHF from Alpaca format."""
    if example.get("input", "").strip():
        query = (
            f"### Instruction:\n{example['instruction']}\n\n"
            f"### Input:\n{example['input']}\n\n"
            f"### Response:\n"
        )
    else:
        query = (
            f"### Instruction:\n{example['instruction']}\n\n"
            f"### Response:\n"
        )
    return {"query": query}

dataset = dataset.map(format_query)
{% else %}

def format_query(example):
    """Extract prompts for RLHF."""
    return {"query": example.get("prompt", example.get("instruction", ""))}

dataset = dataset.map(format_query)
{% endif %}

# Tokenize queries
def tokenize(example):
    tokens = tokenizer(
        example["query"],
        truncation=True,
        max_length={{ max_seq_length }},
        padding=False,
    )
    example["input_ids"] = tokens["input_ids"]
    return example

dataset = dataset.map(tokenize, batched=False)
dataset.set_format(type="torch")

# ============================================================
# PPO Config
# ============================================================

ppo_config = PPOConfig(
    model_name=model_name,
    learning_rate={{ learning_rate }},
    batch_size={{ batch_size * gradient_accumulation_steps }},
    mini_batch_size={{ batch_size }},
    ppo_epochs={{ ppo_epochs | default(4, true) }},
    gradient_accumulation_steps={{ gradient_accumulation_steps }},
    seed=SEED,
    log_with={{ '"' + report_to + '"' if report_to and report_to != "none" else "None" }},
)

# ============================================================
# PPO Trainer
# ============================================================

ppo_trainer = PPOTrainer(
    config=ppo_config,
    model=model,
    tokenizer=tokenizer,
    dataset=dataset,
)

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

print("Starting RLHF (PPO) training...")

generation_kwargs = {
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
    "max_new_tokens": 256,
}

for epoch in range({{ num_epochs }}):
    print(f"\nEpoch {epoch + 1}/{{ num_epochs }}")
    for batch_idx, batch in enumerate(ppo_trainer.dataloader):
        query_tensors = batch["input_ids"]

        # Generate responses
        response_tensors = ppo_trainer.generate(
            query_tensors,
            return_prompt=False,
            **generation_kwargs,
        )
        batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

        # Compute rewards
        texts = [q + r for q, r in zip(batch["query"], batch["response"])]
        reward_inputs = tokenizer(
            texts, return_tensors="pt", padding=True, truncation=True
        ).to(reward_model.device)
        with torch.no_grad():
            rewards = reward_model(**reward_inputs).logits.squeeze(-1)
        rewards_list = [r for r in rewards]

        # PPO step
        stats = ppo_trainer.step(query_tensors, response_tensors, rewards_list)
        ppo_trainer.log_stats(stats, batch, rewards_list)

        if batch_idx % 10 == 0:
            print(f"  Batch {batch_idx}: mean_reward={torch.stack(rewards_list).mean():.4f}")

# Save
output_dir = "./output/{{ model_name | replace('/', '_') }}_rlhf/final"
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
print(f"RLHF training complete! Model saved to {output_dir}")
