(* Calculations for arbitrary atom,
                       including p-orbitals *)
(* Step c - minimization of energy functional *)
(*              by equating partial derivatives to zero *)
(*  Finding initial guess *)

npart=4; (* Number of electrons *)
ndig0=64; (* Accuracy *)
ndig=16; (* Accuracy of FindRoot *)
zmax=10;
zmin=1/2;
zstep0=1/10;
ca1=1/20;cz1=11/10; (* If a changes less than ca1, then increase Zstep cz1 times *)
ca2=1/10;cz2=1/2; (* If a changes more than ca2, then decrease Zstep by cz2 *)
zshift=1/100; (* Shift to imaginary direction to avoid singularities on the real axis *)
ashift=(1+I)/100; (* small perturbing to avoid degeneracy *)

SetDirectory["~/math/boron"];
os=OpenWrite["guess"<>ToString[npart]<>".dat",PageWidth->Infinity];
"Minimization of energy functional for an atom with npart electrons">>watchc;
Definition[npart]>>>watchc;
Clear[cf,mat,l,n,enlas,ens,A,a0,a,pri,x];
Get["cf"<>ToString[npart]<>".dat"];
Get["mat"<>ToString[npart]<>".dat"];
nhalf=Floor[npart/2]; (* Number of bi-particle orbitals *)

(* Summation of contributions from different permutations *)
perm=Permutations[Range[npart]];
lperm=Length[perm];
{mean,meanne,meant,meanee}=
Sum[pl=perm[[l]];
   {l,pl}>>>watchc;
   cf@@pl mat@@pl, {l,lperm}];
Clear[cf,mat];

(* Expressing energy functional and its derivatives through variables a$n=A[n] *)
$ModuleNumber=1;
Do[A[n]=Module[{a},a],{n,npart}];
enla=(meant-meanne+lambda meanee)/mean;
Clear[mean,meanne,meant,meanee];
"---------- Computation of the energy functional completed.">>>watchc;
Do[enlas[n]=D[enla,A[n]],{n,npart}];
"---------- Computation of the derivatives completed.">>>watchc;

(* Initial guess at Z=Zmax *)
If[zmax!=10,
  Print["Warning: Initial guess at Z=Zmax may be inaccurate."]];
If[npart==2,(
  a0[1]=1.080;
  a0[2]=0.857;
 )];
If[npart==3,(
  a0[1]=1.027;
  a0[2]=0.882;
  a0[3]=0.228;
 )];
If[npart==4,(
  a0[1]=1.039;
  a0[2]=0.855;
  a0[3]=0.218+0.088 I;
  a0[4]=0.218-0.088 I;
 )];
If[npart==5,(
  a0[1]=1.046;
  a0[2]=0.850;
  a0[3]=0.207+0.082 I;
  a0[4]=0.207-0.082 I;
  a0[5]=0.386
 )];

Do[a0[n]={a0[n],a0[n]+(-1)^n ashift},{n,npart}];

z=N[zmax,ndig0];
zstep=N[zstep0,ndig0];

(* Printing operator for real numbers*)
(* Printing operator for complex numbers*)
pri[x_,n_]:=(
   xp=SetPrecision[x,n];
   (*xp=Chop[xp,Abs[xp]/10^n];*)
   OutputForm/@{"  \t", Re[xp]//FortranForm, " ", Im[xp]//FortranForm} );

cond=True;
While[cond, (* Start loop *)
(* Finding stationary points of the energy functional *)
   la=1/(z+zshift I);
   en=enla/.lambda->la;
   Do[ens[n]=enlas[n]/.lambda->la,{n,npart}];
   s=FindRoot[Release[Table[ens[n]==0,{n,npart}]],
     Release[Sequence@@Table[{A[n],a0[n]},{n,npart}]],
     WorkingPrecision->ndig,MaxIterations->99];
   en0=en/.s;
   Do[a[n]=A[n]/.s,{n,npart}];
   da=Max@Table[Abs [(a[n]-a0[n][[1]])/a[n]], {n,npart}];
   da1=Max@Table[a1d=a0[2 n-1][[1]]-a0[2 n][[1]];
                 a2d=a[2 n-1]-a[2 n];
                 Abs[(a1d-a2d)/a2d], {n,nhalf}];
   da=Max[da,da1];
   If[da<ca1,zstep=cz1 zstep];
   If[da>ca2,zstep=cz2 zstep];
   Do[a0[n]={a[n],a0[n][[1]]},{n,npart}];
(* Printing results *)
   line=Flatten[{pri[1/la,6], Table[pri[a[n],8],{n,npart}]}];
   line1=Flatten[{pri[1/la,4], pri[en0,4],Table[pri[a[n],3],{n,npart}]}];
   Write[ os, Sequence@@line];
   Print[Sequence@@line1];
   cond=(z>=zmin);
   z=z-zstep
]; (* End of loop *)
Close[os];
"---------- Completed.">>>watchc;