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

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


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


(*Maps between different representations of so(3)*)
Clear[checkMap];
checkMap[A_List]:={A[[3,2]],A[[1,3]],A[[2,1]]}
Clear[hatMap];
hatMap[a_List]:={{0,-a[[3]],a[[2]]},{a[[3]],0,-a[[1]]},{-a[[2]],a[[1]],0}};

(*Definition of the B map*)
inertiaVals={1,2,3};
Inertia=DiagonalMatrix[inertiaVals];
InertiaInv=DiagonalMatrix[1/inertiaVals];
BMap[Mat_]:=hatMap[InertiaInv . checkMap[Mat]]

(*Initial state*)
W0=hatMap[RandomReal[{-1,1},3]];


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


(* ::Text:: *)
(*Notice that the functions use the value the value for h which is defined at the moment. During the time integration we are going to adjust the value o h accordingly.*)


Clear[resIsoSpectralMat];
resIsoSpectralMat[Wk_List,WIn_List] :=With[{W=ArrayReshape[WIn,{3,3}]},
	With[{Bw=hatMap[InertiaInv . checkMap[W]]},With[{R=Wk-(W + h/2*(W . Bw-Bw . W) - h^2/4*(Bw . W . Bw))},
		1/2 Tr[Transpose[R] . R]]]]

Clear[gradIsoSpectralMat];
gradIsoSpectralMat[Wk_List,WIn_List] := With[{W=ArrayReshape[WIn,{3,3}]},
	With[{Bw=hatMap[InertiaInv . checkMap[W]]},With[{R=Wk-(W + h/2*(W . Bw-Bw . W) - h^2/4*(Bw . W . Bw))},
		Flatten[Table[With[{Ones=ReplacePart[ConstantArray[0,Dimensions[Wk]],{i,j}->1]},
		With[{dR=-(Ones + h/2*((Ones . Bw + W . (hatMap[InertiaInv . checkMap[Ones]]))-((hatMap[InertiaInv . checkMap[Ones]]) . W+Bw . Ones)) - h^2/4 (((hatMap[InertiaInv . checkMap[Ones]]) . W + Bw . Ones) . Bw + (Bw . W) . Ones) )},
		1/2 (Tr[Transpose[R] . dR]+Tr[Transpose[dR] . R])]],{i,1,3},{j,1,3}]
		]]]]


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


(*Timestep*)
h=0.01;

(*Arrays and initial conditions setup*)
nsteps=2048;
resHist=ConstantArray[0,nsteps+1];
wHist=ConstantArray[0,{nsteps+1,3,3}];
wHist[[1]]=W0;
Wk = wHist[[1]];

(*Simulation loop*)
Monitor[
For[istep=1,istep<=nsteps,istep++,
  (*Implicit solver*)
  optRes=FindMinimum[{en=resIsoSpectralMat[Wk,W]},{W,Wk},
	Gradient :> gradIsoSpectralMat[Wk,W],
	Method->"QuasiNewton",MaxIterations->1000,
	AccuracyGoal->24];

  resHist[[istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  Wk =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  
  wHist[[istep+1]]=Wk;
  
],{istep}]


(*Residual plot for the implicit solver*)
ListPlot[resHist,PlotRange->Full]


(*Energy conservation*)
ensRigidBody=Table[-(1/2)Tr[BMap[wHist[[k]]] . wHist[[k]]]+1/2 Tr[BMap[wHist[[1]]] . wHist[[1]]],{k,1,Length[wHist]}];
fig=ListPlot[ensRigidBody-ensRigidBody[[1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in energy",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)","\[Delta]E"},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]


(*Eigenvalues*)
evalsRigidBody=Transpose[Table[Im[Eigenvalues[wh]],{wh,wHist}]];
fig=ListPlot[evalsRigidBody-evalsRigidBody[[;;,1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in eigenvalues",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)",""},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]


(*Casimirs*)
CasimirsRigidBody=Transpose[Table[{Tr[wh . wh]},{wh,wHist}]];
fig=ListPlot[CasimirsRigidBody-CasimirsRigidBody[[;;,1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in Casimir function",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)",""},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]


(* ::Section::Closed:: *)
(*2-stage 2nd order (Qin & Zhang Meth22)*)


(*Timestep. Here we divide the timestep by 2 since b1 = b2 = 1/2*)
h=0.01/2;

(*Arrays and initial conditions setup*)
nsteps=2048;
resHist=ConstantArray[0,nsteps+1];
wHist=ConstantArray[0,{nsteps+1,3,3}];
wHist[[1]]=W0;
Wk = wHist[[1]];

(*Simulation loop*)
Monitor[
For[istep=1,istep<=nsteps,istep++,
  (*Implicit solver*) 	
  optRes=FindMinimum[{en=resIsoSpectralMat[Wk,W]},{W,Wk},
	Gradient :> gradIsoSpectralMat[Wk,W],
	Method->"QuasiNewton",MaxIterations->1000,
	AccuracyGoal->24];

  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  WkHalf =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  
  (*Implicit solver*) 
  optRes=FindMinimum[{en=resIsoSpectralMat[WkHalf,W]},{W,WkHalf},
	Gradient :> gradIsoSpectralMat[WkHalf,W],
	Method->"QuasiNewton",MaxIterations->1000,
	AccuracyGoal->24];
	
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)	  
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  Wk =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  
  wHist[[istep+1]]=Wk;
  If[istep==nsteps, h=2*h]
],{istep}]


(*Energy conservation*)
ensRigidBody=Table[-(1/2)Tr[BMap[wHist[[k]]] . wHist[[k]]]+1/2 Tr[BMap[wHist[[1]]] . wHist[[1]]],{k,1,Length[wHist]}];
fig=ListPlot[ensRigidBody-ensRigidBody[[1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in energy",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)","\[Delta]E"},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]


(*Eigenvalues*)
evalsRigidBody=Transpose[Table[Im[Eigenvalues[wh]],{wh,wHist}]];
fig=ListPlot[evalsRigidBody-evalsRigidBody[[;;,1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in eigenvalues",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)",""},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]


(*Casimirs*)
CasimirsRigidBody=Transpose[Table[{Tr[wh . wh]},{wh,wHist}]];
fig=ListPlot[CasimirsRigidBody-CasimirsRigidBody[[;;,1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in Casimir function",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)",""},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]


(* ::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*)
h = 0.01;
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=2048;
wHist=ConstantArray[0,{nsteps+1,3,3}];
wHist[[1]]=W0;
resHist=ConstantArray[0,{7,nsteps+1}];
Wk = wHist[[1]];

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

  resHist[[1,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));
  
  h=h2;
  (*Implicit solver*) 
  optRes=FindMinimum[{en=resIsoSpectralMat[WAux,W]},{W,WAux},
	Gradient :> gradIsoSpectralMat[WAux,W],
	Method->"QuasiNewton",MaxIterations->1000,
	AccuracyGoal->24];

  resHist[[2,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));

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

  resHist[[3,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));

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

  resHist[[4,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));

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

  resHist[[5,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));

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

  resHist[[6,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  WAux =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));

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

  resHist[[7,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = hatMap[InertiaInv . checkMap[WRes]];
  Wk =(WRes - h/2*(WRes . Bw-Bw . WRes) - h^2/4*(Bw . WRes . Bw));

  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}]


(*Energy conservation*)
ensRigidBody=Table[-(1/2)Tr[BMap[wHist[[k]]] . wHist[[k]]]+1/2 Tr[BMap[wHist[[1]]] . wHist[[1]]],{k,1,Length[wHist]}];
fig=ListPlot[ensRigidBody-ensRigidBody[[1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in energy",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)","\[Delta]E"},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]


(*Eigenvalues*)
evalsRigidBody=Transpose[Table[Im[Eigenvalues[wh]],{wh,wHist}]];
fig=ListPlot[evalsRigidBody-evalsRigidBody[[;;,1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in eigenvalues",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)",""},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]


(*Casimirs*)
CasimirsRigidBody=Transpose[Table[{Tr[wh . wh]},{wh,wHist}]];
fig=ListPlot[CasimirsRigidBody-CasimirsRigidBody[[;;,1]],
	Joined->True,PlotRange->Full,
	PlotLabel->"Change in Casimir function",
	AxesLabel->{"\!\(\*SubscriptBox[\(t\), \(n\)]\)",""},
	ImageSize->Large,
	BaseStyle->{FontSize->18}]
