(* 
utility: ratio of the reward acquired, minus effort expended, divided by total time,
alpha: food increment per successful trial,
ns: number of successful trials,
Ts: duration of a trial,
cs: effort cost of a trial,
nL: number of licks during a harvest,
TL: duration of a lick,
cL: cost of a lick. 
*)

F[ns_,nl_,TL_,Ts_]:=(alpha betas ns(1-1/(1+ betaL nl))-nl cL -ns^2 cs)/(nl TL + ns Ts);

(* Fig. 2A plot utility as a function of lick number for various amounts of work period trials *)
(* low licking cost *)
Show[
Table[Plot[F[myns,nL,0.2,1] /. {alpha -> 20, betas-> 0.5, betaL -> 0.3, cL-> .5, cs-> .5}, {nL, 0, 20}, PlotRange -> All], {myns, 10}], Frame-> True, FrameLabel->{"lick","Capture Rate"}]

(* Fig. 2A plot utility as a function of lick number for various amounts of work period trials *)
(* high licking cost *)
Show[
Table[Plot[F[myns,nL,0.2,1] /. {alpha -> 20, betas-> 0.5, betaL -> 0.3, cL-> 2.5, cs-> .5}, {nL, 0, 20}, PlotRange -> All], {myns, 10}], Frame-> True, FrameLabel->{"lick","Capture Rate"}]

(* Find optimal number of licks and work period trials assuming a constant cost of licking.
In these calculations we assume cost of licking does not depent on lick duration *)

Simplify[D[F[ns,nL,TL,Ts],ns]]

Simplify[D[F[ns,nL,TL,Ts],nL]]

nsoptimal[]=Simplify[Flatten[Solve[D[F[ns,nL,TL,Ts],ns]==0, {ns}]][[2]]]

nLoptimal[]=Simplify[Flatten[Solve[D[F[ns,nL,TL,Ts],nL] == 0, {nL}]]][[2]]

Findoptimal[myalpha_,myTs_,myTL_,mycs_,mycL_]:=
Module[{},
nLtemp = 1;For[i=1,i<=50, i++,nstemp = ns/. nsoptimal[] /. {alpha->myalpha, betas-> 0.5, betaL -> 0.3,Ts-> myTs, TL-> myTL,cs->mycs, cL-> mycL, nL-> nLtemp};
nLtemp= nL /. nLoptimal[] /. {alpha->myalpha, betas-> 0.5, betaL -> 0.3,Ts-> myTs, TL-> myTL,cs->mycs, cL->  mycL, ns-> nstemp};
];
{nstemp,nLtemp}]

Graphics[{
Yellow,
Line[Table[{lickcost, Findoptimal[20,1,0.2,0.5,lickcost][[1]]}, {lickcost, 0.2, 3, 0.1}]],
Green,
Line[Table[{lickcost, Findoptimal[25,1,0.2,0.5,lickcost][[1]]}, {lickcost, 0.2, 3, 0.1}]],
Red,
Line[Table[{lickcost, Findoptimal[30,1,0.2,0.5,lickcost][[1]]}, {lickcost, 0.2, 3, 0.1}]]},
Frame -> True, FrameLabel -> {"lick cost","optimal saccade trials"}]

(* plot the maximum capture rate as a function of number of work period saccades *)
Graphics[{
Yellow,
Line[Table[
{lickcost,
Module[{numsacs,numlicks},
{numsacs,numlicks}=Findoptimal[20,1,0.2,0.5,lickcost];
F[numsacs, numlicks, 0.2,1]/. {alpha-> 20, betas-> 0.5, betaL -> 0.3, cL-> lickcost, cs-> .5}
]},{lickcost, 0.2,3,0.1}]],
Green,
Line[Table[
{lickcost,
Module[{numsacs,numlicks},
{numsacs,numlicks}=Findoptimal[25,1,0.2,0.5,lickcost];
F[numsacs, numlicks, 0.2,1]/. {alpha-> 25, betas-> 0.5, betaL -> 0.3, cL-> lickcost, cs-> .5}
]},{lickcost, 0.2,3,0.1}]],
Red,
Line[Table[
{lickcost,
Module[{numsacs,numlicks},
{numsacs,numlicks}=Findoptimal[30,1,0.2,0.5,lickcost];
F[numsacs, numlicks, 0.2,1]/. {alpha-> 30, betas-> 0.5, betaL -> 0.3, cL-> lickcost, cs-> .5}
]},{lickcost, 0.2,3,0.1}]]
}, Frame -> True, FrameLabel -> {"lick cost","Maximum capture rate"}]

(* -------------------------------------------- 
add the ability to change movement vigor 
------------------------------
*)
F[ns_,nl_,Ts_, TL_]:=(alpha betas ns(1-1/(1+ betaL nl))-nl (dL^2/TL+ kL TL) -ns^2 cs)/(nl TL + ns Ts);

Simplify[D[F[ns,nL, Ts,TL], ns]]

nsoptimal[]=FullSimplify[Solve[D[F[ns,nL,Ts,TL],ns]==0, {ns}]][[2]]

nLoptimal[]=FullSimplify[Solve[D[F[ns,nL,Ts,TL],nL]==0, {nL}]][[2]]

Simplify[D[F[ns,nL,Ts,TL],TL]]

TLoptimal[]=FullSimplify[Solve[D[F[ns,nL,Ts,TL],TL]==0, {TL}]][[2]]

Findoptimal[myalpha_,myTs_,mycs_, mydL_]:=
Module[{},
nLtemp = 1;
TLtemp=0.2;
For[i=1,i<=50, i++,nstemp = ns/. nsoptimal[] /. {alpha->myalpha, betas-> 0.5, betaL -> 0.3,Ts-> myTs, kL-> 1, dL-> mydL,TL-> TLtemp,cs->mycs, nL-> nLtemp};
nLtemp= nL /. nLoptimal[] /. {alpha->myalpha, betas-> 0.5, betaL -> 0.3,Ts-> myTs, kL-> 1, dL-> mydL,TL-> TLtemp,cs->mycs, ns-> nstemp};
TLtemp=TL/. TLoptimal[] /. {alpha->myalpha, betas-> 0.5, betaL -> 0.3,Ts-> myTs,kL-> 1, dL-> mydL,cs->mycs, ns-> nstemp, nL-> nLtemp};
If[Im[TL]!= 0 , TLtemp=0.01] ;
];
{nstemp,nLtemp, TLtemp}]

Graphics[{
Yellow,
Line[Table[
{tubedist, Findoptimal[20,1,0.5,tubedist][[1]]}, {tubedist, 0.1,0.3,0.02}]],
Green,
Line[Table[
{tubedist, Findoptimal[22.5,1,0.5,tubedist][[1]]}, {tubedist,  0.1,0.3,0.02}]],
Red,
Line[Table[
{tubedist, Findoptimal[25,1,0.5,tubedist][[1]]}, {tubedist,  0.1,0.3,0.02}]]
},
Frame-> True, AspectRatio->1,FrameLabel -> {"tube dist","optimal saccade trials"}]

Graphics[{
Yellow,
Line[Table[
{tubedist, Findoptimal[20,1,0.5,tubedist][[2]]}, {tubedist, 0.1,0.3,0.02}]],
Green,
Line[Table[
{tubedist, Findoptimal[22.5,1,0.5,tubedist][[2]]}, {tubedist, 0.1,0.3,0.02}]],
Red,
Line[Table[
{tubedist, Findoptimal[25,1,0.5,tubedist][[2]]}, {tubedist, 0.1,0.3,0.02}]]
},Frame-> True, AspectRatio->1,FrameLabel -> {"tube distance","optimal lick number"}]

(* optimal lick duration *)
Graphics[{
Yellow,
Line[Table[
{tubedist, Findoptimal[20,1,0.5,tubedist][[3]]}, {tubedist, 0.1,0.3,0.02}]],
Green,
Line[Table[
{tubedist, Findoptimal[22.5,1,0.5,tubedist][[3]]}, {tubedist, 0.1,0.3,0.02}]],
Red,
Line[Table[
{tubedist, Findoptimal[25,1,0.5,tubedist][[3]]}, {tubedist, 0.1,0.3,0.02}]]
},Frame-> True, AspectRatio->1,FrameLabel -> {"tube distance","optimal lick duration (sec)"}]

Graphics[
Table[
Line[Table[{TL, tubedist^2/TL +  TL}, {TL, .05, 0.6, 0.01}]],
 {tubedist, 0.1, 0.3, 0.05}],
 Frame-> True, AspectRatio->1,FrameLabel -> {"Lick duration (sec)","Metabolic cost (a.u.)"}]

(* metabolic optimal lick duration as a function of tube distance *)Solve[D[dL^2/TL + KL TL, TL]==0, {TL}]