# !pip install mesa-frames mesa

from mesa_frames import ModelDF, AgentSetPolars


class MoneyModelDF(ModelDF):
    def __init__(self, N: int, agents_cls):
        super().__init__()
        self.n_agents = N
        self.agents += agents_cls(N, self)

    def step(self):
        # Executes the step method for every agentset in self.agents
        self.agents.do("step")

    def run_model(self, n):
        for _ in range(n):
            self.step()

import polars as pl


class MoneyAgentPolars(AgentSetPolars):
    def __init__(self, n: int, model: ModelDF):
        super().__init__(model)
        self += pl.DataFrame(
            {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)}
        )

    def step(self) -> None:
        self.do("give_money")

    def give_money(self):
        self.select(self.wealth > 0)
        other_agents = self.agents.sample(
            n=len(self.active_agents), with_replacement=True
        )
        self["active", "wealth"] -= 1
        new_wealth = other_agents.group_by("unique_id").len()
        self[new_wealth["unique_id"], "wealth"] += new_wealth["len"]

class MoneyAgentPolarsConcise(AgentSetPolars):
    def __init__(self, n: int, model: ModelDF):
        super().__init__(model)
        ## Adding the agents to the agent set
        # 1. Changing the agents attribute directly (not recommended, if other agents were added before, they will be lost)
        """self.agents = pl.DataFrame(
            {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)}
        )"""
        # 2. Adding the dataframe with add
        """self.add(
            pl.DataFrame(
                {
                    "unique_id": pl.arange(n, eager=True),
                    "wealth": pl.ones(n, eager=True),
                }
            )
        )"""
        # 3. Adding the dataframe with __iadd__
        self += pl.DataFrame(
            {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)}
        )

    def step(self) -> None:
        # The give_money method is called
        # self.give_money()
        self.do("give_money")

    def give_money(self):
        ## Active agents are changed to wealthy agents
        # 1. Using the __getitem__ method
        # self.select(self["wealth"] > 0)
        # 2. Using the fallback __getattr__ method
        self.select(self.wealth > 0)

        # Receiving agents are sampled (only native expressions currently supported)
        other_agents = self.agents.sample(
            n=len(self.active_agents), with_replacement=True
        )

        # Wealth of wealthy is decreased by 1
        # 1. Using the __setitem__ method with self.active_agents mask
        # self[self.active_agents, "wealth"] -= 1
        # 2. Using the __setitem__ method with "active" mask
        self["active", "wealth"] -= 1

        # Compute the income of the other agents (only native expressions currently supported)
        new_wealth = other_agents.group_by("unique_id").len()

        # Add the income to the other agents
        # 1. Using the set method
        """self.set(
            attr_names="wealth",
            values=pl.col("wealth") + new_wealth["len"],
            mask=new_wealth,
        )"""

        # 2. Using the __setitem__ method
        self[new_wealth, "wealth"] += new_wealth["len"]


class MoneyAgentPolarsNative(AgentSetPolars):
    def __init__(self, n: int, model: ModelDF):
        super().__init__(model)
        self += pl.DataFrame(
            {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)}
        )

    def step(self) -> None:
        self.do("give_money")

    def give_money(self):
        ## Active agents are changed to wealthy agents
        self.select(pl.col("wealth") > 0)

        other_agents = self.agents.sample(
            n=len(self.active_agents), with_replacement=True
        )

        # Wealth of wealthy is decreased by 1
        self.agents = self.agents.with_columns(
            wealth=pl.when(pl.col("unique_id").is_in(self.active_agents["unique_id"]))
            .then(pl.col("wealth") - 1)
            .otherwise(pl.col("wealth"))
        )

        new_wealth = other_agents.group_by("unique_id").len()

        # Add the income to the other agents
        self.agents = (
            self.agents.join(new_wealth, on="unique_id", how="left")
            .fill_null(0)
            .with_columns(wealth=pl.col("wealth") + pl.col("len"))
            .drop("len")
        )