import agentpy as ap

from random import uniform

from worker import WorkerAgent
from firm import FirmAgent


class Env(ap.Model):

    float_market_size: float
    percent_rh: float
    percent_search: float
    n_worker: int
    n_firm: int
    e_revenue: float

    a_lst_worker: ap.AgentList
    a_lst_firm: ap.AgentList

    def setup(self):
        # 工作人员、企业数量、搜寻企业数量赋值
        self.n_worker = self.p.n_worker
        self.n_firm = self.p.n_firm
        self.percent_search = self.p.percent_search
        # 工人、企业列表
        self.a_lst_worker = ap.AgentList(self)
        self.a_lst_firm = ap.AgentList(self)
        self.e_revenue = 119.3

        # 在工人列表中添加工人
        for i in range(self.n_worker):
            # 初始化  workeragent，并把alpha属性传过去
            w = WorkerAgent(self, self.p.alpha)
            self.a_lst_worker.append(w)

        # 在企业列表中添加企业，放入一个is_RH_ratio, 即有多大比例的企业是属于RH类型的
        for i in range(self.n_firm):
            # 对于企业属性true or false 的判断， 影响到firm 板块下， self.s_IsRH = is_RH 语句的判断
            f = FirmAgent(self, self.p.is_RH_ratio >= uniform(0, 1))
            self.a_lst_firm.append(f)

    def update_e_revenue(self):
        self.e_revenue += 0.01 * self.e_revenue

    def step(self):
        self.update_e_revenue()
        # 先清空每次的选择列表
        self.a_lst_firm.empty_apply()
        # 一开始worker要去选择很多firm
        self.a_lst_worker.select_firm()
        # 第二步， firm 去选 worker
        self.a_lst_firm.select_worker()

        if self.t == 100:
            self.stop()
        pass

    def provide_lst_random_firms(self, the_worker: WorkerAgent):
        '''选择企业数量 = 企业总数*百分比
        选择企业的列表 = 随机选择的企业的个数
        如果员工处于被雇佣的状态：
           如果员工工作的企业在随机选定的企业列表中：
             打开列表中的企业
             移除该企业
             返回值：移除后，再重新选择随机选择企业
           否则：
            返回值：选择企业列表
        '''
        n_select_firms = int(self.percent_search * self.n_firm)
        a_lst_select_firms = self.a_lst_firm.random(n_select_firms)
        if the_worker.s_is_hired:
            if the_worker.working_firm in a_lst_select_firms:
                # 转换为 list
                lst_f = list(self.a_lst_firm)
                lst_f.remove(the_worker.working_firm)
                return ap.AgentList(self, lst_f).random(n_select_firms)
        # 假如以上都不满足， 直接返回
        return ap.AgentList(self, a_lst_select_firms)


if __name__ == '__main__':
    dict_para = {'n_worker': 100,
                 'n_firm': 20,
                 'percent_search': 0.2,
                 'alpha': 0.5,
                 'is_RH_ratio': 0.5}
    my_model = Env(dict_para)
    my_model.run()