import numpy as np
from scipy.signal import convolve2d

def part_1():
  handle = open('20a.txt','r')
  #read in the dictionary, convert it to binary
  ref = handle.readline().strip()
  handle.readline() #skip the blank
  vals = {'.':0,'#':1}
  ref = [vals[x] for x in ref]

  #now read in the initial image, conver to binary
  lines = []
  for line in handle:
    row = [vals[x] for x in line.strip()]
    lines.append(row)

  image = np.array(lines)

  #this is the key to the puzzle - you not only have to keep track of what
  #is happening in the image boundary, but also what happens to the infinite background
  #the infinite background will start at all zero, but will potentially change
  #if the first value in the dictionary is not zero (changing to all ones)
  #and then will change back again on the next step (assuming the last entry is a 0)
  pad_val = 0

  #enhance 50 times
  for i in range(50):
    image = enhance(image,ref,pad_val)
    if i == 1 or i == 49: #really 2 and 50
      print(np.sum(image)) #count the ones.

    #update the padding value, if it's currently 0, it goes to ref[0],
    #if it's currently 1, it goes to ref[-1] (which is what you get from all 1s)
    if pad_val == 0:
      pad_val = ref[0]
    else:
      pad_val = ref[-1]

#use the magic of convolution!
def enhance(img,ref,pad_val):
    result = np.zeros_like(img)
    img = np.pad(img,(2,2),constant_values=pad_val)

    kernel = np.array([[1,2,4],[8,16,32],[64,128,256]])
    indices = convolve2d(img,kernel,mode='full',fillvalue=pad_val)

    result = np.array(ref)[indices]
    return result


def enhance_slow(img,ref,pad_val):

  result = np.zeros_like(img)
  #first we need to pad out the image with the current pad value
  #note that we pad by TWO per edge, because we need to specially handle any cell
  #which depends on a cell in the image (so the image itself plus a border of 1)
  #the second layer of padding lets us figure out the new values of that first
  #layer of external cells.
  img = np.pad(img,(2,2),constant_values=pad_val)
  #we also need to figure out what the padded values in the result will be
  #this is a little hacky, but basically we want the result to be the same size
  #as the input image, so that it's easy to match coordinates, and so have some
  #extra values along the edges
  #this could be optimized away, I think
  if pad_val == 0:
    pad_val = ref[0]
  else:
    pad_val = ref[-1]
  result = np.pad(result,(2,2),constant_values=pad_val)

  #now we just loop over the range of cells which need to be calculated
  #this extends one pixel past the original image boundaries (which have
  #been padded by TWO, thus the 1 and -1)

  #if we wanted to go FAST we could convolve with a specially crafted filter:
  # [1,2,4]
  # [8,16,32]
  # [64,128,256]
  #to get each code, and then just batch update
  #but this was faster to code at 12AM.
  for x in range(1,len(img)-1):
    for y in range(1,len(img)-1):
      vals = img[x-1:x+2,y-1:y+2]
      vals = vals.reshape(-1) #luckily np reshape does the thing we want
      code = int(''.join([str(x) for x in list(vals)]),2) #lookup the code
      result[x,y] = ref[code] #fill in the result cell
  return result


part_1()