class Agent {
  constructor(b, r, s, w, h) {
    this.network = b;
    this.rect = r;
    this.speed = s;
    this.width = w;
    this.height = h;
  }
}
class Environment {
  constructor(w, h, r, c, es, as) {
    this.width = w;
    this.height = h;
    this.rows = r;
    this.columns = c;
    this.enemySpeed = es;
    this.agentSpeed = as;
    this.agent = this.resetAgent();
    this.eps = Environment.MAX_EPS;
    this.discount = Environment.DISCOUNT;
  }
}

const ACTIONS = [MOVE_RIGHT, MOVE_DOWN, MOVE_LEFT, MOVE_UP];
chooseAction(state, eps) {
  if (random(0, 1) < eps) {
    return ACTIONS[random([0, 1, 2, 3])]; // рандомный шаг
  } else {
    return tf.tidy(() => {
      // сеть возвращает массив из 4 значений
      const probs = this.network.predict(state).dataSync();
      // шаг с максимальным значением
      return ACTIONS[probs.indexOf(Math.max(...probs))];
    });
  }
}

update(action) {
  switch (action) {
    case MOVE_UP:
    this.rect.top = this.rect.top - this.speed;
    break;
    case MOVE_DOWN:
    this.rect.top = this.rect.top + this.speed;
    break;
    case MOVE_RIGHT:
    this.rect.left = this.rect.left + this.speed;
    break;
    case MOVE_LEFT:
    this.rect.left = this.rect.left - this.speed;
    break;
  }
}

createModel(inputShape) {
  const model = tf.sequential();
  model.add(tf.layers.dense({ inputShape: [inputShape], units: 36, activation: 'relu' }));
  model.add(tf.layers.dense({ units: 36, activation: 'relu' }));
  model.add(tf.layers.dropout({ rate: 0.20 }));
  model.add(tf.layers.dense({ units: Agent.ACTIONS.length }));
  model.compile({ optimizer: 'adam', loss: 'meanSquaredError' });
  return model;
}

calcReward() {
  // находим нормализованную дистанцию до цели (от 0 до 1) и умножаем на -0.2
  let reward = distance(
    this.agent.rect.left,
    this.width,
    this.agent.rect.top,
    this.height) / distance(0, this.width, 0, this.height) * -0.2;
  const agentRect = toRect(this.agent.rect);
  const enemiesRects = this.enemies.map(e => toRect(e));
  const goalRect = toRect(this.goal);
  const intersected = enemiesRects.filter(e => rectsIntersected(e, agentRect));
  reward += intersected.length && -10;
  if (rectsIntersected(agentRect, goalRect)) reward = 100;
  return reward;
}

function setup() {
  mem = new ReplayMemory();
  env = new Environment(450, 300, 4, 6, 4, 2);
  createCanvas(...env.dims);
}

async function draw() {
  CURRENT_STEP++;
  background(220);
  drawGoal(env.goal);
  drawNet(env.net);
  drawEnemies(env.enemies);
  drawAgent(env.agent.rect);
  ...

// environment
updateAgent(STATE = this.getStateTensor()) {
  const action = this.agent.chooseAction(STATE, this.eps);
  this.agent.update(action);
  const nextState = this.getStateTensor();
  return [nextState, action, this.isDone()];
}

// draw in sketch file
const [nextState, action, done] = env.updateAgent(STATE);
const reward = env.calcReward();
mem.append([STATE, action, reward, nextState, done]);
STATE = nextState;
...

async function replay() {
  let miniBatch = mem.sample(500);
  const filtered = miniBatch.filter(Boolean);
  // фильтруем если очень мало шагов сделали
  if (!filtered.length) return;
  let currentStates = filtered.map((dp) => { return dp[0].dataSync() });
  // предсказываем Q для каждого текущего состояния s в памяти
  let currentQs = await env.agent.network.predict(tf.tensor(currentStates)).array();
  let newCurrentStates = filtered.map((dp) => { return dp[3].dataSync() });
  // предсказываем Q для каждого состояния s', в которое мы попали из s
  let futureQs = await env.agent.network.predict(tf.tensor(newCurrentStates)).array();
  let X = [];
  let Y = [];
  for (let index = 0; index < filtered.length; index++) {
    // берем один слайс
    const [state, action, reward, newState, done] = filtered[index];
    let newQ;
    let currentQ;
    // уравнение Беллмана
    if (!done) {
      let maxFutureQ = Math.max(...futureQs[index]);
      // находим максимальный Q для следующего состояния (s')
      // и складываем с моментальным вознаграждением (r)
      newQ = reward + (env.discount * maxFutureQ);
    }
    // если финальный переход, просто учитываем сам переход
    else { newQ = reward }
    currentQ = currentQs[index];
    // корректируем текущее значение Q на то, которое посчитали
    currentQ[action] = newQ;
    X.push(state.dataSync()); // 9 параметров нашего состояния
    Y.push(currentQ); // массив из 4 значений Q для наших шагов (вправо, вниз, влево, вверх)
  }
  // учим нашу сеть скорректированными данными
  await env.agent.network.fit(tf.tensor(X), tf.tensor(Y), { verbose: 0 });
}