clc; clear; close all;

parameters;

P_fas = [1e-7, 5e-7, 1e-6, 5e-6, 1e-5, 5e-5, 1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1, 1e0];

P_d_FARs = zeros(length(P_fas), 1);
P_d_narrows = zeros(length(P_fas), 1);


%% FAR

% Signal Model
% epi = B / f_c;
epi = 0;
[C_n, A] = get_Psi(FAR_N, FAR_M, epi);
fac = A * A';
A = A / sqrt(abs(fac(1, 1)));
f_n = f_c + C_n * B / M;

[Lambda, Lambda_C, x] = get_sparse_vector(FAR_N, FAR_M, betas_wide, true);

% Define filename
filename = ...
    "data2/" + ...
    "FAR" + "_" + ...
    string(FAR_N) + "_" + ...
    string(FAR_M)+ "_" + ...
    method + "_" + ...
    string(noise_sigma) + "_" + ...
    string(length(Lambda)) + "_" + ...
    string(LASSO_lambda) + "_" + ...
    trail_times + ...
    ".mat";


for P_fa_idx = 1: length(P_fas)
    P_fa = P_fas(P_fa_idx);
    if exist(filename, "file")
        load( ...
            filename, ...
            "recovery_results", "sigma_w2", "thresholds",...
            "C_n", "A", "Lambda", "Lambda_C", "x" ...
        );
    else
        % Mento Carlo Recovery
        [recovery_results, sigma_w2, thresholds] = All_Recovery(A, x, P_fa, noise_sigma / FAR_M);
    
        save( ...
            filename, ...
            "recovery_results", "sigma_w2", "thresholds",...
            "C_n", "A", "Lambda", "Lambda_C", "x" ...
        );
    end
    
    % Compare sigmas (from paper) and vars (from simulation)
    % [ss, vs, H0_mean, H0_max, H1_mean, H1_max] = get_sigma_var(recovery_results, sigma_w2, x);
    
    T_opt = zeros(2, 1);
    
    delta_r = c / (2 * B);
    tmp_1 = repelem(f_n, FAR_M);
    tmp_2 = repmat(1: FAR_M, 1, FAR_N)';
    tmp_3 = exp(-1 * 1j * 2 * pi * 2 * delta_r .* tmp_2 .* tmp_1 / c);
    
    idx_1 = 1;
    idx_2 = 1;
    i1 = 1; i2 = 1;

    Lambda_distributes = zeros(2, length(Lambda) * trail_times);
    Lambda_C_distributes = zeros(2, length(Lambda_C) * trail_times);

    for t = 1:trail_times
        % 0 假设
        x_0_hat = squeeze(recovery_results(1, t, :));
        s = x_0_hat .* tmp_3;
        for i = 1: FAR_N
            L = (i - 1) * FAR_M + 1;
            R = L - 1 + FAR_M;
            T_opt(1, idx_1) = foo(s(L:R));
            idx_1 = idx_1 + 1;
        end

        % 1 假设
        x_1_hat = squeeze(recovery_results(2, t, :));
        s = x_1_hat .* tmp_3;
        T_opt(2, idx_2) = foo(s(Lambda));
        idx_2 = idx_2 + 1;

        for i = 1: length(Lambda) + length(Lambda_C)
            if ismember(i, Lambda_C)
                Lambda_C_distributes(1, i1) = x_0_hat(i);
                Lambda_C_distributes(2, i1) = x_1_hat(i);
                i1 = i1 + 1;
            elseif ismember(i, Lambda)
                Lambda_distributes(1, i2) = x_0_hat(i);
                Lambda_distributes(2, i2) = x_1_hat(i);
                i2 = i2 + 1;
            end
        end
    end
    
    % figure;
    % subplot(211); plot(real(sum(squeeze(recovery_results(1, 1:10, :)))));
    % subplot(212); plot(real(sum(squeeze(recovery_results(2, 1:10, :)))));
    
    % figure;
    % subplot(211); histfit(T_opt(1, 1:idx_1 - 1));
    % subplot(212); histfit(T_opt(2, 1:idx_2 - 1));
    
    figure;
    subplot(2, 2, 1); histfit(real(Lambda_C_distributes(1, :)));
    subplot(2, 2, 2); histfit(real(Lambda_distributes(1, :)));
    subplot(2, 2, 3); histfit(real(Lambda_C_distributes(2, :)));
    subplot(2, 2, 4); histfit(real(Lambda_distributes(2, :)));

    figure;
    histfit(real([Lambda_distributes(2, :) Lambda_C_distributes(2, :)]));

    means = [ ...
                 mean(Lambda_C_distributes(1, :)), ...
                 mean(Lambda_distributes(1, :)), ...
                 mean(Lambda_C_distributes(2, :)), ...
                 mean(Lambda_distributes(2, :)), ...
                 mean([Lambda_distributes(1, :) Lambda_C_distributes(1, :)]), ...
                 mean([Lambda_distributes(1, :) Lambda_C_distributes(1, :) Lambda_C_distributes(2, :)]), ...
                 mean([Lambda_distributes(2, :)]) ...
             ];

    stds = [ ...
                std(Lambda_C_distributes(1, :)), ...
                std(Lambda_distributes(1, :)), ...
                std(Lambda_C_distributes(2, :)), ...
                std(Lambda_distributes(2, :)), ...
                 std([Lambda_distributes(1, :) Lambda_C_distributes(1, :)]), ...
                 std([Lambda_distributes(1, :) Lambda_C_distributes(1, :) Lambda_C_distributes(2, :)]), ...
                 std([Lambda_distributes(2, :)]) ...
            ];


    fprintf("H_00: mu = %.15f, std = %.15f\n", means(1), stds(1));
    fprintf("H_01: mu = %.15f, std = %.15f\n", means(2), stds(2));
    fprintf("H_10: mu = %.15f, std = %.15f\n", means(3), stds(3));
    fprintf("H_11: mu = %.15f, std = %.15f\n", means(4), stds(4));
    fprintf("H_0: mu = %.15f, std = %.15f\n", means(5), stds(5));
    fprintf("H_0 + Lambda^C: mu = %.15f, std = %.15f\n", means(6), stds(6));
    fprintf("H_1 Lambda: mu = %.15f, std = %.15f\n\n", means(7), stds(7));
    

    H_0_mean = mean(T_opt(1, 1:idx_1-1));
    H_0_std = std(T_opt(1, 1:idx_1-1));
    H_1_mean = mean(T_opt(2, 1:idx_2-1));
    H_1_std = std(T_opt(2, 1:idx_2-1));


    threshold = norminv(1 - P_fa, H_0_mean, H_0_std);
    P_d = 1 - normcdf(threshold, H_1_mean, H_1_std);
    P_d_FARs(P_fa_idx) = P_d;
end

format long;
fprintf("H_0 mean: %.15f\n", H_0_mean);
fprintf("H_0 std: %.15f\n", H_0_std);
fprintf("H_1 mean: %.15f\n", H_1_mean);
fprintf("H_1 std: %.15f\n", H_1_std);

%% PD
[s_T_narrow, A] = narrow_signal_model_2(FAR_N);
[Lambda, Lambda_C, x] = get_sparse_vector(FAR_N, 1, betas_narrow, false);
filename = ...
    "data2/" + ...
    "PD" + "_" + ...
    method + "_" + ...
    string(noise_sigma) + "_" + ...
    string(length(Lambda)) + "_" + ...
    trail_times + ...
    ".mat";

for P_fa_idx = 1: length(P_fas)
    [s_T, A] = narrow_signal_model_2(FAR_N);
    P_fa = P_fas(P_fa_idx);
    if exist(filename, "file")
        load( ...
            filename, ...
            "recovery_results", "sigma_w2", "thresholds",...
            "f_c", "A", "Lambda", "Lambda_C", "x" ...
        );
    else
        % Mento Carlo Recovery
        [recovery_results, sigma_w2, thresholds] = All_Recovery(A, x, P_fa, noise_sigma);
    
        save( ...
            filename, ...
            "recovery_results", "sigma_w2", "thresholds",...
            "f_c", "A", "Lambda", "Lambda_C", "x" ...
        );
    end

    T_opt = zeros(2, trail_times);
        
    delta_r = T_p * c / 2;
    tmp_1 = (1: FAR_N)';
    tmp_3 = exp(-1 * 1j * 2 * pi * 2 * delta_r * f_c .* tmp_1 / c);

    for t = 1:trail_times
        % 0 假设
        x_hat = recovery_results(1, t, :);
        s = squeeze(x_hat) .* tmp_3;
        for i = 1: FAR_N
            T_opt(1, idx_1) = foo(s(i));
            idx_1 = idx_1 + 1;
        end

        % 1 假设
        x_hat = recovery_results(2, t, :);
        s = squeeze(x_hat) .* tmp_3;
        T_opt(2, idx_2) = foo(s(Lambda));
        idx_2 = idx_2 + 1;
    end

    % figure;
    % subplot(211); plot(real(sum(squeeze(recovery_results(1, 1:10, :)))));
    % subplot(212); plot(real(sum(squeeze(recovery_results(2, 1:10, :)))));

    H_0_mean = mean(T_opt(1, 1:idx_1-1));
    H_0_std = std(T_opt(1, 1:idx_1-1));
    H_1_mean = mean(T_opt(2, 1:idx_2-1));
    H_1_std = std(T_opt(2, 1:idx_2-1));
    
    threshold = norminv(1 - P_fa, H_0_mean, H_0_std);
    P_d = 1 - normcdf(threshold, H_1_mean, H_1_std);
    P_d_narrows(P_fa_idx) = P_d;
end

format long;
fprintf("H_0 mean: %.15f\n", H_0_mean);
fprintf("H_0 std: %.15f\n", H_0_std);
fprintf("H_1 mean: %.15f\n", H_1_mean);
fprintf("H_1 std: %.15f\n", H_1_std);


%% ROC

figure;
hold on;
grid on;
semilogx(P_fas, P_d_FARs);
semilogx(P_fas, P_d_narrows);
xlim([8e-8, 1]);
% ylim([0, 1]);
xlabel("P_{fa}");
ylabel("P_{d}");
title("ROC");