import numpy as np
import matplotlib.pyplot as plt

class Carry:
	def __init__(self, arr):
		if not isinstance(arr, np.ndarray):
			arr = np.array(arr)
		overflow = None
		over_pos = None
		coeff_sum = 0
		#find largest carry
		for i, row in enumerate(arr):
			for j, val in enumerate(row):
				coeff_sum += val
				if val > 0:
					if overflow is not None:
						raise ValueError("array supplied has more than one overflow digit")
					overflow = val
					over_pos = (i, j)

		if coeff_sum < 0:
			raise ValueError("sum of coefficients too small")

		self._arr = arr
		self.over_pos = over_pos
		self.overflow = overflow

	def _update(self, base, update_tuples, zerowall):
		#dim_x, dim_y = base.shape
		pos_x, pos_y = self.over_pos
		for (x,y) in update_tuples:
			idx = base[(x,y)]
			dec = (idx // self.overflow) # - 1

			wall = zerowall.copy()
			sh_x, sh_y = self._arr.shape
			x -= pos_x
			y -= pos_y
			
			wall[x:x + sh_x, y:y + sh_y] = dec * self._arr
			base -= wall

	def update(self, base):
		zerowall = np.zeros(base.shape, dtype=base.dtype)
		while True:
			update_tuples = list(zip(*np.where(base >= self.overflow)))
			if not update_tuples:
				break
			self._update(base, update_tuples, zerowall)

	def times(self, base, val, center=(0,0)):
		#perform multiplication using the carry
		base[center] *= val
		self.update(base)
		return base
		
	def add(self, base, val=1, center=(0,0)):
		base[center] += val
		self.update(base)
		return base

laplace = Carry(np.array([[0,-1,0],[-1,4,-1],[0,-1,0]]))

#test = np.zeros((100, 100))
#test[50,50] = 4