"""Example usage of the Web Gym environment.

GRPO implementation adapted from:
https://github.com/open-thought/tiny-grpo

Usage:

(Mac) Set MPS as the default device:
```
export PYTORCH_ENABLE_MPS_FALLBACK=1
```

Run the script:
```
python examples/agent_training.py
```

Examples using cli args:
```bash
python examples/agent_training.py \
    --model_id google/gemma-3-270m-it \
    --enable_gradient_checkpointing \
    --n_hops 1 \
    --n_tries_per_hop 10 \
    --rollout_min_new_tokens 64 \
    --rollout_max_new_tokens 128 \
    --group_size 4
```

With weights and biases:

```bash
export WANDB_API_KEY=...
python examples/agent_training.py \
    --model_id google/gemma-3-1b-it \
    --enable_gradient_checkpointing \
    --n_hops 1 \
    --n_tries_per_hop 10 \
    --rollout_min_new_tokens 256 \
    --rollout_max_new_tokens 512 \
    --group_size 4 \
    --wandb_project aigym-agent-training
```
"""

import json
import tempfile
from dataclasses import dataclass
from functools import partial
from typing import Literal, Protocol, cast

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from peft import LoraConfig, PeftModelForCausalLM, get_peft_model
from rich import print as rprint
from rich.console import Console
from torch.nn.utils import clip_grad_norm_
from torch.nn.utils.clip_grad import _get_total_norm
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    CompileConfig,
    GenerationConfig,
    PreTrainedModel,
    PreTrainedTokenizer,
    Qwen2ForCausalLM,
)

import aigym.pprint as pprint
import aigym.prompts as prompts
import wandb
from aigym.agent import Agent
from aigym.env import WikipediaGymEnv
from aigym.exceptions import InvalidActionError
from aigym.types import Action, ActionBatch, Observation, RolloutBatch


@dataclass
class TrainingArgs:
    experiment_name: str | None = None
    n_episodes: int = 30
    model_id: str = "google/gemma-3-1b-it"
    model_save_dir: str = "./output_model"
    ref_model_id: str | None = None
    static_env: bool = False  # if True, the environment will not change over each episode
    start_url_anchors: str | None = None
    optim: str = "adamw"
    lr: float = 1e-4
    group_size: int = 4
    advantage_eps: float = 1e-8
    clip_eps: float = 0.2
    kl_weight: float = 0.01
    max_grad_norm: float = 1.0
    use_lora: bool = True
    lora_r: int = 16
    lora_alpha: int = 64
    lora_dropout: float = 0.1
    use_bnb_quantization: bool = False
    enable_gradient_checkpointing: bool = False
    attn_implementation: str = "eager"
    n_hops: int = 1
    n_tries_per_hop: int = 5
    rollout_min_new_tokens: int = 256
    rollout_max_new_tokens: int = 1024
    rollout_temperature: float = 1.25
    rollout_top_k: int = 64
    rollout_top_p: float = 0.95
    rollout_repetition_penalty: float = 1.0
    rollout_no_repeat_ngram_size: int = 0
    wandb_project: str = "aigym-agent-training"


class TrainingLogger(Protocol):
    def log_training_run_info(self, env: WikipediaGymEnv, wandb_url: str | None) -> None: ...
    def log_environment(self, env: WikipediaGymEnv) -> None: ...
    def log_actions(self, actions: list[Action], action_probs: list[float], step_action_index: int | None) -> None: ...
    def log_observation(self, observation: Observation) -> None: ...
    def log_metrics(self, metrics: dict[str, float]) -> None: ...
    def log_rewards(self, rewards: torch.Tensor, returns: torch.Tensor) -> None: ...
    def flush(self, episode: int, step: int) -> None: ...


def masked_mean(
    tensor: torch.Tensor,
    action_mask: torch.Tensor | None,
    dim: int = None,
) -> torch.Tensor:
    if action_mask is None:
        return tensor.mean(axis=dim)
    return (tensor * action_mask).sum(axis=dim) / action_mask.sum(axis=dim)


def copy_model(
    model: PreTrainedModel | PeftModelForCausalLM,
    model_copy: PreTrainedModel | PeftModelForCausalLM,
) -> PreTrainedModel:
    if isinstance(model, PeftModelForCausalLM):
        model = cast(PeftModelForCausalLM, model)
        model_copy.load_state_dict({k: v.clone() for k, v in model.state_dict().items()})
    else:
        model_copy = type(model)(model.config)
        model_copy.load_state_dict({k: v.clone() for k, v in model.state_dict().items()})

    return model_copy


class GRPOLoss(nn.Module):
    """GRPO actor loss"""

    def __init__(self, clip_eps: float, kl_weight: float) -> None:
        super().__init__()
        self.clip_eps = clip_eps
        self.kl_weight = kl_weight

    def forward(
        self,
        log_probs: torch.Tensor,
        log_probs_old: torch.Tensor | None,
        log_probs_ref: torch.Tensor,
        returns: torch.Tensor,
        action_mask: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        kl = approx_kl_divergence(
            log_probs=log_probs,
            log_probs_ref=log_probs_ref,
            action_mask=action_mask,
        )

        if log_probs_old is not None:
            ratio = (log_probs - log_probs_old).exp()
        else:
            ratio = log_probs

        surr1 = ratio * returns
        surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * returns
        loss = -torch.min(surr1, surr2) + self.kl_weight * kl
        loss = masked_mean(loss, action_mask, dim=-1).mean()
        return loss, kl.mean()


def approx_kl_divergence(
    log_probs: torch.Tensor,
    log_probs_ref: torch.Tensor,
    action_mask: torch.Tensor | None,
) -> torch.Tensor:
    """
    Reference: http://joschu.net/blog/kl-approx.html
    """
    log_ratio = log_probs_ref - log_probs
    if action_mask is not None:
        log_ratio = log_ratio * action_mask
    return log_ratio.exp() - log_ratio - 1


def compute_log_probs(
    model: PreTrainedModel,
    sequence_ids: torch.Tensor,
    attention_mask: torch.Tensor,
):
    position_ids = attention_mask.long().cumsum(dim=-1) - 1
    position_ids = position_ids.clone()  # Clone to avoid modifying input
    position_ids.masked_fill_(mask=(attention_mask == 0), value=1)

    torch.cuda.empty_cache()

    output = model.forward(
        input_ids=sequence_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        use_cache=False,
    )
    # truncate logits to only include the action tokens
    logits = output["logits"]

    log_probs = F.log_softmax(logits, dim=-1)
    return log_probs.gather(dim=-1, index=sequence_ids.unsqueeze(-1)).squeeze(-1)


def reward_function(action: Action, observation: Observation) -> float:
    if action.error_type and action.error_type.type == "url_not_in_context":
        return -1.0
    elif action.error_type:
        return -0.5
    elif action.url == observation.next_url:
        return 0.5
    elif action.url == observation.target_url:
        return 1.0
    else:
        return 0.0


@torch.no_grad()
def policy(
    training_args: TrainingArgs,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizer,
    generation_config: GenerationConfig,
    prompt: str,
) -> RolloutBatch:
    # tokenize and prepare inputs for batch generation
    chat_messages = [
        {
            "role": "system",
            "content": prompts.REASONING_TEMPLATE,
        },
        {
            "role": "user",
            "content": prompt,
        },
    ]
    chat_prompt = tokenizer.apply_chat_template(chat_messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer(
        [chat_prompt],
        return_tensors="pt",
        padding=True,
        padding_side="left",
        return_attention_mask=True,
    ).to(model.device)

    model_inputs["attention_mask"] = model_inputs["attention_mask"].repeat(training_args.group_size, 1)
    model_inputs["input_ids"] = model_inputs["input_ids"].repeat(training_args.group_size, 1)

    # generate completions
    model.eval()
    with torch.no_grad():
        outputs = model.generate(
            **model_inputs,
            generation_config=generation_config,
            tokenizer=tokenizer,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            return_dict_in_generate=True,
            output_logits=True,
        )

    sequence_ids = outputs.sequences[:, model_inputs["input_ids"].shape[1] :]
    completions = tokenizer.batch_decode(sequence_ids, skip_special_tokens=True)

    # shape: batch_size, sequence_length, vocab_size
    logits = torch.concat([x.unsqueeze(1) for x in outputs.logits], dim=1)
    log_probs = F.log_softmax(logits, dim=-1)

    # get the log probabilities for selected token
    token_log_probs = torch.gather(log_probs, dim=-1, index=sequence_ids.unsqueeze(-1)).squeeze(-1)
    mask = sequence_ids != tokenizer.eos_token_id
    action_log_probs = (token_log_probs * mask).sum(dim=-1)

    # normalize log probs by the number of tokens in the sequence
    action_log_probs = action_log_probs / mask.sum(dim=-1)
    action_probs = F.softmax(action_log_probs, dim=0)

    return RolloutBatch(
        sequence_ids=sequence_ids,
        input_ids=model_inputs["input_ids"],
        attention_mask=model_inputs["attention_mask"],
        completions=completions,
        logits=logits,
        log_probs=token_log_probs,
        action_probs=action_probs,
    )


def reconstruct_sequence_ids(
    action_batch: ActionBatch,
    tokenizer: PreTrainedTokenizer,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Reconstructs sequence ids from the action batch.

    The action mask is used to mask out the input id from the loss calculation.
    The attention mask is used to mask out the padding tokens from the loss calculation.
    """

    log_probs_old = None

    if action_batch.log_probs is not None:
        log_probs_old = action_batch.log_probs.detach()

    sequence_ids = action_batch.sequence_ids

    # action mask makes sure end of sequence tokens are masked out of the
    # loss calculation
    action_mask = torch.full_like(sequence_ids, fill_value=True, dtype=torch.bool)
    action_mask[sequence_ids == tokenizer.eos_token_id] = False
    attention_mask = sequence_ids != tokenizer.eos_token_id

    return sequence_ids, action_mask, attention_mask, log_probs_old


def update_policy(
    action_batch: ActionBatch,
    returns: torch.Tensor,
    optimizer: optim.Optimizer,
    tokenizer: PreTrainedTokenizer,
    objective: GRPOLoss,
    training_args: TrainingArgs,
    model: PreTrainedModel,
    reference_model: PreTrainedModel,
):
    model.train()
    optimizer.zero_grad()

    sequence_ids, action_mask, attention_mask, log_probs_old = reconstruct_sequence_ids(
        action_batch,
        tokenizer,
    )

    log_probs = compute_log_probs(model, sequence_ids, attention_mask)
    with torch.no_grad():
        log_probs_ref = compute_log_probs(reference_model, sequence_ids, attention_mask)

    loss, kl = objective(
        log_probs=log_probs,
        log_probs_old=log_probs_old,
        log_probs_ref=log_probs_ref,
        returns=returns,
        action_mask=action_mask,
    )

    if not loss.isfinite():
        print(f"Loss not finite, skipping backward, loss={loss}, returns: {returns}")
        del log_probs, log_probs_old, log_probs_ref, sequence_ids, action_mask, attention_mask, loss, kl
        torch.cuda.empty_cache()
        return model, None, None, None

    loss.backward()
    grad_norm_preclip = clip_grad_norm_(model.parameters(), max_norm=training_args.max_grad_norm)
    grad_norm = _get_total_norm([p.grad for p in model.parameters() if p.grad is not None])
    print(f"loss={loss: .10f}, kl={kl: .10f}, grad_norm_preclip={grad_norm_preclip: .10f}, grad_norm={grad_norm: .10f}")
    optimizer.step()
    torch.cuda.empty_cache()
    return model, loss, kl, grad_norm_preclip, grad_norm


def select_action(action_batch: ActionBatch, selection_type: Literal["greedy", "sample"] = "greedy") -> tuple[Action | None, int | None]:
    if (action_batch.action_probs == 0).all():
        return None, None

    # Sample from the action probability distribution
    if selection_type == "greedy":
        index = action_batch.action_probs.argmax().item()
    elif selection_type == "sample":
        index = torch.multinomial(action_batch.action_probs, num_samples=1).item()

    action = action_batch.actions[index]
    return action, index


class PrintLogger:

    def log_training_run_info(self, env: WikipediaGymEnv, wandb_url: str | None) -> None:
        pass

    def log_environment(self, env: WikipediaGymEnv) -> None:
        pprint.print_travel_path(env.travel_path)

    def log_actions(self, actions: list[Action], action_probs: list[float], step_action_index: int | None) -> None:
        for i, action in enumerate(actions):
            pprint.print_action(action, action_prob=action_probs[i], step_action_index=step_action_index, index=i)

    def log_observation(self, observation: Observation) -> None:
        pprint.print_observation(observation)
        pprint.print_context(observation)

    def log_metrics(self, metrics: dict[str, float]) -> None:
        _metrics = {k: float(v) for k, v in metrics.items()}
        pprint.print_metrics(_metrics)

    def log_rewards(self, rewards: torch.Tensor, returns: torch.Tensor) -> None:
        pprint.print_rewards(
            {
                "Raw rewards": rewards.squeeze().tolist(),
                "Returns": returns.squeeze().tolist(),
            }
        )

    def flush(self, episode: int, step: int) -> None: ...


def main(training_args: TrainingArgs, logger: TrainingLogger | None = None):
    if logger is None:
        logger = PrintLogger()

    if training_args.wandb_project is None:
        run = wandb.init(mode="disabled")
    else:
        run = wandb.init(project=training_args.wandb_project)

    if torch.cuda.is_available():
        print(f"Using CUDA: {torch.cuda.get_device_name(torch.cuda.current_device())}")
        print("\nAvailable GPUs:")
        for i in range(torch.cuda.device_count()):
            print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
            print(f"  Memory: {torch.cuda.get_device_properties(i).total_memory / 1024**3:.1f} GB")
        print()

    print("Loading model")
    tokenizer = AutoTokenizer.from_pretrained(training_args.model_id)

    bnb_config = None
    if training_args.use_bnb_quantization:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )

    reference_model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
        training_args.ref_model_id or training_args.model_id,
        dtype="auto",
        device_map="auto",
        attn_implementation=training_args.attn_implementation,
        quantization_config=bnb_config,
    )

    model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
        training_args.model_id,
        dtype="auto",
        device_map="auto",
        attn_implementation=training_args.attn_implementation,
        quantization_config=bnb_config,
    )

    lora_config = None
    if training_args.use_lora:
        target_lora_modules = ["q_proj", "k_proj", "v_proj"]
        lora_config = LoraConfig(
            r=training_args.lora_r,
            lora_alpha=training_args.lora_alpha,
            target_modules=target_lora_modules,
            lora_dropout=training_args.lora_dropout,
            bias="none",
            task_type="CAUSAL_LM",
        )

        with tempfile.TemporaryDirectory() as tmp_dir:
            reference_model = get_peft_model(reference_model, lora_config, adapter_name="default")
            reference_model.save_pretrained(tmp_dir)

            model = get_peft_model(model, lora_config, adapter_name="default")
            model.load_adapter(tmp_dir, adapter_name="default", is_trainable=True)
            model.print_trainable_parameters()

        reference_model.set_adapter("default")
        model.set_adapter("default")

    reference_model.eval()

    if training_args.enable_gradient_checkpointing:
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})

    if training_args.optim == "adamw":
        optimizer = optim.AdamW(model.parameters(), lr=training_args.lr)
    elif training_args.optim == "sgd":
        optimizer = optim.SGD(model.parameters(), lr=training_args.lr)
    else:
        raise ValueError(f"Invalid optimizer: {training_args.optim}")

    objective = GRPOLoss(
        clip_eps=training_args.clip_eps,
        kl_weight=training_args.kl_weight,
    )

    generation_config = GenerationConfig(
        do_sample=True,
        min_new_tokens=training_args.rollout_min_new_tokens,
        max_new_tokens=training_args.rollout_max_new_tokens,
        temperature=training_args.rollout_temperature,
        top_k=training_args.rollout_top_k,
        top_p=training_args.rollout_top_p,
        repetition_penalty=training_args.rollout_repetition_penalty,
        no_repeat_ngram_size=training_args.rollout_no_repeat_ngram_size,
        padding=True,
        padding_size="left",
        return_attention_mask=True,
        pad_token_id=tokenizer.eos_token_id,
        stop_strings=["<end_of_turn>"],
        compile_config=CompileConfig(fullgraph=False),
    )

    print(f"Initializing agent with generation config {generation_config}")
    agent = Agent(
        policy=partial(
            policy,
            training_args,
            model,
            tokenizer,
            generation_config,
        ),
        stream=False,
    )

    if training_args.start_url_anchors is not None:
        start_url_anchors = json.loads(training_args.start_url_anchors)
    else:
        start_url_anchors = None

    env = WikipediaGymEnv(
        n_hops=training_args.n_hops,
        start_url_anchors=start_url_anchors,
    )
    n_tries = training_args.n_hops * training_args.n_tries_per_hop

    logger.log_training_run_info(env, run.url if run is not None else None)

    # Separate rollout step from model updates, i.e. implement an experience
    # replay buffer so that you can sample from it, then update the model.
    # This clarifies the roll of the `old_log_probs` in the loss calculation.
    # You'll have multiple rollouts, you'll sample from this buffer in minibatches
    # so therefore the current model log probs will start to deviate from the
    # old log probs.
    # See: https://www.perplexity.ai/search/can-you-give-me-the-basic-grpo-g3kdNUI4RSmbAtoxi_HcjQ#6

    total_cumulative_reward = 0
    console = Console(highlight=False)
    global_step = 0
    observation, info = env.reset()
    initial_travel_path = env.travel_path
    for episode in range(1, training_args.n_episodes + 1):
        console.rule(f"▶️ Episode {episode}")

        if episode > 1:
            # reset the environment if it's not the first episode
            if training_args.static_env:
                observation, info = env.reset_manual(initial_travel_path)
            else:
                observation, info = env.reset()

        n_steps = 0
        episode_cumulative_reward = 0
        target_found = False

        for step in range(1, n_tries + 1):
            logger.log_environment(env)
            global_step += 1
            n_steps += 1
            console.rule(f"Step {step}")
            logger.log_observation(observation)
            action_batch: ActionBatch = agent.act(observation)

            step_action: Action | None = None
            step_action_index: int | None = None
            rewards: list[float] = []
            for i, action in enumerate(action_batch.actions):
                reward = reward_function(action, observation)
                rewards.append(reward)

            step_action, step_action_index = select_action(action_batch)
            logger.log_actions(action_batch.actions, action_batch.action_probs, step_action_index)

            rewards: torch.Tensor = torch.tensor(rewards, dtype=model.dtype).unsqueeze(1)
            returns = (rewards - rewards.mean()) / (rewards.std() + training_args.advantage_eps)
            returns = returns.to(model.device)
            logger.log_rewards(rewards, returns)

            if (returns == 0).all():
                update_metrics = {}
            else:
                # update the model
                model, loss, kl, grad_norm_preclip, grad_norm = update_policy(
                    action_batch,
                    returns,
                    optimizer,
                    tokenizer,
                    objective,
                    training_args,
                    model,
                    reference_model,
                )
                update_metrics = {
                    "loss": loss,
                    "kl": kl,
                    "grad_norm_preclip": grad_norm_preclip,
                    "grad_norm": grad_norm,
                }

            rewards_sum = rewards.sum()
            total_cumulative_reward += rewards_sum
            episode_cumulative_reward += rewards_sum

            metrics = {
                "max_return": returns.max(),
                "mean_reward": rewards.mean(),
                "n_steps": n_steps,
                **update_metrics,
            }
            logger.log_metrics(metrics)
            wandb.log(metrics, step=global_step)
            logger.flush(episode, step)

            if step_action is not None:
                # If the action batch contains at least one item with the correct
                # next-page target, take a step. Otherwise, don't change the state.
                try:
                    observation, env_reward, terminated, truncated, info = env.step(step_action)
                except InvalidActionError as e:
                    rprint(f"Error stepping with action {step_action.url}")
                    continue

                if terminated or truncated:
                    rprint(f"Episode terminated or truncated at step {step}")
                    target_found = terminated
                    break

        wandb.log(
            {
                "episode": episode,
                "episode_cumulative_reward": episode_cumulative_reward,
                "episode_mean_reward": episode_cumulative_reward / n_steps,
                "target_found": int(target_found),
                "n_steps": n_steps,
                "negative_n_steps": -n_steps,
                "total_cumulative_rewards": total_cumulative_reward.item(),
            },
            step=global_step,
        )


    rprint("Task finished!")
    env.close()

    model.save_pretrained(training_args.model_save_dir)


if __name__ == "__main__":
    from transformers import HfArgumentParser

    parser = HfArgumentParser(TrainingArgs)
    training_args, *_ = parser.parse_args_into_dataclasses()
    main(training_args)
