import random
import numpy as np
import math

from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa
from skimage.io import imsave
from skimage.util import random_noise

maxSlope = 10 # restrict the maximum slope of generated lines for stability
minLength = 20 # restrict the minimum length of line segments

class ICircleDataset:
	'''
	Generator of circle segment images.

	Images will have 1 random circle each, filled with noise and distractor lines.
	Class also offers functionality for drawing line parameters, hypotheses and point predictions.
	'''

	def __init__(self, imgW = 64, imgH = 64, margin = -5, bg_clr = 0.5):
		'''
		Constructor. 

		imgW -- image width (default 64)
		imgH -- image height (default 64)
		margin -- lines segments are sampled within this margin, negative value means that a line segment can start or end outside the image (default -5)
		bg_clr -- background intensity (default 0.5)
		'''
		
		self.imgW = imgW
		self.imgH = imgH
		self.margin = margin
		self.bg_clr = bg_clr

	def draw_circle(self, data, cX, cY, r, clr, alpha=1.0):
		'''
		Draw a circle with the given color and opacity.

		data -- image to draw to
		cX -- x value of circle center
		cY -- y value of circle center
		r -- radius of circle
		clr -- line color, triple of values
		alpha -- opacity (default 1.0)
		'''

		cY = int(cY * self.imgH)
		cX = int(cX * self.imgW)
		r = int(r * self.imgW)

		rr, cc, val =  circle_perimeter_aa(cY, cX, r)
		set_color(data, (rr, cc), clr, val)

	def draw_hyps(self, labels, scores, data=None):
		'''
		Draw a set of line hypothesis for a batch of images.
		
		labels -- line parameters, array shape (NxMx2) where 
			N is the number of images in the batch
			M is the number of hypotheses per image
			2 is the number of line parameters (intercept, slope)
		scores -- hypotheses scores, array shape (NxM), see above, higher score will be drawn with higher opacity
		data -- batch of images to draw to, if empty a new batch wil be created according to the shape of labels
		
		'''

		n = labels.shape[0] # number of images
		m = labels.shape[1] # number of hypotheses

		if data is None: # create new batch of images
			data = np.zeros((n, self.imgH, self.imgW, 3), dtype=np.float32)
			data.fill(self.bg_clr)

		clr = (0, 0, 1)

		for i in range (0, n):
			for j in range (0, m):
				lY1 = int(labels[i, j, 0] * self.imgH)
				lY2 = int(labels[i, j, 1] * self.imgW + labels[i, j, 0] * self.imgH)
				self.draw_line(data[i], 0, lY1, self.imgW, lY2, clr, scores[i, j])

		return data

	def draw_models(self, labels, data=None, correct=None):
		'''
		Draw circles for a batch of images.
	
		labels -- circle parameters, array shape (Nx3) where 
			N is the number of images in the batch
			3 is the number of circles parameters (center x,  center y, radius)
		data -- batch of images to draw to, if empty a new batch wil be created according to the shape of labels 
			and circles will be green, circles will be blue otherwise
		correct -- array of shape (N) indicating whether a circle estimate is correct 
		'''

		n = labels.shape[0]
		if data is None: 
			data = np.zeros((n, self.imgH, self.imgW, 3), dtype=np.float32)
			data.fill(self.bg_clr)
			clr = (0, 1, 0)
		else:
			clr = (0, 0, 1)

		for i in range (0, n):

			self.draw_circle(data[i], labels[i, 0], labels[i, 1], labels[i, 2], clr)

			if correct is not None:
				
				# draw border green if estiamte is correct, red otherwise
				if correct[i]: borderclr = (0, 1, 0)
				else: borderclr = (1, 0, 0)
				
				set_color(data[i], line(0, 0, 0, self.imgW-1), borderclr)			
				set_color(data[i], line(0, 0, self.imgH-1, 0), borderclr)			
				set_color(data[i], line(self.imgH-1, 0, self.imgH-1, self.imgW-1), borderclr)			
				set_color(data[i], line(0, self.imgW-1, self.imgH-1, self.imgW-1), borderclr)			

		return data

	def draw_points(self, points, data, inliers=None):
		'''
		Draw 2D points for a batch of images.

		points -- 2D points, array shape (Nx2xM) where 
			N is the number of images in the batch
			2 is the number of point dimensions (x, y)
			M is the number of points
		data -- batch of images to draw to
		inliers -- soft inlier score for each point, 
			if given and score < 0.5 point will be drawn green, red otherwise
		'''

		n = points.shape[0] # number of images
		m = points.shape[2] # number of points

		for i in range (0, n):
			for j in range(0, m):

				clr = (0.2, 0.2, 0.2) # draw predicted points as dark circles
				if inliers is not None and inliers[i, j] > 0.5:
					clr = (0.7, 0.7, 0.7) # draw inliers as light circles
					
				r = int(points[i, 0, j] * self.imgH)
				c = int(points[i, 1, j] * self.imgW)
				rr, cc = circle(r, c, 2)
				set_color(data[i], (rr, cc), clr)

		return data

	def samples(self, n):
		'''
		Create new input images of random line segments and distractors along with ground truth parameters.

		n -- number of images to create
		'''

		data = np.zeros((n, self.imgH, self.imgW, 3), dtype=np.float32)
		data.fill(self.bg_clr)
		labels = np.zeros((n, 3), dtype=np.float32)

		for i in range (0, n):
			data[i] = random_noise(data[i], mode='speckle')

		return data, labels

[((51, 17, 51, 47), 'skimage.draw.circle_perimeter_aa', 'circle_perimeter_aa', ({(51, 37, 51, 39): 'cY', (51, 41, 51, 43): 'cX', (51, 45, 51, 46): 'r'}, {}), '(cY, cX, r)', False, 'from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa\n'), ((52, 2, 52, 37), 'skimage.draw.set_color', 'set_color', ({(52, 12, 52, 16): 'data', (52, 18, 52, 26): '(rr, cc)', (52, 28, 52, 31): 'clr', (52, 33, 52, 36): 'val'}, {}), '(data, (rr, cc), clr, val)', False, 'from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa\n'), ((158, 9, 158, 65), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((160, 11, 160, 45), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((71, 10, 71, 66), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((98, 10, 98, 66), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((163, 13, 163, 50), 'skimage.util.random_noise', 'random_noise', (), '', False, 'from skimage.util import random_noise\n'), ((146, 13, 146, 28), 'skimage.draw.circle', 'circle', ({(146, 20, 146, 21): 'r', (146, 23, 146, 24): 'c', (146, 26, 146, 27): '2'}, {}), '(r, c, 2)', False, 'from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa\n'), ((147, 4, 147, 37), 'skimage.draw.set_color', 'set_color', ({(147, 14, 147, 21): 'data[i]', (147, 23, 147, 31): '(rr, cc)', (147, 33, 147, 36): 'clr'}, {}), '(data[i], (rr, cc), clr)', False, 'from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa\n'), ((114, 23, 114, 49), 'skimage.draw.line', 'line', ({(114, 28, 114, 29): '(0)', (114, 31, 114, 32): '(0)', (114, 34, 114, 35): '(0)', (114, 37, 114, 48): '(self.imgW - 1)'}, {}), '(0, 0, 0, self.imgW - 1)', False, 'from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa\n'), ((115, 23, 115, 49), 'skimage.draw.line', 'line', ({(115, 28, 115, 29): '(0)', (115, 31, 115, 32): '(0)', (115, 34, 115, 45): '(self.imgH - 1)', (115, 47, 115, 48): '(0)'}, {}), '(0, 0, self.imgH - 1, 0)', False, 'from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa\n'), ((116, 23, 116, 69), 'skimage.draw.line', 'line', ({(116, 28, 116, 39): '(self.imgH - 1)', (116, 41, 116, 42): '(0)', (116, 44, 116, 55): '(self.imgH - 1)', (116, 57, 116, 68): '(self.imgW - 1)'}, {}), '(self.imgH - 1, 0, self.imgH - 1, self.imgW - 1)', False, 'from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa\n'), ((117, 23, 117, 69), 'skimage.draw.line', 'line', ({(117, 28, 117, 29): '(0)', (117, 31, 117, 42): '(self.imgW - 1)', (117, 44, 117, 55): '(self.imgH - 1)', (117, 57, 117, 68): '(self.imgW - 1)'}, {}), '(0, self.imgW - 1, self.imgH - 1, self.imgW - 1)', False, 'from skimage.draw import line, line_aa, circle, set_color, circle_perimeter_aa\n')]