%% Load pre-defined data
load('ue_positions.mat');  % [numSteps x numUEs x 2] Position of each UE
load('task_mat.mat');       % [numSteps x numUEs] Given task for each UE
load('PL_BS.mat');          % [numSteps x numUEs x numBS] Path loss between UEs and BSs
load('PL_D2D.mat');         % [numSteps x numUEs x numUEs] Path loss between UEs

%% System parameters
[numSteps,numUEs] = size(task_mat);
numBS             = size(PL_BS,3);
time_step         = 0.01;        % seconds per slot
UE_procRate       = 175e6;       % bits/sec
BS_procRate       = 10e9;        % bits/sec
BW_BS             = 100e6;       % Hz
BW_D2D            = 5e6;         % Hz
Pt_BS             = 23;          % dBm
Pt_D2D            = 23;          % dBm
noise_dBm         = -174;        % dBm/Hz
noise_BS_dBm      = noise_dBm + 10*log10(BW_BS);
noise_D2D_dBm     = noise_dBm + 10*log10(BW_D2D);

% Convert dBm to linear Watts
Pt_lin_BS    = 10^(Pt_BS/10)/1000;      % W
noise_lin_BS = 10^(noise_BS_dBm/10)/1000; % W

%% State variables
UE_next_free = zeros(numUEs,1); % Matrix to save the time when each UE would get idle
BS_next_free = zeros(numBS,1); % Time when each BS gets idle
all_delays   = [];
cnt_local = 0; cnt_bs = 0; cnt_d2d = 0; cnt_total = 0;

%% Simulation loop
for t = 1:numSteps
    currentTime = (t-1)*time_step;
    tasks       = task_mat(t,:)';              % [numUEs x 1]
    pos         = squeeze(ue_positions(t,:,:));% [numUEs x 2]
    PLbs        = squeeze(PL_BS(t,:,:));       % [numUEs x numBS]
    PLd2d       = squeeze(PL_D2D(t,:,:));      % [numUEs x numUEs]

    % Track who has offloaded to each BS this slot
    txToBS   = false(numUEs,1);
    targetBS = zeros(numUEs,1);
    
    activeUEs = find(tasks>0); % only when the UE has task
    for i = activeUEs'
        cnt_total = cnt_total + 1;
        S = tasks(i);
        waitSelf = max(0, UE_next_free(i) - currentTime); % queue wait time on itself
        delay_local = waitSelf + S/UE_procRate;

        %% BS candidate delays (with interference & BW sharing)
        pl_i = PLbs(i,:);  % 1 x numBS in dB
        
        % compute interference power and number of contending UEs per BS
        interf_bs = zeros(1,numBS);
        nCont     = zeros(1,numBS);
        for bs = 1:numBS
            offs = find(txToBS & targetBS==bs);
            nCont(bs) = numel(offs) + 1;  % include current UE
            if isempty(offs)
                interf_bs(bs) = 0;
            else
                interf_bs(bs) = sum(Pt_lin_BS .* 10.^(-PLbs(offs,bs)/10));
            end
        end
        % signal power
        sig_bs = Pt_lin_BS * 10.^(-pl_i/10);
        % SINR [-]
        sinr_bs_lin = sig_bs ./ (noise_lin_BS + interf_bs);
        % effective bandwidth share per UE
        bw_eff = BW_BS ./ nCont;
        % compute rates and times
        rate_bs = bw_eff .* log2(1 + sinr_bs_lin);
        tx_bs   = S ./ rate_bs;            % 1 x numBS
        waitBS_vec  = max(0, BS_next_free - currentTime);  % 1 x numBS
        waitBS_vec = transpose(waitBS_vec);
        proc_bs = S/BS_procRate;           % scalar
        delay_bs_vec = waitBS_vec + tx_bs + proc_bs; % 1 x numBS
        [delay_bs, bestBS] = min(delay_bs_vec);

        %% D2D candidate delays
        dists = sqrt(sum((pos - pos(i,:)).^2,2));
        dists(i) = Inf;
        peers = find(dists < 200);
        if isempty(peers)
            delay_peer = Inf;
        else
            pl_p       = PLd2d(i,peers);              % 1 x P
            pl_p = transpose(pl_p);
            sinr_p_lin = 10.^((Pt_D2D - pl_p - noise_D2D_dBm)/10);
            rate_p     = BW_D2D * log2(1 + sinr_p_lin);
            tx_p_col   = S ./ rate_p;             % P x 1
            waitP      = max(0, UE_next_free(peers) - currentTime); % P x 1
            proc_p     = S/UE_procRate;               % scalar
            delay_p_vec= waitP + tx_p_col + proc_p;  % P x 1
            [delay_peer, idx] = min(delay_p_vec);     % scalar
            bestPeer   = peers(idx);
        end

        %% Choose action with minimum delay
        [minDelay, choice] = min([delay_local, delay_bs, delay_peer]);
        switch choice
            case 1  % local
                UE_next_free(i) = max(UE_next_free(i),currentTime) + S/UE_procRate;
                cnt_local = cnt_local + 1;
            case 2  % BS offload
                % schedule FIFO service at BS
                arrival = currentTime + tx_bs(bestBS);
                startP  = max(BS_next_free(bestBS), arrival);
                BS_next_free(bestBS) = startP + proc_bs;
                % track for interference/BW sharing
                txToBS(i)   = true;
                targetBS(i) = bestBS;
                cnt_bs = cnt_bs + 1;
            case 3  % D2D offload
                arrival = currentTime + tx_p_col(idx);
                startP  = max(UE_next_free(bestPeer), arrival);
                UE_next_free(bestPeer) = startP + proc_p;
                cnt_d2d = cnt_d2d + 1;
        end
        all_delays(end+1,1) = minDelay;
    end
end

%% Results
avgDelay = mean(all_delays);
fprintf('Avg delay: %.4f s\n', avgDelay);
counts = [cnt_local, cnt_bs, cnt_d2d];
fprintf(' Local Ratio: %.2f%% \n', cnt_local/cnt_total*100);
fprintf(' BS Ratio: %.2f%% \n', cnt_bs/cnt_total*100);
fprintf(' D2D Ratio: %.2f%% \n', cnt_d2d/cnt_total*100);
figure; pie(counts, {'Local','BS','D2D'}); title('Mode share');
figure; histogram(all_delays,50,'Normalization','cdf'); xlabel('Delay (s)'); ylabel('CDF');
