%% Something wrong with the poles from RZIP at the moment script not working
%
% [+MEQ MatlabEQuilibrium Toolbox+] Swiss Plasma Center EPFL Lausanne 2022. All rights reserved.

%% Few simple comparison of RZIP dynamics with FGE
% Compute the free evolution ( contant feed-forward inputs) of RZIP linear state-space system A,B,C,D
% Project the RZIP evolution into eigenspace to see the evelotion of each eigenstate independently 
% Compute the growth rate from the max(eig(A))
% Estimate the grow rate alpha from time evolution of most unstable eigenstate xe~ exp(alpha*t) 
% Estimate the grow rate alpha from time evolution of plasma axis position zA ~ exp(alpha*t)  

clear all
close all
clc

%% Comparison of RZIP and fge

path_meq = '/home/carpanes/meq_GIT';
path_RZIP = '/home/carpanes/rzip'; 

addpath(genpath(path_meq));
addpath(genpath(path_RZIP));

add_rzip_paths

%% Set simulation parameters
shot = 57796;
t = 0.35;
tint = 0.01;
simtime = t:5e-5:t+tint;
nt = numel(simtime);


%% Run RZIP
[LFGS,~,LYFGS] = fgs('TCV', shot, t, 'insrc', 'liu');
[sys,sysinfo] = make_rzip_model('v',256,'fgs',LFGS,LYFGS,'TCV',1:19,shot,t);

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

X0 = zeros(sysinfo.nx,1);
%X0(sysinfo.ia(19)) = 10;
%Xe0 = V\X0;

simtime1 = t:4e-5:t+0.08;
nt1 = numel(simtime1);
U = zeros(nt1,  sysinfo.na);
U(:,19) = 1; % Step perrturbation in G coil
[Y,T,X] = lsim(sys, U , simtime1, X0 );

% Fetch RZIP output into LY structure for direct comparison
LYRZIP = struct();
LYRZIP.t = T;
LYRZIP.zA = (X(:,sysinfo.indZIp) + sysinfo.xeq(sysinfo.indZIp))./(X(:,sysinfo.indIp) + sysinfo.xeq(sysinfo.indIp) );
LYRZIP.rA = (X(:,sysinfo.indRIp) + sysinfo.xeq(sysinfo.indRIp))./(X(:,sysinfo.indIp) + sysinfo.xeq(sysinfo.indIp) );
LYRZIP.Ip = (X(:,sysinfo.indIp) + sysinfo.xeq(sysinfo.indIp));
LYRZIP.Ia =  (X(:,sysinfo.inda) + repmat(sysinfo.xeq(sysinfo.inda)', nt1,1 ));
LYRZIP.Iv = (X(:,sysinfo.indv) + repmat(sysinfo.xeq(sysinfo.indv)', nt1,1 ));
LYRZIP.Xe = zeros(nt1, sysinfo.nx);

for ii= 1:numel(T)
LYRZIP.Xe(ii,:) = V\[X(ii,:)]';
LYRZIP.Xeord(ii,:) = LYRZIP.Xe(ii,mask);
end

%% Solve FGE
[L,LX] = fge('tcv', shot,simtime, 'usepreconditioner', 1, 'algot', 'Euler_implicit');
LX.Va(17,:) = 1;
LY= fget(L,LX);

%% Compare plot between RZIP and FGE
figure
hold on 
plot(LYRZIP.t, LYRZIP.Ia, 'r')
plot(LY.t, LY.Ia, 'b-.');
title('RZIP vs FGE')
ylabel('I_a')
xlabel('t')

figure
hold on 
plot(LYRZIP.t, LYRZIP.Iv, 'r')
plot(LY.t, LY.Iu, 'b-.');
title('RZIP vs FGE')
ylabel('I_v')
xlabel('t')


figure
hold on 
plot(LYRZIP.t, LYRZIP.Ip, 'r')
plot(LY.t, LY.Ip, 'b-.');
title('RZIP vs FGE')
ylabel('I_p')
xlabel('t')
legend('RZIP', 'FGE')


figure
hold on 
plot(LYRZIP.t, abs(LYRZIP.zA-LYRZIP.zA(1)));
plot(LY.t, abs(LY.zA - LY.zA(1)));
legend('RZIP', 'FGE');
ylabel('zA [m]');
xlabel('t [s]');


figure
hold on 
plot(LYRZIP.t, LYRZIP.rA);
plot(LY.t, LY.rA);
legend('RZIP', 'FGE');
ylabel('rA');

%% Estimate growth rate from free evolution of zA
tshiftRZIP = LYRZIP.t - LYRZIP.t(1);
istart = 800;
iend = 1100;
[P] = polyfit(tshiftRZIP(istart:iend), log(abs((LYRZIP.zA(istart:iend)-LYRZIP.zA(1))/LYRZIP.zA(1))+1) , 1 );
figure
hold on 
plot(tshiftRZIP, log(abs((LYRZIP.zA-LYRZIP.zA(1))/LYRZIP.zA(1))+1) );
plot(tshiftRZIP(istart:iend), polyval(P,tshiftRZIP(istart:iend)), 'r.-');
legend(['RZIP max(eig)' num2str(max(real(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)')

%% Estimate the growth rate from only the eigenvector of the max related to the unstable pole
tshiftRZIP = LYRZIP.t - LYRZIP.t(1);
istart = 800;
iend = 1100;
[P] = polyfit(tshiftRZIP(istart:iend), log(abs(LYRZIP.Xeord(istart:iend,1))) , 1 );
figure
hold on 
plot(tshiftRZIP, log(abs(LYRZIP.Xeord(:,1))));
plot(tshiftRZIP(istart:iend), polyval(P,tshiftRZIP(istart:iend)));
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(Xe)')

%% Plot all eigenstate evolution
figure
hold on 
plot(LYRZIP.t, log(abs(LYRZIP.Xeord)), 'b')
plot(LYRZIP.t, log(abs(LYRZIP.Xeord(:,1))), 'r.')
plot(LYRZIP.t, log(abs(LYRZIP.Xeord(:,end-21:end))), 'g.')
plot(LYRZIP.t, log(abs(LYRZIP.Xeord(:,2))), 'y.')
plot(LYRZIP.t, log(abs(LYRZIP.Xeord(:,3))), 'y.')
title('Xe')
xlabel('t [s]')
ylabel('log(abs(Xe)')

%% 
figure
plot(sorted, '.')
ylabel('|eig|')









