mass_symmetry = false;

subject = 'delaware3';
trial = 'ss_walking1';

%subject = 'subject01';
%trial = 'walk1';

gravity = [0 -9.80665 0];

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

model = read_osimFile(sprintf('%s.osim',subject), true);
nbodies = length(model.bodies);

acc = read_motionFile(sprintf('Results/%s_%s_RRA1_BodyKinematics_acc.sto',subject,trial));
acc_t = acc.data(:,1);
acc_xyz = acc.data(:,find_columns_by_label(acc.labels,'center_of_mass'));

grf = read_motionFile(sprintf('%s_%s_ik.mot',subject,trial));
grf_t = grf.data(:,1);
grf_xyzxyz = grf.data(:,find_columns_by_label(grf.labels,'ground_force_v'));
grf_xyz = [sum(grf_xyzxyz(:,[1 4]),2) sum(grf_xyzxyz(:,[2 5]),2) sum(grf_xyzxyz(:,[3 6]),2)];

stored_res = read_motionFile(sprintf('Results/%s_%s_RRA1_Actuation_force.sto',subject,trial));
stored_res_t = res.data(:,1);
stored_res_xyz = res.data(:,find_columns_by_label(res.labels,'F[XYZ]'));

tmin = min(acc_t);
tmax = max(acc_t);
T = tmin:0.001:tmax;
n = length(T);

acc_xyz = interp1(acc_t, acc_xyz, T);
grf_xyz = interp1(grf_t, grf_xyz, T);
stored_res_xyz = interp1(stored_res_t, stored_res_xyz, T);

acc_gravity = reshape(ones(n,1)*gravity,3*n,1) * ones(1,nbodies);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

masses = zeros(nbodies,1);
acc_bodies = [];
for i=1:nbodies
    acc_body = acc.data(:,find_columns_by_label(acc.labels,sprintf('%s_[XYZ]',model.bodies(i).name)));
    if size(acc_body,2) ~= 3
        error(sprintf('Could not find BodyKinematics accelerations for body "%s"', model.bodies(i).name));
    end
    acc_bodies(:,i) = reshape(acc_body, 3*n, 1);
    masses(i) = model.bodies(i).mass;
end

% Set up system to solve
A = acc_bodies - acc_gravity;
b = reshape(grf_xyz, 3*n, 1);

if ~ mass_symmetry
    new_masses = lsqnonneg(A,b)
else
    Asymm = [];

    for i=1:nbodies
        name = model.bodies(i).name;

        if strfind(name, '_r')
            % right body
            left_body_name = strrep(name, '_r', '_l');
            left_body_index = 0;
            for j=1:nbodies
                if strcmp(model.bodies(j).name,left_body_name)
                    left_body_index = j;
                    break;
                end
            end
            if ~ left_body_index
                error(sprintf('Could not locate symmetric body for "%s"', name));
            end

            Asymm(end+1,:) = zeros(1,nbodies);
            Asymm(end,i) = 1;
            Asymm(end,left_body_index) = 1;

        elseif strfind(name, '_l')
            % left body -- skip because right body handled this

        else
            % body with no symmetric counterpart

            Asymm(end+1,:) = zeros(1,nbodies);
            Asymm(end,i) = 1;
        end
    end

    new_masses = Asymm' * lsqnonneg(A*Asymm',b);
end

for i=1:length(model.bodies)
    disp(sprintf('%10s:\tOLD=%10f\tOPT=%10f', model.bodies(i).name, model.bodies(i).mass, new_masses(i)));
end
disp(sprintf('%10s:\tOLD=%10f\tOPT=%10f', 'TOTAL', sum(masses), sum(new_masses)));

% Baseline values
ma = reshape(acc_bodies * masses, n, 3);
res_xyz = ma - reshape(acc_gravity * masses, n, 3) - grf_xyz;

% New values
new_ma = reshape(acc_bodies * new_masses, n, 3);
new_res_xyz = new_ma - reshape(acc_gravity * new_masses, n, 3) - grf_xyz;

disp(sprintf('OLD RESIDUAL NORM: X=%f Y=%f Z=%f', norm(res_xyz(:,1)), norm(res_xyz(:,2)), norm(res_xyz(:,3))));
disp(sprintf('NEW RESIDUAL NORM: X=%f Y=%f Z=%f', norm(new_res_xyz(:,1)), norm(new_res_xyz(:,2)), norm(new_res_xyz(:,3))));

disp(sprintf('Mean Fy = %f, recommended adjustment = %f', mean(res_xyz(:,2)), mean(res_xyz(:,2))/gravity(2)));


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

close all;

plot(T, [grf_xyz(:,2) ma(:,2) res_xyz(:,2) new_ma(:,2) new_res_xyz(:,2) stored_res_xyz(:,2)]);
legend({'grf', 'mass*acc', 'res', 'mass*new_acc', 'new_res', 'stored_res'}, 'Interpreter', 'none');

%plot(T, [total_mass*acc_xyz(:,2) grf_xyz(:,2) res_xyz(:,2) weight(:,2) F_ext(:,2)]);
%legend('mass*acc','grf','res','weight','Fext');
%
%figure;
%plot(T, [total_mass*new_acc_xyz(:,2) grf_xyz(:,2) new_res_xyz(:,2) weight(:,2) new_F_ext(:,2)]);
%legend('mass*acc','grf','res','weight','Fext');
