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

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


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


(*Definition of the B map*)
BMap[Mat_List]:={
{0,Mat[[1,2]],0,-Mat[[1,4]]},
{-Mat[[2,1]],0,Mat[[2,3]],0},
{0,-Mat[[3,2]],0,Mat[[3,4]]},
{Mat[[4,1]],0,-Mat[[4,3]],0}};

(*Hamiltonian function and extended Hamiltonian function *)
H[Mat_]:=2*Tr[Mat.Mat];
extH[Mat_]:=-(1/2)*Tr[Transpose[Mat].BMap[Mat]]+H[Mat];

(*Initial state*)
W0={{-1,-1,0,1},{-1,1,1,0},{0,1,-1,-1},{1,0,-1,1}};


(* ::Subsection::Closed:: *)
(*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,{4,4}]},
	With[{Bw=BMap[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,{4,4}]},
	With[{Bw=BMap[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 . (BMap[Ones]))-((BMap[Ones]) . W+Bw . Ones)) - h^2/4 (((BMap[Ones]) . W + Bw . Ones) . Bw + (Bw . W) . Ones) )},
		1/2 (Tr[Transpose[R] . dR]+Tr[Transpose[dR] . R])]],{i,1,4},{j,1,4}]
		]]]]				


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


(*Timestep*)
h=0.1;

(*Arrays and initial conditions setup*)
nsteps=1024;
resHist=ConstantArray[0,nsteps+1];
wHist=ConstantArray[0,{nsteps+1,4,4}];
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,
	WorkingPrecision->32,
	AccuracyGoal->32];

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


(*Conservation Hamiltonian*)
ListPlot[Table[H[wHist[[k]]]-H[wHist[[1]]],{k,1,Length[wHist]}]]


(*Conservation extended Hamiltonian*)
ListPlot[Table[extH[wHist[[k]]]-extH[wHist[[1]]],{k,1,Length[wHist]}]]


(*Isospectrality*)
Ev=Transpose[Table[Sort[N[Eigenvalues[wHist[[k]]]]],{k,1,Length[wHist]}]];
ListPlot[
{Ev[[1]]-Ev[[1,1]],
Ev[[2]]-Ev[[2,1]],
Ev[[3]]-Ev[[3,1]],
Ev[[4]]-Ev[[4,1]]},Joined->True]


(* ::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.1;
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=1024;
wHist=ConstantArray[0,{nsteps+1,4,4}];
wHist[[1]]=W0;
resHist=ConstantArray[0,{7,nsteps+1}];
Wk = wHist[[1]];

(*Simultion 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,
	WorkingPrecision->32,
	AccuracyGoal->32];
  
  resHist[[1,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = BMap[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,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[2,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];

  (*Explicit step*)  
  Bw = BMap[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,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[3,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];

  (*Explicit step*)    
  Bw = BMap[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,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[4,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];

  (*Explicit step*)  
  Bw = BMap[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,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[5,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = BMap[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,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[6,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  (*Explicit step*)
  Bw = BMap[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,
	WorkingPrecision->32,
	AccuracyGoal->32];

  resHist[[7,istep+1]] = optRes[[1]];
  WRes = N[W/.optRes[[2]]];
  
  Bw = BMap[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}]


(*Conservation Hamiltonian.*)
ListPlot[Table[H[wHist[[k]]]-H[wHist[[1]]],{k,1,Length[wHist]}]]


(*Conservation extended Hamiltonian*)
ListPlot[Table[extH[wHist[[k]]]-extH[wHist[[1]]],{k,1,Length[wHist]}]]


(*Isospectrality*)


Ev=Transpose[Table[Sort[N[Eigenvalues[wHist[[k]]]]],{k,1,Length[wHist]}]];
ListPlot[
{Ev[[1]]-Ev[[1,1]],
Ev[[2]]-Ev[[2,1]],
Ev[[3]]-Ev[[3,1]],
Ev[[4]]-Ev[[4,1]]},Joined->True]
