%  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;

%% Load data and normalize
input_data = load('input_section.txt');
output_data_raw = load('output_nonDD.txt');
target_data = output_data_raw(1,:) ./ max(output_data_raw(1,:)); % Normalize the first row

%% List of network architectures to test
hidden_options = {
    [50 100 50] % Only test one architecture, you can add more here
};

K = 5;
cv = cvpartition(size(input_data,2), 'KFold', K);

fprintf("=== K-FOLD CROSS-VALIDATION ===\n");
best_perf = inf;
best_hidden = [];
mse_history = cell(length(hidden_options), K); % Store MSE per epoch

for i = 1:length(hidden_options)
    hiddenLayerSize = hidden_options{i};
    fprintf('\n[Model %d] Hidden Layers: %s\n', i, mat2str(hiddenLayerSize));
    perf_all = zeros(K,1);
    
    for k = 1:K
        trainIdx = training(cv, k);
        testIdx = test(cv, k);
        
        inputTrain = input_data(:, trainIdx);
        targetTrain = target_data(:, trainIdx);
        
        inputTest = input_data(:, testIdx);
        targetTest = target_data(:, testIdx);

        % Initialize network
        net = fitnet(hiddenLayerSize, 'traingda');
        net.divideFcn = 'dividetrain'; % Use the entire fold for training
        net.trainParam.epochs = 2000;
        net.trainParam.goal = 1e-5;
        net.trainParam.lr = 0.01;
        net.trainParam.lr_inc = 1.05;
        net.trainParam.lr_dec = 0.7;
        net.trainParam.mc = 0.1;
        net.trainParam.min_grad = 1e-5;
        net.trainParam.showWindow = false;
        net.trainParam.showCommandLine = false;

        % Train the network
        [net, tr] = train(net, inputTrain, targetTrain);

        % Predict and compute MSE
        predict = net(inputTest);
        perf = perform(net, targetTest, predict);
        perf_all(k) = perf;

        % Store MSE per epoch
        mse_history{i,k} = tr.perf;
    end

    avg_perf = mean(perf_all);
    fprintf('>> Avg. Performance (MSE): %.6f\n', avg_perf);
    
    if avg_perf < best_perf
        best_perf = avg_perf;
        best_hidden = hiddenLayerSize;
    end
end

fprintf("\n=== BEST HYPERPARAMETER ===\n");
fprintf("Best Hidden Layers: %s\n", mat2str(best_hidden));
fprintf("Best Avg. MSE: %.6f\n", best_perf);

%% PLOT MSE FOR EACH FOLD (per epoch)
fprintf("\n=== PLOT MSE PER FOLD ===\n");
for k = 1:K
    mse_k = mse_history{1, k}; % Assuming using the first model (i = 1)
    figure;
    semilogy(1:length(mse_k), mse_k, 'LineWidth', 2);
    xlabel('Epoch');
    ylabel('MSE (log scale)');
    title(sprintf('Training Performance - Fold %d', k));
    grid on;

    % Mark best MSE
    [best_val, best_epoch] = min(mse_k);
    hold on;
    plot(best_epoch, best_val, 'go', 'MarkerSize', 8, 'LineWidth', 2);
    text(best_epoch, best_val, sprintf('\\leftarrow Best MSE = %.6f', best_val), ...
        'VerticalAlignment', 'bottom', 'FontSize', 10);
end

%% FINAL TRAINING ON FULL DATASET (70/15/15)
fprintf("\n=== FINAL TRAINING ON ALL DATA (70/15/15 split) ===\n");

% Initialize network with best configuration
net_final = fitnet(best_hidden, 'traingda');
net_final.divideParam.trainRatio = 0.7;
net_final.divideParam.valRatio   = 0.15;
net_final.divideParam.testRatio  = 0.15;

% Training parameters
net_final.trainParam.epochs = 2000;
net_final.trainParam.goal = 1e-5;
net_final.trainParam.lr = 0.01;
net_final.trainParam.lr_inc = 1.05;
net_final.trainParam.lr_dec = 0.7;
net_final.trainParam.mc = 0.1;
net_final.trainParam.min_grad = 1e-5;
net_final.trainParam.showWindow = true;
net_final.trainParam.showCommandLine = true;

% Train on full dataset
[net_final, tr_final] = train(net_final, input_data, target_data);

% Predict and final performance
predict_all = net_final(input_data);
performance = perform(net_final, target_data, predict_all);
fprintf('Final performance (All data): %.6f\n', performance);

% Regression plot for all data
figure;
plotregression(target_data, predict_all, 'Final Trained Model (All Data)');

% Plot training/validation/test performance
figure;
plotperform(tr_final);

% Regression plots for each subset
trainInd = tr_final.trainInd;
valInd   = tr_final.valInd;
testInd  = tr_final.testInd;

figure; plotregression(target_data(:,trainInd), net_final(input_data(:,trainInd)), 'Train');
figure; plotregression(target_data(:,valInd),   net_final(input_data(:,valInd)),   'Validation');
figure; plotregression(target_data(:,testInd),  net_final(input_data(:,testInd)),  'Test');
