function result = example_eeg_acquisition_to_matlab(device_name, channel_ids, sample_rate_hz, annotation_count, output_mat_path)
%EXAMPLE_EEG_ACQUISITION_TO_MATLAB Continuous acquisition with annotations.
%
% Usage:
%   example_eeg_acquisition_to_matlab()
%   example_eeg_acquisition_to_matlab('BA MINI 001')
%   result = example_eeg_acquisition_to_matlab('BA MINI 001', [], 250, 10)

project_root = fileparts(fileparts(mfilename('fullpath')));
addpath(fullfile(project_root, 'matlab'));

if nargin < 1
    device_name = '';
end
if nargin < 2
    channel_ids = [];
end
if nargin < 3 || isempty(sample_rate_hz)
    sample_rate_hz = 250;
end
if nargin < 4 || isempty(annotation_count)
    annotation_count = 10;
end
if nargin < 5 || isempty(output_mat_path)
    timestamp = char(datetime("now", Format="yyyyMMdd-HHmmss"));
    output_mat_path = fullfile(tempdir, ...
        sprintf('brainaccess-%s.mat', timestamp));
end

ba_setup();
ba_init();
cleanup_obj = onCleanup(@() safe_shutdown());
devices = ba_scan();
device_name = resolve_device_name(device_name, devices);

fprintf('Connecting to %s\n', device_name);
connect_result = ba_connect(device_name);
if ~connect_result.connected
    error('brainaccess:api', 'Connection failed: %s (%s)', ...
        connect_result.message, connect_result.status_name);
end

device_features = ba_get_device_features();
channel_ids = resolve_channel_ids(channel_ids, device_features);
fprintf('Using %d EEG channels.\n', numel(channel_ids));

ba_configure_eeg(channel_ids, sample_rate_hz, 8, []);
ba_start_stream();

warmup_seconds = 5.0;
fprintf('Stream started. Waiting %.1f seconds before annotations.\n', warmup_seconds);
pause(warmup_seconds);

all_data = [];
annotation_messages = strings(1, annotation_count);
chunk_meta = cell(1, annotation_count);

for idx = 1:annotation_count
    pause(1.0);
    annotation_messages(idx) = string(idx);
    fprintf('Sending annotation %s\n', annotation_messages(idx));
    ba_annotate(char(annotation_messages(idx)));

    [chunk, meta] = ba_read_chunk(1000, true);
    chunk_meta{idx} = meta;
    if isstruct(meta) && isfield(meta, 'available') && meta.available && ~isempty(chunk)
        all_data = append_chunk(all_data, chunk);
    end
end

settle_seconds = 2.0;
fprintf('Waiting %.1f seconds for final chunks.\n', settle_seconds);
pause(settle_seconds);

[final_chunk, final_meta] = ba_read_chunk(0, true);
if isstruct(final_meta) && isfield(final_meta, 'available') ...
        && final_meta.available && ~isempty(final_chunk)
    all_data = append_chunk(all_data, final_chunk);
end

annotations = ba_get_annotations();
ba_stop_stream();
ba_disconnect();
ba_close();
clear cleanup_obj

save(output_mat_path, 'all_data', 'annotations', 'annotation_messages', ...
    'chunk_meta', 'channel_ids', 'device_name', 'final_meta', 'sample_rate_hz');
fprintf('Saved acquisition: %s\n', output_mat_path);

plot_eeg(all_data, sample_rate_hz, device_name, annotations);

result = struct( ...
    'annotation_messages', {cellstr(annotation_messages)}, ...
    'annotations', annotations, ...
    'channel_ids', channel_ids, ...
    'data', all_data, ...
    'device_name', device_name, ...
    'device_features', device_features, ...
    'devices', devices, ...
    'final_meta', final_meta, ...
    'output_mat_path', output_mat_path, ...
    'sample_rate_hz', sample_rate_hz);
end

function channel_ids = resolve_channel_ids(channel_ids, device_features)
if ~isempty(channel_ids)
    return;
end

electrode_count = double(device_features.electrode_count);
channel_ids = 0:(electrode_count - 1);
end

function device_name = resolve_device_name(device_name, devices)
if ~isempty(device_name)
    return;
end

device_name = getenv('BA_DEVICE_NAME');
if ~isempty(device_name)
    return;
end

if isempty(devices)
    error('brainaccess:api', ...
        'No devices found. Set BA_DEVICE_NAME or pass a device name explicitly.');
end

device_name = devices(1).name;
end

function data = append_chunk(data, chunk)
if isempty(data)
    data = chunk;
    return;
end

if size(data, 1) ~= size(chunk, 1)
    error('brainaccess:api', ...
        'Chunk channel count changed during acquisition.');
end

data = [data, chunk];
end

function plot_eeg(data, sample_rate_hz, device_name, annotations)
if isempty(data)
    error('brainaccess:api', 'No samples were collected.');
end

channel_count = size(data, 1);
sample_count = size(data, 2);
time_axis = (0:(sample_count - 1)) / sample_rate_hz;
centered = data - mean(data, 2);
spacing = max(std(centered, 0, 2));
if isempty(spacing) || spacing <= 0
    spacing = 1;
end
offsets = ((channel_count - 1):-1:0)' * spacing * 4;

figure('Name', 'BrainAccess Acquisition With Annotations');
plot(time_axis, (centered + offsets)');
xlabel('Time (s)');
ylabel('Amplitude + offset');
title(sprintf('BrainAccess annotated acquisition: %s', device_name));
grid on;

annotation_count = min(numel(annotations.annotations), numel(annotations.timestamps));
if annotation_count > 0
    hold on;
    for idx = 1:annotation_count
        annotation_time = double(annotations.timestamps(idx)) / sample_rate_hz;
        if annotation_time < time_axis(1) || annotation_time > time_axis(end)
            continue;
        end
        xline(annotation_time, '--', annotations.annotations{idx}, ...
            'Color', [0.85, 0.2, 0.2], ...
            'LineWidth', 1.0, ...
            'LabelVerticalAlignment', 'middle', ...
            'LabelHorizontalAlignment', 'left');
    end
    hold off;
end
end

function safe_shutdown()
status = ba_status();
if status.streaming
    ba_stop_stream();
end
if status.connected
    ba_disconnect();
end
if status.core_initialized || status.manager_present
    ba_close();
end
end
