from matplotlib import pyplot as plt
import numpy as np
import sympy as sp
from matplotlib.animation import FuncAnimation

print('Ready')

class Curve: # class for the initial line and ultimate polynomial
	def __init__(self, init_data, data, degree, zorder = -1):
		self.init_x_data, self.init_y_data = init_data[0], init_data[1]
		self.init_slope, self.init_intercept = np.polyfit(self.init_x_data, self.init_y_data, 1)
		if self.init_slope == 0: # fix reciprocal problem if the slope is 0
			self.init_slope = 1e-6
		self.x_data, self.y_data = data[0], data[1]
		self.data_slope, self.data_intercept = np.polyfit(self.x_data, self.y_data, 1)
		if self.data_slope == 0: # fix reciprocal problem if the slope is 0
			self.data_slope = 1e-6
		self.degree = degree
		self._zorder = zorder
		(self.average_x, self.average_y) = (np.average(x_data), np.average(y_data))
		self.trans_slope, self.trans_intercept = self.init_slope, self.init_intercept
		if degree == 1:
			self.line_distance_diff, self.line_slope_angle_diff = self.get_line_info_and_perp(
				(self.average_x,self.average_y))[0:2]
			self.plot, = ax.plot((MIN_PLOT, MAX_PLOT),
			(MAX_PLOT * self.init_slope + self.init_intercept, MAX_PLOT * self.init_slope + self.init_intercept), 
			lw = 3, color = FIT_LINE_COLOR, zorder = self._zorder)
		else:
			data_poly_fit = np.polyfit(x_data, y_data, self.degree)
			self.data_poly_eval = np.poly1d(data_poly_fit)
			self.plot, = ax.plot(X_SPACED, self.data_poly_eval(X_SPACED), 
			lw = 3, color = FIT_LINE_COLOR, zorder = self._zorder)
			polynomial_spacing = self.degree + 2 # the number of points used to make the polynomial oscillation smooth
			self.poly_x_spaced = np.linspace(MIN_PLOT, MAX_PLOT, polynomial_spacing)
			self.poly_distances = [self.data_poly_eval(a) - (a * self.data_slope + self.data_intercept) 
				for a in self.poly_x_spaced]
	def intercept_oscillate(self, time):
		trans_line_distance = self.line_distance_diff * get_osc_funct(time)
		line_perp_slope = -1 / self.init_slope
		line_perp_angle = np.arctan(line_perp_slope)
		(line_point_x, line_point_y) = (self.average_x + trans_line_distance * np.cos(line_perp_angle), 
			self.average_y + trans_line_distance * np.sin(line_perp_angle))
		self.trans_slope = self.init_slope	
		self.trans_intercept = line_point_y - self.init_slope * line_point_x
	def slope_oscillate(self, time):
		trans_slope_angle_difference = self.line_slope_angle_diff * get_osc_funct(time)
		self.trans_slope = np.tan(trans_slope_angle_difference + np.arctan(self.data_slope))
		self.trans_intercept = self.average_y - self.trans_slope * self.average_x
		return self.trans_slope, self.trans_intercept
	def get_line_info_and_perp(self, data_point):
		# transform line equation into Hess form
		(w1, w2, b) = (-self.trans_slope, 1, -self.trans_intercept)
		# apply the distance formula d = (w·x+b)/||w|| to get the distance
		# between where we are and where we want to be
		line_distance_diff = (np.array((w1, w2)) @ np.array(data_point) + b) / np.linalg.norm((w1, w2))
		# also get the difference in slope between where we are 
		# and where we want to be
		slope_angle_diff = np.arctan2(self.trans_slope, 1) - np.arctan2(self.data_slope, 1)
		line_perp_slope = -1 / self.trans_slope
		line_perp_angle = np.arctan2(line_perp_slope, 1)
		# fix a problem with the connecting lines sometimes pointing in the wrong direction
		line_distance_diff = line_distance_diff * self.trans_slope / abs(self.trans_slope)
		(line_point_x, line_point_y) = (data_point[0] + line_distance_diff * np.cos(line_perp_angle), 
			data_point[1] + line_distance_diff * np.sin(line_perp_angle))
		return line_distance_diff, slope_angle_diff, (line_point_x, line_point_y)
	# write a function to get the best-fitting polynomial as we oscillate around the final polynomial
	def get_polynomial(self, time):
		trans_poly_distances = [x * (1 - get_osc_funct(time)) for x in self.poly_distances]
		trans_poly_x = [a for a, b in zip(self.poly_x_spaced, trans_poly_distances)]
		trans_poly_y = [self.data_slope * a + self.data_intercept + b for a, b in 
		zip(self.poly_x_spaced, trans_poly_distances)]
		self.trans_poly_fit = np.polyfit(trans_poly_x, trans_poly_y, self.degree)
		self.trans_poly_eval = np.poly1d(self.trans_poly_fit)
	# write a function to get the closest point between a data point and a polynomial
	def get_closest_point(self, data_point):
		x = sp.symbols("x")
		list0 = np.ndarray.tolist(self.trans_poly_fit)
		poly_fit_sym = sp.Poly(list0, x)
		# take the derivative of the fitted polynomial
		poly_deriv_sym = poly_fit_sym.diff(x)
		# set the following term to zero to minimize the (square of the) distance between the point and polynomial
		equation = x - data_point[0] + (poly_fit_sym - data_point[1]) * poly_deriv_sym
		coeff = sp.Poly(equation).all_coeffs()
		solved = np.roots(coeff)
		# chop any tiny imaginary component arising from the numerical solution
		solved = solved.real[abs(solved.imag) < 1e-6]
		dist = [(n - data_point[0])**2 + (self.data_poly_eval(n) - data_point[1])**2 for n in solved] # get distances
		solved = solved[np.argmin(dist)] # get shortest distance if there's more than one
		return solved, self.trans_poly_eval(solved)
	def dim(self, time, speedup = 1):
		self.plot.set_alpha(max(1 - speedup * time, 0))
	def brighten(self, time, speedup = 1):
		self.plot.set_alpha(min(speedup * time, 1))

class Connector: # thin perpendicular connector lines between each point and fitting curve
	def __init__(self, data):
		x_data, y_data = data[0], data[1]
		self.connector_line = []
		for i in range(len(x_data)):
			(line_point_x, line_point_y) = line_1.get_line_info_and_perp((x_data[i], y_data[i]))[2]
			self.connector_line_component, = ax.plot((x_data[i], line_point_x), 
				(y_data[i], line_point_y), 
				lw = 1, color = 'r', zorder = -1)
			self.connector_line.append(self.connector_line_component)
	def update(self, object):
		for j in range(len(x_data)):
			if object.degree == 1:
				(line_point_x, line_point_y) = object.get_line_info_and_perp((x_data[j], y_data[j]))[2]
			else:
				(line_point_x, line_point_y) = object.get_closest_point((x_data[j], y_data[j]))
			self.connector_line[j].set_xdata((x_data[j], line_point_x))
			self.connector_line[j].set_ydata((y_data[j], line_point_y))
	def brighten(self, time, speedup = 1):
		for j in range(len(x_data)): # the connector lines fade in
			self.connector_line[j].set_alpha(min(time, 1))
	def dim(self, time, speedup = 1):
		for j in range(len(x_data)):
			self.connector_line[j].set_alpha(max(1 - speedup * time, 0))

class Animate:
	def __init__(self, fig, frame_list, smoothness):
		self.smoothness = smoothness
		self.frame_names = [i[0] for i in frame_list]
		self.frame_durations = [self.smoothness * i[1] for i in frame_list]
		print("Starting to accumulate animation frames")
		anim = FuncAnimation(fig, self.update, frames = sum(self.frame_durations), interval = 100 / self.smoothness, 
			init_func = self.initial, blit = False)
		anim.save('animation.gif', dpi = 72, writer = 'imagemagick')
	def time_past(self, period): # total time before a certain segment
		return sum(self.frame_durations[0:self.frame_names.index(period)])
	def time_future(self, period): # total time including a certain segment
		return sum(self.frame_durations[0:self.frame_names.index(period) + 1])
	def time_present(self, period): # duration of a certain segment
		return self.frame_durations[self.frame_names.index(period)]
	def normalized_time(self, i, period): # how far we are into a certain segment
		return (i - self.time_past(period)) / self.time_present(period)
	def effective_time(self, i, period): # used as the time in the oscillation function
		return (i - self.time_past(period)) / self.smoothness
	def initial(self):
		return (line_1.plot, line_2.plot, connectors.connector_line,)
	def update(self, i):
		print(f"Frame: {i+1}")
		if i == 0:
			line_1.plot.set_alpha(1) # the first line is set as visible
			line_1.intercept_oscillate(0) # initialize the first line (is this necessary?)
			line_2.plot.set_alpha(0) # the second line is set as invisible
			line_2.intercept_oscillate(0) # initialize the second line (is this necessary?)
			poly.plot.set_alpha(0) # the polynomial is set as invisible
			connectors.update(line_1)
		if i <= self.time_future('fadein'): # this is the component line fade-in part
			connectors.brighten(self.normalized_time(i, 'fadein'))
		elif i <= self.time_future('fadein_pause'): # pause if desired
			pass
		elif i <= self.time_future('intercept_oscillate'): # this the line translation part
			line_1.intercept_oscillate(self.effective_time(i, 'intercept_oscillate'))
			line_1.plot.set_ydata((MIN_PLOT * line_1.trans_slope + 
				line_1.trans_intercept, MAX_PLOT * line_1.trans_slope + line_1.trans_intercept))
			connectors.update(line_1)
		elif i <= self.time_future('slope_oscillate'): # this is the line rotation part
			line_1.slope_oscillate(self.effective_time(i, 'slope_oscillate'))
			line_1.plot.set_ydata((MIN_PLOT * line_1.trans_slope + 
				line_1.trans_intercept, MAX_PLOT * line_1.trans_slope + line_1.trans_intercept))
			connectors.update(line_1)
		elif i <= self.time_future('slope_pause'): # pause if desired
			pass
		elif i <= self.time_future('poly_oscillate'): # this is the polynomial part
			line_1.plot.set_alpha(0)
			poly.plot.set_alpha(1)
			poly.get_polynomial(self.effective_time(i, 'poly_oscillate'))
			poly.plot.set_ydata(poly.trans_poly_eval(X_SPACED))
			connectors.update(poly)
		elif i <= self.time_future('poly_pause'): # pause if desired
			pass
		elif i <= self.time_future('fadeout'): 
			# fade from the correctly fitted line to the original line
			# speed up the fade-out a little relative to the fade-in
			speedup = 2
			norm_time = self.normalized_time(i, 'fadeout')
			connectors.dim(norm_time, speedup)
			poly.dim(norm_time, speedup)
			line_2.brighten(norm_time, 1)
		elif i <= self.time_future('fadeout_pause'): # pause if desired 
			pass		
		return (line_1.plot, line_2.plot, connectors.connector_line,)

def get_random_points(no_of_points, mean_x, sd_x, mean_y, sd_y, function):
	x_data = np.random.normal(mean_x, sd_x, no_of_points)
	y_data = np.random.normal(mean_y, sd_y, no_of_points)
	for idx, val in enumerate(y_data): # add an arbitrary cubic function
		y_data[idx] = (y_data[idx] + function.subs(x, x_data[idx]))
	# keep only data points that lie inside the plotted area
	check_limits_array = np.transpose(np.array([x_data, y_data]))
	x_data = [x[0] for x in check_limits_array if 
		MIN_PLOT < x[0] < MAX_PLOT and MIN_PLOT < x[1] < MAX_PLOT]
	y_data = [x[1] for x in check_limits_array if 
		MIN_PLOT < x[0] < MAX_PLOT and MIN_PLOT < x[1] < MAX_PLOT]
	no_of_actual_points = len(x_data)
	return x_data, y_data, no_of_actual_points

# write an oscillation function that overshoots and then settles down
def get_osc_funct(time): # set some spring, mass, damper values that provide pleasing oscillation
	k = 4 # spring value
	m = 1 # mass value
	c = 1.2 # damper value
	wn = np.sqrt(k / m) # natural frequency
	z = c / (2 * m * wn) # damping coefficient
	wd = np.sqrt(1 - z**2) * wn # damped natural frequency
	# two equivalent functions to find the displacement over time; either works
	osc_funct = np.exp(-z * wn * time) * (np.cos(wd * time) + z / np.sqrt(1 - z**2) * np.sin(wd * time))
	return osc_funct

NO_OF_POINTS = 15
(MEAN_X, SD_X) = (5, 5)
(MEAN_Y, SD_Y) = (6, 0.70)
x = sp.symbols("x")
(MIN_PLOT, MAX_PLOT) = (0, 10)
CUBIC = (x - (MAX_PLOT - MIN_PLOT) / 2) / 2 - 1 / 10 * (x - (MAX_PLOT - MIN_PLOT) / 2)**3
x_data, y_data, no_of_actual_points = get_random_points(NO_OF_POINTS, 
	MEAN_X, SD_X, MEAN_Y, SD_Y, CUBIC)

# pick some initial fitting line properties; 
# the line will pop up from this arrangement
INIT_LINE_INTERCEPT = 1
INIT_LINE_SLOPE = 0
FIT_LINE_COLOR = '0.5'
PLOT_POINT_NUMBER = 50
X_SPACED = np.linspace(MIN_PLOT, MAX_PLOT, PLOT_POINT_NUMBER)
POLY_DEGREE = 3

# define figure and axes and set some graphing parameters
fig, ax = plt.subplots(figsize = (4, 4))
ax.set_aspect(1)
plt.rc('font', size = 16)
marker_style = dict(linestyle = 'none', markersize = 10, markeredgecolor = 'k')
plt.box(on = None)
fig.set_facecolor('#c0c0c0')
plt.rcParams['savefig.facecolor'] = '#c0c0c0'
ax.spines['bottom'].set_position('zero')
ax.spines['left'].set_position('zero')

# plot the axes
plt.xticks(np.arange(MIN_PLOT, MAX_PLOT + 1, 2), fontsize = 10)
plt.yticks(np.arange(MIN_PLOT, MAX_PLOT + 1, 2), fontsize = 10)
plt.xlim([MIN_PLOT - 1, MAX_PLOT + 1])
plt.ylim([MIN_PLOT - 1, MAX_PLOT + 1])
ax.annotate('', xy = (MIN_PLOT, MAX_PLOT + 1), xytext = (MIN_PLOT, MIN_PLOT - 0.1), 
	arrowprops = {'arrowstyle': '->', 'lw': 1, 'facecolor': 'k'}, va = 'center')
ax.annotate('', xy = (MAX_PLOT + 1, MIN_PLOT), xytext = (MIN_PLOT - 0.1, MIN_PLOT), 
	arrowprops = {'arrowstyle': '->', 'lw': 1, 'facecolor': 'k'}, va = 'center', zorder = -1) 
plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05, wspace=0.0, hspace=0.0)

# plot the data points
ax.plot(x_data, y_data, marker = "o", 
	markerfacecolor = 'red', **marker_style, zorder = 0)

# now make the fitting lines - the first one moves; 
# the second fades back in at the end
line_1 = Curve(((0, 1),(INIT_LINE_INTERCEPT, INIT_LINE_INTERCEPT + INIT_LINE_SLOPE)), 
	(x_data, y_data), degree = 1)
line_2 = Curve(((0, 1),(INIT_LINE_INTERCEPT, INIT_LINE_INTERCEPT + INIT_LINE_SLOPE)), 
	(x_data, y_data), degree = 1)
poly = Curve(((0, 1),(INIT_LINE_INTERCEPT, INIT_LINE_INTERCEPT + INIT_LINE_SLOPE)), 
	(x_data, y_data), degree = POLY_DEGREE)
connectors = Connector((x_data, y_data))

# define the duration of each segment of the fitting animation
FRAME_LIST = [['fadein', 6],
	['fadein_pause', 2],
	['intercept_oscillate', 8],
	['slope_oscillate', 8],
	['slope_pause', 2],
	['poly_oscillate', 8],
	['poly_pause', 2],
	['fadeout', 6],
	['fadeout_pause', 2]]
SMOOTHNESS = 5

Animate(fig, FRAME_LIST, SMOOTHNESS)
print('Done')