# Optimal putting of a golf ball on a putting green
# MKS system of units

param g := 9.8; # acc due to gravity
param m := 0.01; # mass of a golf ball (in kilograms)

param x0 :=  1; # coords of starting point
param y0 :=  2;

param xn :=  1; # coords of ending point
param yn := -2;

param n := 50; # number of time discretizations

param mu; 

var T >= 0; # total time for the putt

var x{0..n};  # coordinates of the trajectory
var y{0..n}; 

# Here we define the elevation of the green
#var z   {i in 0..n} = flatness*x[i];
#var dzdx{i in 0..n} = flatness;
#var dzdy{i in 0..n} = 0;
#var d2zdx2{i in 0..n} = 0;
#var d2zdy2{i in 0..n} = 0;

# Here is a parabolic elevation function for the green.
# Choose either this function or the one above by uncommenting appropriately
#param flatness := 0.14;  # let mu := 0.1, 0.3; initial arc to the right
#param flatness := 0.14;  # let mu := 0.1; initial path straight
#var z   {i in 0..n} = flatness*(x[i]^2 + y[i]^2)/2;
#var dzdx{i in 0..n} = flatness*x[i];
#var dzdy{i in 0..n} = flatness*y[i];
#var d2zdx2{i in 0..n} = flatness;
#var d2zdy2{i in 0..n} = flatness;

#let mu := 0.1;
#var z   {i in 0..n} = 0.5/(1 + x[i]^2 + y[i]^2);
#var dzdx{i in 0..n} = -1.0*x[i]*z[i]^2;
#var dzdy{i in 0..n} = -1.0*y[i]*z[i]^2;
##var d2zdx2{i in 0..n} = -1.0*(z[i]^2-x[i]^2*2*z[i]^3);
##var d2zdy2{i in 0..n} = -1.0*(z[i]^2-y[i]^2*2*z[i]^3);

#let mu := 0.2;
var z   {i in 0..n} = -0.2*atan(x[i]);
var dzdx{i in 0..n} = -0.2/(1+x[i]^2);
var dzdy{i in 0..n} =  0;
#var d2zdx2{i in 0..n} = 0.4*x[i]/(1+x[i]^2)^2;
#var d2zdy2{i in 0..n} = 0;

#let mu := 0.4;
#var z   {i in 0..n} =  0.2*sin(1.5*(x[i]+y[i]));
#var dzdx{i in 0..n} =  0.2*1.5*cos(1.5*(x[i]+y[i]));
#var dzdy{i in 0..n} =  0.2*1.5*cos(1.5*(x[i]+y[i]));
#var d2zdx2{i in 0..n} = -0.1*1.5*1.5*sin(1.5*(x[i]+y[i]));
#var d2zdy2{i in 0..n} = -0.1*1.5*1.5*sin(1.5*(x[i]+y[i]));

#let mu := 0.3, 0.5;
#var z   {i in 0..n} =  0.5*log(x[i]^2+y[i]^2+1);
#var dzdx{i in 0..n} =  x[i]/(x[i]^2+y[i]^2+1);
#var dzdy{i in 0..n} =  y[i]/(x[i]^2+y[i]^2+1);
#var d2zdx2{i in 0..n} = -0.1*1.5*1.5*sin(1.5*(x[i]+y[i]));
#var d2zdy2{i in 0..n} = -0.1*1.5*1.5*sin(1.5*(x[i]+y[i]));

# The velocity vector.  v[i] denotes the derivative at the midpoint of
# the interval i(T/n) to (i+1)(T/n).
var vx{i in 0..n-1} = (x[i+1]-x[i])*n/T; 
var vy{i in 0..n-1} = (y[i+1]-y[i])*n/T; 
var vz{i in 0..n-1} = (z[i+1]-z[i])*n/T; 

# The acceleration vector.  a[i] denotes the accel at the midpoint of
# the interval (i-0.5)(T/n) to (i+0.5)(T/n), i.e. at i(T/n).
var ax{i in 1..n-1} = (vx[i]-vx[i-1])*n/T;
var ay{i in 1..n-1} = (vy[i]-vy[i-1])*n/T;
var az{i in 1..n-1} = (vz[i]-vz[i-1])*n/T;

var Nz{i in 1..n-1} = m*(
			g - ax[i]*dzdx[i] - ay[i]*dzdy[i] + az[i]
#			g + d2zdx2[i]*vx[i] + d2zdy2[i]*vy[i]
			)
                     /(dzdx[i]^2 + dzdy[i]^2 + 1);
var Nx{i in 1..n-1} = -dzdx[i]*Nz[i]; 
var Ny{i in 1..n-1} = -dzdy[i]*Nz[i]; 
#var Nmag{i in 1..n-1} = sqrt(Nx[i]^2 + Ny[i]^2 + Nz[i]^2);
var Nmag{i in 1..n-1} = m*(
			  g - ax[i]*dzdx[i] - ay[i]*dzdy[i] + az[i]
#			  g + d2zdx2[i]*vx[i] + d2zdy2[i]*vy[i]
			  )
		     /sqrt(dzdx[i]^2 + dzdy[i]^2 + 1);

var vx_avg{i in 1..n-1} = (vx[i]+vx[i-1])/2;
var vy_avg{i in 1..n-1} = (vy[i]+vy[i-1])/2;
var vz_avg{i in 1..n-1} = (vz[i]+vz[i-1])/2;
var speed{i in 1..n-1} = sqrt(vx_avg[i]^2 + vy_avg[i]^2 + vz_avg[i]^2);

var Frx{i in 1..n-1} = -mu*Nmag[i]*vx_avg[i]/speed[i];
var Fry{i in 1..n-1} = -mu*Nmag[i]*vy_avg[i]/speed[i];
var Frz{i in 1..n-1} = -mu*Nmag[i]*vz_avg[i]/speed[i];

#minimize finalspeed: (sqrt(0.001 + vx[n-1]^2 + vy[n-1]^2) - 0.25)^2;
minimize finalspeed: vx[n-1]^2 + vy[n-1]^2;

subject to newt_x {i in 1..n-1}: ax[i] = (Nx[i] + Frx[i])/m;
subject to newt_y {i in 1..n-1}: ay[i] = (Ny[i] + Fry[i])/m;
#subject to newt_z {i in 1..n-1}: az[i] = (Nz[i] + Frz[i] - m*g)/m;

subject to xinit: x[0] = x0;
subject to yinit: y[0] = y0;

subject to xfinal: x[n] = xn;
subject to yfinal: y[n] = yn;

let T := 5;

#param r := sqrt(x0^2 + y0^2);
#param th0 := atan(y0/x0);
#param thn := atan(yn/xn);
#display th0, thn, r;
#let {i in 0..n} x[i] := r*cos((i/n)*thn + (1-i/n)*th0);
#let {i in 0..n} y[i] := r*sin((i/n)*thn + (1-i/n)*th0);
let {i in 0..n} x[i] := (i/n)*xn + (1-i/n)*x0;
let {i in 0..n} y[i] := (i/n)*yn + (1-i/n)*y0;

let mu := 0.2;
solve;

let mu := 0.2;
solve;

display x,y,speed;

display T;

printf: "const puttPath = [\n" > putt_data.js;
printf {i in 0..n}: "  [%10f, %10f, %10f],\n", x[i], z[i]+0.04, y[i] >> putt_data.js;
printf: "  [%10f, %10f, %10f],\n", x[n], z[n]+0.04, y[n] >> putt_data.js;
printf: "  [%10f, %10f, %10f]\n", x[0], z[0]+0.04, y[0] >> putt_data.js;
printf: "];\n\n" >> putt_data.js;

printf: "const greenHeights = [\n" >> putt_data.js;
printf {i in -2.5..2.5 by 0.25, j in -2.5..2.5 by 0.25}: "  %10f,\n", -0.2*atan(j) >> putt_data.js;
printf: "];\n\n" >> putt_data.js;

printf: "const greenInfo = {\n" >> putt_data.js;
printf: "  xDimension: 21,\n" >> putt_data.js;
printf: "  zDimension: 21,\n" >> putt_data.js;
printf: "  xSpacing: 0.25,\n" >> putt_data.js;
printf: "  zSpacing: 0.25,\n" >> putt_data.js;
printf: "  xOffset: -2.5,\n" >> putt_data.js;
printf: "  zOffset: -2.5\n" >> putt_data.js;
printf: "};\n" >> putt_data.js;

printf: "ax, ax - (Nx+Frx)/m | ";
printf: "ay, ay - (Ny+Fry)/m | ";
printf: "az, az - (Nz+Frz-mg)/m \n";
printf {i in 1..n-1}: "%10f %10f | %10f %10f | %10f %10f \n", 
	ax[i], ax[i] - (Nx[i]+Frx[i])/m,
	ay[i], ay[i] - (Ny[i]+Fry[i])/m,
	az[i], az[i] - (Nz[i]+Frz[i]-m*g)/m;

display x0-x[0], y0-y[0], xn-x[n], yn-y[n];