from typing import Optional

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from matplotlib import ticker


def max_sim_agents(scenes, ax: Optional[plt.Axes] = None, bins: int = 20, **kwargs):
    max_sim_agents = []
    for scene in scenes:
        max_sim_agents.append(max([len(agents) for agents in scene.agent_presence]))

    if ax is not None:
        ax = sns.histplot(max_sim_agents, bins=bins, stat="proportion", ax=ax, **kwargs)
        ax.set_xlabel("Max. Simultaneous Agents")
        ax.yaxis.set_major_formatter(ticker.PercentFormatter(1.0, decimals=0))
    else:
        return np.histogram(max_sim_agents, bins=bins, **kwargs)


def sim_agents(scenes, ax: Optional[plt.Axes] = None, bins: int = 20, **kwargs):
    sim_agents = []
    for scene in scenes:
        sim_agents.extend(([len(agents) for agents in scene.agent_presence]))

    if ax is not None:
        ax = sns.histplot(sim_agents, bins=bins, stat="proportion", ax=ax, **kwargs)
        ax.set_xlabel("Simultaneous Agents")
        ax.yaxis.set_major_formatter(ticker.PercentFormatter(1.0, decimals=0))
    else:
        return np.histogram(sim_agents, bins=bins, **kwargs)
