# Outils Python pour le calcul de trajectoire atmospherique
# en phase propulse ou balistique.
#
#
#  Formation IAS 2018
#
# Janvier - Fevrier 2018

import numpy as np
import matplotlib.pyplot as plt

class traj():

	def __init__(self):
		#Position initiale---------------
		self.x=0.
		self.y=0.
		self.theta=np.pi/2.-0.1				 # Direction initiale (rad)
		u0=0.
		self.u=u0*np.cos(self.theta)
		self.v=u0*np.sin(self.theta)
		#Geometrie:----------------------
		self.D=5.452						 # Diametre principal (m)
		self.H=38.5							 # Hauteur (m)
		self.De=2.15						 # Diametre sortie tuyere (m)
		self.V0=0.8*self.H*np.pi*self.D**2/4 # Volume reservoir (m3)
		self.m0=10000.						 # Masse utile (Kg)
		#Aerodynamique-------------------
		self.rho=1.24 
		self.Sd=np.pi*self.D**2/4.
		self.Sl=self.H*self.D
		self.Cd=0.30                         # Coeff trainee
		self.Cl=0.00    					 # Coeff portance
		self.alpha=0.						 # incidence
		self.wind=0.
		self.pa=1.013e5						 # pression atm au sol
		#Propulsion----------------------
		self.Se=np.pi*self.De**2/4.
		self.Me=5.0							 # Mach sortie tuyere
		self.Ti=3500.						 # Temp Combustion (K)
		self.pi=115.*1.e5					 # Pression Combustion (Pa)
		self.gamma=1.27
		self.r=555.
		self.gg=1.+0.5*(self.gamma-1.)*self.Me**2
		self.ggg=-self.gamma/(self.gamma-1)
		self.Te=self.Ti*(self.gg)**(-1)
		self.pe=self.pi*self.gg**self.ggg
		self.rhoe=self.pe/(self.r*self.Te)
		self.rho_c=self.rhoe*self.gg**(1./(self.gamma-1.))
		self.cs=np.sqrt(self.gamma*self.r*self.Te)
		self.ue=self.Me*self.cs
		self.q=self.rhoe*self.ue*self.Se	 # Debit mass sortie (kg/s)
		self.m=self.m0+self.V0*self.rho_c
		#Gravite
		self.g0=9.81
		#Integration temporelle----------
		self.dt=1e-1

	def launch(self):
		self.time=0.
		self.count=0
		self.fic1=file("Trajectoire.dat",'w')
		self.update_p()
		if (self.m*self.g0>=self.p*np.sin(self.theta)):
			print "Low Thrust"
			print "Poussee: "+str(self.p)
			print "Poid   : "+str(self.m*self.g0)
		else:
			print('Poussee(kN) = '+str(self.p/1000.))
			print('ISP (s) = '+str(self.p/(self.q*self.g0)))
			print('Duree Propulsion (s): '+str(self.rho_c*self.V0/(self.q)))
			print('Masse Carburant (kg): '+str(self.V0*self.rho_c))
			print('Masse Totale (kg): '+str(self.m))
			raw_input("Launch ? ")
		while (self.y>=0. and self.y<1e5):
			self.time+=self.dt
			self.count+=1
			self.update_atm()
			self.update_m()
			self.update_p()
			self.update_L()
			self.update_D()
			self.update_u()
			dx=self.u*self.dt
			dy=self.v*self.dt
			self.x+=dx
			self.y+=dy
			if (np.abs(self.v)<=0.1 and self.time!=0.):
				self.theta=-np.pi/2.
				Vlim=np.sqrt(self.m*self.g0/(0.5*self.rho*self.Sd*self.Cd))
				raw_input("Apogee : "+str(self.y/1000.)+" km ....continue ?")
				raw_input("Vitesse limite de chute : "+str(Vlim)+" m/s ...continue ?")
			else:
				self.theta=np.arctan2(self.v,self.u)
			#self.plotpos()
			self.monitoring()
			self.write()
			if self.count>1e6:
				break
		self.fic1.close()

	def write(self):
			self.fic1.write('%.4f %.4f %.4f %.4f %.4f %.4f %.4f\n' %
				(self.time,self.x,self.y,self.u,self.v,self.alpha,self.V0))


	def plot(self):
			data=np.loadtxt('Trajectoire.dat',dtype='f')
			p=np.where(data[:,-1]>0)
			fig=plt.figure()
			fig.add_subplot(511)
			plt.plot(data[:,0],data[:,2])
			plt.ylabel('Altitude')
			plt.plot(data[p,0],data[p,2],'>r')
			fig.add_subplot(512)
			plt.plot(data[:,0],data[:,3])
			plt.ylabel('Ux')
			fig.add_subplot(513)
			plt.plot(data[:,0],data[:,4])
			plt.ylabel('Uy')
			fig.add_subplot(514)
			plt.plot(data[:,0],data[:,5])
			plt.ylabel('$\\alpha$')
			fig.add_subplot(515)
			plt.plot(data[:,0],data[:,6])
			plt.ylabel('Reservoir')
			plt.show()
	
	def plotpos(self):
			data=np.loadtxt('Trajectoire.dat',dtype='f')
			xmin=data[:,1].min()
			xmax=data[:,1].max()
			ymin=data[:,2].min()
			ymax=data[:,2].max()
			fig=plt.figure()
			frame=-1
			for time in data[:,0]:
				frame+=1
				plt.plot(data[frame,1],data[frame,2],'ok')
				plt.xlim(xmin,xmax);plt.ylim(ymin,ymax)
				plt.savefig(str(frame)+'.png')
				plt.close()

	def monitoring(self):
		if np.mod(int(100*self.time),100.)==0.:
		#	print self.time,self.y,self.v,self.theta,self.alpha,self.D,self.p
			print self.p/self.g0,self.m


	def update_m(self):
		self.update_reserv()
		self.m=self.m0+self.V0*self.rho_c
		self.J=self.m*self.H**2/12.

	def	update_atm(self):
		self.pa=1.013e5*(1.-0.0065*self.y/288.15)**5.255
		self.rho=352.995*(1.-2.25577e-5*self.y)**5.25516
		self.rho=self.rho/(288.15-0.0065*self.y)
		self.g0=3.9858e14/(6.371e6+self.y)**2

	def update_reserv(self):
		if self.V0>0:
			self.V0-=self.dt*self.q/self.rho_c
			self.seg='.r'
		else:
			self.V0=0.
			self.seg='.k'
			self.q=0.
			self.ue=0.
			self.pe=self.pa

	def update_p(self):
		self.p=self.Se*(self.rhoe*self.ue**2+(self.pe-self.pa))

	def update_L(self):
		self.L=0.5*self.rho*(self.u**2+self.v**2)*self.Sl*self.Cl

	def update_D(self):
		self.D=0.5*self.rho*(self.u**2+self.v**2)*self.Sd*self.Cd

	def update_u(self):
		ca=np.cos(self.theta)
		sa=np.sin(self.theta)
		self.u=self.m*self.u+self.dt*(self.p*ca-self.D*ca-self.L*sa)
		self.v=self.m*self.v+self.dt*(self.p*sa-self.D*sa+self.L*ca-self.m*self.g0)
		self.u=self.u/self.m
		self.v=self.v/self.m

t=traj()
t.launch()
t.plot()
plt.show()
t.plotpos()