%% Calculate system nonlinearity measure
%
% [+MEQ MatlabEQuilibrium Toolbox+] Swiss Plasma Center EPFL Lausanne 2022. All rights reserved.

shot = 63783;
tstart = 0.3;
tend = tstart +  0.005;
dt = 1e-5;
simtime = tstart:dt:tend;


% Linearize fge about equilibrium point
[M, init0, yind, Uind] = fgelin('fge', shot, tstart, 'observables', {'rA', 'zA'} , 'ifequilinit', 1, 'insrc', 'liu');

[V,D] = eig(M.A);
[sorted,mask] = sort(diag(D), 'descend');

D = D(mask);
V = V(:,mask);

%%
X0 = 1e-7*V(:,1);

sys = ss(M.A, M.B, M.C, M.D);

% Simulate linear system free evolution
nt = numel(simtime);
U = zeros(nt,  init0.L.G.na + init0.L.ng + 1 + numel(init0.xdot));
[Y,T,X] = lsim(sys, U , simtime, X0);

%%

% Simulate FGE
% Generate structure, and compute the preconditioner
[L,LX] = fge('tcv', shot, simtime(1), 'preconditioner', 1, 'debug', 0 );

% Set feedforward traces and initital state for equilibrium point
LX.Va = repmat(L.G.Ra.*LX.Ia, 1, numel(simtime)); % Va = Ra*Ia
LX.Iu = init0.LYfgs.Iu; % Iu = 0
LX.Fe = L.Mee*[init0.LYfgs.Ia;init0.LYfgs.Iu] + L.Mey*init0.LYfgs.Iy(:); % Use Iy computed with Iu = 0
Fe0 = LX.Fe;
LX.Fe = LX.Fe + X0 ;
LX.ag = repmat(LX.ag, 1, numel(simtime));
LX.Ip = repmat(LX.Ip, 1, numel(simtime));
LX.rBt = repmat(LX.rBt, 1, numel(simtime));
LX.t = simtime;
% Run fge
LY = fget(L,LX);


%%
% Plot all obervables in observables list
for ii = fieldnames(yind).'
  hf = figure;
  ax = axes(hf);
  hold(ax, 'on')
  plot(ax, LY.t, LY.(ii{:})- LY.(ii{:})(:,1) + Y(1,yind.(ii{:})),'r');
  plot(ax, T, Y(:,yind.(ii{:})),'b');
  xlabel(ax,'t [s]')
  ylabel(ax, ['\delta' ii{:}])
  p(1) = plot(ax,NaN,NaN,'r'); % trick to add single line in legend
  p(2) = plot(ax,NaN,NaN,'b'); % Trck to add single line in legend
  legend(p, {'FGE', 'fgelin'})
  title(['shot' num2str(shot)])
end

%%
figure 
hold on 
plot(X0)
plot(LY.t, LY.zA)


%%
Xeord = [];
Xe = [];
LY.Xe = [];
for ii= 1:numel(T)
Xe(ii,:) = V\[X(ii,:)]';
LY.Xe(:,ii) = V\[LY.Fe(:,ii)-Fe0];
end

desort = 1:20:500;
hf = figure
ax = axes(hf)
hold(ax,'on')
plot(ax,LY.t(desort), LY.Xe(1:end,desort), 'r')
plot(ax,T(desort),Xe(desort,1:end), 'b*')
ylabel(ax,'Xe')
xlabel(ax,'t')
title(ax, 'Evolution of eigenstates')
p(1) = plot(ax,NaN,NaN,'b*'); % trick to add single line in legend
p(2) = plot(ax,NaN,NaN,'r'); % Trck to add single line in legend
legend(p, 'linFge', 'FGE')



figure
plot(T,X, 'b')
plot(LY.t, LY.Fe-Fe0, 'r*')
legend('fgelin', 'FGE')


%%
tshiftRZIP = LY.t - LY.t(1);
istart = 400;
iend = 500;
[P] = polyfit(tshiftRZIP(istart:iend), log(abs((LY.zA(istart:iend)-LY.zA(1))/LY.zA(1))+1) , 1 );
figure
hold on 
plot(tshiftRZIP, log(abs((LY.zA-LY.zA(1))/LY.zA(1))+1) );
plot(tshiftRZIP(istart:iend), polyval(P,tshiftRZIP(istart:iend)), 'r.-');
legend(['RZIP max(eig)' num2str(max(eig(sys.A)))], ['fit: growth rate '  num2str(P(1)) 'Hz'])
title(['Data vs exponential fit RZIP ' num2str(shot)])
xlabel('t [s]')
ylabel('log(|zA-zA0|/|zA0| + 1)')


