#!/usr/bin/env wolframscript
(* ::Package:: *)

(* ::Section::Closed:: *)
(*Setup*)


(* ::Subsection:: *)
(*General*)


Nmax = 7;
s=(Nmax-1)/2;

(*Basis elements*)
TMats = Table[Table[Table[
	(-1)^(s-m1)*Sqrt[2*l+1]*ThreeJSymbol[{s,-m1},{l,m},{s,m2}]
	,{m1,-s,s},{m2,-s,s}]
	,{m,-l,l}],{l,0,Nmax-1}];

(*Build Laplacian*)
TMatsTab = Flatten[TMats,1];
THat = Table[Flatten[TMatsTab[[i]]],{i,1,Nmax^2}];
LapEvals = DiagonalMatrix[
	Flatten[Table[Table[l*(l+1),{m,-l,l}],{l,0,Nmax-1}]]
	];
LapMat=Transpose[N[ConjugateTranspose[THat] . LapEvals . THat]];

(*Build inverse Laplacian*)
InvLapEvals=DiagonalMatrix[
	Flatten[Table[Table[Piecewise[{{1/(l*(l+1)),l!=0}}],{m,-l,l}],{l,0,Nmax-1}]]
];
InvLapMat=Transpose[N[ConjugateTranspose[THat] . InvLapEvals . THat]];
InvLaplaceOp[W_] := ArrayReshape[InvLapMat . Flatten[W],{Nmax,Nmax}];

(*Hamiltonian*)
H[Mat_]:=Tr[ConjugateTranspose[InvLaplaceOp[Mat]] . Mat];

(*Enstrophy*)
Ens[Mat_]:=Tr[ConjugateTranspose[Mat] . Mat]



(* ::Subsection:: *)
(*Implicit problem function and gradient*)


(* ::Text:: *)
(*Have to work with extended versions of the matrix since the optimization has problems with complex inputs*)


Clear[resIsoSpectralMat];
resIsoSpectralMat[WkExt_List,WInExt_List] :=With[{Wk=WkExt[[1;;Nmax]]+I*WkExt[[Nmax+1;;2*Nmax]]},With[{W=WInExt[[1;;Nmax]]+I*WInExt[[Nmax+1;;2*Nmax]]},
	With[{Bw=InvLaplaceOp[W]},With[{R=Wk-(W + h/2*(W . Bw-Bw . W) - h^2/4*(Bw . W . Bw))},
		Re[1/2 Tr[ConjugateTranspose[R] . R]]]]]]
		
GradInverseLaplacianRe=Table[With[{Ones=ReplacePart[ConstantArray[0,Dimensions[TMatsTab[[1]]]],{i,j}->1]},InvLaplaceOp[Ones]],{i,1,Nmax},{j,1,Nmax}];
GradInverseLaplacianIm=Table[With[{Ones=ReplacePart[ConstantArray[0,Dimensions[TMatsTab[[1]]]],{i,j}->1]},InvLaplaceOp[I*Ones]],{i,1,Nmax},{j,1,Nmax}];

Clear[gradIsoSpectralMat];
gradIsoSpectralMat[WkExt_List,WInExt_List] :=With[{Wk=WkExt[[1;;Nmax]]+I*WkExt[[Nmax+1;;2*Nmax]]}, With[{W=WInExt[[1;;Nmax]]+I*WInExt[[Nmax+1;;2*Nmax]]},
	With[{Bw=InvLaplaceOp[W]},With[{R=Wk-(W + h/2*(W . Bw-Bw . W) - h^2/4*(Bw . W . Bw))},
		Join[
		Flatten[Table[With[{Ones=ReplacePart[ConstantArray[0,Dimensions[Wk]],{i,j}->1],InvLap=GradInverseLaplacianRe[[i,j]]},
With[{dRre=-(Ones + h/2*((Ones . Bw + W . InvLap)-(InvLap . W+Bw . Ones)) - h^2/4 ((InvLap . W + Bw . Ones) . Bw + (Bw . W) . InvLap) )},
		1/2 (Tr[ConjugateTranspose[R] . dRre]+Tr[ConjugateTranspose[dRre] . R])]],{i,1,Nmax},{j,1,Nmax}]
		],
		Flatten[Table[With[{Ones=ReplacePart[ConstantArray[0,Dimensions[Wk]],{i,j}->1],InvLap=GradInverseLaplacianIm[[i,j]]},
With[{dRim=-(I*Ones + h/2*(((I*Ones) . Bw + W . InvLap)-(InvLap . W+Bw . (I*Ones))) - h^2/4 ((InvLap . W + Bw . (I*Ones)) . Bw + (Bw . W) . InvLap) )},
		1/2 (Tr[ConjugateTranspose[R] . dRim]+Tr[ConjugateTranspose[dRim] . R])]],{i,1,Nmax},{j,1,Nmax}]
		]
		]
		]]]]


(* ::Subsection:: *)
(*Initial conditions*)


(*Initial conditions from spherical harmonics coefficients for real signal*)

alm=Table[Join[{RandomReal[{-1,1}]},RandomComplex[{-1-I,1+I},Max[0,l]]],{l,0,Nmax-1}];
ylm=Table[With[{lalm=alm[[l+1]]},Join[Table[Conjugate[lalm[[-m+1]]]*If[Mod[m,2]==1,-1,1],{m,-l,-1}],lalm]],{l,0,Nmax-1}];
ylm[[1]]=0.;
ylmFlat=Flatten[ylm];

W0 = Sum[ylmFlat[[i]]*(I*TMatsTab[[i]]),{i,1,Length[TMatsTab]}];
W0Ext = Join[Re[W0],Im[W0]];


(* ::Section::Closed:: *)
(*Implicit midpoint method(1-stage 2-order Meth12)*)


(*Timestep. Here we consider the time scaling present in Modin & Viviani [15]*)
NormW0 = Re[Sqrt[Tr[ConjugateTranspose[W0].W0]]];
h=0.0025*Nmax^(3/2)*NormW0/(4*Sqrt[Pi]);

(*Arrays and initial conditions setup*)
nsteps=1024;
resHist=ConstantArray[0,nsteps+1];
wHist=ConstantArray[0,{nsteps+1,Nmax,Nmax}];
wHist[[1]]=W0/NormW0;
Wk = wHist[[1]];
WkExt = Join[Re[Wk],Im[Wk]];

(*Simulation loop*)
Monitor[
For[istep=1,istep<=nsteps,istep++,

  WkB=InvLaplaceOp[Wk];
  WkNextEuler=Wk+h/2*(WkB . Wk-Wk . WkB);
  WkExtEuler=Join[Re[WkNextEuler],Im[WkNextEuler]];
  
  (*Implicit solver*)
  optRes=FindMinimum[{en=resIsoSpectralMat[WkExt,WExt]},{WExt,WkExtEuler},
	Gradient :> Re[gradIsoSpectralMat[WkExt,WExt]],
	Method->"QuasiNewton",MaxIterations->1000,
	AccuracyGoal->24];
  
  resHist[[istep+1]] = optRes[[1]];
  WExtRes = N[WExt/.optRes[[2]]];
  WRes = WExtRes[[1;;Nmax]]+I*WExtRes[[Nmax+1;;2*Nmax]];
  
  (*Explicit step*)
  Bw = InvLaplaceOp[WRes];
  Wk =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  WkExt=Join[Re[Wk],Im[Wk]];
  wHist[[istep+1]]=Wk;
],{istep, en}]


Wcoeffs=Monitor[Chop[Table[Table[-Tr[Transpose[wHist[[k]]] . (I*TMatsTab[[i]])],{i,1,Length[TMatsTab]}],{k,1,Length[wHist]}]],{k}];

LaplaceInv0=Flatten[Join[{0},Table[Table[1/(l*(l+1)),{m,-l,l}],{l,1,Nmax-1}]]];
En=Table[Dot[Wcoeffs[[k]]*LaplaceInv0,Conjugate[Wcoeffs[[k]]]],{k,1,Length[Wcoeffs]}];
EnBands=Table[Table[With[{cc=Wcoeffs[[k]][[l^2+1;;(l+1)^2]]},Dot[cc*(1/(l(l+1))),Conjugate[cc]]],{l,1,Nmax-1}],{k,1,Length[Wcoeffs]}];


ListPlot[resHist,PlotRange->Full]


ListPlot[(En-En[[1]])/En[[1]],Joined->True,PlotRange->Full,
	PlotLabel->"Relative change in energy",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)","\[Delta]E"},
	ImageSize->Large,
	BaseStyle->{FontSize->16}]


ListPlot[Transpose[EnBands]-EnBands[[1]],Joined->True,PlotRange->Full,
	PlotLabel->"Change in energy per band (in (l,m) space)",
	PlotLegends->Table[Style["l="<>ToString[l],18],{l,1,Nmax}],
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)","\!\(\*SubscriptBox[\(\[Delta]E\), \(l\)]\)"},
	ImageSize->Large,
	BaseStyle->{FontSize->16}]


ListPlot[Transpose[EnBands],Joined->True,PlotRange->Full,
	PlotLabel->"Energy per band (in (l,m) space)",
	PlotLegends->Table[Style["l="<>ToString[l],18],{l,1,Nmax}],
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)","\!\(\*SubscriptBox[\(E\), \(l\)]\)"},
	ImageSize->Large,
	BaseStyle->{FontSize->16}]


EigenvaluesEuler = Chop[Transpose[Table[Sort[N[Im[Eigenvalues[wHist[[k]]]]]],{k,1,Length[wHist]}]]];
ListPlot[EigenvaluesEuler[[;;]]-EigenvaluesEuler[[;;,1]],Joined->True,PlotRange->Full,
	PlotLabel->"Eigenvalues",ImageSize->Large,
	ImageSize->Large,
	BaseStyle->{FontSize->16}]


(*Conservation Hamiltonian*)
ListPlot[Table[Re[H[wHist[[k]]]-H[wHist[[1]]]],{k,1,Length[wHist]}],
	Joined->True,PlotMarkers->False,
	PlotLabel->"Hamiltonian",ImageSize->Large]


(*Conservation Enstrophy*)
ListPlot[Table[Re[Ens[wHist[[k]]]-Ens[wHist[[1]]]],{k,1,Length[wHist]}],Joined->True,
	PlotLabel->"Enstrophy",
	Joined->True,PlotMarkers->Automatic,
	ImageSize->Large]


ListLogPlot[resHist,Joined->True,PlotRange->Full,ImageSize->Large]


(* ::Section::Closed:: *)
(*7-stage 4th order (symmetric and symplectic Meth74)*)


(*RK coefficients*)
b1=1.4944291445422252;
b2=\[Minus]2.3484252147655893;
b3=2.8982950506987870;
b4=1-2*b1-2*b2-2*b3;
b5=b3;
b6=b2;
b7=b1;


(*Timesteps*)
NormW0 = Re[Sqrt[Tr[ConjugateTranspose[W0].W0]]];
h=0.0025*Nmax^(3/2)*NormW0/(4*Sqrt[Pi]);
h1=h*b1;
h2=h*b2;
h3=h*b3;
h4=h*b4;
h5=h*b5;
h6=h*b6;
h7=h*b7;


(*Arrays and initial conditions setup*)
nsteps=512;
wHist=ConstantArray[0,{nsteps+1,Nmax,Nmax}];
wHist[[1]]=W0;
resHist=ConstantArray[0,{7,nsteps+1}];
Wk = wHist[[1]];
WkExt = Join[Re[Wk],Im[Wk]];

(*Simultion loop*)
Monitor[
For[istep=1,istep<=nsteps,istep++,
  h=h1;
  (*Implicit solver*)  
  optRes=FindMinimum[{en=resIsoSpectralMat[WkExt,WExt]},{WExt,WkExt},
	Gradient :> gradIsoSpectralMat[WkExt,WExt],
	Method->"QuasiNewton",MaxIterations->1000,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[1,istep+1]] = optRes[[1]];
  WExtRes = N[WExt/.optRes[[2]]];
  
  WRes = WExtRes[[1;;Nmax]]+I*WExtRes[[Nmax+1;;2*Nmax]];
  
  (*Explicit step*)
  Bw = InvLaplaceOp[WRes];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  WAuxExt = Join[Re[WAux],Im[WAux]];
  
  h=h2;
  (*Implicit solver*) 
  optRes=FindMinimum[{en=resIsoSpectralMat[WAuxExt,WExt]},{WExt,WAuxExt},
	Gradient :> gradIsoSpectralMat[WAuxExt,WExt],
	Method->"QuasiNewton",MaxIterations->1000,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[2,istep+1]] = optRes[[1]];
  WExtRes = N[WExt/.optRes[[2]]];
  
  WRes = WExtRes[[1;;Nmax]]+I*WExtRes[[Nmax+1;;2*Nmax]];
  
  (*Explicit step*)
  Bw = InvLaplaceOp[WRes];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  WAuxExt = Join[Re[WAux],Im[WAux]];

  h=h3;
  (*Implicit solver*) 
  optRes=FindMinimum[{en=resIsoSpectralMat[WAuxExt,WExt]},{WExt,WAuxExt},
	Gradient :> gradIsoSpectralMat[WAuxExt,WExt],
	Method->"QuasiNewton",MaxIterations->1000,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[2,istep+1]] = optRes[[1]];
  WExtRes = N[WExt/.optRes[[2]]];
  
  WRes = WExtRes[[1;;Nmax]]+I*WExtRes[[Nmax+1;;2*Nmax]];
  
  (*Explicit step*)
  Bw = InvLaplaceOp[WRes];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  WAuxExt = Join[Re[WAux],Im[WAux]];

  h=h4;
  (*Implicit solver*) 
  optRes=FindMinimum[{en=resIsoSpectralMat[WAuxExt,WExt]},{WExt,WAuxExt},
	Gradient :> gradIsoSpectralMat[WAuxExt,WExt],
	Method->"QuasiNewton",MaxIterations->1000,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[2,istep+1]] = optRes[[1]];
  WExtRes = N[WExt/.optRes[[2]]];
  
  WRes = WExtRes[[1;;Nmax]]+I*WExtRes[[Nmax+1;;2*Nmax]];
  
  (*Explicit step*)
  Bw = InvLaplaceOp[WRes];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  WAuxExt = Join[Re[WAux],Im[WAux]];

  h=h5;
  (*Implicit solver*) 
  optRes=FindMinimum[{en=resIsoSpectralMat[WAuxExt,WExt]},{WExt,WAuxExt},
	Gradient :> gradIsoSpectralMat[WAuxExt,WExt],
	Method->"QuasiNewton",MaxIterations->1000,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[2,istep+1]] = optRes[[1]];
  WExtRes = N[WExt/.optRes[[2]]];
  
  WRes = WExtRes[[1;;Nmax]]+I*WExtRes[[Nmax+1;;2*Nmax]];

  (*Explicit step*)
  Bw = InvLaplaceOp[WRes];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  WAuxExt = Join[Re[WAux],Im[WAux]];

  h=h6;
  (*Implicit solver*) 
  optRes=FindMinimum[{en=resIsoSpectralMat[WAuxExt,WExt]},{WExt,WAuxExt},
	Gradient :> gradIsoSpectralMat[WAuxExt,WExt],
	Method->"QuasiNewton",MaxIterations->1000,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[2,istep+1]] = optRes[[1]];
  WExtRes = N[WExt/.optRes[[2]]];
  
  WRes = WExtRes[[1;;Nmax]]+I*WExtRes[[Nmax+1;;2*Nmax]];
  
  (*Explicit step*)
  Bw = InvLaplaceOp[WRes];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  WAuxExt = Join[Re[WAux],Im[WAux]];

  h=h7;
  (*Implicit solver*) 
  optRes=FindMinimum[{en=resIsoSpectralMat[WAuxExt,WExt]},{WExt,WAuxExt},
	Gradient :> gradIsoSpectralMat[WAuxExt,WExt],
	Method->"QuasiNewton",MaxIterations->1000,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[2,istep+1]] = optRes[[1]];
  WExtRes = N[WExt/.optRes[[2]]];
  
  WRes = WExtRes[[1;;Nmax]]+I*WExtRes[[Nmax+1;;2*Nmax]];
  
  (*Explicit step*)
  Bw = InvLaplaceOp[WRes];
  Wk =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  WkExt = Join[Re[Wk],Im[Wk]];

  wHist[[istep+1]]=Wk;
  If[istep==nsteps, h=0.01]
],{istep,en}]


ListPlot[
{resHist[[1]],resHist[[2]],resHist[[3]],resHist[[4]],resHist[[5]],resHist[[6]],resHist[[7]]}
,Joined->True,PlotRange->All,PlotLegends->{1,2,3,4,5,6,7}]


Casimirs = Chop[Transpose[Table[Sort[N[Im[Eigenvalues[wHist[[k]]]]]],{k,1,Length[wHist]}]]];
ListPlot[Casimirs[[;;]]-Casimirs[[;;,1]],Joined->True,PlotRange->Full,PlotLabel->"Casimirs"]


(*Conservation Hamiltonian*)
ListPlot[Table[Re[H[wHist[[k]]]-H[wHist[[1]]]],{k,1,Length[wHist]}],Joined->True,PlotLabel->"Hamiltonian"]


(*Conservation Enstrophy*)
ListPlot[Table[Re[Ens[wHist[[k]]]-Ens[wHist[[1]]]],{k,1,Length[wHist]}],Joined->True,PlotLabel->"Enstrophy"]
