from typing import Dict
from typing import Optional
from typing import Any


def map_losses_per_agent_ac(
    critic_losses: Dict, policy_losses: Dict, total_losses: Optional[Dict] = None
) -> Dict:
    """Map seperate losses dict to loss per agent.

    Args:
        critic_losses : critic loss per agent.
        policy_losses : policy loss per agent.
        total_losses : optional total (critic + policy loss) per agent.

    Returns:
        dict with losses grouped per agent.
    """
    assert len(policy_losses) > 0 and (
        len(critic_losses) == len(policy_losses)
    ), "Invalid System Checkpointer."
    logged_losses: Dict[str, Dict[str, Any]] = {}
    for agent in policy_losses.keys():
        logged_losses[agent] = {
            "critic_loss": critic_losses[agent],
            "policy_loss": policy_losses[agent],
        }
        if total_losses is not None:
            logged_losses[agent]["total_loss"] = total_losses[agent]

    return logged_losses