% Implements the nonlinear model and creates state-space matrices A and B
% using MATLAB's symbolic differentiation.
%
% The model physics is taken from Matt Rich's "Model Development..."
% sections 4.1 and 4.2
%
% This linearization assumes the rotor orientation in the simulink model:
% 1     2
%   +dx
%      +dy
%
% 3     4
% In some of the internal simulink files, these are indexed 0, ..., 3
% instead of 1, .., 4. The order is unchanged. 
% Note that the final result should be the same regardless of rotor
% orientation, but everything internal to this function and its parameters
% need to be consistent.
function [A, B] = linearization(m, g, Kt, Kh, Kd, delta_T_z, rx, ry, rz, J, Jreq, Vb, Rm, Kv, Kq, If, signal_mixer)
    % Symbolic state variables
    syms u v w p q r x y z ph th ps
    
    % Rotation matrices. After linearization, these are equivalent to I
    L_EB = [...
        cos(th)*cos(ps), sin(ph)*sin(th)*cos(ps)-cos(ph)*sin(ps), cos(ph)*sin(th)*cos(ps)+sin(ph)*sin(ps)
        cos(th)*sin(ps), sin(ph)*sin(th)*sin(ps)+cos(ph)*cos(ps), cos(ph)*sin(th)*sin(ps)-sin(ph)*cos(ps)
        -sin(th),        sin(ph)*cos(th),                         cos(ph)*cos(th)];
    A_EB = [...
        1,  0,       -sin(th)
        0,  cos(ph),  sin(ph)*cos(th)
        0, -sin(ph),  cos(ph)*cos(th)];
    
    % Rotor positions
    r_h = [...
        -rx,  rx, -rx,  rx
         ry,  ry, -ry, -ry
        -rz, -rz, -rz, -rz
    ];

    % Gamma matrices select force dimensions
    Gamma_H = diag([1, 1, 0]);
    Gamma_Omega = [zeros(2,4); 1, -1, -1, 1];
    
    % uT, uA, uE, uR
    u_c = sym('u_c_%d', [1 4]);
    % Individual rotor commands
    u_p = (signal_mixer * u_c')';
    
    % Rotor angular velocities
    % These values assume no motor transient
    % (i.e. speeds reached immediately)
    omega = sym('w_%d', [1 4]);
    for i = 1:4
        omega(i) = (-1 + sqrt(1 - 4*Rm*Kv*Kq*Kd*(Kv*Rm*If - Kv*u_p(i)*Vb)))/(2*Rm*Kv*Kq*Kd);
    end
    
    % Variable vectors
    v_B_0   = [u; v; w];
    Omega_B = [p; q; r];
    
    % Rotor velocity
    for i = 1:4
        v_h(:,i) = v_B_0 + cross(Omega_B, r_h(:,i));
    end
    
    % Gravity
    G_B = m*g*[-sin(th); cos(th)*sin(ph); cos(th)*cos(ph)];
    
    % Force
    % Rotor Thrust (3 directions x 4 rotors matrix)
    T_B = [zeros(2, 4); -(Kt*omega.^2 + delta_T_z*v_h(3,:).*omega)];
    % Rotor In-plane drag (3 directions x 4 rotors matrix)
    H_B = -(Kh*ones(3,1)*omega).*(Gamma_H*v_h);
    % Total Force (Gravity + rotor force summed over all rotors)
    F_B = G_B + sum(T_B + H_B, 2);

    % Torque
    for i = 1:4
        % Force Lever-arm Torque
        Q_B_F(:,i) = cross(r_h(:,i), T_B(:,i) + H_B(:,i));
        % Change of Angular Momentum Torque (assuming no motor transient)
        Q_B_L(:,i) = -Jreq*cross(Omega_B, omega(i)*Gamma_Omega(:,i));
        % Drag/Induced Torque
        Q_B_d(:,i) = -Kd*omega(i)^2*Gamma_Omega(:,i);
    end
    % Total Torque
    Q_B = sum(Q_B_F + Q_B_d + Q_B_L, 2);
    
    % Define state vector Lambda and its time derivative, Lambda_dot
    Lambda = [u, v, w, p, q, r, x, y, z, ph, th, ps];
    Lambda_dot = [...
        1/m * F_B - cross(Omega_B, v_B_0)
        J^-1 * Q_B - cross(J^-1 * Omega_B, J * Omega_B)
        L_EB * v_B_0
        A_EB * Omega_B
    ];

    % A is partial of Lambda_dot wrt Lambda
    for i = 1:12
        for j = 1:12
            symA(i, j) = diff(Lambda_dot(i), Lambda(j));
        end
    end
    
    % B is partial of Lambda_dot wrt u_T, u_A, u_E, u_R
    for i = 1:12
        for j = 1:4
            symB(i, j) = diff(Lambda_dot(i), u_c(j)); 
        end
    end
    
    % Set equilibrium values
    u = 0; v = 0; w = 0; p = 0; q = 0; r = 0; x = 0; y = 0; z = 0; ph = 0; th = 0; ps = 0;
    u_c_1 = 0; u_c_2 = 0; u_c_3 = 0; u_c_4 = 0;
    
    % Binary search for equilibrium throttle value
    throttle_max = 1; throttle_min = 0;
    while (abs(double(subs(F_B(3)))) > 1e-7)
        u_c_1 = (throttle_max + throttle_min)/2;
        if (double(subs(F_B(3))) > 0)
            throttle_min = double(u_c_1);
        else
            throttle_max = double(u_c_1);
        end
    end
    
    % Convert symbolic linearization matrices to double values
    A = double(subs(symA));
    B = double(subs(symB));
end
%#ok<*AGROW>