from math import *
import matplotlib.pyplot as plt

#Python is in rad

#Fx = f(k, Fz) = Fz * D * sin(C * atan[{Bk - E * [Bk - atan(Bk)]}])

class Wheel:
	def __init__(self):
		self.globalFrictionCoefficient = 1.0

		#Pacejka parameters

		#Lower value means that the curve will be flatter before the peak (and thus have a "wider" peak)
		#Value above 2 means a change of sign of the curve, meaning it has to be <= 2.0 for longitudal grip
		self.Cx = 1.65
		self.Cz = 2.8
		self.Cy = 1.3

		#Essentially a max grip, might as well be constant 1 and only change the global friction coefficient
		self.D = 0.5

		#Shapes the curve as an exponential, which means lower value will mean flatter curve
		#Compared to C it also moves the very beginning of the curve, rather than "only" affecting the area before the peak
		self.Ex = 0.9
		self.Ez = 0.98
		self.Ey = 0.95

		#These two numbers move the peak acceleration during slip, per "radian" (slip) and per Newton of load, respectively
		#Moving the peak "back", i.e. having the peak being at a higher slip rate, also causes the "fallback" (the return to 0 acceleration as the slip approaches zero) to be slower (and thus more smooth)

		#Higher a3 means later peak
		#Newtons per radians
		self.a3 = 55 * 1000 * 0.01

		#Higher a4 moves the peak backwards (earlier fast acceleration, and decreased fallback rate)
		#Fluctuates the total time to reach zero slip depending on the value, so can in some cases increase fallback rate
		#Newtons
		self.a4 = 7000.0 * 1.0

		#How much the global friction coefficient should change per newton of vertical load
		self.frictionCoefficientChange = -0.05 / 1000.0

		#A constant we add to the longitudal velocity to avoid division by zero
		self.longitudalVelocityOffset = 0.1

		self.radius = 0.25

		self.angular_velocity = 0.0
		self.linear_velocity = 0.0

	def getLongitudalSlip(self):
		return (self.radius * self.angular_velocity - self.linear_velocity) / (abs(self.linear_velocity) + self.longitudalVelocityOffset)

	def getGlobalFrictionCoefficient(self, Fz):
		return self.globalFrictionCoefficient + Fz * self.frictionCoefficientChange

	def getLongitudalForce(self, Fz, k):
		u = self.getGlobalFrictionCoefficient(Fz)
		D = self.D * u * Fz
		C = self.Cx

		BCD = self.a3 * sin(2.0 * atan(Fz / self.a4))
		B = BCD / (C * D)
		E = self.Ex

		return D * sin(C * atan(B * k - E * (B * k - atan(B * k))))


	def getVerticalMoment(self, Fz, k):
		u = self.getGlobalFrictionCoefficient(Fz)
		D = self.D * u * Fz
		C = self.Cz
		BCD = self.a3 * sin(2.0 * atan(Fz / self.a4))
		B = BCD / (C * D)

		E = self.Ez

		return D * sin(C * atan(B * k - E * (B * k - atan(B * k))))

	def getLateralForce(self, Fz, k):
		u = self.getGlobalFrictionCoefficient(Fz)
		D = self.D * u * Fz
		C = self.Cy
		BCD = self.a3 * sin(2.0 * atan(Fz / self.a4))
		B = BCD / (C * D)

		E = self.Ey

		return D * sin(C * atan(B * k - E * (B * k - atan(B * k))))

def plotXY(x, y, title, xlabel, ylabel):
	plt.clf()
	plt.plot(x,y)
	plt.margins(0.1)
	plt.xlabel(xlabel)
	plt.ylabel(ylabel)
	plt.title(title)
	plt.show()


#unsigned gravity acceleration
g = 9.82

#Total sprung mass in kilograms
carMass = 1120

simTime = 25.0
simTimeStep = 0.15

wheels = [Wheel(), Wheel(), Wheel(), Wheel()]

tyreLoad = carMass / len(wheels) * g
ang_acc = 0.0 / wheels[0].radius

printTime = -0.4

time = 0.0
timeSincePrint = printTime

for wheel in wheels:
	wheel.linear_velocity = 60.0
	wheel.angular_velocity = wheel.linear_velocity / wheel.radius
	wheel.linear_velocity = 0.0

maxAngVel = -2.0# / wheels[0].radius


timePlot = []
linVelPlot = []
angVelPlot = []
accInGPlot = []
tetsPlot = []
longSlipPlot = []

simIteration = 0
while time < simTime:
	force = 0.0
	avgLongSlip = 0.0
	avgAngVel = 0.0

	for wheel in wheels:
		if simIteration > 0:
			wheel.angular_velocity += ang_acc * simTimeStep
		if maxAngVel >= 0.0:
			wheel.angular_velocity = min(wheel.angular_velocity, maxAngVel)
		longSlip = wheel.getLongitudalSlip()
		longF = wheel.getLongitudalForce(tyreLoad, longSlip)
		force += longF
		avgLongSlip += longSlip / len(wheels)
		avgAngVel += wheel.angular_velocity / len(wheels)
	acc = force * simTimeStep / carMass

	longSlipPlot.append(avgLongSlip)
	timePlot.append(time)
	linVelPlot.append(wheels[0].linear_velocity)
	angVelPlot.append(avgAngVel)
	accInGPlot.append(force / (g * carMass))

	for wheel in wheels:
		wheel.linear_velocity += acc
	#if abs(wheels[0].angular_velocity) > 0.0 and abs(wheels[0].linear_velocity / (wheels[0].angular_velocity * wheels[0].radius) > 0.99):
	#	break

	timeSincePrint += simTimeStep
	if printTime >= 0.0 and timeSincePrint >= printTime:
		print "Time: " + str(time)
		print "Ang. Vel: " + str(avgAngVel) + " (" + str(avgAngVel * wheels[0].radius) + " m/s)"
		print "Lin. Vel: " + str(wheels[0].linear_velocity)
		print "Average Slip: " + str(avgLongSlip)
		print "Acceleration: " + str(force / (g * carMass)) + " g"
		print "----------------------------------"
		timeSincePrint = 0.0
	simIteration += 1
	time += simTimeStep

#plotXY(longSlipPlot, accInGPlot, "Acc-Slip", "Slip", "Acceleration")
#plotXY(angVelPlot, linVelPlot, "Linear-Angular Velocity", "Ang vel", "Lin vel")
#plotXY(timePlot, longSlipPlot, "Slip-Time", "Time", "Slip")
#plotXY(timePlot, angVelPlot, "Angular Velocity-Time", "Time", "Angular Velocity")
#plotXY(timePlot, linVelPlot, "Linear Velocity-Time", "Time", "Linear Velocity")
plotXY(timePlot, accInGPlot, "Acceleration-Time", "Time", "Acceleration (g)")



MzX = []
MzY = []
FyY = []
FxY = []
FxX = []

MzC = 0.25
i = -MzC

while i <= MzC:
	MzX.append(i / 3.1415926535 * 180.0)
	MzY.append(wheels[0].getVerticalMoment(tyreLoad, i))
	FyY.append(wheels[0].getLateralForce(tyreLoad, i))
	i += 0.01
FxC = 16.0
i = -FxC
while i <= FxC:
	longSlip = i
	FxX.append(longSlip)
	FxY.append(wheels[0].getLongitudalForce(tyreLoad, longSlip) / tyreLoad)
	i += 0.01

#plotXY(MzX, MzY, "Self-aligning Torque", "Lateral Slip", "Torque")
plotXY(MzX, FyY, "Lateral Force", "Lateral Slip", "Force")
plotXY(FxX, FxY, "Longitudal Force", "Longitudal Slip", "Force")