% This script demonstrates the use of MAMA-EMD on spike-contaminated
% duffing wave
% This cript is best run under Matlab R2016a and is tested on R2017a
% Code writer: Hui-Wen Yang 2019 @GICE, NTU and LIBSA, NCU
% Please cite:
% Hui-Wen Yang ; Shyh-Kang Jeng ; Hsu-Wen V. Young ; Chen Lin ; Yung-Hung Wang ; 
% Kun Hu ; Men-Tzung Lo. (2019). A Minimum Arclength Method for Removing Spikes 
% in Empirical Mode Decomposition. IEEE Access PP(99):1-1 ?January 2019
clear; close all;

cd .\MA_EMD_code
% Create dufing wave
% cycle length: 10; sample rate: 200; unit: hr
t0=0;
tf=200;
t=[t0:1/200:tf];
y0=[1,1];
[t2,y]=ode113('MTdufng',[t0:1/200:tf],y0);
s1=y(1:10:end,1)';
t2=t2(1:10:end);
fs = 20;

%% EMD on the clean signal (as a reference)
IMF22 = eemd(s1,0,1,4);


%% Add triangular spike
% spikeID = [563 1785 2119]; 
% spikeH = [0.7535 -0.6512 1.4802]*2;
slopes = [5, -8, 10]; centers = [563, 1.5; 1785,-1.3; 2119, 2.9];
spikeIDM = [563 2119]; spikeIDm = [1785];
s2 = s1;
for i=1:3
    outsig = create_triangle(s2,centers(i,:),slopes(i),fs);
   	s2 = outsig + s2;
end

%% original EMD on spike disturbed signal
IMF32 = eemd(s2,0,1,4);



%% Masking + MinArcLength EMD
nPh = 2; ad = std(s1)*0.1;
allmodemarc = cell(nPh,1);  % to save the  results IMF of each nPh
allnewYsM = cell(nPh,1); allnewYsm = cell(nPh,1);
% to save the result Y (adjusted spike position) of each nPh

for ii = 1:nPh
    maskinf=cos(2*pi*(3)*(0:length(s2)-1)/fs+(ii-1)*2*pi/nPh);
    z=s2+ad*maskinf;
    [allmode_minArc,newYsM,newYsm] = MA_EMD(z,spikeIDM,spikeIDm);
    allmodemarc{ii} = allmode_minArc;
    allnewYsM{ii} = newYsM;
    allnewYsm{ii} = newYsm;
end

% average IMFs of all nPh
allmodeMarc = allmodemarc{1};
for ii = 2:nPh
    allmodeMarc = allmodeMarc + allmodemarc{ii};
end
allmodeMarc = allmodeMarc/nPh;

%% Figure 1: Plot the input signal, duffing and spikes
fig1 = figure;
a1 = subplot(6,1,[1:2]);
    plot(t2,s1,'k','LineWidth',1.5)
    title('(a) s1: Duffing wave')
    ylim([-3, 3])
a2 = subplot(6,1,[3:4]);
    plot(t2,s2-s1,'k','LineWidth',1.5)
    ylim([-3, 3])
    title('(b) triangular spikes')
a3 = subplot(6,1,[5:6]);
    plot(t2,s2,'k','LineWidth',1.5);
    ylim([-3, 3])
    xlabel('Time (sec)')
    title('(c) s2: Duffing wave (a) + spike (b)')
    set([a1,a2],'xticklabel',{[]}) 

%% Figure 2: Plot EMD on s1 (original duffing) and s2 (contaminated)
fig2 = figure; 
a2 = subplot(6,1,[1,2]);
    plot(t2,IMF22(1,:),'k','LineWidth',1.5); hold on;
    plot(t2,IMF32(1,:),'r','LineWidth',1);
    axis tight
    ylabel('IMF 1')
    ylim([-2.2, 2.2])
    title('EMD on s1 (black) and s2 (red)')
    legend('EMD on s1','EMD on s2')
a3 = subplot(6,1,[3,4]);
    plot(t2,IMF22(2,:),'k','LineWidth',1.5); hold on;
    plot(t2,IMF32(2,:),'r','LineWidth',1);
    set([a2,a3],'xticklabel',{[]})
    axis tight
    ylabel('IMF 2')
    ylim([-0.3, 0.3])
a4 = subplot(6,1,[5,6]);
    plot(t2,IMF22(3,:),'k','LineWidth',1.5); hold all;
    plot(t2,IMF32(4,:),'r','LineWidth',1);
    axis tight
    ylabel('IMF 3')
    ylim([-0.4 0.4])
% plot(t2,IMF22(4,:)); %residual of EMD
linkaxes([a2,a3,a4],'x')
% xlim([0,150])
xlabel('Time (sec)')

%% Plot MAMA-EMD on s2
fig3 = figure; 
a1 = subplot(8,1,[1,2]);
    plot(t2,allmodeMarc(:,2),'b','LineWidth',1.5)
    axis tight
    title('MAMA-EMD on s2')
    legend('extracted spike')
    ylabel('IMF 0')
a2 = subplot(8,1,[3,4]);
    plot(t2,IMF22(1,:),'k','LineWidth',1.5); hold on;
    plot(t2,allmodeMarc(:,3),'r','LineWidth',1);
    ylabel('IMF 1')
    legend('EMD on s1','MAMA-EMD on s2')
a3 = subplot(8,1,[5,6]);
    plot(t2,IMF22(2,:),'k','LineWidth',1.5); hold on;
    plot(t2,allmodeMarc(:,4),'r','LineWidth',1);
    ylabel('IMF 2')
a4 = subplot(8,1,[7,8]);
    plot(t2,IMF22(3,:),'k','LineWidth',1.5); hold all;
    plot(t2,allmodeMarc(:,5),'r','LineWidth',1);
    set([a1,a2,a3],'xticklabel',{[]})
        ylabel('IMF 3')
linkaxes([a2,a3,a4,a1],'x')
xlim([0,200])
xlabel('Time (sec)')


%% Figure 4: Demo extrema re-adjust
fig4 = figure;
plot(t2,s2,'k','LineWidth',1.5);  % the input signal
hold on;
plot(t2,allmode_minArc(:,1),'b'); % the input + masking
plot(t2(spikeIDM),allnewYsM{2}(1,:,1),'*r')
plot(t2(spikeIDm),allnewYsm{2}(1,:,1),'^r')
[spmax, spmin, flag]=extrema(allmode_minArc(:,1));
spmax(85,2) = allnewYsM{2}(1,1,1);
spmax(317,2) = allnewYsM{2}(1,2,1);
spmin(267,2) = allnewYsm{2}(1,1,1);
upper= spline(spmax(:,1),spmax(:,2),spmax(1,1):spmax(end,1));
lower= spline(spmin(:,1),spmin(:,2),spmin(1,1):spmin(end,1));
plot(t2(spmax(1,1):spmax(end,1)),upper,'c')
plot(t2(spmin(1,1):spmin(end,1)),lower,'g')
legend('input signal','input+masking','adjusted maximum','adjusted minimum','upper envelop','lower envelop')

% close all;
% save('example_dufingspike_trig')
