From fc469f8ec4e32746177232470b7200b458de0455 Mon Sep 17 00:00:00 2001
From: Michael Coppola <michael.n.coppola@gmail.com>
Date: Thu, 18 Feb 2021 01:36:05 -0500
Subject: [PATCH] Added trend line option

---
 ChronoPlotter.py | 45 ++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/ChronoPlotter.py b/ChronoPlotter.py
index 582a9e9..1ec9b04 100644
--- a/ChronoPlotter.py
+++ b/ChronoPlotter.py
@@ -202,6 +202,10 @@ def __init__(self):
 		self.INCREASING = 0
 		self.DECREASING = 1
 
+		# line style
+		self.SOLID_LINE = 0
+		self.DASHED_LINE = 1
+
 		self.series = []
 		self.scroll_area = None
 
@@ -340,6 +344,23 @@ def initUI(self):
 		vd_layout.addWidget(self.vd_location)
 		self.options_layout.addLayout(vd_layout)
 
+		trend_layout = QHBoxLayout()
+		self.trend_checkbox = QCheckBox()
+		self.trend_checkbox.setChecked(False)
+		self.trend_checkbox.stateChanged.connect(self.trendCheckBoxChanged)
+		trend_layout.addWidget(self.trend_checkbox, 0)
+		self.trend_label = QLabel("Show trend line")
+		self.trend_label.setStyleSheet("color: #878787")
+		trend_layout.addWidget(self.trend_label, 1)
+		self.trend_linetype = QComboBox()
+		self.trend_linetype.addItem("solid line")
+		self.trend_linetype.addItem("dashed line")
+		self.trend_linetype.setCurrentIndex(1)
+		self.trend_linetype.setEnabled(False)
+		trend_layout.addWidget(self.trend_linetype)
+		self.options_layout.addLayout(trend_layout)
+
+
 		# Don't resize row heights if window height changes
 		self.options_layout.addStretch(0)
 
@@ -687,15 +708,19 @@ def vdCheckBoxChanged(self):
 		print("vdCheckBoxChanged called")
 		return self.optionCheckBoxChanged(self.vd_checkbox, self.vd_label, self.vd_location)
 
-	def optionCheckBoxChanged(self, checkbox, label, location):
+	def trendCheckBoxChanged(self):
+		print("trendCheckBoxChanged called")
+		return self.optionCheckBoxChanged(self.trend_checkbox, self.trend_label, self.trend_linetype)
+
+	def optionCheckBoxChanged(self, checkbox, label, combo):
 		if checkbox.isChecked():
 			action = "checked"
 			label.setStyleSheet("")
-			location.setEnabled(True)
+			combo.setEnabled(True)
 		else:
 			action = "unchecked"
 			label.setStyleSheet("color: #878787")
-			location.setEnabled(False)
+			combo.setEnabled(False)
 
 		print("checkbox was %s" % action)
 
@@ -743,6 +768,8 @@ def showGraph(self, save_without_showing=False):
 		# by chage weight so it draws correctly if the user-inputted charge weights aren't in order.
 		sorted_series = sorted(self.series.copy(), key=lambda x: x[3].value())
 
+		all_points = []
+
 		enabled_i = 0
 
 		for i, val in enumerate(sorted_series):
@@ -780,6 +807,7 @@ def showGraph(self, save_without_showing=False):
 
 				for v in m_velocs:
 					points.append((enabled_i, v))
+					all_points.append((enabled_i, v))
 
 				# Draw scatter plot
 				scatter_x, scatter_y = list(zip(*points))
@@ -915,6 +943,17 @@ def showGraph(self, save_without_showing=False):
 			plt.plot(line_x, line_y, "-", linewidth=1.5, color="#1c57eb")
 			ax.xaxis.grid(False)
 
+		# Draw trend line
+		if self.trend_checkbox.isChecked():
+			points_x, points_y = list(zip(*all_points))
+			z = numpy.polyfit(points_x, points_y, 1)
+			p = numpy.poly1d(z)
+			if self.trend_linetype.currentIndex() == self.SOLID_LINE:
+				line_type = "-"
+			else:
+				line_type = "--"
+			plt.plot(points_x, p(points_x), "r%s" % line_type, alpha=0.65)
+
 		# Titles and axis labels
 		graph_title = self.graph_title.text()
 		rifle = self.rifle.text()