import numpy as np
import matplotlib.pyplot as plt
import math
from PIL import Image

# Define constants for program
# Radial stepper is connected to Z
# Linear stepper is connected to Y

FEED_RATE = 900
Z_STEP =  0.3
MAX_COLOR = 255.0 # black
NUM_RAD_STEPS = 471 # steps in a full rotation for stepper motor with F FEED_RATE Z Z_STEP
PHYS_RADIUS = 138 # radius of canvas in mm
LIN_STEP = 0.6 # step size for linear oscillation
LIN_RES = 8 # how many different sized oscillations you can make
SAMPLING_RADIUS = 1 # look at all the neighboring pixels within SAMPLING_RADIUS when determining oscillation magnitude

# Open file for writing
f = open("lion.txt", "w")

# Import Image and Convert to Grayscale

im = Image.open('lion.jpg').convert('L')
im_radius = min(im.size) / 2
a = np.asarray(im)

def polar_to_cart(r, theta):
	return (r * math.cos(theta), r * math.sin(theta))

def origin_mid_to_origin_ul(x, y):
	return x + im_radius, -1 * y + im_radius

def avg_pix_value(x, y, pix_rad):
	count = 0
	pix_tot = 0
	for i in range(-pix_rad, pix_rad + 1, 1):
		for j in range(-pix_rad, pix_rad + 1, 1):
			if x + i < im.size[1] and y + j < im.size[0]:
				pix_tot += a[x + i][y + j]
				count+=1
	return pix_tot / count

# Generate Archimedean Spiral Points

loops = math.floor(PHYS_RADIUS / (2 * LIN_STEP * LIN_RES)) # in SPIRAL
step_size = 2 * math.pi / NUM_RAD_STEPS
b = im_radius / (2 * math.pi * loops) # Given that radius of outermost point of spiral = b * 2 * pi * loops
max_theta = math.floor(loops * math.pi * 2) # max value for parameter t which generates points
spiral_points_xy = []
spiral_points_polar = []
t = 0
while t < max_theta:
	r = b * t
	x, y = polar_to_cart(r, t)
	spiral_points_xy += [(x,y)]
	spiral_points_polar += [(r, t)]
	t += step_size

# For each point, perturb deviation from spiral to match color intensity of input image

max_deviation = im_radius / (2 * loops)
final_points_xy = []
final_points_polar = []
num_points = 0
prev_r = 0
for i in range(len(spiral_points_polar)):
	direction  = (-1) ** i

	# Get point on simple archimedian curve
	r, theta = spiral_points_polar[i]
	spi_x, spi_y = polar_to_cart(r, theta)
	im_x, im_y = origin_mid_to_origin_ul(spi_x, spi_y)
	im_x, im_y = int(math.floor(im_x)), int(math.floor(im_y))

	# Calculate how much to oscillate and produce new point
	dev = max_deviation * direction * (1 - avg_pix_value(im_x, im_y, SAMPLING_RADIUS) / MAX_COLOR)
	r = (r + dev) * (PHYS_RADIUS / float(im_radius))
	r = round(r * (1.0 / LIN_STEP)) / (1.0 / LIN_STEP) # round r value to nearest multiple of step size

	# If no need to oscillate, still add baby oscillation so the radial movement doesn't get thrown off
	if prev_r - r <= 0.2:
		r += direction * 0.2

	# Oscillate in a certain direction
	y_move = prev_r - r
	f.write("G91 F900 y%.1f\n" % y_move)

	# Store point for visualization
	x, y = polar_to_cart(r, theta)
	final_points_xy += [(x, y)]
	final_points_polar += [(r, theta)]

	# Move one radial step
	theta += step_size
	f.write("G91 F900 z0.3\n")

	# Store point for visualization
	x, y = polar_to_cart(r, theta)
	final_points_xy += [(x, y)]
	final_points_polar += [(r, theta)]

	# Compute number of points made
	num_points += 2

	# Store previous r
	prev_r = r

print("b: " + str(b))
print("loops: " + str(loops))
print("num_points: " + str(num_points))
print("max_deviation: " + str(max_deviation))

# Plot spiral image
xs = [p[0] for p in final_points_xy]
ys = [p[1] for p in final_points_xy]
plt.plot(xs, ys)
plt.show()