from math import pi

import matplotlib.pyplot as plt
import numpy as np

from das.legacy.Numerics import Propagator
from das.legacy.OptimalControlProblem import *

problem = OptimalControlProblem(name='OrbitTransfer')

# Define the states and rates
problem.State(name='r', expr='vr', units='m')
problem.State(name='theta', expr='vtheta/r', units='m/s')
problem.State(name='vr', expr='(vtheta^2)/r-mu/(r^2)+(T/(m))*1.3*sin(u1)', units='m/s')
problem.State(name='vtheta', expr='-(vr*vtheta)/r+(T/(m))*1.3*sin(u2)', units='s')
problem.State(name='m', expr='delta_m')
# Define the independent variable
problem.Independent(name='t', units='s')

# Define constants
problem.Constant(name='m0', value=1, units='?')
problem.Constant(name='T', value=0.02, units='?')
problem.Constant(name='delta_m', value=-0.01, units='?')
problem.Constant(name='mu', value=1, units='')

# Define constants of motion
problem.ConstantOfMotion(name='Hamiltonian', expr='lambda_r*vr + lambda_theta*(vtheta/r) + lambda_vr*((vtheta^2)/r-mu/(r^2)+(T/(m))*1.3*sin(u1)) + lambda_vtheta*(-(vr*vtheta)/r+(T/(m))*1.3*sin(u2)) + lambda_m*delta_m', units='J')
problem.ConstantOfMotion(name='Control_u1', expr='1.3*T*lambda_vr*cos(u1)/m', units='J')
problem.ConstantOfMotion(name='Control_u2', expr='1.3*T*lambda_vtheta*cos(u2)/m', units='J')


# Define the control variables
problem.Control(name='u1', units='?')
problem.Control(name='u2', units='?')

# Define the cost to be minimized
problem.Cost(expr='0')

# Derive the necessary conditions for optimality
problem.necessaryConditions(controlType='dae', verbose=2)

# Add a propagator
# problem.Propagator(name='Main', propagator=Propagator(name='SIEp', MechanicalSystem=problem.system[0], ConstantsOfMotion=[str(problem.constantsofmotionfilled[0]), str(problem.constantsofmotionfilled[1]), str(problem.constantsofmotionfilled[2]), str(problem.constantsofmotionfilled[3]), str(problem.constantsofmotionfilled[4]), str(problem.constantsofmotionfilled[5])]))
problem.Propagator(name='Main', propagator=Propagator(name='ForwardEuler', MechanicalSystem=problem.system[0], ConstantsOfMotion=[]))
problem.Propagator(name='Main', propagator=Propagator(name='SIEp', MechanicalSystem=problem.system[0], ConstantsOfMotion=[]))
problem.Propagator(name='Main', propagator=Propagator(name='SIEp', MechanicalSystem=problem.system[0], ConstantsOfMotion=[str(problem.constantsofmotionfilled[0])]))
problem.propagators[0].genStepFunc()
problem.propagators[1].genStepFunc()

problem.propagators[0].bake(verbose=1)
problem.propagators[1].bake(verbose=1)
initial = [1.00000000, 0.00000000, 0.00000000, 1.00000000, 1.00000000, -0.03045064, 0.99953627, -5.77629993, -0.00000000, 0.11991710, -7.27665053, 2.47410610, 0, 0]
C0 = problem.propagators[0].solve(tSpan=np.linspace(start=0, stop=20, num=200), x0=initial)
C1 = problem.propagators[1].solve(tSpan=np.linspace(start=0, stop=20, num=200), x0=initial)
# C2 = problem.propagators[2].solve(tSpan=np.linspace(start=0, stop=3.32, num=200), x0=initial)
# C0 = problem.propagators[0].solve(tSpan=np.linspace(start=0, stop=3.32, num=500), x0=[])
# pickle.dump(C0, open("SE_500_100.p", "wb"))

# problem.Propagator(name='Main', propagator=Propagator(name='ForwardEuler', MechanicalSystem=problem.system[0], ConstantsOfMotion=[str(problem.constantsofmotionfilled[0])]))
#
# problem.propagators[1].genStepFunc()
#
# problem.propagators[1].bake(verbose=1)
#
# C1 = problem.propagators[1].solve(tSpan=np.linspace(start=0, stop=626.8731811642932*1.0, num=500), x0=[6418000, -2.122494903, 0.580321976, 2065, -0.090757121, -1.570796327, 0.218484568553597, 0.838619365091245, -0.000000281, 0.423817876, -1.284076759, -0.000076985, -0.048655246, -0.053908033, 0, 0])
#
#
X = np.linspace(start=0, stop=2*pi, num=64)
# plt.plot(C0.x[0,0]*np.sin(X), C0.x[0,0]*np.cos(X), 'g-', linewidth=3)
# plt.plot(C0.x[-1,0]*np.sin(X), C0.x[-1,0]*np.cos(X), 'g-', linewidth=3)

plt.plot(C0.x[:,0]*np.sin(C0.x[:,1]),C0.x[:,0]*np.cos(C0.x[:,1]),'g-')
plt.plot(C1.x[:,0]*np.sin(C1.x[:,1]),C1.x[:,0]*np.cos(C1.x[:,1]),'r-')
# plt.plot(C2.x[:,0]*np.sin(C2.x[:,1]),C2.x[:,0]*np.cos(C2.x[:,1]),'b-')
plt.legend(['FE','SE'], loc=4)
# plt.plot(C1.x[:,3],C1.x[:,0]-6378000)
#
# pickle.dump(C1, open("SE2000_95.p", "wb" ) )
plt.show()