function [t, q, p, q0, qd, log] = fly_2dmatlab_simulator()
% 3DOF robofly simulator and renderer (extensible to 6DOF)
% Sawyer B. Fuller, 2016.09.28 
% 
% Simulates robofly dynamics including stroke-averaged aerodynamic drag.
% 
% The states that are computed are q = [thetay, omegay_body, x_world,
% z_world, xdot_body, zdot_body], that is, a 3DOF state vector. The
% position x_world is given in world coordinates and the angle thetay is
% the Euler angle around the y-axis. Velocity vectors (xdot_body and omega) are
% calculated in the coordinates of the body frame (but note that, as is
% convention, computed velocities are not _relative_ to the moving body
% frame, just given with respect the current orientation of its coordinate
% frame). To compute velocities in the world frame, use v_world = R *
% v_body, where R is the rotation matrix from the rot_matrix function.)
% Internally, dynamics are computed with a full 6DOF, and the fly_dynamics
% function can be easily augmeted to accomodate corresponding three-element
% force and torque vectors inputs and produce the full 12-element state
% vector. This simulation loosely follows the treatment given in RB McGhee,
% ER Bachmann, MJ Zyda, "Rigid body dynamics, inertial reference frames,
% and graphics coordinate systems: A resolution of conflicting conventions
% and terminology," MOVES Academic Group Technical Report NPS-MV-01-002,
% Naval Postgraduate School, Monterey, CA (2000).
%
% useful simulation parameters: 
% * TFINAL: time duration of simulation  
% * PLOT: whether to plot trajectory results
% * set the controller under 'controllers' comment. There 
%   a separate controller for each of the the altitude and torque dynamnics

% reset persistent variables in functions
clear all

global ctrl_dt log

TFINAL = 1; % time duration of simulation
PLOT = 1; % whether to plot trajectory results
figure(1)
clf()
CTRL_HZ = 500; % simulation update rate

% robofly params
p.winglength = .6 * 2.54 / 100; 
p.l = .013; % length of body
p.h = .0025; % thickness of body
p.J = 1.4e-9; % from science paper/solidworks
p.b_w = 2.0e-4; % aero drag on wings from wind tunnel tests, Ns/m
p.c_w = (p.h/2 + p.winglength * 2./3)^2 * p.b_w; % rot drag coeff around z [Nsm]
p.r_w = .008; % z-dist from ctr of wings to ctr of mass 
p.m = 80e-6; %mass of robofly
p.g = 9.81; 
p.max_f_c = 1.5 * p.m * p.g; 
p.max_tau_c = 1e-6;  

% rendering
diagnl = 1/2*sqrt(p.l^2+p.h^2);
angl = atan(p.l/p.h);
plotbox = inline('plot([x+diag*cos(th-angl) x+diag*cos(th+angl) x+diag*cos(th-angl+pi) x+diag*cos(th+angl+pi) x+diag*cos(th-angl)], [y+diag*sin(th-angl) y+diag*sin(th+angl) y+diag*sin(th-angl+pi) y+diag*sin(th+angl+pi) y+diag*sin(th-angl)])','x','y','th','angl','diag');

ctrl_dt = 1/CTRL_HZ; 
t = 0:ctrl_dt:TFINAL; 

%q (the state) = theta, thetadot, xworld, zworld, xdot, zdot
q0 = [.2, 0, 0, p.l/2, 0, 0]'; 
qd = [0, 0, .04, .05, 0, 0]; % desired state

% initialize state
q = zeros(length(t), length(q0)); 
u = zeros(length(t), 2);
qode = q0'; 

% controller selection
%thrust_controller = @open_loop_thrust;
%thrust_controller = @proportional_thrust;
thrust_controller = @PD_thrust;
%roll_controller = @open_loop_torque;
%roll_controller = @PD_torque;
roll_controller = @PD_and_lateral_PD_torque;

for i = 1:length(t)
    control_torque = roll_controller(t(i), qode(end,:), qd, p);
    control_thrust = thrust_controller(t(i), qode(end,:), qd, p); 
    control_torque = clip(control_torque, p.max_tau_c); % limit to reasonable range
    control_thrust = clip(control_thrust, p.max_f_c); 
    [tode,qode] = ode45(@fly_dynamics, [t(i), t(i)+ctrl_dt], qode(end,:), ...
        [], [control_thrust, control_torque], p); 
    if size(qode, 1) > 100
        disp('unstable simulation detected, aborting');
        break
    end
    q(i,:) = qode(end, :); 
    u(i,:) = [control_thrust, control_torque]; 

    % animate
    if mod(i, 3) == 0
        plot([-.05, .05], [0, 0])
        hold on
        plotbox(q(i, 3), q(i, 4), -q(i, 1), angl, diagnl);
        axis equal
        plot(qd(3), qd(4), '+')
        hold off
        title(['t=' num2str(t(i),'%1.3f')])
        xlim([-.06, .06])
        drawnow
    end
end

if PLOT
    figure(1)
    % floor
    plot([-.05, .05], [0, 0])
    hold on
    %    trajectory of cm
    plot(q(:,3), q(:,4), 'k')
    for i = 1:10:length(t)
        plotbox(q(i, 3), q(i, 4), -q(i, 1), angl, diagnl);
    end
    axis equal
    % setpoint location
    plot(qd(3), qd(4), '+')
    xlim([-.06, .06])
    xlabel('x position (m)')
    ylabel('y position (m)')
    hold off
    
    figure(2)
    subplot(221)
    plot(t, q(:,4))
    xlabel('time (s)')
    ylabel('z (m)')
    subplot(223)
    plot(t, u(:, 1))
    xlabel('time (s)')
    ylabel('fc (N)')
    subplot(222)
    plot(t, q(:,1), 'r')
    hold on 
    plot(t, q(:,3))
    hold off
    xlabel('time (s)')
    ylabel('theta (rad, red), x (m, blue)')
    subplot(224)
    plot(t, u(:, 2))
    xlabel('time (s)')
    ylabel('tauc (Nm)')
    
    % bias corrections
    if isfield(log, 'correction_torque')
        figure(3)
        subplot(311)
        plot(t, log.correction_torque); 
        ylabel('correction torque (Nm)')
        subplot(312)
        plot(t, log.correction_theta); 
        ylabel('theta correction (rad)')
        subplot(313)
        plot(t, log.thetad); 
        ylabel('theta desired (rad)')
        xlabel('time (s)')
    end
end

end

% controller functions ---------------------

function f_c = open_loop_thrust(t, q, qd, p)
f_c = p.m * p.g * 1.1; % feedforward gravity cancellation
end

function f_c = proportional_thrust(t, q, qd, p)
error = qd(4) - q(4);
k = 1e-2;
f_c = p.m * p.g + k * error; % feedforward gravity cancellation
end

function f_c = PD_thrust(t, q, qd, p)
% initialize
global ctrl_dt
persistent last_error % remember last error
if isempty(last_error)
    last_error = 0; % need an initial value
end

% compute
% q = [thetay, omegay_body, x_world, z_world, xdot_body, zdot_body]
error = qd(4) - q(4);
edot = (error - last_error)/ctrl_dt;
k = 1e-2; % proportional gain
kd = 1e-3; % derivative gain
f_c = k * error + kd * edot; 
f_c = f_c + p.m * p.g; % feedforward gravity cancellation
last_error = error; % save for next function call
end

function tau_c = open_loop_torque(t, q, qd, p)
tau_c = 0; 
end

function tau_c = PD_torque(t, q, qd, p)
% q = [thetay, omegay_body, x_world, z_world, xdot_body, zdot_body]
error = qd(1) - q(1);
edot = qd(2) - q(2);
k = 1e-6; % proportional gain
kd = 1e-7; % derivative gain
tau_c = k * error + kd * edot; 
end

function tau_c = PD_and_lateral_PD_torque(t, q, qd, p)
% initialize
global ctrl_dt
persistent last_theta_error % remember last error
if isempty(last_theta_error)
    last_theta_error = 0; % need an initial value
end

% q = [thetay, omegay_body, x_world, z_world, xdot_body, zdot_body]

% outer loop: lateral control
% choose a setpoint attitude
k_lateral = 20; % proportional gain
kd_lateral = 2; % derivative gain
error_lateral = qd(3) - q(3);
errordot_lateral = qd(5) - q(5); 
theta_desired = k_lateral * error_lateral + kd_lateral * errordot_lateral; 

theta_error = theta_desired - q(1);
theta_errordot = (theta_error - last_theta_error)/ctrl_dt;
k = 1e-6; % proportional gain
kd = 1e-7; % derivative gain
tau_c = k * theta_error + kd * theta_errordot;
last_theta_error = theta_error; 
end




function qdot = fly_dynamics(t, q, u, p)
% given state and inputs u from controller, calculate derivatives for ode45
% integrator. 
v = [q(5), 0, q(6)]'; 
theta = [0, q(1), 0]'; 
omega = [0, q(2), 0]'; 
R = rot_matrix(theta);
W = omega2thetadot_matrix(theta);

[f_d, tau_d] = fly_aerodynamics(v, omega, p);

% forces (f)
f_g = R' * [0, 0, - p.g * p.m]'; % gravity
f_c = [0, 0, u(1)]'; % control force
f = f_c + f_d + f_g; %- p.g * p.m * [-sin(theta), 0, cos(theta)]'; 

% moments/torues (tau)
tau_c = [0, u(2), 0]'; 
tau = tau_c + tau_d; 

fictitious_f = cross(omega, v); %[0; omega; 0], [q(4); 0; q(6)]); % omega x v
fictitious_tau = cross(omega, p.J * omega); 

% geometric
xdotworld = R * v; % world coordinates
vdot = 1/p.m * f - fictitious_f; % body coordinates

thetadot = W * omega; % world euler angles
omegadot = 1/p.J * (tau - fictitious_tau); % body coordinates 

qdot = [thetadot(2); 
        omegadot(2); 
        xdotworld(1); 
        xdotworld(3);
        vdot(1); 
        vdot(3)];
end

function c = cross(a,b) % default cross func is slow
c = [a(2)*b(3)-a(3)*b(2); a(3)*b(1)-a(1)*b(3); a(1)*b(2)-a(2)*b(1)];
end


function R = rot_matrix(euler_theta)
% R matrix to convert 3-vectors in body coords to world coords
% v = Rv' where v is in world frame and v' is in body frame
% convention here is ZYX convention: the three coordinates in euler_theta  
% mean rotate around world Z, then world Y, then world X 
% alternately, rot around body x, then new body y, then new body z
% (the latter two are equivalent)

c1 = cos(euler_theta(1)); 
c2 = cos(euler_theta(2));
c3 = cos(euler_theta(3));
s1 = sin(euler_theta(1));
s2 = sin(euler_theta(2));
s3 = sin(euler_theta(3));
R = [c2.*c3, -c2.*s3, s2;
     c1.*s3 + c3.*s1.*s2, c1.*c3 - s1.*s2.*s3, -c2.*s1;
     s1.*s3 - c1.*c3.*s2, c3.*s1 + c1.*s2.*s3, c1.*c2];
end

function W = omega2thetadot_matrix(euler_theta)
% transform euler angle rates to angular rot rate vector omega
% thetadot = W * omega
% this needs tobe corrected for 3d case for the current euler angle convention. 
st1 = sin(euler_theta(1));
ct1 = cos(euler_theta(1));
tt2 = tan(euler_theta(2));
ct2 = cos(euler_theta(2));
W = [1, st1*tt2, ct1*tt2; 
    0, ct1, -st1;
    0, st1/ct2, ct1/ct2];
end

function [f_d, tau_d] = fly_aerodynamics(v, omega, p)
% calculate stroke-averaged forces due to aerodynamic drag on flapping 
% wings. assumes force is applied at point r_w away from center of mass in
% z-direction in body-coordinates. v_w is the vel at that point. 
% assumes drag force f_d = -b_w * v_w. this has been tested in a wind
% tunnel for x- and y-directions but not z-direction
r_w = [0, 0, p.r_w]';
v_w = v + cross(omega, r_w); % vel at midpoint of wings
f_d = - p.b_w * v_w;  
tau_d = cross(r_w, f_d); %p.r_w * f_d(1); % r_w cross f_d
end

function clipped = clip(x, absmax)
    clipped = max(min(x, absmax), -absmax);
end