clear 
close all


% Load data
edfFile = 'r10.edf';
info = edfinfo(edfFile);
data1 = edfread(edfFile, 'SelectedSignals', 'Direct_1');
d = cell2mat(table2cell(data1));
data2 = edfread(edfFile, 'SelectedSignals', 'Abdomen_1');
B2 = cell2mat(table2cell(data2));
data3 = edfread(edfFile, 'SelectedSignals', 'Abdomen_2');
B3 = cell2mat(table2cell(data3));
data4 = edfread(edfFile, 'SelectedSignals', 'Abdomen_3');
B4 = cell2mat(table2cell(data4));
data5 = edfread(edfFile, 'SelectedSignals', 'Abdomen_4');
B5 = cell2mat(table2cell(data5));


% Setup parameters
Fs = 1000;
Ts = 1/Fs;
Tm = (0:length(d)-1) / Fs;


% First-stage signal pre-processing for noise reduction
% Apply bandpass filter to eliminate baseline wander and high-frequency noise
b_notch = designfilt('bandstopiir', 'FilterOrder', 2, ...
                  'HalfPowerFrequency1', 48, 'HalfPowerFrequency2', 52, ...
                  'DesignMethod', 'butter', 'SampleRate', Fs);


b_bp = designfilt('bandpassiir', 'FilterOrder', 4, ...
                'HalfPowerFrequency1', 1, 'HalfPowerFrequency2', 100, ...
                'DesignMethod', 'butter', 'SampleRate', Fs);


% Apply filters to abdominal signals
B2_filtered = filtfilt(b_notch, B2);
B2_filtered = filtfilt(b_bp, B2_filtered);
B3_filtered = filtfilt(b_notch, B3);
B3_filtered = filtfilt(b_bp, B3_filtered);
B4_filtered = filtfilt(b_notch, B4);
B4_filtered = filtfilt(b_bp, B4_filtered);
B5_filtered = filtfilt(b_notch, B5);
B5_filtered = filtfilt(b_bp, B5_filtered);


% Apply filters to direct signal (reference)
d_filtered = filtfilt(b_notch, d);
d_filtered = filtfilt(b_bp, d_filtered);


% Differential approach using multiple abdominal channels
% Creates a better reference signal for extracting fetal ECG
noisy_sig = B2_filtered;


% Setup Cuckoo Search Algorithm for optimizing LMS parameters
% Cuckoo Search parameters
n_nests = 25;           % Number of nests (population size)
max_iterations = 25;    % Maximum number of iterations
pa = 0.25;              % Probability of abandoning worse nests


% Search space boundaries
min_order = 80;         % Minimum filter order
max_order = 300;        % Maximum filter order (reduced for efficiency)
min_mu = 0.01;          % Min learning rate
max_mu = 0.99;          % Max learning rate


% Initialize nests [order, step_size]
nests = zeros(n_nests, 2);
fitness = zeros(n_nests, 1);


% Generate initial nests with Lévy flight distribution
for i = 1:n_nests
    nests(i, :) = [randi([min_order, max_order]), rand * (max_mu - min_mu) + min_mu];
end


% Select shorter segment for optimization to speed up process
start_time = 0;
end_time = 10; 
start_idx = find(Tm >= start_time, 1);
end_idx = find(Tm >= end_time, 1);
d_opt = d_filtered(start_idx:end_idx);
noisy_opt = noisy_sig(start_idx:end_idx);


% Track best solution
best_fitness = -inf;
best_solution = [0, 0];
fitness_history = zeros(max_iterations, 1);


% Define Lévy flight function
function s = levy_flight(beta)
    % Mantegna's algorithm for Lévy flights
    sigma_u = (gamma(1+beta)*sin(pi*beta/2)/(gamma((1+beta)/2)*beta*2^((beta-1)/2)))^(1/beta);
    u = randn(1) * sigma_u;
    v = randn(1);
    s = u/abs(v)^(1/beta);
end


% Cuckoo Search main loop
for iter = 1:max_iterations
    % Evaluate fitness for each nest
    for i = 1:n_nests
        order = round(nests(i, 1));
        mu = nests(i, 2);


        % Run LMS filter with these parameters
        lms = dsp.LMSFilter(order + 1, 'StepSize', mu, 'Method', 'Normalized LMS', 'WeightsOutputPort', true);
        [y, e, w] = step(lms, noisy_opt, d_opt);


        % Use wavelet-based SNR computation to better evaluate signal quality
        [~, fetal_signal] = extract_fetal_wavelet(e);
        curr_snr = wSNR(d_opt, fetal_signal);


        % Balance SNR with computational complexity
        % Higher weight on SNR (0.9) vs order reduction (0.1)
        fitness(i) = 0.9 * curr_snr - 0.1 * (order / max_order);
    end


    % Find current best nest
    [current_best_fitness, best_idx] = max(fitness);


    % Update best solution if improved
    if current_best_fitness > best_fitness
        best_fitness = current_best_fitness;
        best_solution = nests(best_idx, :);
    end


    % Keep track of progress
    fitness_history(iter) = best_fitness;
    fprintf('Iteration %d: Best Order = %d, Step Size = %.4f, Fitness = %.4f\n', ...
        iter, round(best_solution(1)), best_solution(2), best_fitness);


    % Get a new cuckoo by Lévy flight
    for i = 1:n_nests
        % Generate step size using Lévy flight
        beta = 1.5;  % Parameter for Lévy distribution (typically between 1 and 2)
        step_size_order = levy_flight(beta) * 20;  % Scale for order parameter
        step_size_mu = levy_flight(beta) * 0.1;    % Scale for step size parameter


        % Select a random nest
        j = randi(n_nests);


        % Generate new nest using Lévy flight
        new_order = round(nests(i, 1) + step_size_order * (nests(i, 1) - nests(j, 1)));
        new_mu = nests(i, 2) + step_size_mu * (nests(i, 2) - nests(j, 2));


        % Ensure bounds
        new_order = max(min_order, min(max_order, new_order));
        new_mu = max(min_mu, min(max_mu, new_mu));


        % Evaluate new solution
        lms = dsp.LMSFilter(new_order + 1, 'StepSize', new_mu, 'Method', 'Normalized LMS', 'WeightsOutputPort', true);
        [y, e, w] = step(lms, noisy_opt, d_opt);
        [~, fetal_signal] = extract_fetal_wavelet(e);
        new_snr = wSNR(d_opt, fetal_signal);
        new_fitness = 0.9 * new_snr - 0.1 * (new_order / max_order);


        % Replace if better
        if new_fitness > fitness(i)
            nests(i, :) = [new_order, new_mu];
            fitness(i) = new_fitness;
        end
    end


    % Abandon worst nests and build new ones
    [sorted_fitness, idx] = sort(fitness, 'descend');
    for i = floor(n_nests * (1-pa))+1:n_nests
        % Replace worst nests with new random solutions
        % But biased towards the best solution (exploitation)
        if rand < 0.5
            % Generate completely new nest
            nests(idx(i), :) = [randi([min_order, max_order]), rand * (max_mu - min_mu) + min_mu];
        else
            % Generate solution biased towards best solution
            alpha = 0.3 * rand;  % Small random factor
            nests(idx(i), :) = best_solution + alpha * (rand(1, 2) .* 2 - 1) .* [30, 0.1];


            % Ensure bounds
            nests(idx(i), 1) = max(min_order, min(max_order, round(nests(idx(i), 1))));
            nests(idx(i), 2) = max(min_mu, min(max_mu, nests(idx(i), 2)));
        end
    end
end


% Apply optimal parameters to full dataset
optimal_order = round(best_solution(1));
optimal_mu = best_solution(2);
fprintf('\nOptimal Filter Order: %d\n', optimal_order);
fprintf('Optimal Step Size: %.4f\n', optimal_mu);


% Process full signal with optimized parameters
lms_optimal = dsp.LMSFilter(optimal_order + 1, 'StepSize', optimal_mu, 'Method', 'Normalized LMS', 'WeightsOutputPort', true);
[y_optimal, e_optimal, w_optimal] = step(lms_optimal, noisy_sig, d_filtered);


% Apply advanced post-processing for better signal quality
[fECG_filtered_optimal, fetal_ecg_final] = advanced_fecg_extraction(e_optimal, Fs);


% Calculate SNR with optimized parameters
Tm_full = (0:length(y_optimal)-1) / Fs;
start_time = 0;
end_time = 30;
start_idx = find(Tm_full >= start_time, 1);
end_idx = find(Tm_full >= end_time, 1);
selected_signal_opt = y_optimal(start_idx:end_idx);
selected_original = noisy_sig(start_idx:end_idx);
selected_tm = Tm_full(start_idx:end_idx);
selected_fecg = fetal_ecg_final(start_idx:end_idx);


% Calculate SNR using wavelet-based method
SNR_optimal = wSNR(d_filtered(start_idx:end_idx), selected_fecg);
fprintf('SNR with Optimized Parameters: %.4f\n', SNR_optimal);


% Compare with a baseline LMS filter (fixed parameters)
baseline_order = 200;
baseline_mu = 0.2;
lms_baseline = dsp.LMSFilter(baseline_order + 1, 'StepSize', baseline_mu, 'Method', 'Normalized LMS', 'WeightsOutputPort', true);
[y_baseline, e_baseline, w_baseline] = step(lms_baseline, noisy_sig, d_filtered);
[~, fetal_ecg_baseline] = advanced_fecg_extraction(e_baseline, Fs);
SNR_baseline = wSNR(d_filtered(start_idx:end_idx), fetal_ecg_baseline(start_idx:end_idx));
fprintf('SNR with Baseline Parameters: %.4f\n', SNR_baseline);
fprintf('Improvement: %.4f dB\n', SNR_optimal - SNR_baseline);


% Plot results
figure;
subplot(3,1,1);
plot(selected_tm, selected_original);
title('Original Abdominal ECG Signal');
xlabel('Time (s)');
ylabel('Amplitude');
ylim([-250 400]);
xlim([0 6]);


subplot(3,1,2);
plot(selected_tm, -e_optimal(start_idx:end_idx));
title(['After LMS Filtering (Order: ' num2str(optimal_order) ', μ: ' num2str(optimal_mu, '%.4f') ')']);
xlabel('Time (s)');
ylabel('Amplitude');
ylim([-250 400]);
xlim([0 6]);


subplot(3,1,3);
plot(selected_tm, selected_fecg);
title(['Final Fetal ECG Signal (SNR: ' num2str(SNR_optimal, '%.2f') ' dB)']);
xlabel('Time (s)');
ylabel('Amplitude');
ylim([-250 400]);
xlim([0 6]);




% Required helper functions
function [denoised_signal, fetal_signal] = extract_fetal_wavelet(signal)
    % Wavelet-based extraction of fetal ECG
    % Decompose signal using wavelet transform
    [c, l] = wavedec(signal, 5, 'db4');


    % Threshold coefficients to remove noise
    thr = median(abs(c))/0.6745 * sqrt(2*log(length(signal)));
    c_t = wthresh(c, 's', thr/2.5);  % Soft thresholding with reduced threshold for better fetal preservation


    % Reconstruct signal
    denoised_signal = waverec(c_t, l, 'db4');


    % Extract fetal components (detail coefficients in certain bands)
    d3 = wrcoef('d', c_t, l, 'db4', 3);
    d4 = wrcoef('d', c_t, l, 'db4', 4);
    d5 = wrcoef('d', c_t, l, 'db4', 5);


    % Combine relevant detail coefficients for fetal signal
    fetal_signal = d3 + d4 + d5;
end


function snr_val = wSNR(reference, extracted)
    % Wavelet-based SNR calculation for better evaluation of fetal ECG
    % Apply wavelet denoising to both signals
    [~, ref_fetal] = extract_fetal_wavelet(reference);
    [~, ext_fetal] = extract_fetal_wavelet(extracted);


    % Normalize both signals
    ref_fetal = ref_fetal / std(ref_fetal);
    ext_fetal = ext_fetal / std(ext_fetal);


    % Calculate correlation-based SNR
    r = corrcoef(ref_fetal, ext_fetal);
    correlation = r(1,2);
    snr_val = 20 * log10(abs(correlation) / sqrt(1 - correlation^2));


    % Ensure reasonable SNR range
    snr_val = max(0, min(snr_val, 30));  % Cap between 0-30 dB
end


function [denoised_signal, enhanced_fetal] = advanced_fecg_extraction(signal, fs)
    % Multi-stage fetal ECG extraction


    % Stage 1: Wavelet denoising
    [denoised_signal, fetal_wavelet] = extract_fetal_wavelet(signal);


    % Stage 2: Bandpass filtering to isolate fetal frequency range (typically 1-35 Hz)
    bp_fetal = designfilt('bandpassiir', 'FilterOrder', 4, ...
                'HalfPowerFrequency1', 1, 'HalfPowerFrequency2', 35, ...
                'DesignMethod', 'butter', 'SampleRate', fs);
    fetal_filtered = filtfilt(bp_fetal, denoised_signal);


    % Stage 3: Non-linear energy operator to enhance QRS complexes
    y = fetal_filtered;
    psi = y(2:end-1).^2 - y(1:end-2).*y(3:end);
    psi = [0; psi; 0];  % Pad to maintain original length


    % Stage 4: Moving average smoothing
    window_size = round(fs * 0.025);  % 25ms window
    smoothed = movmean(psi, window_size);


    % Stage 5: Combine enhanced signal with wavelet output
    enhanced_fetal = fetal_wavelet + 0.6 * smoothed;


    % Normalize output
    enhanced_fetal = enhanced_fetal / max(abs(enhanced_fetal)) * max(abs(fetal_wavelet));
end