require_relative 'util'

# --- Day 23: Safe Cracking ---

# This is one of the top floors of the nicest tower in EBHQ. The
# Easter Bunny's private office is here, complete with a safe hidden
# behind a painting, and who wouldn't hide a star in a safe behind a
# painting?

# The safe has a digital screen and keypad for code entry. A sticky
# note attached to the safe has a password hint on it: "eggs". The
# painting is of a large rabbit coloring some eggs. You see 7.

# When you go to type the code, though, nothing appears on the
# display; instead, the keypad comes apart in your hands, apparently
# having been smashed. Behind it is some kind of socket - one that
# matches a connector in your prototype computer! You pull apart the
# smashed keypad and extract the logic circuit, plug it into your
# computer, and plug your computer into the safe.

# Now, you just need to figure out what output the keypad would have
# sent to the safe. You extract the assembunny code from the logic
# chip (your puzzle input).

# The code looks like it uses almost the same architecture and
# instruction set that the monorail computer used! You should be able
# to use the same assembunny interpreter for this as you did there,
# but with one new instruction:

# tgl x toggles the instruction x away (pointing at instructions like
# jnz does: positive means forward; negative means backward):

# - For one-argument instructions, inc becomes dec, and all other
#   one-argument instructions become inc.
# - For two-argument instructions, jnz becomes cpy, and all other
#   two-instructions become jnz.
# - The arguments of a toggled instruction are not affected.
# - If an attempt is made to toggle an instruction outside the
#   program, nothing happens.
# - If toggling produces an invalid instruction (like cpy 1 2) and an
#   attempt is later made to execute that instruction, skip it
#   instead.
# - If tgl toggles itself (for example, if a is 0, tgl a would target
#   itself and become inc a), the resulting instruction is not
#   executed until the next time it is reached.

# For example, given this program:

#     cpy 2 a
#     tgl a
#     tgl a
#     tgl a
#     cpy 1 a
#     dec a
#     dec a

# - cpy 2 a initializes register a to 2.
# - The first tgl a toggles an instruction a (2) away from it, which
#   changes the third tgl a into inc a.
# - The second tgl a also modifies an instruction 2 away from it,
#   which changes the cpy 1 a into jnz 1 a.
# - The fourth line, which is now inc a, increments a to 3.
# - Finally, the fifth line, which is now jnz 1 a, jumps a (3)
#   instructions ahead, skipping the dec a instructions.

# In this example, the final value in register a is 3.

# The rest of the electronics seem to place the keypad entry (the
# number of eggs, 7) in register a, run the code, and then send the
# value left in register a to the safe.

# What value should be sent to the safe?

input = File.open('23.txt') { |f| f.readlines.map(&:chomp) }
test_input = "cpy 2 a\ntgl a\ntgl a\ntgl a\ncpy 1 a\ndec a\ndec a".split("\n")

class Computer
  attr_accessor :registers

  def initialize(input)
    @registers = { a: 0, b: 0, c: 0, d: 0 }
    @program = input.map { |line| parse(line) }
    @pc = 0
  end

  def parse_arg(arg)
    if arg[/^[-+]?\d+$/]
      arg.to_i
    else
      arg.to_sym
    end
  end

  def parse(line)
    op, *ops = line.split
    [op.to_sym] + ops.map { |o| parse_arg(o) }
  end

  def lookup(reg_or_arg)
    if reg_or_arg.is_a?(Symbol)
      @registers[reg_or_arg]
    else
      reg_or_arg
    end
  end

  def fetch
    assert(@pc < @program.length)
    result = @program[@pc]
    @pc += 1
    result
  end

  def cpy(src, dest)
    @registers[dest] = lookup(src)
  end

  def inc(reg)
    @registers[reg] += 1
  end

  def dec(reg)
    @registers[reg] -= 1
  end

  def jnz(arg1, arg2)
    # HACK: account for fetch always incrementing @pc
    @pc += lookup(arg2) - 1 if lookup(arg1) != 0
  end

  def tgl(arg)
    # HACK: account for fetch always incrementing @pc
    target = @pc - 1 + lookup(arg)
    return unless target >= 0 && target < @program.length

    op = @program[target][0]
    @program[target][0] = case op
                          when :inc then :dec
                          when :dec, :tgl then :inc
                          when :jnz then :cpy
                          else :jnz
                          end
  end

  def step
    op, *args = fetch
    send(op, *args)
  end

  def done?
    @pc < 0 || !@program[@pc]
  end

  def run
    step until done?
  end
end

def easy(input)
  computer = Computer.new(input)
  computer.registers[:a] = 7
  computer.run
  computer.registers[:a]
end

assert(easy(test_input) == 3)
puts "easy(input): #{easy(input)}"

# --- Part Two ---

# The safe doesn't open, but it does make several angry noises to
# express its frustration.

# You're quite sure your logic is working correctly, so the only other
# thing is... you check the painting again. As it turns out, colored
# eggs are still eggs. Now you count 12.

# As you run the program with this new input, the prototype computer
# begins to overheat. You wonder what's taking so long, and whether
# the lack of any instruction more powerful than "add one" has
# anything to do with it. Don't bunnies usually multiply?

# Anyway, what value should actually be sent to the safe?

class OptimizingComputer < Computer
  attr_accessor :profile

  def initialize(input)
    super(input)
    @profile = Hash.new { |h, k| h[k] = 0 }
  end

  ## program snapshot
  #    0000 [:cpy, :a, :b]  # b = a
  #    0001 [:dec, :b]      # b -= 1
  #                         # loop
  #    0002 [:cpy, :a, :d]  #   d = a
  #    0003 [:cpy, 0, :a]   #   a = 0
  #                         #   until d.zero? do
  #    0004 [:cpy, :b, :c]  #   c = b
  #                         #     until c.zero? do
  #    0005 [:inc, :a]      #       a += 1
  #    0006 [:dec, :c]      #       c -= 1
  #    0007 [:jnz, :c, -2]  #     end
  #    0008 [:dec, :d]      #     d -= 1
  #    0009 [:jnz, :d, -5]  #   end
  #    0010 [:dec, :b]      #   b -= 1
  #    0011 [:cpy, :b, :c]  #   c = b
  #    0012 [:cpy, :c, :d]  #   d = c
  #                         #   until d.zero? do
  #    0013 [:dec, :d]      #     d -= 1
  #    0014 [:inc, :c]      #     c += 1
  #    0015 [:jnz, :d, -2]  #   end
  #    0016 [:tgl, :c]      #   ???
  #    0017 [:cpy, -16, :c] #   c = -16
  #    0018 [:jnz, 1, :c]   # end
  #    0019 [:cpy, 89, :c]  # c = 89
  #                         # until c.zero? do
  #    0020 [:jnz, 90, :d]  #   jro d
  #                         #   until d.zero? do
  #    0021 [:inc, :a]      #     a += 1
  #    0022 [:inc, :d]      #     d += 1
  #    0023 [:jnz, :d, -2]  #   end
  #    0024 [:inc, :c]      #   c += 1
  #    0025 [:jnz, :c, -5]  # end

  ## the hotspot
  #    0005    [:inc, :a]      29.235587627557
  #    0006    [:dec, :c]      29.235587627557
  #    0007    [:jnz, :c, -2]  29.235587627557
  #    0008    [:dec, :d]      4.095730288844516
  #    0009    [:jnz, :d, -5]  4.095730288844516

  ## its code
  #    0002 [:cpy, :a, :d]  #   d = a
  #    0003 [:cpy, 0, :a]   #   a = 0
  #                         #   until d.zero? do
  #    0004 [:cpy, :b, :c]  #   c = b
  #                         #     until c.zero? do
  #    0005 [:inc, :a]      #       a += 1
  #    0006 [:dec, :c]      #       c -= 1
  #    0007 [:jnz, :c, -2]  #     end
  #    0008 [:dec, :d]      #     d -= 1
  #    0009 [:jnz, :d, -5]  #   end

  ## what it does
  # outer loop runs d times (set to a)
  # inner loop runs c times (set to b)
  # a is used as accumulator (set to 0)
  # inner loop increments a for c times (a += c)
  # outer loop increments a by c again, for d times
  # therefore a += c * d

  ## the optimization
  #    0005 [:incmul, :a, :c, :d]  a += c * d
  #    0006 [:cpy, 0, :c]   #   c = 0
  #    0007 [:cpy, 0, :d]   #   d = 0
  #    0008 [:nop]
  #    0009 [:nop]

  def at_hotspot?(target)
    (target + 5 < @program.length && # 5 lines of context
      @program[target][0] == :inc &&
      @program[target + 1][0] == :dec &&
      @program[target + 2][0] == :jnz && @program[target + 2][2] == -2 &&
      @program[target + 3][0] == :dec &&
      @program[target + 4][0] == :jnz && @program[target + 4][2] == -5)
  end

  def patch_hotspot(target)
    acc = @program[target][1]
    arg1 = @program[target + 1][1]
    arg2 = @program[target + 3][1]
    @program[target] = [:incmul, acc, arg1, arg2]
    @program[target + 1] = [:cpy, 0, arg1]
    @program[target + 2] = [:cpy, 0, arg2]
    @program[target + 3] = [:nop]
    @program[target + 4] = [:nop]
  end

  def optimize
    @program.length.times do |i|
      patch_hotspot(i) if at_hotspot?(i)
    end
  end

  def nop; end

  def incmul(reg, arg1, arg2)
    @registers[reg] += lookup(arg1) * lookup(arg2)
  end

  def print_profile
    tally = @profile.values.reduce(&:+)
    top = @profile.to_a.sort { |a, b| a[1] <=> b[1] }.last(5)
    top.reverse.each do |k, v|
      line_no = k.to_s.rjust(4, '0')
      line = @program[k]
      puts "#{line_no}\t#{line}\t#{100 * v.to_f / tally}"
    end
  end

  def step
    @profile[@pc] += 1
    op, *args = fetch
    send(op, *args)
  end

  def tgl(arg)
    super(arg)
    # p @program
    optimize
    # p @program
    # print_profile
  end
end

def hard(input)
  computer = OptimizingComputer.new(input)
  computer.registers[:a] = 12
  computer.run
  computer.registers[:a]
end

puts "hard(input): #{hard(input)}"