import os
from unsloth import FastLanguageModel, PatchFastRL, is_bfloat16_supported
PatchFastRL("GRPO", FastLanguageModel)

import torch
import re
from datasets import load_dataset, Dataset, concatenate_datasets
from transformers import AutoConfig, AutoTokenizer
from trl import GRPOConfig, GRPOTrainer

# -------------------------------
# Model setup
# -------------------------------
max_seq_length = 256
lora_rank = 64

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Phi-4-mini-reasoning-unsloth-bnb-4bit",
    cache_dir="/mnt/d/AI/models/phi-4-mini-reasoning-unsloth-bnb-4bit",
    max_seq_length=max_seq_length,
    max_lora_rank=lora_rank,
    gpu_memory_utilization=0.85,
    device_map="auto",
    fast_inference=False,
)

print(type(tokenizer))
print(model.config.model_type)

model = FastLanguageModel.get_peft_model(
    model,
    r=lora_rank,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_alpha=lora_rank,
    use_gradient_checkpointing="unsloth",
    random_state=3407,
)

# -------------------------------
# Prompt format
# -------------------------------
SYSTEM_PROMPT = """You are Villager. Respond in the following format:
<think>
...
</think>
..."""

XML_COT_FORMAT = """<think>
{reasoning}
</think>
{answer}"""

# -------------------------------
# Extraction helpers
# -------------------------------
def extract_xml_answer(text: str) -> str:
    if "</think>" in text:
        return text.split("</think>")[-1].strip()
    return text.strip()

def extract_think(text: str) -> str:
    if "<think>" in text and "</think>" in text:
        return text.split("<think>")[-1].split("</think>")[0].strip()
    return ""

def extract_hash_answer(text: str) -> str | None:
    if "####" not in text:
        return None
    return text.split("####")[1].strip()

# -------------------------------
# Dataset loaders
# -------------------------------
def get_gsm8k_questions(split="train") -> Dataset:
    data = load_dataset("openai/gsm8k", "main")[split]
    data = data.map(lambda x: {
        "prompt": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": x["question"]}
        ],
        "answer": extract_hash_answer(x["answer"])
    })
    return data

def get_mcwiki(split="train") -> Dataset:
    data = load_dataset("lparkourer10/minecraft-wiki")[split]
    data = data.map(lambda x: {
        "prompt": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": x["question"]}
        ],
        "answer": x["answer"]
    })
    return data

gsm8k = get_gsm8k_questions()
mcwiki = get_mcwiki()
dataset = concatenate_datasets([gsm8k, mcwiki])

# -------------------------------
# Reward functions
# -------------------------------
def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
    responses = [completion[0]['content'] for completion in completions]
    q = prompts[0][-1]['content']
    extracted_responses = [extract_xml_answer(r) for r in responses]
    print('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}",
          f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
    return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]

def int_reward_func(completions, **kwargs) -> list[float]:
    responses = [completion[0]['content'] for completion in completions]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]

def strict_format_reward_func(completions, **kwargs) -> list[float]:
    """Check for strict <think>...</think>{answer} format."""
    pattern = r"^<think>\n.*?\n</think>\n.*$"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r, re.DOTALL) for r in responses]
    return [0.5 if match else 0.0 for match in matches]

def soft_format_reward_func(completions, **kwargs) -> list[float]:
    """Check for loose <think>...</think>{answer} format."""
    pattern = r"<think>.*?</think>.*"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r, re.DOTALL) for r in responses]
    return [0.5 if match else 0.0 for match in matches]

def count_xml(text) -> float:
    count = 0.0
    if text.count("<think>\n") == 1:
        count += 0.25
    if text.count("\n</think>\n") == 1:
        count += 0.25
    return count

def xmlcount_reward_func(completions, **kwargs) -> list[float]:
    contents = [completion[0]["content"] for completion in completions]
    return [count_xml(c) for c in contents]

# -------------------------------
# Training setup
# -------------------------------
training_args = GRPOConfig(
    use_vllm=True,
    learning_rate=2e-5,
    adam_beta1=0.9,
    adam_beta2=0.95,
    weight_decay=0.01,
    warmup_ratio=0.05,
    lr_scheduler_type="cosine",
    optim="adamw_8bit",
    logging_steps=5,
    bf16=is_bfloat16_supported(),
    fp16=not is_bfloat16_supported(),
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    num_generations=2,
    max_prompt_length=256,
    max_completion_length=256,
    max_steps=500,
    save_steps=100,
    max_grad_norm=1.0,
    report_to="wandb",
    output_dir="outputs_phi4",
    run_name="Villager"
)

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[
        xmlcount_reward_func,
        soft_format_reward_func,
        strict_format_reward_func,
        int_reward_func,
        correctness_reward_func,
    ],
    args=training_args,
    train_dataset=dataset,
)

trainer.train()
model.save_lora("grpo_saved_lora")