clc
clear all
close all

% This code is used to perform nonlinear system identification using
% Hammerstein-Wiener model

load suspension.mat
Ts = 0.001; N = 100000;

roadx = {hxA, hxB, hxC, hxD, hxE};
roads = {hsA, hsB, hsC, hsD, hsE};

bodyx = {bodyxsuppressedA, bodyxsuppressedB, bodyxsuppressedC, bodyxsuppressedD, bodyxsuppressedE};
bodys = {bodysuppressedA, bodysuppressedB, bodysuppressedC, bodysuppressedD, bodysuppressedE};

tyrex = {tirexsuppressedA, tirexsuppressedB, tirexsuppressedC, tirexsuppressedD, tirexsuppressedE};
tyres = {tiresuppressedA, tiresuppressedB, tiresuppressedC, tiresuppressedD, tiresuppressedE};

bodynx = {4,6,10,10,2};
tyrenx = {5,8,10,8,10};

% Find excited harmonics
X = abs(fft(hsA)); c = find(X > 0.00001);

% body
for i = 1:5
    u_train = [roadx{i};roadx{i}]; y_train = bodyx{i}(N+1:3*N); y_train = y_train - mean(y_train);
    u_val_same = [roadx{i}]; y_val_same = bodyx{i}(3*N+1:4*N); y_val_same = y_val_same - mean(y_val_same);
    u_val_diff = [roads{i};roads{i};roads{i}]; y_val_diff = bodys{i}(N+1:4*N); y_val_diff = y_val_diff - mean(y_val_diff);
    Z2 = fft(bodys{i}(N+1:2*N)); Y2 = zeros(N,1); Y2(c) = Z2(c); y2 = ifft(Y2);
    Z3 = fft(bodys{i}(2*N+1:3*N)); Y3 = zeros(N,1); Y3(c) = Z3(c); y3 = ifft(Y3);
    Z4 = fft(bodys{i}(3*N+1:4*N)); Y4 = zeros(N,1); Y4(c) = Z4(c); y4 = ifft(Y4);
    y_val_lin = [y2;y3;y4];

    dat = iddata(y_train,u_train,Ts);
    nl1 = idPolynomial1D; nl1.Degree = 2; inputNL = [nl1]; outputNL = [nl1];
    initialModel = idnlhw([bodynx{i} bodynx{i} 1],inputNL,outputNL); initialModel.Ts = Ts;
    tic
    sysbodyNL{i} = nlhw(dat,initialModel);
    time_taken = toc;

    [yh,fit,X0] = compare(iddata(bodyx{i}(1:3*N),[roadx{i};roadx{i};roadx{i}],Ts),sysbodyNL{i});
    yest_train = yh.y; yest_train(1:N) = []; yest_train = yest_train - mean(yest_train);
    FIT_train = 100*(1-sqrt(sum((y_train - yest_train).^2)/sum((y_train - mean(y_train)).^2)));
    MSE_train = mean((y_train - yest_train).^2);

    yest_val_same = yest_train; yest_val_same(1:N) = [];
    FIT_val_same = 100*(1 - sqrt(sum((y_val_same-yest_val_same).^2)/sum((y_val_same - mean(y_val_same)).^2)));
    MSE_val_same = mean((y_val_same - yest_val_same).^2);

    [yh,fit,X0] = compare(iddata(bodys{i}(1:4*N),[roads{i};roads{i};roads{i};roads{i}],Ts),sysbodyNL{i});
    yest_val_diff = yh.y; yest_val_diff(1:N) = []; yest_val_diff = yest_val_diff - mean(yest_val_diff);
    FIT_val_diff = 100*(1 - sqrt(sum((y_val_diff - yest_val_diff).^2)/sum((y_val_diff - mean(y_val_diff)).^2)));
    MSE_val_diff = mean((y_val_diff - yest_val_diff).^2);

    [FIT_train FIT_val_same FIT_val_diff time_taken]

    figure(1)
    loglog([0.01:0.01:3000],(abs(fft(yh.y)).^2)/(3*N),'.','LineWidth',2)
    hold on
    grid on
    xlabel('Spatial Frequency [cycles/m]')
    ylabel('Frequency Response Magnitude')
    legend('Class A','Class B','Class C','Class D','Class E')
end


% tyre
for i = 1:5
    u_train = [roadx{i};roadx{i}]; y_train = tyrex{i}(N+1:3*N); y_train = y_train - mean(y_train);
    u_val_same = [roadx{i}]; y_val_same = tyrex{i}(3*N+1:4*N); y_val_same = y_val_same - mean(y_val_same);
    u_val_diff = [roads{i};roads{i};roads{i}]; y_val_diff = tyres{i}(N+1:4*N); y_val_diff = y_val_diff - mean(y_val_diff);
    Z2 = fft(tyres{i}(N+1:2*N)); Y2 = zeros(N,1); Y2(c) = Z2(c); y2 = ifft(Y2);
    Z3 = fft(tyres{i}(2*N+1:3*N)); Y3 = zeros(N,1); Y3(c) = Z3(c); y3 = ifft(Y3);
    Z4 = fft(tyres{i}(3*N+1:4*N)); Y4 = zeros(N,1); Y4(c) = Z4(c); y4 = ifft(Y4);
    y_val_lin = [y2;y3;y4];

    dat = iddata(y_train,u_train,Ts);
    nl1 = idPolynomial1D; nl1.Degree = 2; inputNL = [nl1]; outputNL = [nl1];
    initialModel = idnlhw([tyrenx{i} tyrenx{i} 1],inputNL,outputNL); initialModel.Ts = Ts;
    tic
    systyreNL{i} = nlhw(dat,initialModel);
    time_taken = toc;

    [yh,fit,X0] = compare(iddata(tyrex{i}(1:3*N),[roadx{i};roadx{i};roadx{i}],Ts),systyreNL{i});
    yest_train = yh.y; yest_train(1:N) = []; yest_train = yest_train - mean(yest_train);
    FIT_train = 100*(1 - sqrt(sum((y_train - yest_train).^2)/sum((y_train - mean(y_train)).^2)));
    MSE_train = mean((y_train - yest_train).^2);

    yest_val_same = yest_train; yest_val_same(1:N) = [];
    FIT_val_same = 100*(1 - sqrt(sum((y_val_same - yest_val_same).^2)/sum((y_val_same - mean(y_val_same)).^2)));
    MSE_val_same = mean((y_val_same - yest_val_same).^2);

    [yh,fit,X0] = compare(iddata(tyres{i}(1:4*N),[roads{i};roads{i};roads{i};roads{i}],Ts),systyreNL{i});
    yest_val_diff = yh.y; yest_val_diff(1:N) = []; yest_val_diff = yest_val_diff - mean(yest_val_diff);
    FIT_val_diff = 100*(1 - sqrt(sum((y_val_diff - yest_val_diff).^2)/sum((y_val_diff - mean(y_val_diff)).^2)));
    MSE_val_diff = mean((y_val_diff - yest_val_diff).^2);

    [FIT_train FIT_val_same FIT_val_diff time_taken]

    figure(2)
    loglog([0.01:0.01:3000],(abs(fft(yh.y)).^2)/(3*N),'.','LineWidth',2)
    hold on
    grid on
    xlabel('Spatial Frequency [cycles/m]')
    ylabel('Frequency Response Magnitude')
    legend('Class A','Class B','Class C','Class D','Class E')
end
