#!/usr/bin/env python3
"""
Genesis-2.0 — Phase 2: GRPO with Verifiable Rewards

Run on RunPod RTX PRO 6000 Blackwell (96 GB) using TRL + Unsloth + vLLM.

Steps:
1. Load Phase 1 checkpoint (or Genesis-1.0 base + all-linear LoRA)
2. Configure GRPOTrainer with vLLM for rollouts
3. Run GRPO with rule-based reward functions
4. Save + upload adapter

Usage:
    python3 train_grpo.py [--from-dpo /workspace/genesis2-dpo]
"""

import os
import sys
import json
import torch
import argparse
from typing import Optional

os.environ["UNSLOTH_VLLM_STANDBY"] = "1"


class Config:
    # Model
    base_model = "Qwen/Qwen3.6-35B-A3B"
    # If --from-dpo is set, load that adapter. Otherwise load Genesis-1.0 SFT adapter
    initial_adapter = "jacobeen06/Genesis-1.0-SFT-adapter"
    dpo_adapter_path = None  # Override with --from-dpo

    # Output
    output_dir = "/workspace/genesis2-grpo"
    hf_repo = "jacobeen06/Genesis-2.0-GRPO-adapter"

    # QLoRA
    load_in_4bit = True
    bnb_4bit_quant_type = "nf4"
    bnb_4bit_compute_dtype = torch.bfloat16

    # All-linear LoRA (Config C)
    lora_r = 32
    lora_alpha = 64
    lora_dropout = 0.0
    lora_target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
        "gate",
    ]
    use_rslora = True

    # GRPO
    num_generations = 4  # G=4 responses per prompt
    max_length = 4096
    max_prompt_length = 3072
    beta = 0.04  # KL penalty (start)
    beta_decay = True  # Decay to 0.01 during training
    clip_high = 0.28  # DAPO-style asymmetric
    clip_low = 0.20

    # Training
    learning_rate = 3e-6
    lr_scheduler_type = "cosine"
    warmup_ratio = 0.05
    per_device_train_batch_size = 1
    gradient_accumulation_steps = 4
    num_train_epochs = 1
    logging_steps = 5
    save_steps = 100
    save_total_limit = 2

    # vLLM
    vllm_gpu_memory_utilization = 0.90

    # Data
    prompts_data = "/workspace/training_prompts.jsonl"  # Prompts for GRPO rollout


def load_model(config: Config):
    """Load base model + merge Genesis-1.0 SFT → attach all-linear LoRA."""
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import PeftModel, LoraConfig, get_peft_model

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=config.load_in_4bit,
        bnb_4bit_quant_type=config.bnb_4bit_quant_type,
        bnb_4bit_compute_dtype=config.bnb_4bit_compute_dtype,
    )

    model = AutoModelForCausalLM.from_pretrained(
        config.base_model,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )

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

    # Load adapter
    if config.dpo_adapter_path and os.path.exists(config.dpo_adapter_path):
        print(f"Loading DPO adapter from {config.dpo_adapter_path}")
        model = PeftModel.from_pretrained(model, config.dpo_adapter_path)
    else:
        print(f"Loading Genesis-1.0 SFT adapter from {config.initial_adapter}")
        model = PeftModel.from_pretrained(model, config.initial_adapter)

    model = model.merge_and_unload()

    # Attach new LoRA for GRPO
    lora_config = LoraConfig(
        r=config.lora_r,
        lora_alpha=config.lora_alpha,
        target_modules=config.lora_target_modules,
        lora_dropout=config.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        use_rslora=config.use_rslora,
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    return model, tokenizer


def reward_func(completions, **kwargs):
    """
    Rule-based reward function for GRPO.
    Runs the rewards module on each generated completion.
    """
    # Import our reward functions
    sys.path.insert(0, "/workspace/genesis-rlhf")
    from rewards import combined_reward

    rewards = []
    for completion in completions:
        score = combined_reward(completion)
        rewards.append(score)
    return rewards


def load_prompts(config: Config):
    """Load training prompts for GRPO rollout."""
    prompts = []
    with open(config.prompts_data) as f:
        for line in f:
            line = line.strip()
            if line:
                item = json.loads(line)
                prompts.append(item.get("prompt", item.get("text", "")))
    print(f"Loaded {len(prompts)} prompts for GRPO")
    return prompts


def train_grpo(config: Config):
    from trl import GRPOTrainer
    from transformers import TrainingArguments
    from datasets import Dataset

    model, tokenizer = load_model(config)
    prompts = load_prompts(config)

    # Create dataset with just prompts
    dataset = Dataset.from_list([{"prompt": p} for p in prompts])

    # Training args
    training_args = TrainingArguments(
        output_dir=config.output_dir,
        per_device_train_batch_size=config.per_device_train_batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        learning_rate=config.learning_rate,
        lr_scheduler_type=config.lr_scheduler_type,
        warmup_ratio=config.warmup_ratio,
        num_train_epochs=config.num_train_epochs,
        logging_steps=config.logging_steps,
        save_steps=config.save_steps,
        save_total_limit=config.save_total_limit,
        bf16=True,
        tf32=True,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        report_to="wandb" if os.environ.get("WANDB_API_KEY") else "none",
        run_name="genesis2-grpo",
        dataloader_num_workers=2,
    )

    # GRPO Trainer
    trainer = GRPOTrainer(
        model=model,
        reward_funcs=[reward_func],
        args=training_args,
        train_dataset=dataset,
        tokenizer=tokenizer,
        num_generations=config.num_generations,
        max_length=config.max_length,
        max_prompt_length=config.max_prompt_length,
        beta=config.beta,
        clip_high=config.clip_high,
        clip_low=config.clip_low,
    )

    # Train
    print("Starting GRPO training...")
    trainer.train()

    # Save
    trainer.save_model(config.output_dir)
    tokenizer.save_pretrained(config.output_dir)

    # Upload
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        api.create_repo(config.hf_repo, exist_ok=True)
        api.upload_folder(
            folder_path=config.output_dir,
            repo_id=config.hf_repo,
            commit_message="Genesis-2.0 GRPO adapter",
        )
        print(f"Uploaded to {config.hf_repo}")
    except Exception as e:
        print(f"Upload failed (non-fatal): {e}")

    print(f"DONE! Adapter saved to {config.output_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--from-dpo", type=str, default=None,
                        help="Path to DPO-trained adapter to start from")
    args = parser.parse_args()

    config = Config()
    config.dpo_adapter_path = args.from_dpo

    print("=" * 60)
    print("Genesis-2.0 — Phase 2: GRPO Training")
    print("=" * 60)
    print(f"Base model: {config.base_model}")
    print(f"Start from DPO: {config.dpo_adapter_path or 'No (from SFT)'}")
    print(f"Generations per prompt: {config.num_generations}")
    print(f"Output: {config.output_dir}")
    print()

    train_grpo(config)