public class Person extends Point {
  private int state = State.NORMAL;
}

import numpy as np
import matplotlib.pyplot as plt

class People(object):
  def __init__(self, count=1000, first_infected_count=3):
    self.count = count
    self.first_infected_count = first_infected_count
    self.init()

  def init(self):
    self._people = np.random.normal(0, 100, (self.count, 2))
    self.reset()

def reset(self):
    self._round = 0
    self._status = np.array([0] * self.count)
    self._timer = np.array([0] * self.count)
    self.random_people_state(self.first_infected_count, 1)

def random_people_state(self, num, state=1):
    """隨機(jī)挑選人設(shè)置狀態(tài)
    """
    assert self.count > num
    # TODO：極端情況下會(huì)出現(xiàn)無(wú)限循環(huán)
    n = 0
    while n < num:
      i = np.random.randint(0, self.count)
      if self._status[i] == state:
        continue
      else:
        self.set_state(i, state)
        n += 1

  def set_state(self, i, state):
    self._status[i] = state
    # 記錄狀態(tài)改變的時(shí)間
    self._timer[i] = self._round

 @property
  def healthy(self):
    return self._people[self._status == 0]

  @property
  def infected(self):
    return self._people[self._status == 1]

  def update(self):
    """每一次迭代更新"""
    self.change_state()
    self.affect()
    self.move()
    self._round += 1
    self.report()

  def change_state(self):
    dt = self._round - self._timer
    # 必須先更新時(shí)鐘再更新?tīng)顟B(tài)
    d = np.random.randint(3, 5)
    self._timer[(self._status == 1) & ((dt == d) | (dt > 14))] = self._round
    self._status[(self._status == 1) & ((dt == d) | (dt > 14))] += 1

  def infect_possible(self, x=0., safe_distance=3.0):
    """按概率感染接近的健康人
    x 的取值參考正態(tài)分布概率表，x=0 時(shí)感染概率是 50%
    """
    for inf in self.infected:
      dm = (self._people - inf) ** 2
      d = dm.sum(axis=1) ** 0.5
      sorted_index = d.argsort()
      for i in sorted_index:
        if d[i] >= safe_distance:
          break # 超出范圍，不用管了
        if self._status[i] > 0:
          continue
        if np.random.normal() > x:
          continue
        self._status[i] = 1
        # 記錄狀態(tài)改變的時(shí)間
        self._timer[i] = self._round

  def move(self, width=1, x=.0):
    movement = self.random_movement(width=width)
    # 限定特定狀態(tài)的人員移動(dòng)
    switch = self.random_switch(x=x)
    movement[switch == 0] = 0
    self._people = self._people + movement

  def random_switch(self, x=0.):
    """隨機(jī)生成開(kāi)關(guān)，0 - 關(guān)，1 - 開(kāi)

    x 大致取值范圍 -1.99 - 1.99；
    對(duì)應(yīng)正態(tài)分布的概率， 取值 0 的時(shí)候?qū)?yīng)概率是 50%
    :param x: 控制開(kāi)關(guān)比例
    :return:
    """
    normal = np.random.normal(0, 1, self.count)
    switch = np.where(normal < x, 1, 0)
    return switch

 def report(self):
    plt.cla()
    # plt.grid(False)
    p1 = plt.scatter(self.healthy[:, 0], self.healthy[:, 1], s=1)
    p2 = plt.scatter(self.infected[:, 0], self.infected[:, 1], s=1, c='pink')
    p3 = plt.scatter(self.confirmed[:, 0], self.confirmed[:, 1], s=1, c='red')

    plt.legend([p1, p2, p3], ['healthy', 'infected', 'confirmed'], loc='upper right', scatterpoints=1)
    t = "Round: %s, Healthy: %s, Infected: %s, Confirmed: %s" % \
      (self._round, len(self.healthy), len(self.infected), len(self.confirmed))
    plt.text(-200, 400, t, ha='left', wrap=True)

if __name__ == '__main__':
  np.random.seed(0)
  plt.figure(figsize=(16, 16), dpi=100)
  plt.ion()
  p = People(5000, 3)
  for i in range(100):
    p.update()
    p.report()
    plt.pause(.1)
  plt.pause(3)