import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, List, Optional

import mlx.core as mx
import mlx.nn as nn
import numpy as np
from mlx.utils import tree_flatten, tree_map
from mlx_lm.generate import batch_generate
from mlx_lm.sample_utils import make_sampler
from mlx_lm.tuner.callbacks import TrainingCallback
from tqdm import tqdm
from functools import partial

from .grpo_reward_functions import (
    RewardFunctions,
    r1_accuracy_reward_func,
    r1_count_xml,
    r1_int_reward_func,
    r1_soft_format_reward_func,
    r1_strict_format_reward_func,
)
from .sft_trainer import SFTTrainingArgs, average_gradients, grad_checkpoint


@dataclass
class GRPOTrainingArgs(SFTTrainingArgs):
    group_size: int = field(
        default=4,
        metadata={"help": "Number of responses per prompt."},
    )
    beta: float = field(default=0.1, metadata={"help": "KL penalty coefficient."})
    epsilon: float = field(
        default=1e-4, metadata={"help": "The Epsilon for numerical stability."}
    )
    epsilon_high: float = field(
        default=None,
        metadata={
            "help": "For DAPO Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound specified in argument epsilon."
        },
    )
    max_completion_length: int = field(
        default=512, metadata={"help": "Number of Generations."}
    )
    reference_model_path: str = field(
        default=None,
        metadata={
            "help": "Path to reference model weights. If None, uses the same model."
        },
    )
    temperature: float = field(
        default=0.8,
        metadata={
            "help": "Temperature for sampling. The higher the temperature, the more random the completions."
        },
    )
    top_p: float = field(
        default=0.95,
        metadata={"help": "Top-p sampling parameter."},
    )
    top_k: int = field(default=20, metadata={"help": "Top-k sampling parameter."})
    min_p: float = field(
        default=0.0, metadata={"help": "Minimum probability for sampling."}
    )
    grpo_loss_type: str = field(
        default="grpo",
        metadata={
            "help": "Type of loss to use for GRPO. Supported: 'grpo', 'bnpo', 'dr_grpo'."
        },
    )
    reward_weights: Optional[List[float]] = field(
        default=None,
        metadata={
            "help": "Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are weighted equally with weight `1.0`."
        },
    )
    importance_sampling_level: str = field(
        default=None,
        metadata={
            "help": "importance_sampling_level (`str`, *optional*, defaults to None): "
            "Controls whether importance sampling ratios are computed at the 'token' or 'sequence' level. "
            "keeps the raw per-token log-probability ratios (one weight per token).  'sequence' averages the "
            "log-probability ratios across valid tokens to produce a single ratio per sequence. The "
            "GSPO paper https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more "
            "stable training and better alignment with  sequence-level rewards.."
        },
    )


def get_per_token_logps(model: nn.Module, inputs, lengths):
    logits = model(inputs).astype(mx.float16)
    logits = logits[:, :-1, :]
    targets = inputs[:, 1:]
    per_token_logps = []
    for i in range(logits.shape[0]):
        seq_len = int(lengths[i]) - 1
        if seq_len <= 0:
            # If sequence is too short, return empty log probs
            per_token_logps.append(mx.array([]))
            continue
        seq_logits = logits[i, :seq_len]
        seq_targets = targets[i, :seq_len]
        log_probs = nn.log_softmax(seq_logits, axis=-1)
        token_log_probs = mx.take_along_axis(
            log_probs, seq_targets.reshape(seq_len, 1), axis=-1
        ).squeeze(-1)
        per_token_logps.append(token_log_probs)
    return per_token_logps


def generate_grpo(
    model: nn.Module,
    tokenizer,
    prompt_tokens,
    max_tokens: int,
    group_size: int,
    batch_size: int,
    end_token: str,
    temperature: float,
    top_p: float,
    top_k: int,
    min_p: float,
):
    was_training = model.training
    model.eval()
    try:
        all_completions = []
        all_completion_texts = []
        batch_indices = []

        total_samples = len(prompt_tokens)

        use_eos_token = False
        if end_token:
            try:
                tokenizer.add_eos_token(end_token)
                use_eos_token = True
            except ValueError:
                use_eos_token = False

        for i in range(0, total_samples, batch_size):
            current_batch_size = min(batch_size, total_samples - i)
            batch_prompts = prompt_tokens[i : i + current_batch_size]

            batched_prompts = []
            batched_indices = []
            for j, prompt in enumerate(batch_prompts):
                for k in range(group_size):
                    batched_prompts.append(prompt)
                    batched_indices.append(i + j)

            sampler = make_sampler(
                temperature,
                top_p=top_p,
                min_p=min_p,
                top_k=top_k,
            )

            results = batch_generate(
                model=model,
                tokenizer=tokenizer,
                prompts=batched_prompts,
                max_tokens=max_tokens,
                sampler=sampler,
                verbose=False,
            )

            for idx, completion_text in enumerate(results.texts):
                completion_ids = tokenizer.encode(completion_text)

                if not use_eos_token and end_token:
                    end_sequence = tokenizer.encode(end_token)
                    if (
                        len(completion_ids) >= len(end_sequence)
                        and completion_ids[-len(end_sequence) :] == end_sequence
                    ):
                        completion_ids = completion_ids[: -len(end_sequence)]

                if len(completion_ids) == 0:
                    completion_ids = mx.array([], dtype=mx.int32)
                else:
                    completion_ids = mx.array(completion_ids)
                all_completions.append(mx.stop_gradient(completion_ids))
                all_completion_texts.append(completion_text)
                batch_indices.append(batched_indices[idx])

            # PATCH: Clear memory after each batch
            del results
            mx.eval(all_completions[-len(batched_prompts) :])
            mx.clear_cache()

        if not all_completions:
            raise ValueError(
                "No valid completions generated. Check that prompts are not empty "
                "and end_token configuration is correct."
            )
        return all_completions, all_completion_texts, batch_indices
    finally:
        mx.clear_cache()
        if was_training:
            model.train()


def calculate_rewards_and_advantages(
    reward_funcs: List[RewardFunctions],
    expanded_prompts: List[str],
    all_completion_texts: List[str],
    expanded_answers: List[str],
    expanded_types: List,
    batch_indices: List[int],
    unique_prompt_indices: List[int],
    reward_weights: Optional[List[float]] = None,
):
    """Calculate rewards and advantages for completions."""
    # Calculate rewards from all reward functions
    all_func_rewards = []
    for reward_func in reward_funcs:
        raw_rewards = reward_func(
            prompts=expanded_prompts,
            completions=all_completion_texts,
            answer=expanded_answers,
            types=expanded_types,
        )
        if raw_rewards is None:
            processed_rewards = [float("nan")] * len(all_completion_texts)
        else:
            processed_rewards = [
                float(r) if r is not None else float("nan") for r in raw_rewards
            ]
        func_rewards = mx.array(processed_rewards)
        all_func_rewards.append(func_rewards)

    rewards = mx.stack(all_func_rewards, axis=1)

    # Check for all NaN rows
    all_nan_rows = mx.all(mx.isnan(rewards), axis=1)
    if mx.any(all_nan_rows):
        nan_row_idx = mx.argmax(all_nan_rows).item()
        warning_msg = (
            f"All reward functions returned None for prompt: {expanded_prompts[nan_row_idx]}, "
            f"completion: {all_completion_texts[nan_row_idx]}, "
            f"answer: {expanded_answers[nan_row_idx]}. "
            "Please ensure that at least one reward function returns a valid reward."
        )
        raise RuntimeError(warning_msg)

    # Apply reward weights
    if reward_weights is not None:
        if len(reward_weights) != len(reward_funcs):
            raise ValueError(
                f"Number of reward weights ({len(reward_weights)}) must match number of reward "
                f"functions ({len(reward_funcs)})"
            )
        reward_weights = mx.array(reward_weights, dtype=mx.float32)
    else:
        reward_weights = mx.ones(len(reward_funcs), dtype=mx.float32)

    # Handle NaN values and compute weighted sum
    valid_reward_mask = ~mx.isnan(rewards)
    rewards_no_nan = mx.where(valid_reward_mask, rewards, mx.zeros_like(rewards))
    rewards = (rewards_no_nan * mx.expand_dims(reward_weights, 0)).sum(axis=1)

    # Group rewards by prompt
    num_unique_prompts = len(unique_prompt_indices)
    rewards_by_prompt = [[] for _ in range(num_unique_prompts)]
    for i, prompt_idx in enumerate(batch_indices):
        prompt_position = unique_prompt_indices.index(prompt_idx)
        rewards_by_prompt[prompt_position].append(rewards[i])

    # Calculate advantages
    advantages = mx.zeros_like(rewards)
    for i, prompt_rewards in enumerate(rewards_by_prompt):
        if len(prompt_rewards) > 1:
            prompt_rewards = mx.array(prompt_rewards)
            mean_reward = mx.mean(prompt_rewards)
            std_reward = mx.std(prompt_rewards)
            indices = [
                j
                for j, idx in enumerate(batch_indices)
                if idx == unique_prompt_indices[i]
            ]
            for j, idx in enumerate(indices):
                advantages[idx] = (prompt_rewards[j] - mean_reward) / (
                    std_reward + 1e-4
                )
        else:
            idx = batch_indices.index(unique_prompt_indices[i])
            advantages[idx] = 0.0

    # Calculate reward metrics
    reward_metrics = {}
    for i, reward_func in enumerate(reward_funcs):
        func_name = reward_func.__name__
        raw_rewards = reward_func(
            prompts=expanded_prompts,
            completions=all_completion_texts,
            answer=expanded_answers,
        )
        valid_mask = ~mx.isnan(
            mx.array(
                [
                    reward if reward is not None else float("nan")
                    for reward in raw_rewards
                ]
            )
        )
        valid_rewards = mx.array(
            [
                reward
                for reward in raw_rewards
                if reward is not None and not mx.isnan(reward)
            ]
        )
        if len(valid_rewards) > 0:
            reward_metrics[f"{func_name}_mean"] = mx.mean(valid_rewards)
            reward_metrics[f"{func_name}_std"] = (
                mx.std(valid_rewards) if len(valid_rewards) > 1 else mx.zeros(1)
            )
            reward_metrics[f"{func_name}_coverage"] = valid_mask.sum() / len(
                raw_rewards
            )
        else:
            reward_metrics[f"{func_name}_mean"] = float("nan")
            reward_metrics[f"{func_name}_std"] = float("nan")
            reward_metrics[f"{func_name}_coverage"] = 0.0

    # Calculate grouped rewards statistics
    grouped_rewards_mean = mx.array(
        [mx.mean(mx.array(rewards)) for rewards in rewards_by_prompt]
    )
    grouped_rewards_std = mx.array(
        [
            mx.std(mx.array(rewards)) if len(rewards) > 1 else mx.zeros(1)
            for rewards in rewards_by_prompt
        ]
    )

    # Prepare reward-specific metrics
    reward_specific_metrics = {
        "total_rewards_mean": mx.mean(rewards),
        "total_rewards_std": mx.std(rewards),
        "grouped_rewards_mean": mx.mean(grouped_rewards_mean),
        "grouped_rewards_std": mx.mean(grouped_rewards_std),
        **reward_metrics,
    }

    return advantages, reward_specific_metrics


def grpo_loss(
    model,
    ref_model,
    batch,
    completions=None,
    completion_texts=None,
    batch_indices=None,
    advantages=None,
    reward_metrics=None,
    beta: float = 0.1,
    epsilon: float = 1e-4,
    epsilon_high: float = None,
    max_tokens: int = 64,
    importance_sampling_level: str = "token",
    grpo_loss_type: str = "grpo",
):
    all_completions = completions
    batch_indices = batch_indices

    if not all_completions:
        raise ValueError(
            "No completions were generated. Please check your model and inputs."
        )

    # Prepare padded completions
    max_length = max(ids.shape[0] for ids in all_completions)
    padded_completions = []
    attention_masks = []

    for completion_ids in all_completions:
        completion_tensor = completion_ids
        padding_length = max_length - completion_tensor.shape[0]
        if padding_length > 0:
            padding = mx.zeros((padding_length,), dtype=completion_tensor.dtype)
            padded_ids = mx.concatenate([completion_tensor, padding])
            mask = mx.concatenate(
                [mx.ones_like(completion_tensor), mx.zeros_like(padding)]
            )
        else:
            padded_ids = completion_tensor
            mask = mx.ones_like(completion_tensor)
        padded_completions.append(padded_ids)
        attention_masks.append(mask)

    inputs = mx.stack(padded_completions)
    attention_mask = mx.stack(attention_masks)
    lengths = attention_mask.sum(axis=1)

    # Calculate log probabilities
    token_log_probs = get_per_token_logps(model, inputs, lengths)

    if ref_model is None:
        ref_token_log_probs = token_log_probs
    else:
        ref_token_log_probs = get_per_token_logps(ref_model, inputs, lengths)

    del inputs, attention_mask
    mx.clear_cache()

    # Pad log probabilities
    max_len = max(x.shape[0] for x in token_log_probs)
    padded_log_probs = []
    padded_ref_log_probs = []

    for i in range(len(token_log_probs)):
        seq_len = token_log_probs[i].shape[0]
        padding = mx.zeros((max_len - seq_len,))

        padded_log_probs.append(mx.concatenate([token_log_probs[i], padding]))
        padded_ref_log_probs.append(mx.concatenate([ref_token_log_probs[i], padding]))

    token_log_probs = mx.stack(padded_log_probs)
    ref_token_log_probs = mx.stack(padded_ref_log_probs)

    # Create mask for valid tokens
    length_mask = mx.arange(token_log_probs.shape[1])[None, :] < (lengths[:, None] - 1)

    # Compute log ratio for importance sampling
    log_ratio = token_log_probs - mx.stop_gradient(ref_token_log_probs)

    # Apply importance sampling based on level
    if importance_sampling_level == "token":
        log_importance_weights = log_ratio
    elif importance_sampling_level == "sequence":
        # Average log ratio over sequence length for each sequence
        sequence_log_ratio = (log_ratio * length_mask).sum(axis=1) / mx.maximum(
            length_mask.sum(axis=1), 1.0
        )
        log_importance_weights = mx.expand_dims(sequence_log_ratio, axis=1)
    elif importance_sampling_level is None or importance_sampling_level == "none":
        log_importance_weights = mx.zeros_like(log_ratio)
    else:
        raise ValueError(
            f"Unknown importance sampling level: {importance_sampling_level}. "
            "Possible values are 'token', 'sequence', or None."
        )

    # Calculate importance weights
    coef_1 = mx.exp(log_importance_weights)

    # Apply PPO like clipping
    epsilon_high = epsilon_high if epsilon_high else epsilon
    coef_2 = mx.clip(coef_1, 1 - epsilon, 1 + epsilon_high)

    # Track clipping metrics
    is_low_clipped = (coef_1 < 1 - epsilon) & (advantages.reshape(-1, 1) < 0)
    is_high_clipped = (coef_1 > 1 + epsilon_high) & (advantages.reshape(-1, 1) > 0)
    is_region_clipped = is_low_clipped | is_high_clipped

    # Calculate both unclipped and clipped objectives
    unclipped_obj = coef_1 * advantages.reshape(-1, 1)
    clipped_obj = coef_2 * advantages.reshape(-1, 1)

    # Take the minimum (pessimistic bound)
    per_token_loss = -mx.minimum(unclipped_obj, clipped_obj)

    # Add KL penalty if beta is non-zero
    if beta != 0.0:
        # r_i,t = π_θ / π_old (already computed as coef_1)
        # KL = r_i,t * (π_ref / π_θ) - log(π_ref / π_θ) - 1
        log_ratio_ref_theta = token_log_probs - ref_token_log_probs
        ratio_ref_theta = mx.exp(log_ratio_ref_theta)

        # Unbiased KL estimator
        kl_div = coef_1 * ratio_ref_theta - log_ratio_ref_theta - 1

        # Add KL penalty
        per_token_loss = per_token_loss + beta * kl_div
    else:
        # Compute KL divergence using Schulman's approximator
        log_ratio = ref_token_log_probs - token_log_probs
        kl_div = mx.exp(log_ratio) - log_ratio - 1

    if grpo_loss_type == "grpo":
        loss = (per_token_loss * length_mask).sum() / length_mask.sum()
    elif grpo_loss_type == "bnpo":
        loss = (per_token_loss * length_mask).sum() / mx.maximum(length_mask.sum(), 1.0)
    elif grpo_loss_type == "dr_grpo":
        loss = (per_token_loss * length_mask).sum() / (
            per_token_loss.shape[0] * max_tokens
        )
    else:
        raise ValueError(f"Unknown loss type: {grpo_loss_type}")

    # Calculate mean KL divergence for metrics
    mean_kl = ((kl_div * length_mask).sum(axis=1) / length_mask.sum(axis=1)).mean()

    # Calculate token generation statistics
    completion_lengths = [comp.shape[0] for comp in all_completions]
    max_generated = max(completion_lengths) if completion_lengths else 0
    min_generated = min(completion_lengths) if completion_lengths else 0
    avg_generated = (
        sum(completion_lengths) / len(completion_lengths) if completion_lengths else 0
    )
    # Count how many hit the max token limit
    hit_max_tokens = sum(1 for length in completion_lengths if length >= max_tokens)
    hit_max_ratio = (
        hit_max_tokens / len(completion_lengths) if completion_lengths else 0
    )

    metrics = {
        "kl": mean_kl,
        "average_generated_tokens": avg_generated,
        "max_generated_tokens": max_generated,
        "min_generated_tokens": min_generated,
        "hit_max_tokens_ratio": hit_max_ratio,
        "clip_ratio_low": (
            (is_low_clipped * length_mask).sum() / length_mask.sum()
            if length_mask.sum() > 0
            else mx.zeros(1)
        ),
        "clip_ratio_high": (
            (is_high_clipped * length_mask).sum() / length_mask.sum()
            if length_mask.sum() > 0
            else mx.zeros(1)
        ),
        "clip_ratio_total": (
            (is_region_clipped * length_mask).sum() / length_mask.sum()
            if length_mask.sum() > 0
            else mx.zeros(1)
        ),
        **reward_metrics,  # Include reward-specific metrics
    }

    mx.clear_cache()

    return loss, length_mask.sum(axis=1).sum(), metrics


def iterate_grpo_batches(dataset, batch_size, max_seq_length, train=False):
    has_types = isinstance(dataset[0], tuple) and len(dataset[0]) == 5

    if (
        not dataset
        or not isinstance(dataset[0], tuple)
        or (not has_types and len(dataset[0]) != 4)
    ):
        raise ValueError(
            "Dataset must be list of (prompt_tokens, answer_tokens, prompt_str, answer_str[, type]) tuples"
        )

    def length_key(i):
        return len(dataset[i][0]) + len(dataset[i][1])

    idx = sorted(range(len(dataset)), key=length_key)

    if len(dataset) < batch_size:
        raise ValueError(
            f"Dataset must have at least batch_size={batch_size} "
            f"examples but only has {len(dataset)}."
        )

    step = mx.distributed.init().size()
    if batch_size % step != 0:
        raise ValueError("The batch size must be divisible by the number of workers")

    def batch_index_generator():
        for i in range(0, len(idx) - batch_size + 1, batch_size):
            yield idx[i : i + batch_size : step]

    while True:
        indices = (
            np.random.permutation(list(batch_index_generator()))
            if train
            else batch_index_generator()
        )

        for batch_idx in indices:
            current_batch = [dataset[j] for j in batch_idx]

            prompts_tokens = [item[0] for item in current_batch]
            answers_tokens = [item[1] for item in current_batch]
            prompts_text = [item[2] for item in current_batch]
            answers_text = [item[3] for item in current_batch]
            types = [item[4] for item in current_batch] if has_types else None

            yield prompts_tokens, answers_tokens, prompts_text, answers_text, types

        if not train:
            break


def evaluate_grpo(
    model: nn.Module,
    ref_model: Optional[nn.Module],
    dataset,
    tokenizer,
    batch_size,
    num_batches,
    beta: float,
    epsilon: float,
    epsilon_high: float,
    group_size: int,
    max_seq_length: int,
    max_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
    min_p: float,
    reward_funcs: Optional[List[RewardFunctions]] = [
        r1_accuracy_reward_func,
        r1_int_reward_func,
        r1_strict_format_reward_func,
        r1_soft_format_reward_func,
        r1_count_xml,
    ],
    reward_weights: Optional[List[float]] = None,
    loss_fn: callable = grpo_loss,
    iterate_batches: callable = iterate_grpo_batches,
    grpo_loss_type: str = "grpo",
    importance_sampling_level: str = "token",
    end_answer_token: str = "</answer>",
):
    model.eval()
    all_losses = 0
    ntokens = 0
    all_metrics = None

    index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)

    for _, batch in zip(
        index_iterator,
        iterate_batches(
            dataset=dataset,
            batch_size=batch_size,
            max_seq_length=max_seq_length,
        ),
    ):
        prompt_tokens, answer_tokens, prompt_text, answer_text, type_info = batch

        all_completions, all_completion_texts, batch_indices = generate_grpo(
            model=model,
            tokenizer=tokenizer,
            prompt_tokens=prompt_tokens,
            max_tokens=max_tokens,
            group_size=group_size,
            batch_size=batch_size,
            end_token=end_answer_token,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            min_p=min_p,
        )

        # Prepare expanded data for reward calculation
        expanded_answers = []
        expanded_prompts = []
        expanded_types = []
        unique_prompt_indices = sorted(set(batch_indices))
        grouped_completions = {idx: [] for idx in unique_prompt_indices}

        for i, completion_idx in enumerate(batch_indices):
            grouped_completions[completion_idx].append(i)

        ordered_completions = []
        ordered_completion_texts = []
        ordered_batch_indices = []

        for prompt_idx in unique_prompt_indices:
            completion_indices = grouped_completions[prompt_idx]
            for idx in completion_indices:
                ordered_completions.append(all_completions[idx])
                ordered_completion_texts.append(all_completion_texts[idx])
                ordered_batch_indices.append(prompt_idx)
                expanded_answers.append(answer_text[prompt_idx])
                expanded_prompts.append(prompt_text[prompt_idx])
                expanded_types.append(
                    type_info[prompt_idx] if type_info is not None else None
                )

        # Calculate rewards and advantages outside of the loss function
        advantages, reward_metrics = calculate_rewards_and_advantages(
            reward_funcs=reward_funcs,
            expanded_prompts=expanded_prompts,
            all_completion_texts=ordered_completion_texts,
            expanded_answers=expanded_answers,
            expanded_types=expanded_types,
            batch_indices=ordered_batch_indices,
            unique_prompt_indices=unique_prompt_indices,
            reward_weights=reward_weights,
        )

        # Update the loss function call to use the new signature
        losses, toks, metrics = loss_fn(
            model=model,
            ref_model=ref_model,
            batch=(prompt_tokens, answer_tokens, prompt_text, answer_text, type_info),
            completions=ordered_completions,
            completion_texts=ordered_completion_texts,
            batch_indices=ordered_batch_indices,
            advantages=advantages,
            reward_metrics=reward_metrics,
            beta=beta,
            epsilon=epsilon,
            epsilon_high=epsilon_high,
            importance_sampling_level=importance_sampling_level,
            grpo_loss_type=grpo_loss_type,
            max_tokens=max_tokens,
        )

        del all_completions, all_completion_texts, batch_indices
        del ordered_completions, ordered_completion_texts, ordered_batch_indices
        del advantages, reward_metrics
        mx.clear_cache()

        all_losses += losses * toks
        ntokens += toks

        if all_metrics is None:
            all_metrics = {k: v * toks for k, v in metrics.items()}
        else:
            for k, v in metrics.items():
                all_metrics[k] += v * toks

    mx.eval(all_losses, ntokens)

    all_losses = mx.distributed.all_sum(all_losses, stream=mx.cpu)
    ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
    all_metrics = {k: mx.distributed.all_sum(v) for k, v in all_metrics.items()}

    avg_metrics = {k: (v / ntokens).item() for k, v in all_metrics.items()}
    avg_loss = (all_losses / ntokens).item()

    return avg_loss, ntokens, avg_metrics


def train_grpo(
    model: nn.Module,
    ref_model: Optional[nn.Module],
    tokenizer,
    optimizer,
    train_dataset,
    val_dataset: Optional[Any] = None,
    reward_funcs: Optional[List[RewardFunctions]] = [
        r1_accuracy_reward_func,
        r1_int_reward_func,
        r1_strict_format_reward_func,
        r1_soft_format_reward_func,
        r1_count_xml,
    ],
    args: GRPOTrainingArgs = GRPOTrainingArgs(),
    loss_fn: callable = grpo_loss,
    iterate_batches: callable = iterate_grpo_batches,
    training_callback: TrainingCallback = None,
    end_answer_token: str = "</answer>",
):
    mx.set_wired_limit(mx.device_info()["max_recommended_working_set_size"])
    world = mx.distributed.init()
    world_size = world.size()
    rank = world.rank()
    if world_size > 1:
        tqdm.write(f"Node {rank} of {world_size}")

    if args.grad_checkpoint:
        grad_checkpoint(model.layers[0])

    grad_accum_steps = args.gradient_accumulation_steps
    if grad_accum_steps < 1:
        raise ValueError("gradient_accumulation_steps must be at least 1")

    state = [model.state, optimizer.state, mx.random.state]

    @partial(mx.compile, inputs=state, outputs=state)
    def step(
        batch,
        ordered_completions,
        ordered_completion_texts,
        ordered_batch_indices,
        advantages,
        reward_metrics,
        prev_grad,
        do_update,
    ):
        prompt_tokens, answer_tokens, prompt_text, answer_text, type_info = batch

        # Update the loss function call to use the new signature
        (lvalue, toks, metrics), grad = loss_value_and_grad(
            model,
            batch=(prompt_tokens, answer_tokens, prompt_text, answer_text, type_info),
            completions=ordered_completions,
            completion_texts=ordered_completion_texts,
            batch_indices=ordered_batch_indices,
            advantages=advantages,
            reward_metrics=reward_metrics,
            beta=args.beta,
            epsilon=args.epsilon,
            epsilon_high=args.epsilon_high,
            ref_model=ref_model,
            grpo_loss_type=args.grpo_loss_type,
            importance_sampling_level=args.importance_sampling_level,
            max_tokens=args.max_completion_length,
        )

        del ordered_completions, ordered_completion_texts, ordered_batch_indices
        del advantages, reward_metrics
        mx.clear_cache()

        if prev_grad is not None:
            grad = tree_map(lambda x, y: x + y, grad, prev_grad)

        if do_update:
            grad = average_gradients(grad)
            if grad_accum_steps > 1:
                grad = tree_map(lambda x: x / grad_accum_steps, grad)
            optimizer.update(model, grad)
            grad = None
            mx.clear_cache()

        return lvalue, toks, metrics, grad

    loss_value_and_grad = nn.value_and_grad(model, loss_fn)

    model.train()
    losses = 0
    n_tokens = 0
    steps = 0
    trained_tokens = 0
    accumulated_metrics = {
        "total_rewards_mean": 0,
        "total_rewards_std": 0,
        "grouped_rewards_mean": 0,
        "grouped_rewards_std": 0,
        "kl": 0,
        "average_generated_tokens": 0,
        "max_generated_tokens": 0,
        "min_generated_tokens": 0,
        "hit_max_tokens_ratio": 0,
        "clip_ratio_low": 0,
        "clip_ratio_high": 0,
        "clip_ratio_total": 0,
    }
    grad_accum = None
    for reward_func in reward_funcs:
        func_name = reward_func.__name__
        accumulated_metrics[f"{func_name}_mean"] = 0
        accumulated_metrics[f"{func_name}_std"] = 0
        accumulated_metrics[f"{func_name}_coverage"] = 0

    start = time.perf_counter()
    pbar = tqdm(range(1, args.iters + 1), desc="Training", disable=rank != 0)
    for it in pbar:
        batch = next(
            iterate_batches(
                dataset=train_dataset,
                batch_size=args.batch_size,
                max_seq_length=args.max_seq_length,
                train=True,
            )
        )

        if (
            val_dataset is not None
            and len(val_dataset) > 0
            and (it == 1 or it % args.steps_per_eval == 0 or it == args.iters)
        ):
            stop = time.perf_counter()
            val_loss, val_ntokens, val_metrics = evaluate_grpo(
                model=model,
                dataset=val_dataset,
                loss_fn=loss_fn,
                ref_model=ref_model,
                reward_funcs=reward_funcs,
                tokenizer=tokenizer,
                group_size=args.group_size,
                batch_size=args.batch_size,
                num_batches=args.val_batches,
                max_seq_length=args.max_seq_length,
                max_tokens=args.max_completion_length,
                beta=args.beta,
                epsilon=args.epsilon,
                epsilon_high=args.epsilon_high,
                iterate_batches=iterate_batches,
                grpo_loss_type=args.grpo_loss_type,
                end_answer_token=end_answer_token,
                temperature=args.temperature,
                top_p=args.top_p,
                top_k=args.top_k,
                min_p=args.min_p,
            )
            val_time = time.perf_counter() - stop
            if rank == 0:
                tqdm.write(
                    f"Iter {it}: "
                    f"Val loss {val_loss:.3f}, "
                    f"Val took {val_time:.3f}s"
                )

            if training_callback is not None:
                val_info = {
                    "iteration": it,
                    "val_loss": val_loss,
                    "val_time": val_time,
                }
                training_callback.on_val_loss_report(val_info)

            model.train()
            start = time.perf_counter()

        prompt_tokens, answer_tokens, prompt_text, answer_text, type_info = batch

        all_completions, all_completion_texts, batch_indices = generate_grpo(
            model=model,
            tokenizer=tokenizer,
            prompt_tokens=prompt_tokens,
            max_tokens=args.max_completion_length,
            group_size=args.group_size,
            batch_size=args.batch_size,
            end_token=end_answer_token,
            temperature=args.temperature,
            top_p=args.top_p,
            top_k=args.top_k,
            min_p=args.min_p,
        )

        # Prepare expanded data for reward calculation
        expanded_answers = []
        expanded_prompts = []
        expanded_types = []
        unique_prompt_indices = sorted(set(batch_indices))
        grouped_completions = {idx: [] for idx in unique_prompt_indices}

        for i, completion_idx in enumerate(batch_indices):
            grouped_completions[completion_idx].append(i)

        ordered_completions = []
        ordered_completion_texts = []
        ordered_batch_indices = []

        for prompt_idx in unique_prompt_indices:
            completion_indices = grouped_completions[prompt_idx]
            for idx in completion_indices:
                ordered_completions.append(all_completions[idx])
                ordered_completion_texts.append(all_completion_texts[idx])
                ordered_batch_indices.append(prompt_idx)
                expanded_answers.append(answer_text[prompt_idx])
                expanded_prompts.append(prompt_text[prompt_idx])
                expanded_types.append(
                    type_info[prompt_idx] if type_info is not None else None
                )

        advantages, reward_metrics = calculate_rewards_and_advantages(
            reward_funcs=reward_funcs,
            expanded_prompts=expanded_prompts,
            all_completion_texts=ordered_completion_texts,
            expanded_answers=expanded_answers,
            expanded_types=expanded_types,
            batch_indices=ordered_batch_indices,
            unique_prompt_indices=unique_prompt_indices,
            reward_weights=(
                args.reward_weights if hasattr(args, "reward_weights") else None
            ),
        )

        del all_completions, all_completion_texts, batch_indices

        lvalue, toks, metrics, grad_accum = step(
            batch,
            ordered_completions,
            ordered_completion_texts,
            ordered_batch_indices,
            advantages,
            reward_metrics,
            grad_accum,
            it % grad_accum_steps == 0,
        )
        losses += lvalue
        n_tokens += toks
        steps += 1

        for k, v in metrics.items():
            accumulated_metrics[k] += v

        mx.eval(state, losses, n_tokens, grad_accum)

        if it % args.steps_per_report == 0 or it == args.iters:
            stop = time.perf_counter()

            train_loss = mx.distributed.all_sum(losses).item() / (steps * world_size)
            avg_metrics = {
                k: v / (steps * world_size) for k, v in accumulated_metrics.items()
            }
            n_tokens = mx.distributed.all_sum(n_tokens).item()
            learning_rate = optimizer.learning_rate.item()
            it_sec = args.steps_per_report / (stop - start)
            tokens_sec = float(n_tokens) / (stop - start)
            trained_tokens += n_tokens
            peak_mem = mx.get_peak_memory() / 1e9

            if rank == 0:
                avg_metrics = {}
                for k, v in accumulated_metrics.items():
                    accumulated_v = v / (steps * world_size)
                    if isinstance(accumulated_v, mx.array):
                        avg_metrics[k] = float(accumulated_v.item())
                    else:
                        avg_metrics[k] = float(accumulated_v)

                pbar.set_postfix(
                    {
                        "loss": f"{train_loss:.3f}",
                        "it/s": f"{it_sec:.3f}",
                    }
                )
                reward_metrics_str = ""
                for reward_func in reward_funcs:
                    func_name = reward_func.__name__
                    mean_key = f"{func_name}_mean"
                    std_key = f"{func_name}_std"
                    cov_key = f"{func_name}_coverage"

                    if mean_key in avg_metrics:
                        display_name = func_name.replace("_reward_func", "").replace(
                            "r1_", ""
                        )
                        reward_metrics_str += (
                            f"  • {display_name}: "
                            f"μ={avg_metrics[mean_key]:.3f}, "
                            f"σ={avg_metrics[std_key]:.3f}, "
                            f"cov={avg_metrics[cov_key]:.2%}\n"
                        )
                tqdm.write(
                    f"\n{'='*80}\n"
                    f"Iter {it}:\n"
                    f"{'-'*80}\n"
                    f"Loss: {train_loss:.3f}\n"
                    f"Total Rewards:  μ={avg_metrics['total_rewards_mean']:.3f}, "
                    f"σ={avg_metrics['total_rewards_std']:.3f}\n"
                    f"Group Rewards:  μ={avg_metrics['grouped_rewards_mean']:.3f}, "
                    f"σ={avg_metrics['grouped_rewards_std']:.3f}\n"
                    f"KL Divergence: {avg_metrics['kl']:.12f}\n"
                    f"{'-'*80}\n"
                    f"Generation Stats:\n"
                    f"  • Avg tokens: {avg_metrics['average_generated_tokens']:.1f}\n"
                    f"  • Min tokens: {avg_metrics['min_generated_tokens']:.0f}\n"
                    f"  • Max tokens: {avg_metrics['max_generated_tokens']:.0f} "
                    f"(limit: {args.max_completion_length})\n"
                    f"  • Hit limit: {avg_metrics['hit_max_tokens_ratio']:.1%}\n"
                    f"{'-'*80}\n"
                    f"Individual Reward Functions:\n"
                    f"{reward_metrics_str}"
                    f"{'-'*80}\n"
                    f"Clipping:  low={avg_metrics['clip_ratio_low']:.3f}, "
                    f"high={avg_metrics['clip_ratio_high']:.3f}, "
                    f"total={avg_metrics['clip_ratio_total']:.3f}\n"
                    f"Learning Rate: {learning_rate:.4e}\n"
                    f"Speed: {it_sec:.3f} it/s, {tokens_sec:.1f} tok/s\n"
                    f"Memory: {peak_mem:.3f}GB\n"
                    f"{'='*80}\n"
                )

            if training_callback is not None:
                train_info = {
                    "iteration": it,
                    "train_loss": train_loss,
                    **{f"train_{k}": v for k, v in avg_metrics.items()},
                    "learning_rate": learning_rate,
                    "iterations_per_second": it_sec,
                    "tokens_per_second": tokens_sec,
                    "trained_tokens": trained_tokens,
                    "peak_memory": peak_mem,
                }
                training_callback.on_train_loss_report(train_info)

            losses = 0
            n_tokens = 0
            steps = 0
            accumulated_metrics = {k: 0 for k in accumulated_metrics}
            start = time.perf_counter()

        if it % args.steps_per_save == 0:
            adapter_weights = dict(tree_flatten(model.trainable_parameters()))
            mx.save_safetensors(str(args.adapter_file), adapter_weights)
            checkpoint = (
                Path(args.adapter_file).parent / f"{it:07d}_adapters.safetensors"
            )
            mx.save_safetensors(str(checkpoint), adapter_weights)
            tqdm.write(
                f"\n"
                f"Iter {it}: Saved adapter weights to "
                f"{args.adapter_file} and {checkpoint}."
            )

    adapter_weights = dict(tree_flatten(model.trainable_parameters()))
    mx.save_safetensors(str(args.adapter_file), adapter_weights)
    tqdm.write(f"Saved final weights to {args.adapter_file}.")
