import numpy as np
from part1 import load_diagram, print_diagram

debug = False

def quantum_propagate_diagram(in_diagram):
	diagram = in_diagram.copy()
	ylen = len(diagram)

	# Find start location
	y0 = 0
	x0, = np.where(diagram[y0] == "S")[0]

	# Emit starting beam
	diagram[y0 + 1, x0] = "|"
	
	# Propagate
	worlds_count = quantum_propagate(x0, diagram[y0 + 2:])

	return worlds_count

def quantum_propagate(x, subdiagram):
	ylen = len(subdiagram)
	# print_diagram(subdiagram)
	# print()
	for y in range(0, ylen - 1):
		if subdiagram[y, x] == ".":
			# subdiagram[y, x] = "|"
			if debug:
				print("--- no split ---")
				print_diagram(subdiagram)
				print()
		elif subdiagram[y, x] == "^":
			# subdiagram[y, (x - 1, x + 1)] = "|"
			# Split into two worlds, ignore what has happened before.
			d1 = subdiagram[y + 1:]#.copy()
			if debug:
				print("--- split left ---")
				print_diagram(d1)
				print()
			worlds_count = quantum_propagate(x - 1, d1)
			d2 = subdiagram[y + 1:]#.copy()
			if debug:
				print("--- split right ---")
				print_diagram(d2)
				print()
			worlds_count += quantum_propagate(x + 1, d2)
			# print(worlds_count)
			return worlds_count
	# Reached the end, so this must be a world
	if debug:
		print("Reached end of world!")
	return 1

def sninkogate(diagram):
	ylen = len(diagram)

	# Find start location
	y0 = 0
	x0, = np.where(diagram[y0] == "S")[0]

	# Save worlds per column
	I = np.zeros_like(diagram, dtype=int)

	# Emit starting beam
	I[y0 + 1, x0] = 1

	for y in range(2, ylen):
		for x in np.where((diagram[y] == "^")&(I[y - 1] > 0))[0]:
			# NOTE: I do not ever check whether carets are on the boundary.
			I[y, (x - 1, x + 1)] += I[y - 1, x]
		for x in np.where((diagram[y] == ".")&(I[y - 1] > 0))[0]:
			I[y, x] += I[y - 1, x]
		if debug:
			print_diagram(diagram)
			print()
	return I[-1].sum()


if __name__ == "__main__":
	test_diagram = load_diagram("testinput")

	assert quantum_propagate_diagram(test_diagram) == 40

	diagram = load_diagram("input")
	print(sninkogate(diagram))