BeginPackage["harm`"]

harmxy::usage =
   "harmxy[omx0, omy0, qx0, qy0, omx, omy, en]\n
solves equations {H == en, Grad W == la Grad H}\n
where W = 1/2 (px^2/omx0 + py^2/omy0 + omx0(qx - qx0)^2 + \n
      omy0(qy - qy0)^2),\n
H = 1/2(px^2 + py^2 + omx^2 qx^2 + omy^2 qy^2),\n
and finds the solution with minimal W.\n
Returnes {px, py, qx, qy, la, W}"

harm0::usage =
   "harm0[a, x0, en]\n
solves equations {H == en, Grad W == la Grad H}\n
where W = 1/2 (a1 (x1-x10)^2 + a2 (x2-x20)^2 + ... + aN (xN-xN0)^2),\n
H = 1/2(x1^2 + x2^2 + ... + xN^2),\n
and finds the solution with minimal W.\n
Returnes {{x1, x2, ..., xN}, la, W}"

Begin["`private`"]

harmxy[omx0_, omy0_, qx0_, qy0_, omx_, omy_, en_]:=
(
(* Scaling to the problem with
H = 1/2(x[1]^2 + x[2]^2 + x[3]^2 + x[4]^2)
W = 1/2(a[1]x[1]^2 + a[2]x[2]^2 + a[3]x[3]^2 + a[4]x[4]^2) *)

a={1/omx0, 1/omy0, omx0/omx^2, omy0/omy^2};
x0={0, 0, qx0 omx, qy0 omy};

{{px, py, qx, qy}, la, wig}=harm0[a,x0,en];

(* Returning to the original scale *)
{qx, qy}={qx/omx, qy/omy};

{px, py, qx, qy, la, wig}
)


harm0[a_List, x0_List, en_]:=
(
nn=Length[a]; 
nn1=Length[x0]; 
If[nn1=!=nn,Print["harm0: error: a and x0 have different lengths - ",
   nn," and ",nn1,"! Aborting..."];Abort[]];
If[en<0,Print["harm0: error: E<0! E = ",en,". Aborting..."];Abort[]];
(* N = 1 *)
If[nn==1,
   {a1,x01}={a[[1]],x0[[1]]};
   sq=Sqrt[2 en];
   If[x01==0,
      la=a1;
      x={sq};
      w=a1 en,
      la=a1 (1-x01/sq);
      x={a1 x01/(a1-la)};
      wig=a1/2 (la x01/(a1-la))^2
     ];
   Return[{x,la,wig}],
(* N > 1 *)
   {nsort,asort,x0sort}=Transpose[Sort[Transpose[{Range[nn],a,x0}],#1[[2]]<#2[[2]]&]];
   {a1,a2}=asort[[{1,2}]];
   {n1,x01}=First/@{nsort,x0sort};
   If[a1==a2,Print["harm0: warning: Special case of a1 = a2! The minimum may be degenerate."]];
(* Special case *)
   If[x01==0,
      fmax=1/2 Sum[(asort[[i]] x0sort[[i]]/(asort[[i]]-a1))^2,{i,2,nn}];
      If[en>=fmax,
         la=a1;
         x=Table[If[i==n1,Sqrt[2(en-fmax)],a[[i]] x0[[i]]/(a[[i]]-a1)],{i,nn}];
         wig=a1(en-fmax)+
            1/2 Sum[asort[[i]] (a1 x0sort[[i]]/(asort[[i]]-a1))^2,{i,2,nn}];
         Return[{x,la,wig}];
        ]
     ];
(* Generic case *)
(*
   pol=Numerator[Together[
      1/2 Sum[(a[[i]] x0[[i]]/(a[[i]]-lambda))^2,{i,nn}]-en]];
   lan=Chop[lambda/.Solve[pol==0]];
   mlan=Length[lan];
   la=10.^99;
   Do[la1=lan[[n]];
      If[Im[la1]==0,If[la1<la,la=la1]]
      ,{n,mlan}];
*)
   f=1/2 Sum[(a[[i]] x0[[i]]/(a[[i]]-lambda))^2,{i,nn}];
   {lastart,lamin,lamax}={a1-10.^-1,a1-10.^1,a1};
   la=lambda/.FindRoot[f==en,{lambda,a1-10.^-6},MaxIterations->300];
   x=Table[a[[i]] x0[[i]]/(a[[i]]-la),{i,nn}];
   wig=1/2 Sum[a[[i]] (la x0[[i]]/(a[[i]]-la))^2,{i,nn}];
   Return[{x,la,wig}];
];
)

End[]

EndPackage[]