import random
from mesa import Agent, Model
from mesa.time import RandomActivation
from mesa.space import MultiGrid

class HouseholdAgent(Agent):
    """ 住房市场主体：高、中、低收入群体 """
    def __init__(self, unique_id, model, income_level):
        super().__init__(unique_id, model)
        self.income_level = income_level
        self.has_house = income_level != "low"
        self.wants_to_buy = not self.has_house
        self.wants_to_sell = self.has_house and random.random() < 0.3
        self.house_quality = random.uniform(1, 5)

    def step(self):
        """ 代理行为：购房、卖房、升级 """
        if self.wants_to_buy:
            self.has_house = True
            self.wants_to_buy = False
            self.house_quality += 0.5

class HousingMarketModel(Model):
    def __init__(self, N):
        self.num_agents = N
        self.grid = MultiGrid(10, 10, True)
        self.schedule = RandomActivation(self)

        for i in range(self.num_agents):
            income_group = random.choice(["high", "middle", "low"])
            agent = HouseholdAgent(i, self, income_group)
            self.schedule.add(agent)
            self.grid.place_agent(agent, (random.randint(0, 9), random.randint(0, 9)))

    def step(self):
        """ 执行市场更新 """
        self.schedule.step()
