using JuMP, Ipopt

tol_fun=1e-11
mu=1e2

mod = Model(with_optimizer(Ipopt.Optimizer, max_iter=3000,tol=tol_fun,acceptable_tol=tol_fun,print_level=5))

# Parameters

N=100
hor_pos=1.0;
pi=3.141592;

Ldim=1;
m=86;
alphaval=10

Tmax=260;
rpmmax=240;
Dev=8.5;
omegamax=rpmmax/60;
c = 0.22;
g=9.81;
Fmax=2*pi*Tmax/Dev;
Vmax=Dev*omegamax;
tau=sqrt(m*Ldim/Fmax)
vmax=Vmax*tau/Ldim;
beta=m*g/Fmax
rho=1.225;
A=1;
dtmin=0.0001;
dtmax=1;

C=rho*c*A*Ldim/(2*m)
theta0=0
thetaN=pi/4
Rmin=0.2
R0=10

###################################
# Initial guess
###################################

r_old=ones(1,N)
theta_old=ones(1,N)
rdot_old=ones(1,N)
thetadot_old=ones(1,N)
Fr_old=ones(1,N)
Ftheta_old=ones(1,N)
dt_old=0.01

###################################
# Compile variables
###################################

@variable(mod, Fr[i=1:N], start=Fr_old[i])
@variable(mod, r[i=1:N]>=Rmin, start=r_old[i])
@variable(mod, rdot[i=1:N], start=rdot_old[i])
@variable(mod, theta0<=theta[i=1:N]<=thetaN, start=theta_old[i])
@variable(mod, thetadot[i=1:N]>=0, start=thetadot_old[i])

@variable(mod, dtmin<=dt<=dtmax , start = dt_old[1])

@NLexpression(mod, t_f, dt*N)

@NLexpression(mod, n[i=1:N],   r[i]-1 )
@NLexpression(mod, s[i=1:N],   theta[i] )

@NLexpression(mod, alpha[i=1:N],  tan(pi/180*alphaval) )
@NLexpression(mod, alphadash[i=1:N],  0)

@NLexpression(mod, dsdr[i=1:N], 0)
@NLexpression(mod, dsdtheta[i=1:N], 1)
@NLexpression(mod, dndr[i=1:N], 1)
@NLexpression(mod, dndtheta[i=1:N], 0)

@NLexpression(mod, dfdr[i=1:N], alpha[i]*dndr[i]+alphadash[i]*dsdr[i]*n[i])
@NLexpression(mod, dfdtheta[i=1:N], alpha[i]*dndtheta[i]+alphadash[i]*dsdtheta[i]*n[i])

@NLexpression(mod, v[i=1:N],  sqrt(rdot[i]^2*(1+tan(pi/180*alphaval)^2)+r[i]^2*thetadot[i]^2))
@NLexpression(mod, Ftheta[i=1:N], 0)

@NLexpression(mod, Frproj[i=1:N],  Ftheta[i]*rdot[i]/v[i] + cos(pi/180*alphaval)*Fr[i]*r[i]*thetadot[i]/v[i])

@NLexpression(mod, Fthetaproj[i=1:N], Ftheta[i]*r[i]*thetadot[i]/v[i] - 1/cos(pi/180*alphaval)*Fr[i]*rdot[i]/v[i])


@NLexpression(mod, ODE1[i=2:N],  (theta[i]-theta[i-1])/dt-thetadot[i] )
@NLexpression(mod, ODE2[i=2:N],  ((r[i]-r[i-1])/dt-rdot[i]) )

@NLexpression(mod, ODE3[i=3:N],  ((rdot[i]-rdot[i-1])/dt-r[i]*thetadot[i]^2) - (Frproj[i]-cos(pi/180*alphaval)*beta*dfdr[i]/sqrt(1+dfdr[i]^2+1/r[i]^2*dfdtheta[i]^2)) )

@NLexpression(mod, ODE4[i=3:N],  (r[i]*(thetadot[i]-thetadot[i-1])/dt + 2*rdot[i]*thetadot[i] ) - (Fthetaproj[i]-beta*1/r[i]*dfdtheta[i]/sqrt(1+dfdr[i]^2+1/r[i]^2*dfdtheta[i]^2) -C*thetadot[i]^2*r[i]^2 ) )


###################################
# Solution
###################################


@NLexpression(mod, con, (r[1]-R0)^2+(theta[1]-theta0)^2+(theta[N]-thetaN)^2+ (thetadot[2])^2+(rdot[2])^2+sum(ODE1[i]^2+ODE2[i]^2 for i=2:N)+sum(ODE3[i]^2+ODE4[i]^2 for i=3:N))

@NLexpression(mod, my_obj, t_f + mu*(con) )

@NLobjective(mod, Min, my_obj )

optimize!(mod)

println(sqrt((objective_value(mod)-value(t_f))/mu)/(N))

######################

writedlm(string("data_passive/Fmax"),Fmax)
writedlm(string("data_passive/r_",Int64(round(alphaval))),value.(r))
writedlm(string("data_passive/theta_",Int64(round(alphaval))),value.(theta))