markov_chain

#!/usr/bin/env ruby
# frozen_string_literal: true

require 'csv'
require 'json'
require 'narray'
require 'tsort'

PROBABILITY_OF_4 = 0.1
OUTPUT_FOLDER = 'data/markov_chain'

#
# Merge adjacent tiles. This assumes that the board is sorted.
#
def move(board)
  result = []
  last = nil
  board.each do |value|
    if value == last
      result[-1] = 2 * last
      last = nil
    else
      result << value
      last = value
    end
  end
  result
end

#
# Order states by sum, then lexically.
#
def sort_states(states)
  states.sort_by do |state|
    [state.sum, state]
  end
end

def find_state_index(states, state)
  key = [state.sum, state]
  states.bsearch_index { |x| key <=> [x.sum, x] }
end

#
# Simulate the markov chain model (well, an equivalent one).
#
def simulate_place(board)
  new_tile = 2 + 2 * (rand < PROBABILITY_OF_4 ? 1 : 0)
  board.unshift new_tile
  board.sort!
  board
end

def simulate(max_exponent)
  max_value = 2**max_exponent
  board = []
  2.times { simulate_place(board) }
  moves = 0
  max_size = 0
  while board.max < max_value
    board = move(board)
    simulate_place(board)

    moves += 1
    max_size = board.size if board.size > max_size
  end
  [moves, max_size, board]
end

def build_simulation_histograms(max_exponent, num_trials)
  moves_histogram = Hash.new { |h, k| h[k] = 0 }
  max_size_histogram = Hash.new { |h, k| h[k] = 0 }
  final_state_histogram = Hash.new { |h, k| h[k] = 0 }
  num_trials.times do
    moves, max_size, final_state = simulate(max_exponent)
    moves_histogram[moves] += 1
    max_size_histogram[max_size] += 1
    final_state_histogram[final_state] += 1
  end
  [moves_histogram, max_size_histogram, final_state_histogram]
end

def run_simulations(num_trials = 1_000_000)
  moves_file = File.join(OUTPUT_FOLDER, 'moves_histogram.csv')
  max_size_file = File.join(OUTPUT_FOLDER, 'max_size_histogram.csv')
  final_state_file = File.join(OUTPUT_FOLDER, 'final_state_histogram.csv')
  CSV.open(moves_file, 'w') do |moves_csv|
    moves_csv << %w[max_exponent moves frequency]
    CSV.open(max_size_file, 'w') do |max_size_csv|
      max_size_csv << %w[max_exponent max_size frequency]
      CSV.open(final_state_file, 'w') do |final_state_csv|
        final_state_csv << %w[max_exponent final_state frequency]
        (3..11).each do |max_exponent|
          moves_histogram, max_size_histogram, final_state_histogram =
            build_simulation_histograms(max_exponent, num_trials)
          moves_histogram.keys.sort.each do |moves|
            moves_csv << [max_exponent, moves, moves_histogram[moves]]
          end
          max_size_histogram.keys.sort.each do |max_size|
            max_size_csv << [
              max_exponent, max_size, max_size_histogram[max_size]
            ]
          end
          sort_states(final_state_histogram.keys).each do |final_state|
            final_state_csv << [
              max_exponent, final_state, final_state_histogram[final_state]
            ]
          end
        end
      end
    end
  end
end

# Takes a while...
# run_simulations

#
# Build explicit Markov Chain model.
#

#
# P[:s][:t] is the probability of transitioning from state s to state t.
#
def make_transition_hash
  Hash.new do |h0, state0|
    h0[state0] = Hash.new do |h1, state1|
      h1[state1] = 0.0
    end
  end
end

NEW_TILES = { 2 => 0.9, 4 => 0.1 }.freeze

def build_markov_chain(max_exponent)
  transitions = make_transition_hash
  prestart_state = []

  open_states = []
  NEW_TILES.each do |value0, pr0|
    NEW_TILES.each do |value1, pr1|
      start_state = [value0, value1].sort
      open_states << start_state
      transitions[prestart_state][start_state] += pr0 * pr1
    end
  end

  until open_states.empty?
    state = open_states.pop
    next if transitions.key?(state)
    move_state = move(state)
    NEW_TILES.each do |value, pr|
      new_state = ([value] + move_state).sort
      transitions[state][new_state] += pr
      open_states << new_state unless new_state.max >= 2**max_exponent
    end
  end

  check_transitions(transitions, max_exponent)

  transitions
end

def check_transitions(transitions, max_exponent)
  transitions.each do |state0, successors|
    total_pr = 0
    successors.each do |state1, pr|
      total_pr += pr
      raise "terminal state #{state1}" unless
          transitions.key?(state1) || state1.max >= 2**max_exponent
    end
    raise "pr does not sum to 1: #{state0}" unless (total_pr - 1).abs < 1e-6
  end
end

def print_transitions(transitions)
  transitions.each do |state0, successors|
    successors.each do |state1, pr|
      p [state0, state1, pr]
    end
  end
end

def setup_topological_sort(transitions)
  class <<transitions
    include TSort
    alias_method :tsort_each_node, :each_key
    def tsort_each_child(node0)
      fetch(node0, {}).keys.each do |node1|
        yield node1
      end
    end
  end
end

#
# Find expected hitting times directly, following the method in
# http://www.statslab.cam.ac.uk/~james/Markov/s13.pdf (and many others).
#
def find_hitting_times(max_exponent)
  transitions = build_markov_chain(max_exponent)

  setup_topological_sort transitions
  successor_states = transitions.tsort

  win_states, other_states = successor_states.partition do |state|
    !state.empty? && state.max >= 2**max_exponent
  end
  hitting_times = {}
  win_states.each do |win_state|
    hitting_times[win_state] = 0
  end
  other_states.each do |state|
    hitting_times[state] = 1.0
    transitions[state].each do |state1, pr|
      hitting_times[state] += pr * hitting_times[state1]
    end
  end

  hitting_times
end
# p find_hitting_times(11)

#
# Find win state absorbing probabilities and check the hitting time
# distributions using the canonical form equations from
# https://en.wikipedia.org/wiki/Absorbing_Markov_chain
#

def find_absorbing_states(transitions)
  result = Set.new
  transitions.each do |_state0, successors|
    successors.each do |state1, _pr|
      result << state1 unless transitions.key?(state1)
    end
  end
  result.to_a
end

def make_transition_matrix(transitions, row_states, col_states)
  matrix = NMatrix.float(col_states.size, row_states.size)
  transitions.each do |state0, successors|
    i = find_state_index(row_states, state0)
    next unless i
    successors.each do |state1, pr|
      j = find_state_index(col_states, state1)
      next unless j
      matrix[j, i] = pr
    end
  end
  matrix
end

def make_fundamental_matrices(max_exponent)
  transitions = build_markov_chain(max_exponent)

  transient_states = sort_states(transitions.keys)
  transient_n = transient_states.size
  transient_q = make_transition_matrix(
    transitions, transient_states, transient_states
  )

  absorbing_states = sort_states(find_absorbing_states(transitions))
  absorbing_n = absorbing_states.size
  absorbing_r = make_transition_matrix(
    transitions, transient_states, absorbing_states
  )

  identity = NMatrix.float(absorbing_n, absorbing_n).diagonal!(1)

  n = transient_n + absorbing_n
  fundamental = NMatrix.float(n, n)
  fundamental[0...transient_n, 0...transient_n] = transient_q
  fundamental[transient_n...n, 0...transient_n] = absorbing_r
  fundamental[transient_n...n, transient_n...n] = identity

  [
    transient_states, absorbing_states,
    transient_q, absorbing_r, identity,
    fundamental
  ]
end
# p make_fundamental_matrices(11)

def find_expected_steps(max_exponent)
  transient_states, _, transient_q, = make_fundamental_matrices(max_exponent)
  transient_n = transient_states.size
  identity = NMatrix.float(transient_n, transient_n).diagonal!(1)
  ones = NVector.float(transient_n).fill!(1)

  # Expectation: t = N1 for N = (I - Q)^{-1}
  t = ones / (identity - transient_q)

  # Variance: (2N - I)t - t_sq
  # If (I-Q)v = 2t, then the variance is v - It - t_sq
  v = (2 * t) / (identity - transient_q)
  t_sq = NVector[NArray[t] * NArray[t]][nil, 0, 0]
  vt = v - t - t_sq

  [transient_states, t, vt]
end
# The means should match:
# p find_hitting_times(11)
# p find_expected_steps(11)

def save_expected_steps
  CSV.open(File.join(OUTPUT_FOLDER, 'expected_steps.csv'), 'w') do |csv|
    csv << %w[max_exponent state expected_steps variance_steps]
    (3..11).each do |max_exponent|
      states, expected_steps, variance_steps = find_expected_steps(max_exponent)
      states.each.with_index do |state, i|
        csv << [
          max_exponent,
          state,
          expected_steps[i],
          variance_steps[i]
        ]
      end
    end
  end
end
# save_expected_steps

def find_absorbing_probabilities(max_exponent)
  transient_states, absorbing_states, transient_q, absorbing_r =
    make_fundamental_matrices(max_exponent)
  transient_n = transient_states.size
  identity = NMatrix.float(transient_n, transient_n).diagonal!(1)
  pr = absorbing_r / (identity - transient_q)
  Hash[absorbing_states.zip(pr[nil, 0].to_a.flatten)]
end
# p find_absorbing_probabilities(11)

def save_absorbing_probabilities
  csv_file = File.join(OUTPUT_FOLDER, 'absorbing_probabilities.csv')
  CSV.open(csv_file, 'w') do |csv|
    csv << %w[max_exponent state probability]
    (3..11).each do |max_exponent|
      find_absorbing_probabilities(max_exponent).each do |state, pr|
        csv << [max_exponent, state, pr]
      end
    end
  end
end
# save_absorbing_probabilities

#
# Find the fastest (always 4s) and slowest (always 2s).
#

def find_minmax_moves(max_exponent, new_tile)
  board = [new_tile, new_tile]
  moves = 0
  while board.max < 2**max_exponent
    board = move(board)
    board = ([new_tile] + board).sort
    moves += 1
  end
  moves
end

def save_minmax_moves
  CSV.open(File.join(OUTPUT_FOLDER, 'minmax_moves.csv'), 'w') do |csv|
    csv << %w[max_exponent min_moves max_moves]
    (3..11).each do |max_exponent|
      csv << [
        max_exponent,
        find_minmax_moves(max_exponent, 4),
        find_minmax_moves(max_exponent, 2)
      ]
    end
  end
end
# save_minmax_moves

#
# Find the minimum and maximum number of cells used. It is interesting to know
# whether, under ideal conditions, we can win with a given number of cells to
# work with.
#

def find_min_cells(max_exponent)
  # ... how to do this?
  # the max is easy, because all states in the model have nonzero probability
  # the min seems harder... how to guarantee that there is a path using only
  # a given number of cells... that seems easy enough. increase a threshold
  # and see whether there's a way to get to a win state using states with
  # only at most that number of states.
  transitions = build_markov_chain(max_exponent)
  (2..16).each do |threshold|
    allowed_transitions = filter_states(transitions) do |state|
      state.size <= threshold
    end
    return threshold if winnable?(allowed_transitions, max_exponent)
  end
  nil
end

def find_max_cells(max_exponent)
  transitions = build_markov_chain(max_exponent)
  each_state(transitions).map(&:size).max
end

def each_state(transitions)
  return to_enum(:each_state, transitions) unless block_given?
  transitions.each do |state0, successors|
    yield state0
    successors.each do |state1, pr|
      yield state1 if pr > 0
    end
  end
end

#
# Keep only those states that meet a given condition.
#
def filter_states(transitions)
  Hash[transitions.map do |state0, successors|
    next unless yield(state0)
    [state0, Hash[successors.map do |state1, pr|
      next unless yield(state1)
      [state1, pr]
    end.compact]]
  end.compact]
end

#
# Breadth-first search to determine whether we can reach a winning state from
# the pre-start state.
#
def winnable?(transitions, max_exponent)
  queue = [[]]
  closed = Set.new
  until queue.empty?
    state0 = queue.shift
    return true if !state0.empty? && state0.max >= 2**max_exponent
    next if closed.member?(state0)
    next unless transitions.key?(state0)
    transitions[state0].each do |state1, pr|
      queue << state1 if pr > 0
    end
    closed << state0
  end
  false
end

def save_minmax_cells
  CSV.open(File.join(OUTPUT_FOLDER, 'minmax_cells.csv'), 'w') do |csv|
    csv << %w[max_exponent min_cells max_cells]
    (3..11).each do |max_exponent|
      csv << [
        max_exponent,
        find_min_cells(max_exponent),
        find_max_cells(max_exponent)
      ]
    end
  end
end
# save_minmax_cells

def save_states
  CSV.open(File.join(OUTPUT_FOLDER, 'states.csv'), 'w') do |csv|
    csv << %w[max_exponent state]
    (3..11).each do |max_exponent|
      transitions = build_markov_chain(max_exponent)
      sort_states(each_state(transitions).uniq).each do |state|
        csv << [max_exponent, state]
      end
    end
  end
end
# save_states

def save_canonical_matrix(max_exponent)
  file = "canonical_matrix_#{max_exponent}.csv"
  transient_states, absorbing_states, =
    make_fundamental_matrices(max_exponent)
  all_states = transient_states + absorbing_states
  CSV.open(File.join(OUTPUT_FOLDER, file), 'w') do |csv|
    csv << [nil] + all_states
    (0...all_states.size).each do |i|
      csv << [all_states[i]] + fundamental[nil, i].to_a[0]
    end
  end
end
# save_canonical_matrix(4)

def save_canonical_matrix_sparse(max_exponent)
  matrix_file = "canonical_matrix_sparse_#{max_exponent}.csv"
  states_file = "canonical_matrix_states_#{max_exponent}.csv"
  transient_states, absorbing_states, _, _, _, canonical =
    make_fundamental_matrices(max_exponent)
  all_states = transient_states + absorbing_states
  CSV.open(File.join(OUTPUT_FOLDER, states_file), 'w') do |csv|
    csv << %w[i state]
    all_states.each.with_index do |state, i|
      csv << [i, state]
    end
  end
  CSV.open(File.join(OUTPUT_FOLDER, matrix_file), 'w') do |csv|
    csv << %w[i j probability]
    (0...all_states.size).each do |i|
      (0...all_states.size).each do |j|
        next if canonical[j, i] == 0
        csv << [i, j, canonical[j, i]]
      end
    end
  end
end
save_canonical_matrix_sparse(11)

def save_dot(max_exponent, max_sum, groups)
  pathname = File.join(
    OUTPUT_FOLDER,
    "chain_#{max_exponent}_#{max_sum}_#{groups}.dot"
  )
  transient_states, absorbing_states, _, _, _, fundamental =
    make_fundamental_matrices(max_exponent)
  all_states = transient_states + absorbing_states
  File.open(pathname, 'w') do |f|
    f.puts 'digraph {'
    f.puts '  rankdir=LR;'
    f.puts '  ranksep=1;'

    if groups
      all_states.group_by(&:sum).each do |sum, states|
        next if sum < 4
        next if sum > max_sum
        cluster_name = format('%d', sum)
        f.puts "  subgraph cluster_#{cluster_name} {"
        f.puts "    label=\"#{cluster_name}\";"
        f.puts '    style=filled; color=grey95; margin=16;'
        states.each do |state|
          f.puts "    s_#{state.join('_')};"
        end
        f.puts '  }'
      end
    end

    all_states.each do |state|
      next if state.sum > max_sum
      style = "label=\"{#{state.join(', ')}}\""
      style += ', color=red' if !state.empty? && state.max >= 2048
      f.puts "  s_#{state.join('_')} [#{style}];"
    end

    (0...all_states.size).each do |i|
      next if all_states[i].sum > max_sum
      (0...all_states.size).each do |j|
        next if fundamental[j, i] == 0
        next if fundamental[j, i] == 1
        next if all_states[j].sum > max_sum

        pr = format('%.2f', fundamental[j, i])
        style = []
        style += ['style=dashed'] if pr == '0.10'
        style += [format('label="%s"', pr)] if all_states[i].sum < 6

        f.puts "  s_#{all_states[i].join('_')} ->" \
          " s_#{all_states[j].join('_')} [#{style.join(', ')}]"
      end
    end

    absorbing_states = all_states.select do |state|
      !state.empty? && state.max >= 2048
    end
    absorbing_states.each do |state|
      style = 'label="1.0", dir=back, color=red'
      f.puts "  s_#{state.join('_')} -> s_#{state.join('_')} [#{style}];"
    end

    f.puts '}'
  end
end
# save_dot(11, 8, false)
# save_dot(11, 12, true)
save_dot(11, 4096, true)

def save_json(max_exponent)
  transitions = build_markov_chain(max_exponent)
  pathname = File.join(OUTPUT_FOLDER, "transitions_#{max_exponent}.json")
  File.open(pathname, 'w') do |f|
    f.puts JSON.pretty_generate(transitions)
  end
end
# save_json(11)