# this function computes the polarization currents received by the receptive columns
# and feeds them into neurons
import cv2
from cv_bridge import CvBridge


@nrp.MapSpikeSource("pol_r1", nrp.brain.polar[0],  nrp.dc_source)
@nrp.MapSpikeSource("pol_g2", nrp.brain.polar[2],  nrp.dc_source)
@nrp.MapSpikeSource("pol_b1", nrp.brain.polar[4],  nrp.dc_source)
@nrp.MapSpikeSource("pol_r2", nrp.brain.polar[6],  nrp.dc_source)
@nrp.MapSpikeSource("pol_g1", nrp.brain.polar[8],  nrp.dc_source)
@nrp.MapSpikeSource("pol_b2", nrp.brain.polar[10],  nrp.dc_source)

@nrp.MapSpikeSource("pol_r1_g2", nrp.brain.polar[1],  nrp.dc_source)
@nrp.MapSpikeSource("pol_g2_b1", nrp.brain.polar[3],  nrp.dc_source)
@nrp.MapSpikeSource("pol_b1_r2", nrp.brain.polar[5],  nrp.dc_source)
@nrp.MapSpikeSource("pol_r2_g1", nrp.brain.polar[7],  nrp.dc_source)
@nrp.MapSpikeSource("pol_g1_b2", nrp.brain.polar[9],  nrp.dc_source)
@nrp.MapSpikeSource("pol_b2_r1", nrp.brain.polar[11],  nrp.dc_source)

@nrp.MapCSVRecorder("recorder", filename="polarization_cam_data.csv", headers=["r","g","b"])
@nrp.MapRobotSubscriber("cam_data", Topic('/spikebot/camera_pol/camera', sensor_msgs.msg.Image))
@nrp.Robot2Neuron()

def csv_polarization_camera(t, cam_data, recorder, pol_r1, pol_r2, pol_g1, pol_g2, pol_b1, pol_b2, pol_r1_g2, pol_r2_g1, pol_g1_b2, pol_g2_b1, pol_b1_r2, pol_b2_r1):
	if not isinstance(cam_data.value, type(None)):
		width = 12
		height = 12
		cv_image = CvBridge().imgmsg_to_cv2(cam_data.value, "rgb8") / 256.
		clientLogger.info('image:' + str(cv_image))
		cv_image = cv2.resize(cv_image, (width, height))
		image_result = cv_image
		r1 = r2 = g1 = g2 = b1 = b2 = 0.0
		for x in range(len(image_result)):
			for y in range(len(image_result[x])):
				if x % 3 == 0:
					if image_result[x][y][0] >= 0.5 and x < 6 and y < 6:
						r1 = r1 + 30.0
					elif image_result[x][y][0] >= 0.5 and x >= 6 and y >= 6:
						r1 = r1 + 30.0
					elif image_result[x][y][0] >= 0.5:
						r2 = r2 + 30.0
				if x % 3 -1 == 0:
					if image_result[x][y][1] >= 0.5 and x < 6 and y < 6:
						g1 = g1 + 30.0
					elif image_result[x][y][1] >= 0.5 and x >= 6 and y >= 6:
						g1 = g1 + 30.0
					elif image_result[x][y][1] >= 0.5:
						g2 = g2 + 30.0
				if x % 3 -2 == 0:
					if image_result[x][y][2] >= 0.5 and x < 6 and y < 6:
						b1 = b1 + 30.0
					elif image_result[x][y][2] >= 0.5 and x >= 6 and y >= 6:
						b1 = b1 + 30.0
					elif image_result[x][y][2] >= 0.5:
						b2 = b2 + 30.0
							
		pol_r1.amplitude = r1
		pol_r2.amplitude = r2
		pol_g1.amplitude = g1
		pol_g2.amplitude = g2
		pol_b1.amplitude = b1
		pol_b2.amplitude = b2
		
		pol_r1_g2.amplitude = (r1 + g2)/2
		pol_g2_b1.amplitude = (b1 + g2)/2
		pol_b1_r2.amplitude = (b1 + r2)/2
		pol_r2_g1.amplitude = (g1 + r2)/2
		pol_g1_b2.amplitude = (g1 + b2)/2
		pol_b2_r1.amplitude = (r1 + b2)/2