#计算轨道数据

import numpy as np
import color

#轨道数据表
sheet_measure = sheet("measure_data")

def linear_regression(x, y):
	n = x.size
	sx = x.sum()
	sx2 = (x ** 2).sum()
	sy = y.sum()
	sxy = (x * y).sum()
	delta = n * sx2 - sx ** 2
	a = (n * sxy - sx * sy) / delta
	b = (sx2 * sy - sx * sxy) / delta
	return a, b

def calc_measure_data():
	diff_param = param("height_diff_param").strip()
	if not diff_param:
		print(f"{color.yellow}警告:{color.nc}缺少高差参数")
		diff_param = 0
	else:
		diff_param = float(diff_param)

	with sheet_measure:
		row = 0
		left = []
		right = []
		while not cell(row, 1).empty():
			left.append(cell(row, 1).value)
			right.append(cell(row, 3).value)
			row += 1
		left = np.array(left, dtype = np.float64)
		right = np.array(right, dtype = np.float64)
		n = left.size
		index = np.array(range(left.size), dtype = np.float64)
		h_diff = (left - right) * 1000 + diff_param
		# left_a, left_b = linear_regression(index, left)
		# right_a, right_b = linear_regression(index, right)
		left_a = (left[-1] - left[0]) / (n - 1)
		left_b = left[0]
		right_a = (right[-1] - right[0]) / (n - 1)
		right_b = right[0]

		left_diff = (left - left_a * index - left_b) * 1000
		right_diff = (right - right_a * index - right_b) * 1000

		left_diff2 = left_diff[1: -1] - (left_diff[2:] + left_diff[:-2]) / 2
		right_diff2 = right_diff[1: -1] - (right_diff[2:] + right_diff[:-2]) / 2

		for row, data in enumerate(h_diff):
			cell(row, 4).value = data
		for row, data in enumerate(left_diff):
			cell(row, 5).value = data
		for row, data in enumerate(right_diff):
			cell(row, 6).value = data
		for row, data in enumerate(left_diff2):
			cell(row + 1 , 7).value = data
		for row, data in enumerate(right_diff2):
			cell(row + 1 , 8).value = data

	graph("road_graph").plot(index, left, "左轨", index, right, "右轨")
	graph("road_relative_graph").plot(index, left_diff, "左轨", index, right_diff, "右轨")
	graph("road_horizon_graph").plot(index, h_diff, "轨道水平")
	graph("road_concave_graph").plot(index[1:-1], left_diff2, "左轨", index[1:-1], right_diff2, "右轨")

calc_measure_data()
sheet_measure.fit_content()