function  [cstate,Va,ctrlpar,ctdebug] = meqctrl(mode,cstate,ctrlpar,L,LY,varargin)
% [cstate,Va,ctrlpar,ctdebug] = meqctrl(mode,cstate,ctrlpar,L,LY,varargin)
%
% Generic magnetic controller for meq - implements a simple R,Z,Ip
% controller for anamak and provides extension options for other tokamaks.
% See tutorials/anamak_rzp_control.m for example of its use for anamak.
% Also serves as template for possible design of other controllers.
%
% Inputs:
%   mode: 'setup','init','step' or 'reset' (see later)
%   cstate: vector with the controller state at the previous time step.
%   ctrlpar: controller parameter structure as input. For details, inspect this function
%   L,LY: MEQ L and LY (equilibrium) structures
%   varargin: optional parameter-value pairs to pass into ctrlpar
%
% Outputs:
%   cstate: vector with the controller state at the next time step
%   Va is the controller output voltage (including feedforward part
%   ctrlpar are the controller parameters, returned in the 'setup' mode
%   ctdebug are optional debugging information.
%
% 
% MODES:
% 'setup' mode:
%   Performs controller design and returns ctrlpar. Calls a tokamak-specific function 
%   meqctrlp<tok> to load the controller parameters.
%   * First, a mutual decoupling and resistive compensation term of the form 
%     MR=Maa+Ra/s is placed in front of the open-loop plant.
%   * A vertical control loop (KzQ) is closed on the voltages directly. By default
%     the control direction is chosen to maximize the coupling to the
%     most unstable eigenmode, but custom directions can also be chosen.
%   * A coil current control loop (KIo) orthogonal to the vertical control loop 
%     is closed as well. This gives a stable plant with inputs being the
%     ortogonal coil currents and the vertical position reference.
%   * A radial and current control direction is chosen using a pseudoinverse
%     of the plant gain at low frequencies, to control rIp,Ip while not
%     perturbing zIp
%   * The rIp, Ip loops are closed with their own controllers (KrQ, KIQ)
%     
%   A note on nomenclature, the suffix 'Q' refers to [Piras, PhD thesis EPFL 2010] 
%   and denotes the 'common' mode (mnemonic: 'Q'ommon), to distinguish if 
%   from the symmetric mode 'S' for doublets.
%
% 'init' mode: Initializes the state-space matrices for the controller as 
%   persistent variables and returns the initial state and initial command voltage 
%   (called by fget in the initialization phase)
%
% 'step' mode: Executes controller step, receives previous state, parameters, 
%   and LY from current time step and runs one step of the controller, 
%   returning the control voltages and next controller state.
%   (Called by fget in the time step)
% 
% 'reset' mode: Clears the persistent variables
%
% Setting debug=true plots some transfer functions in the design process

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

persistent A B C D % controller matrices
persistent hp ho % for plotting
persistent TIo    % other useful matrices

switch mode
  case 'setup'
    % input parameters to ctrlpar
    for k = 1:2:numel(varargin)
      assert(ischar(varargin{k}),'parameter name must be as string')
      ctrlpar.(varargin{k}) = varargin{k+1};
    end
    doplot = (~isempty(ctrlpar) && isfield(ctrlpar,'debug') && ctrlpar.debug);
    
    % Call controller design function
    [KzQ,KrQ,KIQ,KIo] = controller_design(L,LY,doplot);
    
    %%
    ctpar_def = struct(...
      'dt',1e-4,...     % Controller sampling time
      'KzQ',KzQ,...     % Vertical position controller
      'KIo',KIo,...     % Plasma orthogonal coil current controler
      'KrQ',KrQ,...     % Radial position controller
      'KIQ',KIQ,...     % Plasma current controller
      'zQref',timeseries(LY.zIp/LY.Ip(1),0),... % z position reference   [m]
      'rQref',timeseries(LY.rIp/LY.Ip(1),0),... % r position reference   [m]
      'IQref',timeseries(LY.Ip(1)       ,0),... % Ip reference           [A]
      'Iaref',timeseries(LY.Ia'      ,LY.t),... % Coil current reference [A]
      'Vff'  ,timeseries(LY.Va',LY.t      ),... % Voltage feedforward    [V]
      'Maa',L.G.Maa,'Ra',L.G.Ra,...  % Coil mutuals and resistances
      'Kmain',1,... % global controller gain switch for testing
      'debug',0 ... % debug flag to display controller information
      );
    
    % replace fields with defaults if not defined
    for field = fieldnames(ctpar_def)'
      if ~isfield(ctrlpar,field{:}), ctrlpar.(field{:}) = ctpar_def.(field{:}); end
    end
    
    cstate = []; Va = []; 
    
  case 'init'
    hp = [];
    % initialize controllers and their state
    
    % construct controller
    KKd = build_position_controller(ctrlpar.Maa,ctrlpar.Ra,ctrlpar.KzQ,ctrlpar.KrQ,ctrlpar.KIQ,ctrlpar.KIo,ctrlpar.dt);
    KKd = ctrlpar.Kmain*KKd;
    ctdebug.KKd = KKd;
    A = KKd.A; B=KKd.B; C=KKd.C; D=KKd.D;
    TIo = ctrlpar.KIo.T;
    cstate = zeros(size(A,2),1); % controller state init
    Va = ctrlpar.Vff.Data(1,:)'; % init feedforward
  case 'step'
    % Step call
    
    % Interpolate references and feedforwards using ZOH
    izQref = ifloor(ctrlpar.zQref.Time,LY.t); % time index of zIp/Ip ref
    zQref  = ctrlpar.zQref.Data(izQref,:);
    
    irQref = ifloor(ctrlpar.rQref.Time,LY.t); % time index of rIp/Ip ref
    rQref  = ctrlpar.rQref.Data(irQref,:);
    
    iIQref = ifloor(ctrlpar.IQref.Time,LY.t); % time index of Ip ref
    IQref  = ctrlpar.IQref.Data(iIQref,:);
    
    iIaref = ifloor(ctrlpar.Iaref.Time,LY.t); % time index of Ia ref
    Iaref  = ctrlpar.Iaref.Data(iIaref,:)';
    
    iVff   = ifloor(ctrlpar.Vff.Time,  LY.t); % time index of Vff ref
    Vff    = ctrlpar.Vff.Data(iVff,:)';
    
    % Get plasma parameters (this is cheating, as formally a
    % reconstruction/estimation is needed
    
    yrQ = LY.rIp/LY.Ip; % centroid position
    yzQ = LY.zIp/LY.Ip; % centroid position
    ypQ = LY.Ip;        % Plasma current
    
    % control errors
    ezIp = (zQref-yzQ).*LY.Ip; % z error times Ip
    erIp = (rQref-yrQ).*LY.Ip; % r error times Ip
    eIp  = (IQref-ypQ);        % plasma current error
    eIa  = (Iaref - LY.Ia);    % coil current error
    
    err  =  [eIa;ezIp;erIp;eIp];
    % control law - state space form
    xc   = A*cstate + B*err;
    uc   = C*cstate + D*err;
    
    cstate = xc;
    
    % assign control output
    Va = uc + Vff; % feedback + feedforward term
    ctdebug = [];
    
    %% Optional plotting
    if ctrlpar.debug
      eIo  = TIo'*eIa;           % orthogonal coil current error
      if isempty(hp) || any(~ishandle(hp(:)))
            
        figure(1); clf;
        hax(1) = subplot(411);
        hax(2) = subplot(412);
        hax(3) = subplot(413);
        hax(4) = subplot(414);
        
        ax=hax(1);
        hp(:,1) = plot(ax,LY.t,yzQ,'ob',LY.t,zQref,'.k'); hold(ax,'on');
        title(ax,'z control')
        ax=hax(2);
        hp(:,2) = plot(ax,LY.t,yrQ,'ob',LY.t,rQref,'.k'); hold(ax,'on');
        title(ax,'r control')
        ax=hax(3);
        hp(:,3) = plot(ax,LY.t,ypQ,'ob',LY.t,IQref,'.k'); hold(ax,'on');
        title(ax,'Ip control')
        ax=hax(4);
        ho = plot(ax,LY.t,eIo,'.'); hold(ax,'on');
        title(ax,'I_{a\perp} error norm')
        drawnow;
      else
        adddatapoint = @(t,y,h) set(h,'XData',[get(h,'XData'),t],'YData',[get(h,'YData'),y]);
        adddatapoint(LY.t,yzQ  ,hp(1,1));
        adddatapoint(LY.t,zQref,hp(2,1));
        adddatapoint(LY.t,yrQ  ,hp(1,2));
        adddatapoint(LY.t,rQref,hp(2,2));
        adddatapoint(LY.t,ypQ  ,hp(1,3));
        adddatapoint(LY.t,IQref,hp(2,3));
        for io=1:numel(eIo)
          adddatapoint(LY.t,eIo(io)  ,ho(io,1));
        end
        if rem(LY.t,20*ctrlpar.dt)==0
          drawnow;
        end
      end
    end
    
  case 'reset'
    clear A B C D Iaref
  otherwise
    error('unknown mode %s',mode);
end
end

function  [KzQ,KrQ,KIQ,KIo] = controller_design(L,LY,doplot)
% Controller design function
%
% Returns structures with controller parameters
% for KzQ (vertical), KrQ (radial), KIQ(plasma current), 
% KIo (coil current orthogonal to z control).%
%
% All controllers are PI(I)D controllers with rolloff:
% K.Kp: proportional gain
% K.Ki: integral gain
% K.Kd: derivative gain
% K.Ki2: Optional second integral gain
% K.tauro: Rolloff pole time constant
% K.T : control direction
%
% K(s) = Kp*(1+Ki/s)(1+Ki2/s)(1+Kd*s)/(tauro*s/(2*pi)+1));

%% Linear system
Ts=0; % continuous time
sys = fgess(L,Ts,{'Ia','zIp','rIp','Ip'});
sys = sys(:,'Va'); % Va only
sys.InputDelay = L.G.Vadelay;

%% Mutual decoupling / resistive compensation
s = tf('s');
MR = ss(L.G.Maa + diag(L.G.Ra)/s);

%% Control parameters
switch lower(L.P.tokamak)
  case 'anamak'
    parfun = @meqctrlpana;
  case 'tcv'
    parfun = @meqctrlptcv;
  otherwise
    error('unrecognized tokamak %s',L.P.tokamak)
end
[KzQ,KrQ,KIQ,KIo,Tconstr] = parfun(L,LY,sys);

%% Control directions
switch lower(L.P.tokamak)
  case 'anamak'
  case 'tcv'
    % Commented for now, later copy to TCV case
    %     % Select FPS only for vertical control (for now)
    %     Tz = zeros(L.G.na,1); Tz(contains(L.G.dima,'G'))=-1;
    %     Tconstr = zeros(L.G.na,1); Tconstr(contains(L.G.dima,'OH'))=[1 -1];
  case 'sparc'
  otherwise
    error('unrecognized tokamak %s',L.P.tokamak)
end

%% Z control
% SISO PID controller
KzQPID = mkPID(KzQ);

if doplot
  %% debugging plots for control design
  figure(1); clf; set(gcf,'name','Z control design')
  drawnow;
  Pz = sys(sys.OutputGroup.zIp,sys.InputGroup.Va)*MR*KzQ.T; % open-loop SISO plant
  % open-loop
  G = Pz*KzQPID;
  
  clf;
  subplot(221)
  bode(Pz,G,KzQPID,{10^1,10^5}); grid on
  subplot(222)
  nyquist(Pz,G), grid on; set(gca,'xlim',[-3 1],'ylim',[-2 2])
  subplot(223)
  cltest = feedback(G,1);
  step(cltest,0.1); y=step(cltest,0.1);
  assert(max(y)<2,'unstable loop or excessive overshoot');
  subplot(224);
  tt = 0:1e-4:0.05; testramp = min(tt,0.02)/0.02;
  lsim(cltest,testramp,tt); title('Ramp response')
  drawnow;
end

iVa = sys.InputGroup.Va;
izIp = sys.OutputGroup.zIp;
sysclz = feedback(sys,MR*KzQ.T*KzQPID,iVa,izIp); % system after having closed vertical loop
% inputs: Va disturbances. Outputs: all sys outputs

%% Ia control - control orthogonal coil directions

if ~isfield(KIo,'T')
  if all(all(KzQ.T==1|KzQ.T==0,2)) % only ones or zeros on each row
    % Specific coils were chosen for each control direction
    iused = any( KzQ.T == 1 ,2 );
    II = eye(L.G.na);
    Torth = II(:,~iused);
    InputName = L.G.dima(~iused);
  else % general method in null space
    % Orthogonal coil directions
    Torth = null(KzQ.T');
    InputName = cellstr(num2str((1:L.G.na-1)','eIo%03d'));
  end
  KIo.T = Torth;
end

KIo.InputName = InputName;
KIoPID = mkPID(KIo);

if doplot
  figure(4); set(gcf,'name','Coil current + zIp closed-loop check')
  
  PIa = sysclz(sysclz.OutputGroup.Ia,sysclz.InputGroup.Va)*MR; % open-loop SISO plant
  KKa = KIo.T*(eye(L.G.na-1)*KIoPID);
  G   = PIa*KKa; % project coil errors onto Torth
  
  clf;
  subplot(221)
  bode(PIa(1,1),G(1,1),KIoPID(1,1),{10^1,10^4}); grid on
  subplot(222)
  subplot(223)
  cltest = Torth'*feedback(G,KIo.T');
  step(cltest(1,1),0.1)
  subplot(224);
  tt = 0:1e-4:0.1; testramp = min(tt,0.02)/0.02;
  lsim(cltest(1,1),testramp,tt); title('Ramp response')
  drawnow;
end

%% Close Ia and z control
% combined Kz and KIo controller. Input: [ezIp,eIo]
Kzo = MR*[Torth*KIoPID,KzQ.T*KzQPID];
Kzo.InputName = [InputName;'ezIp'];
sysVa = sys(:,sysclz.InputGroup.Va);
TFB = blkdiag(Torth',[1,0,0]); % from sys outputs to [Io,zIp]

clza = feedback(sysVa*Kzo,TFB);
clza.InputGroup.eIo = find(contains(Kzo.InputName,InputName))';

% system from orthogonal Va inputs to [Io;rIp;zIp;Ip]
% add integral to Ip output to cancel plant zero at origin

if doplot
  clf;
  step(clza(clza.OutputGroup.zIp,'ezIp'),0.1)
end

%% Now close R,Ip control around these
iIp = clza.OutputGroup.Ip;
KK  = zpk(eye(size(clza,1))); KK(iIp,iIp) = 1/s;
DD  = real(freqresp(KK*clza,0.1,'Hz'));
iii = [clza.OutputGroup.zIp,clza.OutputGroup.rIp,clza.OutputGroup.Ip];

DD  = DD(iii,~Tconstr); % exclude constrained coil combinations from being used for rIp,Ip control

DI(~Tconstr,:) = pinv(DD); % direction to actuate rIp,zIp,Ip separately

na = L.G.na;
Tr = DI(1:na-1,2);
Tp = DI(1:na-1,3);

%% R control
iR = ismember(clza.OutputName,'rIp');

KrQ.T  = Tr;
KrQPID = mkPID(KrQ);
tfin = 0.2;
% SISO PID controller
if doplot
  figure(2); clf;
  set(gcf,'name','R control design')
  Pr = clza(iR,'eIo')*Tr; % open-loop SISO plant
  G  = Pr*KrQPID;
  
  clf;
  subplot(221)
  bode(Pr,G,KrQPID,{10^1,10^4}); grid on
  subplot(222)
  nyquist(Pr,G), grid on; set(gca,'xlim',[-3 1],'ylim',[-2 2])
  subplot(223)
  cltest = feedback(G,1);
  step(cltest,tfin)
  subplot(224);
  tt = 0:1e-4:tfin; testramp = min(tt,tfin/2)/tfin;
  lsim(cltest,testramp,tt); title('Ramp response')
  drawnow;
end

%% Ip control
KIQPID = mkPID(KIQ);
KIQ.T  = Tp;

iIp = ismember(clza.OutputName,'Ip');

if doplot
  figure(3); clf; set(gcf,'name','Ip control design')
  Pp = clza(iIp,'eIo')*Tp; % open-loop SISO plant
  G  = Pp*KIQPID;
  
  clf;
  subplot(221)
  bode(Pp,G,KIQPID,{10^1,10^4}); grid on
  subplot(222)
  nyquist(Pp,G), grid on; set(gca,'xlim',[-3 1],'ylim',[-2 2])
  subplot(223)
  cltest = feedback(G,1);
  step(cltest,0.1)
  subplot(224);
  tt = 0:1e-4:0.05; testramp = min(tt,0.02)/0.02;
  lsim(cltest,testramp,tt); title('Ramp response')
  drawnow;
end

%% Close these two loops too
II  = eye(size(clza,1));
Frp = II(na+[2,3,1],:); % from clza outputs to [rIp,Ip,zIp]

Krp = [Tr*KrQPID,Tp*KIQPID];
Krp.InputName = {'erIp','eIp'};

if doplot
  cl = feedback(clza*blkdiag(Krp,1),Frp);
  clf;
  step(cl(end-2:end,:),0.1)
end
end

function KKd = build_position_controller(Maa,Ra,KzQ,KrQ,KIQ,KIo,Ts)
% Builds combined position controller
% KKd receives [eIa,erIp,ezIp,eIp] inputs and returns Va commands
s = tf('s');
MR = ss(Maa + diag(Ra)/s);
na = size(Maa,2);
% NB: this is improper, but is handled by the rolloff poles in the PID controllers

KzQPID = mkPID(KzQ);
KrQPID = mkPID(KrQ);
KIQPID = mkPID(KIQ);
KIoPID = mkPID(KIo);

%% Formulate complete controller
%  u = Kz*ezIp + KIo*(TIo'*eIa)
%    = Kz*ezIp + KIo*(TIo'*(eIa  +  Tp*Kp*eIp + Tr*Kr*erIp)
%    = [ KIo*[TIo',Tr*Kr,Tp*Kp, TIr] ,Kz ] * [eIa;erIp;eIp;ezIp])

% Continuous time
KKc = ss(MR*[KIo.T*KIoPID*[KIo.T',KrQ.T*KrQPID,KIQ.T*KIQPID] KzQ.T*KzQPID]);
KKc.InputName(1:na)     = cellstr(num2str((1:na)','Ia %03d'));
KKc.InputName(na+(1:3)) = {'erIp';'eIp';'ezIp'};

% Exchange columns such that input order is [eIa,zIp,rIp,Ip]
KKc = KKc(:,[1:na,na+[3,1,2]]);

% Discrete-time
KKd = c2d(KKc,Ts,'tustin'); % discretize using Tustin;

% Assign output names
KKd.OutputName(1:na) = cellstr(num2str((1:na)','Va %03d'));
end

function Tz = vertical_control_direction(method,varargin)
% Tz = vertical_control_direction(method,varargin)
% Computes vertical control direction using various methods
% 
% method: 'eigenvector' - Computes vertical control direction aligned with eigenvector.
%          See F. Pesamosca thesis [EPFL Thesis 8316, Section 5.2.3] 
% no other methods are defined yet

switch method
  case 'eigenvector'
    sys=varargin{1};
    % Compute vertical control direction aligned with eigenvector.
    % See F. Pesamosca thesis [EPFL Thesis 8316, Section 5.2.3]    sys = varargin{1}
    % left eigenvectors
    [V,D] = eig(sys.A');
    [~,jj] = esort(diag(D));
    assert(D(jj(1),jj(1))>0,'Expected unstable eigenvalue, but none found');
    q = V(:,jj(1));    % left eigenvector
    up = sys.B'*q;     % input pole vector
    Tz = up/norm(up);  % scale
    
    % choose sign such that it points opposite to the eigenvector
    Tz = Tz*sign(sum(sys.B*Tz));
    
    % NB: this can also be obtained by SVD of G(pole)
    % Gp = evalfr(ss(sys.A,sys.B(:,iVa),sys.C,sys.D(:,iVa)),p1(1));
    % [U,S,V] = svd(Gp,0);
    % Tz = V(:,1);
  otherwise
    error('undefined method %s',method);
end
end

function PID = mkPID(P)
% Make PID transfer function given controller parameters

s=tf('s');
PID = ss(P.Kp*(1 + P.Ki/s + P.Kd*s)/(P.tauro*s/(2*pi)+1));
if isfield(P,'Ki2')
  % additional integrator
  PID = PID*(1+P.Ki2/s);
end
end

function [KzQ,KrQ,KIQ,KIo,Tconstr] = meqctrlpana(L,~,sys) %#ok<DEFNU>
% Get controller parameters for anamak
Tconstr = false(L.G.na,1); % Directions for extra constraints on coils

s=tf('s');
MR = ss(L.G.Maa + diag(L.G.Ra)/s);
Tz = vertical_control_direction('eigenvector',sys*MR);

tauro = 1/1000; % rolloff time constant

% Vertical control
KzQ.T     = Tz;
KzQ.Kp    = 1000;
KzQ.Kd    = 5e-3;
KzQ.Ki    = 1/10e-3;
KzQ.Ki2   = 0*1/100e-3; 
KzQ.tauro = tauro; % rolloff tau

% Orthogonal coil current direction
KIo.Kp    = 1000;
KIo.Ki    = 0;
KIo.Ki2   = 0; 
KIo.Kd    = 0;
KIo.tauro = tauro; % rolloff tau


% Radial control
KrQ.Kp    = 5;
KrQ.Kd    = 0;
KrQ.Ki    = 1/10e-3;
KrQ.tauro = tauro; % rolloff tau

% Plasma current
KIQ.Kp    = 1;
KIQ.Ki    = 1/5e-3;
KIQ.Ki2   = 0/50e-3;
KIQ.Kd    = 0;
KIQ.tauro = tauro; % rolloff tau

end
