plot_bat_trajectory.m

function plot_bat_trajectory(bat_trajectory, mic_data, batID, date_data, save_folder)
%% Plot bat trajectory with maze and optional mic positions
%% Written by Rie Kaneko on 7/1/2026
%
% Inputs:
%   bat_trajectory - struct from extract_bat_pos()
%   mic_data       - (optional) struct from extract_mic_location()
%   batID          - (optional) bat ID string, e.g., 'bat674B'
%   date_data      - (optional) date string, e.g., '20260501'
%   save_folder    - (optional) ROOT plot folder. Figures are saved as PNG
%                    into <save_folder>\<batID>\<date>\ (subject then
%                    session subfolder, created if missing), dumped flat.
%                    Default root: Z:\Rie\Analysis\Position_processing\plot\
%                    Files: <batID>_<date>_trial<NN>_<view>.png.

%% ===== CHECK INPUTS =====
if ~exist('mic_data', 'var')
    mic_data = [];
end

if ~exist('batID', 'var') || isempty(batID)
    batID = 'unknown';
end

if ~exist('date_data', 'var') || isempty(date_data)
    date_data = 'unknown';
end

if ~exist('save_folder', 'var') || isempty(save_folder)
    save_folder = 'Z:\Rie\Analysis\Position_processing\plot\';
end

%% ===== GET FILENAME AND TRIAL NUMBER =====
[~, filename, ~] = fileparts(bat_trajectory.source_file);

% Try to extract trial number from filename (e.g., "trial_001" -> "001")
trial_num = '';
trial_match = regexp(filename, '(\d+)', 'tokens');
if ~isempty(trial_match)
    trial_num = trial_match{end}{1};  % Get last number in filename
end

% Build title prefix
if ~isempty(trial_num)
    title_prefix = sprintf('%s %s trial#%s', batID, date_data, trial_num);
else
    title_prefix = sprintf('%s %s %s', batID, date_data, filename);
end

% Filesystem-safe base name for saved figures
if ~isempty(trial_num)
    name_base = sprintf('%s_%s_trial%s', batID, date_data, trial_num);
else
    name_base = sprintf('%s_%s_%s', batID, date_data, filename);
end
name_base = regexprep(name_base, '[^\w\-]', '_');

% Subject / session subfolders inside the plot root
subject_safe = regexprep(batID, '[^\w\-]', '_');
session_safe = regexprep(date_data, '[^\w\-]', '_');

% Figure handles (the last two are conditional, so pre-declare as empty)
fig_time  = [];
fig_speed = [];

%% ===== CALCULATE MISSING FIELDS IF NEEDED =====
% Frame rate
if ~isfield(bat_trajectory, 'frame_rate') || isempty(bat_trajectory.frame_rate)
    bat_trajectory.frame_rate = 100;  % Default
end

% Duration
if ~isfield(bat_trajectory, 'duration_sec') || isempty(bat_trajectory.duration_sec)
    if ~isempty(bat_trajectory.x)
        bat_trajectory.duration_sec = length(bat_trajectory.x) / bat_trajectory.frame_rate;
    else
        bat_trajectory.duration_sec = 0;
    end
end

% Total distance
if ~isfield(bat_trajectory, 'total_distance_mm') || isempty(bat_trajectory.total_distance_mm)
    if length(bat_trajectory.x) > 1
        dx = diff(bat_trajectory.x);
        dy = diff(bat_trajectory.y);
        dz = diff(bat_trajectory.z);
        bat_trajectory.total_distance_mm = sum(sqrt(dx.^2 + dy.^2 + dz.^2));
    else
        bat_trajectory.total_distance_mm = 0;
    end
end

% Average speed
if ~isfield(bat_trajectory, 'avg_speed_mm_per_sec') || isempty(bat_trajectory.avg_speed_mm_per_sec)
    if bat_trajectory.duration_sec > 0
        bat_trajectory.avg_speed_mm_per_sec = bat_trajectory.total_distance_mm / bat_trajectory.duration_sec;
    else
        bat_trajectory.avg_speed_mm_per_sec = 0;
    end
end

% Number of frames
if ~isfield(bat_trajectory, 'n_frames') || isempty(bat_trajectory.n_frames)
    bat_trajectory.n_frames = length(bat_trajectory.x);
end

if ~isfield(bat_trajectory, 'n_frames_total') || isempty(bat_trajectory.n_frames_total)
    if isfield(bat_trajectory, 'x_raw')
        bat_trajectory.n_frames_total = length(bat_trajectory.x_raw);
    else
        bat_trajectory.n_frames_total = bat_trajectory.n_frames;
    end
end

%% ===== 3D FIGURE =====
fig_3d = figure('Name', ['3D Flight: ' title_prefix], 'Position', [100, 100, 1200, 900]);
hold on;

% ----- Plot maze -----
if isfield(bat_trajectory, 'maze') && ~isempty(bat_trajectory.maze)
    maze = bat_trajectory.maze;

    if isfield(maze, 'left_wall') && ~isempty(maze.left_wall)
        plot3(maze.left_wall(:,1), maze.left_wall(:,2), maze.left_wall(:,3), ...
            'k-', 'LineWidth', 4, 'DisplayName', 'Maze Wall');
    end

    if isfield(maze, 'right_wall') && ~isempty(maze.right_wall)
        plot3(maze.right_wall(:,1), maze.right_wall(:,2), maze.right_wall(:,3), ...
            'k-', 'LineWidth', 4, 'HandleVisibility', 'off');
    end

    if isfield(maze, 'start_line') && ~isempty(maze.start_line)
        plot3(maze.start_line(:,1), maze.start_line(:,2), maze.start_line(:,3), ...
            'g--', 'LineWidth', 2, 'DisplayName', 'Start Line');
    end
end

% ----- Plot mic positions -----
if ~isempty(mic_data) && isfield(mic_data, 'mic_pos')
    scatter3(mic_data.mic_pos(:,1), mic_data.mic_pos(:,2), mic_data.mic_pos(:,3), ...
        60, 'm', 'filled', 'd', 'DisplayName', 'Microphones');

    % Plot mic pointing directions
    if isfield(mic_data, 'mic_pointing_direction')
        quiver3(mic_data.mic_pos(:,1), mic_data.mic_pos(:,2), mic_data.mic_pos(:,3), ...
            mic_data.mic_pointing_direction(:,1)*50, ...
            mic_data.mic_pointing_direction(:,2)*50, ...
            mic_data.mic_pointing_direction(:,3)*50, ...
            0, 'c', 'LineWidth', 1, 'HandleVisibility', 'off');
    end

    % Label mics (optional - can be noisy with many mics)
    if isfield(mic_data, 'mic_names') && length(mic_data.mic_names) <= 32
        for m = 1:length(mic_data.mic_names)
            if ~any(isnan(mic_data.mic_pos(m,:)))
                text(mic_data.mic_pos(m,1), mic_data.mic_pos(m,2), mic_data.mic_pos(m,3), ...
                    ['  ' mic_data.mic_names{m}], 'FontSize', 7, 'Interpreter', 'none', ...
                    'Color', [0.5 0.5 0.5], 'HandleVisibility', 'off');
            end
        end
    end
end

% ----- Plot flight trajectory -----
if ~isempty(bat_trajectory.x)
    plot3(bat_trajectory.x, bat_trajectory.y, bat_trajectory.z, ...
        'b-', 'LineWidth', 2, 'DisplayName', 'Bat Flight');

    % Start point = green
    plot3(bat_trajectory.x(1), bat_trajectory.y(1), bat_trajectory.z(1), ...
        'go', 'MarkerSize', 16, 'MarkerFaceColor', 'g', 'LineWidth', 2, 'DisplayName', 'Flight Start');

    % End point = red
    plot3(bat_trajectory.x(end), bat_trajectory.y(end), bat_trajectory.z(end), ...
        'ro', 'MarkerSize', 16, 'MarkerFaceColor', 'r', 'LineWidth', 2, 'DisplayName', 'Flight End');
end

% ----- Plot TP and LP from Vicon data -----
if isfield(bat_trajectory, 'tp_position') && ~isempty(bat_trajectory.tp_position)
    plot3(bat_trajectory.tp_position(1), bat_trajectory.tp_position(2), bat_trajectory.tp_position(3), ...
        'c^', 'MarkerSize', 20, 'MarkerFaceColor', 'c', 'LineWidth', 2, 'DisplayName', 'Take Off Perch (Vicon)');
end

Lpts = []; Lnames = {};
if isfield(bat_trajectory,'perch_pos') && isfield(bat_trajectory.perch_pos,'landing') && ~isempty(bat_trajectory.perch_pos.landing)
    Pl = bat_trajectory.perch_pos.landing;
    for kk = 1:numel(Pl)
        if isfield(Pl(kk),'marker1') && numel(Pl(kk).marker1) >= 3 && ~any(isnan(Pl(kk).marker1(1:3)))
            Lpts(end+1,:) = Pl(kk).marker1(1:3); %#ok<AGROW>
            Lnames{end+1} = sprintf('Landing Perch %d (Vicon)', kk); %#ok<AGROW>
        end
    end
end
if isempty(Lpts) && isfield(bat_trajectory,'lp_position') && ~isempty(bat_trajectory.lp_position)
    Lpts = bat_trajectory.lp_position(1:3); Lnames = {'Landing Perch (Vicon)'};
end
lp_cols = {'m','r',[1 0.5 0],'g'};
for kk = 1:size(Lpts,1)
    cc = lp_cols{mod(kk-1,numel(lp_cols))+1};
    plot3(Lpts(kk,1), Lpts(kk,2), Lpts(kk,3), '^', 'MarkerSize', 20, 'MarkerFaceColor', cc, 'MarkerEdgeColor', cc, 'LineWidth', 2, 'DisplayName', Lnames{kk});
end

% ----- Format 3D plot -----
xlabel('X (mm)'); ylabel('Y (mm)'); zlabel('Z (mm)');
title(sprintf('3D Flight: %s', title_prefix), 'Interpreter', 'none');
legend('Location', 'bestoutside');
axis equal; grid on;
view(3);
rotate3d on;

%% ===== 2D TOP-DOWN VIEW (XY) =====
fig_top = figure('Name', ['Top View: ' title_prefix], 'Position', [150, 150, 1200, 900]);
hold on;

% ----- Plot maze (2D) -----
if isfield(bat_trajectory, 'maze') && ~isempty(bat_trajectory.maze)
    maze = bat_trajectory.maze;

    if isfield(maze, 'left_wall') && ~isempty(maze.left_wall)
        plot(maze.left_wall(:,1), maze.left_wall(:,2), 'k-', 'LineWidth', 4, 'DisplayName', 'Maze Wall');
    end

    if isfield(maze, 'right_wall') && ~isempty(maze.right_wall)
        plot(maze.right_wall(:,1), maze.right_wall(:,2), 'k-', 'LineWidth', 4, 'HandleVisibility', 'off');
    end

    if isfield(maze, 'start_line') && ~isempty(maze.start_line)
        plot(maze.start_line(:,1), maze.start_line(:,2), 'g--', 'LineWidth', 2, 'DisplayName', 'Start Line');
    end

    if isfield(maze, 'takeoff_perch') && ~isempty(maze.takeoff_perch)
        plot(maze.takeoff_perch(1), maze.takeoff_perch(2), 'yd', 'MarkerSize', 15, ...
            'MarkerFaceColor', 'y', 'DisplayName', 'Take Off Perch');
    end

    if isfield(maze, 'landing_perch') && ~isempty(maze.landing_perch)
        plot(maze.landing_perch(1), maze.landing_perch(2), 'md', 'MarkerSize', 15, ...
            'MarkerFaceColor', 'm', 'DisplayName', 'Landing Perch');
    end
end

% ----- Plot mics (2D) -----
if ~isempty(mic_data) && isfield(mic_data, 'mic_pos')
    scatter(mic_data.mic_pos(:,1), mic_data.mic_pos(:,2), 40, 'm', 'filled', 'd', 'DisplayName', 'Microphones');
end

% ----- Plot bat flight (2D) -----
if ~isempty(bat_trajectory.x)
    plot(bat_trajectory.x, bat_trajectory.y, 'b-', 'LineWidth', 2, 'DisplayName', 'Bat Flight');
    plot(bat_trajectory.x(1), bat_trajectory.y(1), 'go', 'MarkerSize', 16, ...
        'MarkerFaceColor', 'g', 'LineWidth', 2, 'DisplayName', 'Flight Start');
    plot(bat_trajectory.x(end), bat_trajectory.y(end), 'ro', 'MarkerSize', 16, ...
        'MarkerFaceColor', 'r', 'LineWidth', 2, 'DisplayName', 'Flight End');
end

% ----- TP and LP (2D) -----
if isfield(bat_trajectory, 'tp_position') && ~isempty(bat_trajectory.tp_position)
    plot(bat_trajectory.tp_position(1), bat_trajectory.tp_position(2), 'c^', ...
        'MarkerSize', 20, 'MarkerFaceColor', 'c', 'LineWidth', 2, 'DisplayName', 'Take Off Perch (Vicon)');
end

Lpts = []; Lnames = {};
if isfield(bat_trajectory,'perch_pos') && isfield(bat_trajectory.perch_pos,'landing') && ~isempty(bat_trajectory.perch_pos.landing)
    Pl = bat_trajectory.perch_pos.landing;
    for kk = 1:numel(Pl)
        if isfield(Pl(kk),'marker1') && numel(Pl(kk).marker1) >= 3 && ~any(isnan(Pl(kk).marker1(1:3)))
            Lpts(end+1,:) = Pl(kk).marker1(1:3); %#ok<AGROW>
            Lnames{end+1} = sprintf('Landing Perch %d (Vicon)', kk); %#ok<AGROW>
        end
    end
end
if isempty(Lpts) && isfield(bat_trajectory,'lp_position') && ~isempty(bat_trajectory.lp_position)
    Lpts = bat_trajectory.lp_position(1:3); Lnames = {'Landing Perch (Vicon)'};
end
lp_cols = {'m','r',[1 0.5 0],'g'};
for kk = 1:size(Lpts,1)
    cc = lp_cols{mod(kk-1,numel(lp_cols))+1};
    plot(Lpts(kk,1), Lpts(kk,2), '^', 'MarkerSize', 20, 'MarkerFaceColor', cc, 'MarkerEdgeColor', cc, 'LineWidth', 2, 'DisplayName', Lnames{kk});
end

xlabel('X (mm)'); ylabel('Y (mm)');
title(sprintf('Top-Down View: %s', title_prefix), 'Interpreter', 'none');
legend('Location', 'bestoutside');
axis equal; grid on;

%% ===== 2D SIDE VIEW (XZ) =====
fig_side = figure('Name', ['Side View: ' title_prefix], 'Position', [200, 100, 1200, 600]);
hold on;

% ----- Plot maze (XZ) -----
if isfield(bat_trajectory, 'maze') && ~isempty(bat_trajectory.maze)
    maze = bat_trajectory.maze;

    if isfield(maze, 'left_wall') && ~isempty(maze.left_wall)
        plot(maze.left_wall(:,1), maze.left_wall(:,3), 'k-', 'LineWidth', 4, 'DisplayName', 'Maze Wall');
    end

    if isfield(maze, 'right_wall') && ~isempty(maze.right_wall)
        plot(maze.right_wall(:,1), maze.right_wall(:,3), 'k-', 'LineWidth', 4, 'HandleVisibility', 'off');
    end

    if isfield(maze, 'takeoff_perch') && ~isempty(maze.takeoff_perch)
        plot(maze.takeoff_perch(1), maze.takeoff_perch(3), 'yd', 'MarkerSize', 15, ...
            'MarkerFaceColor', 'y', 'DisplayName', 'Take Off Perch');
    end
end

% ----- Plot mics (XZ) -----
if ~isempty(mic_data) && isfield(mic_data, 'mic_pos')
    scatter(mic_data.mic_pos(:,1), mic_data.mic_pos(:,3), 40, 'm', 'filled', 'd', 'DisplayName', 'Microphones');
end

% ----- Plot bat flight (XZ) -----
if ~isempty(bat_trajectory.x)
    plot(bat_trajectory.x, bat_trajectory.z, 'b-', 'LineWidth', 2, 'DisplayName', 'Bat Flight');
    plot(bat_trajectory.x(1), bat_trajectory.z(1), 'go', 'MarkerSize', 16, ...
        'MarkerFaceColor', 'g', 'LineWidth', 2, 'DisplayName', 'Flight Start');
    plot(bat_trajectory.x(end), bat_trajectory.z(end), 'ro', 'MarkerSize', 16, ...
        'MarkerFaceColor', 'r', 'LineWidth', 2, 'DisplayName', 'Flight End');
end

% ----- TP and LP (XZ) -----
if isfield(bat_trajectory, 'tp_position') && ~isempty(bat_trajectory.tp_position)
    plot(bat_trajectory.tp_position(1), bat_trajectory.tp_position(3), 'c^', ...
        'MarkerSize', 20, 'MarkerFaceColor', 'c', 'LineWidth', 2, 'DisplayName', 'Take Off Perch');
end

Lpts = []; Lnames = {};
if isfield(bat_trajectory,'perch_pos') && isfield(bat_trajectory.perch_pos,'landing') && ~isempty(bat_trajectory.perch_pos.landing)
    Pl = bat_trajectory.perch_pos.landing;
    for kk = 1:numel(Pl)
        if isfield(Pl(kk),'marker1') && numel(Pl(kk).marker1) >= 3 && ~any(isnan(Pl(kk).marker1(1:3)))
            Lpts(end+1,:) = Pl(kk).marker1(1:3); %#ok<AGROW>
            Lnames{end+1} = sprintf('Landing Perch %d (Vicon)', kk); %#ok<AGROW>
        end
    end
end
if isempty(Lpts) && isfield(bat_trajectory,'lp_position') && ~isempty(bat_trajectory.lp_position)
    Lpts = bat_trajectory.lp_position(1:3); Lnames = {'Landing Perch'};
end
lp_cols = {'m','r',[1 0.5 0],'g'};
for kk = 1:size(Lpts,1)
    cc = lp_cols{mod(kk-1,numel(lp_cols))+1};
    plot(Lpts(kk,1), Lpts(kk,3), '^', 'MarkerSize', 20, 'MarkerFaceColor', cc, 'MarkerEdgeColor', cc, 'LineWidth', 2, 'DisplayName', Lnames{kk});
end

xlabel('X (mm)'); ylabel('Z (mm)');
title(sprintf('Side View: %s', title_prefix), 'Interpreter', 'none');
legend('Location', 'bestoutside');
axis equal; grid on;

%% ===== TRAJECTORY OVER TIME =====
if ~isempty(bat_trajectory.x)
    fig_time = figure('Name', ['Position vs Time: ' title_prefix], 'Position', [250, 100, 1200, 800]);

    % Create time vector
    if bat_trajectory.frame_rate > 0
        t = (0:length(bat_trajectory.x)-1) / bat_trajectory.frame_rate;
        time_label = 'Time (sec)';
    else
        t = 1:length(bat_trajectory.x);
        time_label = 'Frame';
    end

    % X position
    subplot(3, 1, 1);
    plot(t, bat_trajectory.x, 'b-', 'LineWidth', 1.5);
    ylabel('X (mm)');
    title(sprintf('Position vs Time: %s', title_prefix), 'Interpreter', 'none');
    grid on;

    % Y position
    subplot(3, 1, 2);
    plot(t, bat_trajectory.y, 'r-', 'LineWidth', 1.5);
    ylabel('Y (mm)');
    grid on;

    % Z position
    subplot(3, 1, 3);
    plot(t, bat_trajectory.z, 'g-', 'LineWidth', 1.5);
    ylabel('Z (mm)');
    xlabel(time_label);
    grid on;
end

%% ===== SPEED PROFILE =====
if length(bat_trajectory.x) > 1
    fig_speed = figure('Name', ['Speed Profile: ' title_prefix], 'Position', [300, 100, 1200, 400]);

    % Calculate instantaneous speed
    dx = diff(bat_trajectory.x);
    dy = diff(bat_trajectory.y);
    dz = diff(bat_trajectory.z);
    dist = sqrt(dx.^2 + dy.^2 + dz.^2);

    if bat_trajectory.frame_rate > 0
        speed = dist * bat_trajectory.frame_rate;  % mm/sec
        t_speed = (0:length(speed)-1) / bat_trajectory.frame_rate;
        time_label = 'Time (sec)';
        speed_label = 'Speed (mm/sec)';
    else
        speed = dist;  % mm/frame
        t_speed = 1:length(speed);
        time_label = 'Frame';
        speed_label = 'Speed (mm/frame)';
    end

    plot(t_speed, speed, 'b-', 'LineWidth', 1.5);
    hold on;
    yline(mean(speed), 'r--', 'LineWidth', 1.5);
    yline(median(speed), 'g--', 'LineWidth', 1.5);

    xlabel(time_label);
    ylabel(speed_label);
    title(sprintf('Speed Profile: %s', title_prefix), 'Interpreter', 'none');
    legend('Speed', sprintf('Mean: %.1f', mean(speed)), sprintf('Median: %.1f', median(speed)), ...
        'Location', 'best');
    grid on;
end

%% ===== SAVE FIGURES =====
% Organize the plot root by subject then session; dump the figures flat
% inside the session folder: <save_folder>\<batID>\<date>\
out_dir = fullfile(save_folder, subject_safe, session_safe);
if ~exist(out_dir, 'dir')
    mkdir(out_dir);
end

figs     = {fig_3d, fig_top,   fig_side, fig_time,           fig_speed};
suffixes = {'3D',   'topdown', 'side',   'position_vs_time', 'speed'};
n_saved  = 0;
for i = 1:numel(figs)
    if ~isempty(figs{i}) && ishandle(figs{i})
        out_png = fullfile(out_dir, sprintf('%s_%s.png', name_base, suffixes{i}));
        try
            exportgraphics(figs{i}, out_png, 'Resolution', 200);
        catch
            saveas(figs{i}, out_png);   % fallback for older MATLAB (< R2020a)
        end
        n_saved = n_saved + 1;
    end
end
if n_saved > 0
    fprintf('\nSaved %d figure(s) to:\n  %s\n', n_saved, out_dir);
end

%% ===== PRINT SUMMARY =====
fprintf('\n===== PLOT SUMMARY =====\n');
fprintf('Trial: %s\n', title_prefix);
fprintf('File: %s\n', filename);
if isfield(bat_trajectory, 'bat_marker_idx') && isfield(bat_trajectory, 'bat_marker_label')
    fprintf('Marker: #%d (%s)\n', bat_trajectory.bat_marker_idx, bat_trajectory.bat_marker_label);
end
fprintf('Frames: %d valid / %d total\n', bat_trajectory.n_frames, bat_trajectory.n_frames_total);
fprintf('Duration: %.2f sec\n', bat_trajectory.duration_sec);
fprintf('Distance: %.1f mm (%.2f m)\n', bat_trajectory.total_distance_mm, bat_trajectory.total_distance_mm/1000);
fprintf('Avg speed: %.1f mm/sec (%.2f m/sec)\n', bat_trajectory.avg_speed_mm_per_sec, bat_trajectory.avg_speed_mm_per_sec/1000);
fprintf('========================\n');

end