%  This code was originally developed by The Center for Engineering Applications &
%  Technology Solution (CEATS), Ho Chi Minh City, Open University, Vietnam.
%  https://ceats.ou.edu.vn/us/ 
%  programming: Hoang-Le Minh
%
%  paper:Optimization of high-rise concrete structure using a sequence design based on a proposed surrogate model
%        (PSM) and Finite element (FE) model updating
%
%    
%  Hoang-Le Minh, Thanh Sang-To, Binh Le-Van, Thanh Cuong-Le
%
%  Email: hoang.lm@ou.edu.vn, cuong.lt@ou.edu.vn
%
%  Journal: Archives of Civil and Mechanical Engineering

clear; clc;
close all
%% Load data and normalize
input_data = load('input_section.txt')';         % N samples x D features
output_nonDD = load('output_nonDD.txt');
target_data = output_nonDD(1,:) ./ max(output_nonDD(1,:))'; % N x 1

X = input_data;
Y = target_data(:);

%% Split data: 70% training, 15% validation, 15% testing
cv = cvpartition(length(Y), 'HoldOut', 0.3);
X_train = X(cv.training, :);
Y_train = Y(cv.training);

X_rest = X(cv.test, :);
Y_rest = Y(cv.test);

cv_val_test = cvpartition(length(Y_rest), 'HoldOut', 0.5);
X_val = X_rest(cv_val_test.training, :);
Y_val = Y_rest(cv_val_test.training);

X_test = X_rest(cv_val_test.test, :);
Y_test = Y_rest(cv_val_test.test);

%% Build a custom KNN regression model
k = 5;

% KNN regression prediction function
predict_knn = @(Xq, Xtrain, Ytrain, k) mean(Ytrain(knnsearch(Xtrain, Xq, 'K', k)), 2);

% Make predictions
Y_pred_train = predict_knn(X_train, X_train, Y_train, k);
Y_pred_val   = predict_knn(X_val, X_train, Y_train, k);
Y_pred_test  = predict_knn(X_test, X_train, Y_train, k);
Y_pred_all   = predict_knn(X, X_train, Y_train, k);

% Compute MSE
mse_train = mean((Y_train - Y_pred_train).^2);
mse_val   = mean((Y_val - Y_pred_val).^2);
mse_test  = mean((Y_test - Y_pred_test).^2);
mse_all   = mean((Y - Y_pred_all).^2);

fprintf('MSE Train: %.6f\n', mse_train);
fprintf('MSE Validation: %.6f\n', mse_val);
fprintf('MSE Test: %.6f\n', mse_test);
fprintf('MSE All: %.6f\n', mse_all);

%% Plot regression results
figure;
plotregression(Y_train, Y_pred_train, 'Train Data');
    ax = gca;
    ax.FontSize = 15;
figure;
plotregression(Y_val, Y_pred_val, 'Validation Data');
    ax = gca;
    ax.FontSize = 15;
figure;
plotregression(Y_test, Y_pred_test, 'Test Data');
    ax = gca;
    ax.FontSize = 15;
figure;
plotregression(Y, Y_pred_all, 'All Data');
    ax = gca;
    ax.FontSize = 15;

%% Plot MSE bar chart
figure;
bar([mse_train, mse_val, mse_test, mse_all]);
set(gca, 'XTickLabel', {'Train', 'Validation', 'Test', 'All'});
ylabel('Mean Squared Error (MSE)');
title('MSE of KNN on different data sets');
grid on;
set(gca, 'FontSize', 14);
