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

coverage_grid = 0.5:.05:0.95; % go from 0.5 to 0.95

% to store empirical coverage values
coverageval_gcv_grid = zeros(length(coverage_grid), length(n_grid));
coverageval_loo_grid = zeros(length(coverage_grid), length(n_grid));

% to store length of prediction intervals
predintval_length_gcv_grid = zeros(length(coverage_grid), length(n_grid));
predintval_length_loo_grid = zeros(length(coverage_grid), length(n_grid));

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

    %% 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, feat_cov);
    y_test = generate_response_vector(X_test, beta0, eps_test, mod_type, feat_cov);

    %% 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));
        
    
    %% form prediction intervals

    % to store end points of the interval for gcv
    predintval_gcv_grid = zeros(n_test, 2, length(coverage_grid));

    % to store end points of the interval for loo
    predintval_loo_grid = zeros(n_test, 2, length(coverage_grid));

    % to store upper and lower qunatiles for gcv and loo
    quanthigh_gcv_grid = zeros(length(coverage_grid), 1);
    quantlow_gcv_grid = zeros(length(coverage_grid), 1);
    quanthigh_loo_grid = zeros(length(coverage_grid), 1);
    quantlow_loo_grid = zeros(length(coverage_grid), 1);

    % to store coverage checks for gcv and loo
    coverageind_gcv_grid = zeros(n_test, length(coverage_grid));
    coverageind_loo_grid = zeros(n_test, length(coverage_grid));

    for j = 1:length(coverage_grid)

        quanthigh_gcv_grid(j) = quantile(resid_gcv, 1 - (1 - coverage_grid(j))/2);
        quanthigh_loo_grid(j) = quantile(resid_loo, 1 - (1 - coverage_grid(j))/2);

        quantlow_gcv_grid(j) = quantile(resid_gcv, (1- coverage_grid(j))/2);
        quantlow_loo_grid(j) = quantile(resid_loo, (1- coverage_grid(j))/2);

        predintval_gcv_grid(:, 1, j) = hy_test + quantlow_gcv_grid(j);
        predintval_gcv_grid(:, 2, j) = hy_test + quanthigh_gcv_grid(j);

        predintval_loo_grid(:, 1, j) = hy_test + quantlow_loo_grid(j);
        predintval_loo_grid(:, 2, j) = hy_test + quanthigh_loo_grid(j);

        coverageind_gcv_grid(:, j) = bitand( ...
            (y_test <= predintval_gcv_grid(:, 2, j)), ...
            (y_test >= predintval_gcv_grid(:, 1, j)));

        coverageind_loo_grid(:, j) = bitand( ...
            (y_test <= predintval_loo_grid(:, 2, j)), ...
            (y_test >= predintval_loo_grid(:, 1, j)));

        coverageval_gcv_grid(i, j) = mean(coverageind_gcv_grid(:, j));
        coverageval_loo_grid(i, j) = mean(coverageind_loo_grid(:, j));

        predintval_length_gcv_grid(i, j) = ...
            quanthigh_gcv_grid(j) - quantlow_gcv_grid(j);
        predintval_length_loo_grid(i, j) = ...
            quanthigh_loo_grid(j) - quantlow_loo_grid(j);
    end

end

%% plots

% coveage
figure
hold on
grid on

for j = 1:length(coverage_grid)
    
    plot(n_grid, coverageval_gcv_grid(:, j), ...
        '+-', ...
        'DisplayName', num2str(coverage_grid(j)));
    hold on
    grid on
    plot(n_grid, coverageval_loo_grid(:, j), ...
        '*-', ...
        'DisplayName', '');
end

legend show
legend box off
legend('Location', 'Best')
xlabel('n')
ylabel('empirical coverage')

save_str = strcat(...
    'coverage_vs_n', ...
    '10gamma', num2str(10 * gamma), '_', ...
    'nt', num2str(n_test), '_', ...
    '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(100 * lambda), '_', ...
    'seed', num2str(seed));

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

% length
figure
hold on
grid on

for j = 1:length(coverage_grid)
    
    plot(n_grid, predintval_length_gcv_grid(:, j), ...
        '+-', ...
        'DisplayName', num2str(coverage_grid(j)));
    hold on
    grid on
    plot(n_grid, predintval_length_loo_grid(:, j), ...
        '*-', ...
        'DisplayName', '');
end

legend show
legend box off
legend('Location', 'Best')
xlabel('n')
ylabel('empirical coverage')

save_str = strcat(...
    'length_vs_n', ...
    '10gamma', num2str(10 * gamma), '_', ...
    'nt', num2str(n_test), '_', ...
    '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(100 * lambda), '_', ...
    'seed', num2str(seed));

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