% Simulate a step responses of ducted-fan VTOL micro-UAV.
%
% Copyright (C) 2024, Naoki Sean Pross, ETH Zürich
% This work is distributed under a permissive license, see LICENSE.txt

function [simout] = uav_sim_step(params, model, ctrl, uncert, nsamp, T, do_plots, do_noise)

% Load closed loop model and add controller
% more or less equivalent to doing usys = lft(Pnom, K)
h = load_system('uav_model_uncertain_clp');
hws = get_param('uav_model_uncertain_clp', 'modelworkspace');

if isfield(ctrl, 'K')
  hws.assignin('K', ctrl.K);
else
  error('You need to provide a controller ctrl.K');
end

% There is an uncertainty block
if isfield(uncert, 'Delta')
  fprintf(' - Running worst case with provided Delta.\n');
  hws.assignin('Delta', uncert.Delta);
else
  fprintf(' - No uncertainty error Delta was provided, setting to zero.\n');
  Delta = tf(0) * zeros(sum(model.uncertain.BlockStructure));
  hws.assignin('Delta', Delta);
end

ulmod_clp = linmod('uav_model_uncertain_clp');
P_clp = minreal(ss(ulmod_clp.a, ulmod_clp.b, ulmod_clp.c, ulmod_clp.d), [], false);

% Check that closed loop is actually stable
nx = size(P_clp.A, 1);
fprintf(' - Closed loop dynamics have %d states\n', nx);
if nx < 60
  [~, nsta, nctrb, nustab, nobsv, nudetb] = pbhtest(P_clp);
  if nx ~= nsta
    error('Closed loop is not stable!');
  end

  fprintf('      %d stable modes, %d unstable modes\n', nsta, nx - nsta);
  fprintf('      %d controllable modes, %d unstabilizable\n', nctrb, nustab);
  fprintf('      %d observable modes, %d undetectable\n', nobsv, nudetb);
else
  eigvals = eig(P_clp.A);
  for i = 1:nx
    if real(eigvals(i)) >= 0
      fprintf('   Closed loop is not stable!\n');
      break;
    end
  end
end

% Generate indices for closed loop plant
o = model.uncertain.dims.output;
i = model.uncertain.dims.input;

dims = struct(...
  'ErrorAlpha', o.OutputErrorAlpha, ...
  'ErrorOmega', o.OutputErrorOmega, ...
  'ErrorPosition', o.OutputErrorPosition, ...
  'ErrorVelocity', o.OutputErrorVelocity, ...
  'ErrorEulerAngles', o.OutputErrorEulerAngles, ...
  'Position', o.OutputNominalPosition, ...
  'Velocity', o.OutputNominalVelocity, ...
  'EulerAngles', o.OutputNominalEulerAngles, ...
  'AngularRates', o.OutputNominalAngularRates, ...
  'PlotAlpha', o.OutputPlotAlpha, ...
  'PlotOmega', o.OutputPlotOmega, ...
  'PlotReference', o.OutputPlotReference, ...
  'PlotPosition', o.OutputPlotPosition, ...
  'InputAlpha', i.InputNominalAlpha, ...
  'InputOmega', i.InputNominalOmega ...
);

idx = struct();
start = 1; fields = fieldnames(dims);
for f = fields'
  dim = getfield(dims, f{1});
  i = (start:(start + dim - 1))';
  idx = setfield(idx, f{1}, i);
  start = start + dim;
end

% Input indices
idx.InputDisturbanceWind = (1:3)';
idx.InputDisturbanceFlaps = (4:7)';
idx.InputReference = (8:10)';

% Create noise
noise = zeros(7, nsamp);
if do_noise
  % Noise (normalized)
  noise_alpha_amp = (.5 * (pi / 180)) / params.actuators.ServoAbsMaxAngle;
  noise_wind_amp = .01;
  noise = [noise_wind_amp * randn(3, nsamp);
    noise_alpha_amp * randn(4, nsamp)];
end

% Step size
step_size_h = .5  / params.normalization.HPosition;
step_size_v = .5  / params.normalization.VPosition;

% Create step inputs (normalized)
ref_step = ones(1, nsamp); % 1d step function

in_step_x = [ noise; step_size_h * ref_step; zeros(2, nsamp) ];
in_step_y = [ noise; zeros(1, nsamp); step_size_h * ref_step; zeros(1, nsamp) ];
in_step_z = [ noise; zeros(2, nsamp); -step_size_v * ref_step ]; % z points down

% Simulation time
t = linspace(0, T, nsamp);

% Scale simulation outputs
S = blkdiag(...
  model.uncertain.scaling.OutputErrorScaling, ...
  model.uncertain.scaling.OutputNominalScaling, ...
  model.uncertain.scaling.OutputPlotScaling, ...
  model.uncertain.scaling.InputNominalScaling ...
);

% Simulate step responses
out_step_x_norm = lsim(P_clp, in_step_x, t, 'foh');
out_step_y_norm = lsim(P_clp, in_step_y, t, 'foh');
out_step_z_norm = lsim(P_clp, in_step_z, t, 'foh');

out_step_x = out_step_x_norm * S;
out_step_y = out_step_y_norm * S;
out_step_z = out_step_z_norm * S;

% Return simulation
simout = struct(...
  'Time', t, ...
  'StepX', out_step_x, ...
  'StepY', out_step_y, ...
  'StepZ', out_step_z, ...
  'StepXNorm', out_step_x_norm, ...
  'StepYNorm', out_step_y_norm, ...
  'StepZNorm', out_step_z_norm, ...
  'Simulink', ulmod_clp, ...
  'StateSpace', P_clp, ...
  'index', idx);

if do_plots
  % Conversion factors
  to_deg = 180 / pi; % radians to degrees
  to_rpm = pi / 30; % rad / s to RPM

  % Figure for flaps and Euler angles
  figure;
  sgtitle(sprintf(...
    '\\bfseries Step Response of Flap and Euler Angles (%s)', ...
    ctrl.Name), 'Interpreter', 'latex');

  % Plot limits
  alpha_max_deg = params.actuators.ServoAbsMaxAngle * to_deg;
  euler_lim_deg = .5;
  omega_max_rpm = (params.actuators.PropellerMaxAngularVelocity ...
    - params.linearization.Inputs(5)) * to_rpm;
  omega_min_rpm = -params.linearization.Inputs(5) * to_rpm;

  % Plot step response from x to alpha
  subplot(2, 3, 1);
  hold on;
  plot(t, out_step_x(:, idx.PlotAlpha(1)) * to_deg);
  plot(t, out_step_x(:, idx.PlotAlpha(2)) * to_deg);
  plot(t, out_step_x(:, idx.PlotAlpha(3)) * to_deg);
  plot(t, out_step_x(:, idx.PlotAlpha(4)) * to_deg);
  plot([0, T], [1, 1] * alpha_max_deg, 'r--');
  plot([0, T], [-1, -1] * alpha_max_deg, 'r--');
  grid on;
  xlim([0, T]);
  ylim([-alpha_max_deg * 1.1, alpha_max_deg * 1.1]);
  title('Horizontal $x$ to Flaps', 'Interpreter', 'latex');
  ylabel('Flap Angle (degrees)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\alpha_1(t)$', '$\alpha_2(t)$', '$\alpha_3(t)$', '$\alpha_4(t)$', ...
    'Interpreter', 'latex');

  % Plot step response from y to alpha
  subplot(2, 3, 2); hold on;
  plot(t, out_step_y(:, idx.PlotAlpha(1)) * to_deg);
  plot(t, out_step_y(:, idx.PlotAlpha(2)) * to_deg);
  plot(t, out_step_y(:, idx.PlotAlpha(3)) * to_deg);
  plot(t, out_step_y(:, idx.PlotAlpha(4)) * to_deg);
  plot([0, T], [1, 1] * alpha_max_deg, 'r--');
  plot([0, T], [-1, -1] * alpha_max_deg, 'r--');
  grid on;
  xlim([0, T]);
  ylim([-alpha_max_deg * 1.1, alpha_max_deg * 1.1]);
  title('Horizontal $y$ to Flaps', 'Interpreter', 'latex');
  ylabel('Flap Angle (degrees)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\alpha_1(t)$', '$\alpha_2(t)$', '$\alpha_3(t)$', '$\alpha_4(t)$', ...
    'Interpreter', 'latex');

  % Plot step response from z to alpha
  subplot(2, 3, 3); hold on;
  plot(t, out_step_z(:, idx.PlotAlpha(1)) * to_deg);
  plot(t, out_step_z(:, idx.PlotAlpha(2)) * to_deg);
  plot(t, out_step_z(:, idx.PlotAlpha(3)) * to_deg);
  plot(t, out_step_z(:, idx.PlotAlpha(4)) * to_deg);
  plot([0, T], [1, 1] * alpha_max_deg, 'r--');
  plot([0, T], [-1, -1] * alpha_max_deg, 'r--');
  grid on;
  xlim([0, T]);
  ylim([-alpha_max_deg * 1.1, alpha_max_deg * 1.1]);
  title('Vertical $z$ to Flaps', 'Interpreter', 'latex');
  ylabel('Flap Angle (degrees)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\alpha_1(t)$', '$\alpha_2(t)$', '$\alpha_3(t)$', '$\alpha_4(t)$', ...
    'Interpreter', 'latex');

  % Plot step response from x to Theta
  subplot(2, 3, 4); hold on;
  plot(t, out_step_x(:, idx.EulerAngles(1)) * to_deg);
  plot(t, out_step_x(:, idx.EulerAngles(2)) * to_deg);
  plot(t, out_step_x(:, idx.EulerAngles(3)) * to_deg);
  grid on;
  xlim([0, T]);
  ylim([-euler_lim_deg, euler_lim_deg]);
  title('Horizontal $x$ to Euler Angles', 'Interpreter', 'latex');
  ylabel('Euler Angle (degrees)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\phi(t)$ Roll ', '$\theta(t)$ Pitch ', '$\psi(t)$ Yaw ', ...
    'Interpreter', 'latex');

  % Plot step response from y to Theta
  subplot(2, 3, 5); hold on;
  plot(t, out_step_y(:, idx.EulerAngles(1)) * to_deg);
  plot(t, out_step_y(:, idx.EulerAngles(2)) * to_deg);
  plot(t, out_step_y(:, idx.EulerAngles(3)) * to_deg);
  grid on;
  xlim([0, T]);
  ylim([-euler_lim_deg, euler_lim_deg]);
  title('Horizontal $y$ to Euler Angles', 'Interpreter', 'latex');
  ylabel('Euler Angle (degrees)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\phi(t)$ Roll ', '$\theta(t)$ Pitch ', '$\psi(t)$ Yaw ', ...
    'Interpreter', 'latex');

  % Plot step response from z to Theta
  subplot(2, 3, 6); hold on;
  plot(t, out_step_z(:, idx.EulerAngles(1)) * to_deg);
  plot(t, out_step_z(:, idx.EulerAngles(2)) * to_deg);
  plot(t, out_step_z(:, idx.EulerAngles(3)) * to_deg);
  grid on;
  xlim([0, T]);
  ylim([-euler_lim_deg, euler_lim_deg]);
  title('Vertical $z$ to Euler Angles', 'Interpreter', 'latex');
  ylabel('Euler Angle (degrees)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\phi(t)$ Roll ', '$\theta(t)$ Pitch ', '$\psi(t)$ Yaw ', ...
    'Interpreter', 'latex');

  % Plot step response from z to omega
  figure;
  sgtitle(sprintf(...
    '\\bfseries Step Response to Propeller (%s)', ...
    ctrl.Name), 'Interpreter', 'latex');

  hold on;
  step(P_clp(idx.PlotOmega, idx.InputReference(3)) * to_rpm, T);
  % plot([0, T], [1, 1] * omega_min_rpm, 'r--');
  % plot([0, T], [1, 1] * omega_max_rpm, 'r--');
  grid on;
  % ylim([omega_min_rpm - 1, omega_max_rpm + 1]);
  title('Vertical $z$ to Thruster $\omega$', 'Interpreter', 'latex');
  ylabel('Angular Velocity (RPM)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\omega(t)$', 'Interpreter', 'latex');

  % Figure for position and velocity
  figure;
  sgtitle(sprintf(...
    '\\bfseries Step Response of Position and Speed (%s)', ...
    ctrl.Name), 'Interpreter', 'latex');

  % Plot step response from horizontal reference to horizontal position
  subplot(2, 2, 1); hold on;
  plot(t, out_step_x(:, idx.PlotPosition(1)));
  plot(t, out_step_y(:, idx.PlotPosition(2)));
  plot(t, out_step_x(:, idx.PlotReference(1)), '--');
  grid on;
  title('Horizontal Position', 'Interpreter', 'latex');
  ylabel('Distance (meters)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$x(t)$', '$y(t)$', 'Interpreter', 'latex');

  % Plot step response horizontal reference to horizontal speed
  subplot(2, 2, 2); hold on;
  plot(t, out_step_x(:, idx.Velocity(1)));
  plot(t, out_step_y(:, idx.Velocity(2)));
  grid on;
  title('Horizontal Velocity', 'Interpreter', 'latex');
  ylabel('Velocity (m / s)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\dot{x}(t)$', '$\dot{y}(t)$', 'Interpreter', 'latex');

  % Plot step response from vertical reference to vertical position
  subplot(2, 2, 3); hold on;
  plot(t, out_step_z(:, idx.PlotPosition(3)));
  plot(t, out_step_z(:, idx.PlotReference(3)), '--');
  grid on;
  title('Vertical Position', 'Interpreter', 'latex');
  ylabel('Distance (meters)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$z(t)$', 'Interpreter', 'latex');

  % Plot step response vertical reference to vertical speed
  subplot(2, 2, 4); hold on;
  plot(t, out_step_z(:, idx.ErrorVelocity(3)));
  grid on;
  title('Vertical Velocity', 'Interpreter', 'latex');
  ylabel('Velocity (m / s)', 'Interpreter', 'latex');
  xlabel('Time (seconds)', 'Interpreter', 'latex');
  legend('$\dot{z}(t)$', 'Interpreter', 'latex');
end

end
% vim:ts=2 sw=2 et: