%% clear all variables and start fresh
clear variables

%% set seed

seed = 5;
rng(seed)
%% problem parameters

gamma = 1; % start with a fixed gamma

n_test = 1e5; % test data to evaluate oos distribution

%% data distribution parameters

% feature parameters
feat_cov_type = 'iso';

feat_type = 't';
feat_param = 5;

% signal parameters
sig_type = 'iso';
sig_energy = 1;
        
% noise parameters
noi_type = 't';
noi_param = 5;
noi_energy = 1;
        
% reponse type
mod_type = 'nlsq';

lambda = 1;
%% vectors to store the ks distances

n_max = 1500;
n_min = 100;
n_step = 100;

n_grid = n_min:n_step:n_max;

n_runs = 100;

ks_oosgcv_grid = zeros(length(n_grid), n_runs);
ks_oosloo_grid = zeros(length(n_grid), n_runs);

for i = 1:length(n_grid)
    
    disp(i)
    
    % get n and p for a given gamma
    n = n_grid(i);
    p = round(n * gamma);

    % fix feature covariance and signal for different runs
    feat_cov = generate_feature_covariance(p, feat_cov_type);        
    beta0 = sqrt(sig_energy) * generate_signal_vector(p, sig_type);

    for j = 1:n_runs
        %% generate train and test features
        
        X = generate_feature_matrix(n, p, feat_cov, feat_type, feat_param);
        X_test = generate_feature_matrix(n_test, p, feat_cov, feat_type, feat_param);

        %% generate train and test responses
       
        eps = sqrt(noi_energy) * generate_noise_vector(n, noi_type, noi_param);
        eps_test = sqrt(noi_energy) * generate_noise_vector(n_test, noi_type, noi_param);

        % response generation
        y = generate_response_vector(X, beta0, eps, mod_type);
        y_test = generate_response_vector(X_test, beta0, eps_test, mod_type);

        %% fit estimator and get smoother
        
        ridge_estim = fit_ridge(X, y, lambda);

        ridge_smoother = get_ridge_smoother(X, lambda);

        %% get predicted values and compute residuals

        hy = predict_linestim(ridge_estim, X);
        resid_orig = y - hy;

        hy_test = predict_linestim(ridge_estim, X_test);
        resid_test = y_test - hy_test;

        resid_gcv = resid_orig / (1 - trace(ridge_smoother) / n);
        resid_loo = resid_orig ./ (1 - diag(ridge_smoother));

        [~, ~, ks_oosgcv] = kstest2(resid_test, resid_gcv);
        [~, ~, ks_oosloo] = kstest2(resid_test, resid_loo);

        ks_oosgcv_grid(i, j) = ks_oosgcv;
        ks_oosloo_grid(i, j) = ks_oosloo;
    end
end

%% plot the ks distance as a function of n

plot(n_grid, mean(ks_oosgcv_grid, 2)', ...
    'DisplayName', 'oos vs. gcv')
hold on
grid on
plot(n_grid, mean(ks_oosloo_grid, 2), ...
    'DisplayName', 'oos vs. loo')
legend show
xlabel('n')
ylabel('KS distance')

% save the ks distance plot

save_str = strcat(...
    'ks_', ...
    '10gamma', num2str(10 * gamma), '_', ...
    'nt', num2str(n_test), '_', ...
    'nr', num2str(n_runs), '_', ...
    'featcov', feat_cov_type, '_', ...
    'feat', feat_type, '_', ...
    '10featparam', num2str(10 * feat_param), '_', ...
    'sig', sig_type, '_', ...
    'noi', noi_type, '_', ...
    '10noiparam', num2str(10 * noi_param), '_', ...
    'mod', mod_type, '_', ...
    '10lam', num2str(10 * lambda), '_', ...
    'seed', num2str(seed));
% 
% save_folder = '../figures/';
% 
% saveas(h_oos, fullfile(save_folder, save_str), 'epsc')

save_fig_str = strcat('../figures/ks/', save_str);
saveas(gcf, save_fig_str)
saveas(gcf, save_fig_str, 'epsc')
saveas(gcf, save_fig_str, 'png')

%% save data

save_data_str = strcat('../data/ks/', save_str, '.mat');

save(save_data_str, ...
    'feat_cov', 'feat_cov_type','feat_param', ...
    'n_grid', 'gamma', 'lambda', 'mod_type', 'n_test', ...
    'noi_energy', 'noi_param', 'noi_type', 'sig_energy', ...
    'seed', 'sig_type', ...
    'ks_oosgcv_grid', 'ks_oosloo_grid')