% ------------------------
% PARAMETERS
% ------------------------
components_keep = [1:10];    % components to keep for denoising
component_idx   = 4;          % for choosing top neurons by loading (style)
top_percent     = 10;        % % neurons to use in plot/glyph
arrow_scale     = 0.6;        % shorten arrows (visual)
K_first_last    = 5;          % use mean of first/last K presentations per stimulus
time_idx        = [];         % e.g., 20:60 to average only response window; [] = all time
P = M;

% ------------------------
% INPUTS (expected in workspace)
% ------------------------
% P      : ktensor from Tensor Toolbox (cp_als / cp_opt)
% labels : trials x 1 vector with values 1 (A) or 2 (B)

assert(exist('P','var')==1 && isa(P,'ktensor') && ndims(P)==3, ...
    'P must be a 3-way ktensor (neurons x time x trials).');
assert(exist('labels','var')==1, 'labels (trials x 1) required.');

% ------------------------
% FULL DENOISED RECONSTRUCTION FROM components_keep
% ------------------------
[Xden, stats_keep] = tt_cp_denoise_keep(P, components_keep); %#ok<NASGU>
% Xden size: N x T x Tr

% ------------------------
% LOADINGS & TOP NEURONS (by component_idx)
% ------------------------
U1 = double(P.U{1});                    % neurons x rank
loadings  = U1(:, component_idx);
n_neurons = size(U1,1);
n_keep    = max(1, round(n_neurons * top_percent / 100));
[~, idx_top] = maxk(loadings, n_keep);

% ------------------------
% TRIAL LABELS & FIRST/LAST K SELECTION
% ------------------------
labels = labels(:);
trial_type_A = find(labels == 1);
trial_type_B = find(labels == 2);
assert(~isempty(trial_type_A) && ~isempty(trial_type_B), 'Need both A and B trials.');

Ka = min(K_first_last, numel(trial_type_A));
Kb = min(K_first_last, numel(trial_type_B));
va = 15;
firstA = trial_type_A(1:Ka);
lastA  = trial_type_A(va-Ka+1:va);
firstB = trial_type_B(1:Kb);
lastB  = trial_type_B(va-Kb+1:va);

% ------------------------
% RESPONSE EXTRACTION FROM DENOISED FULL TENSOR
% ------------------------
% For each neuron: mean over time (optionally restricted by time_idx), then mean across trials set
if isempty(time_idx), time_idx = 1:size(Xden,2); end

% Helper to get mean over time, then mean over a set of trials
mean_over_time = @(X) squeeze(mean(X(:, time_idx, :), 2));   % N x Tr
X_nt = mean_over_time(Xden);                                  % N x Tr

% Take selected neurons and compute early/late means for A/B
X_sel = X_nt(idx_top, :);                                     % n_keep x Tr

x1 = mean(X_sel(:, firstA), 2);   % Stim A early mean
x2 = mean(X_sel(:, lastA ), 2);   % Stim A late  mean
y1 = mean(X_sel(:, firstB), 2);   % Stim B early mean
y2 = mean(X_sel(:, lastB ), 2);   % Stim B late  mean

% Ensure column vectors & arrow deltas
x1 = x1(:); y1 = y1(:); x2 = x2(:); y2 = y2(:);
u = (x2 - x1) * arrow_scale;   % ΔA
v = (y2 - y1) * arrow_scale;   % ΔB

% ------------------------
% SCATTER + ARROWS (old style)
% ------------------------
figure('Color','w'); hold on;
scatter(x1, y1, 60, 'filled', 'MarkerFaceColor', [0.2 0.6 0.9]);
quiver(x1, y1, u, v, 0, 'Color', [0.1 0.1 0.6], 'LineWidth', 1.5);
xlabel(sprintf('Stim A (mean of first %d)', Ka));
ylabel(sprintf('Stim B (mean of first %d)', Kb));
title(sprintf('Top %.0f%% neurons — CP comp %d | keep %s (arrows → last %d)', ...
      top_percent, component_idx, mat2str(components_keep), max(Ka,Kb)));
legend({'Early mean', '→ Late mean'}, 'Location','best');
grid on; axis equal;

% ------------------------
% GLYPH 1 — PROPORTIONS OVER TOTAL (NE/NO/SE/SO)
% ------------------------
N = numel(u);  assert(N>0, 'No arrows to analyze.');
is_NE = (u > 0) & (v > 0);
is_NO = (u < 0) & (v > 0);
is_SE = (u > 0) & (v < 0);
is_SO = (u < 0) & (v < 0);
on_axes = (u == 0) | (v == 0);

cNE = sum(is_NE); cNO = sum(is_NO); cSE = sum(is_SE); cSO = sum(is_SO);
pNE = cNE / N;    pNO = cNO / N;    pSE = cSE / N;    pSO = cSO / N;

glyph_max_len = 1.0;
theta = [pi/4, 3*pi/4, -pi/4, -3*pi/4];
L  = glyph_max_len * [pNE, pNO, pSE, pSO];
dx = L .* cos(theta); dy = L .* sin(theta);

figure('Color','w'); ax = axes; hold(ax,'on');
r = glyph_max_len;
rectangle('Position',[-r -r 2*r 2*r], 'Curvature',[1 1], 'LineStyle',':', 'EdgeColor',[0 0 0], 'LineWidth',0.75);
plot([-r r], [-r r], 'k:', 'LineWidth', 0.75);
plot([-r r], [ r -r], 'k:', 'LineWidth', 0.75);
quiver(0,0,dx(1),dy(1),0,'LineWidth',2); % NE
quiver(0,0,dx(2),dy(2),0,'LineWidth',2); % NO
quiver(0,0,dx(3),dy(3),0,'LineWidth',2); % SE
quiver(0,0,dx(4),dy(4),0,'LineWidth',2); % SO
axis equal; xlim([-r r]); ylim([-r r]); box on; ax.XTick=[]; ax.YTick=[];
title(sprintf('Glyph (counts/TOTAL) — keep %s | first %d / last %d', ...
      mat2str(components_keep), max(Ka,Kb), max(Ka,Kb)));

pct = 100*[pNE pNO pSE pSO];
txt = sprintf('NE %d (%.1f%%) | NO %d (%.1f%%) | SE %d (%.1f%%) | SO %d (%.1f%%) | On-axes: %d/%d', ...
              cNE,pct(1), cNO,pct(2), cSE,pct(3), cSO,pct(4), sum(on_axes), N);
annotation('textbox',[0.0 0.0 1.0 0.07],'String',txt,'EdgeColor','none','HorizontalAlignment','center','FontSize',9);

% ========================
% 20-BIN DIRECTIONAL GLYPH (5 settori per quadrante)
% ========================
assert(exist('u','var')==1 && exist('v','var')==1, 'Servono u e v.');
N = numel(u);  assert(N>0, 'Nessuna freccia.');

% Parametri
nbins_per_quad = 5;
nbins = 4 * nbins_per_quad;   % 20
glyph_max_len = 1.0;
include_axes_in_denominator = true;  % true: normalizza su TUTTE le frecce, incluse quelle sugli assi

% Angoli e lunghezze
theta = atan2(v, u);          % (-pi, pi]
w = hypot(u, v);
on_axes = (u == 0) | (v == 0);

% Binning uniforme in [-pi, pi] con 20 settori
edges   = linspace(-pi, pi, nbins+1);                  % 21 edges
centers = (edges(1:end-1) + edges(2:end)) / 2;        % 20 centri (rad)

% Escludi frecce "on-axes" dalla direzione (coerente con i glifi precedenti)
theta_valid = theta(~on_axes);
w_valid     = w(~on_axes);

% Indici di bin
bin_idx = discretize(theta_valid, edges);              % 1..20
% Conteggi per bin
c = accumarray(bin_idx(~isnan(bin_idx)), 1, [nbins 1], @sum, 0);
% Somma pesi (lunghezze) per bin
cw = accumarray(bin_idx(~isnan(bin_idx)), w_valid(~isnan(bin_idx)), [nbins 1], @sum, 0);

% Denominatori
if include_axes_in_denominator
    Nden = N;
    Wden = sum(w);
else
    Nden = numel(theta_valid);
    Wden = sum(w_valid);
end
if Nden == 0, Nden = 1; end
if Wden == 0, Wden = 1; end

% Lunghezze dei raggi (proporzionali a proporzione sul totale)
L_counts   = glyph_max_len * (c  / Nden);   % conteggi / totale
L_weighted = glyph_max_len * (cw / Wden);   % quota movimento / totale

% ------------------------
% Plot: 20-bin (conteggi)
% ------------------------
figure('Color','w'); ax = axes; hold(ax,'on');
r = glyph_max_len;

% anello guida e diagonali
rectangle('Position',[-r -r 2*r 2*r], 'Curvature',[1 1], 'LineStyle',':', 'EdgeColor',[0 0 0], 'LineWidth',0.75);
plot([-r r], [-r r], 'k:', 'LineWidth', 0.75);
plot([-r r], [ r -r], 'k:', 'LineWidth', 0.75);

% raggi per ogni settore
dx = L_counts .* cos(centers(:));
dy = L_counts .* sin(centers(:));
quiver(zeros(nbins,1), zeros(nbins,1), dx, dy, 0, 'LineWidth', 1.8);

axis equal; xlim([-r r]); ylim([-r r]); box on; ax.XTick=[]; ax.YTick=[];
title(sprintf('20-bin glyph (counts/TOTAL) — %d×(5 per quadrant), on-axes=%d', nbins, sum(on_axes)));

% legenda compatta
pct_counts = 100 * sum(c) / Nden;  %#ok<NASGU> % (per completezza; è sempre <=100%)
annotation('textbox',[0.0 0.0 1.0 0.07], 'String', ...
    sprintf('On-axes: %d/%d  |  Norm=%s', sum(on_axes), N, ternary(include_axes_in_denominator,'TOTAL','quad-defined')), ...
    'EdgeColor','none','HorizontalAlignment','center','FontSize',9);

% ------------------------
% Plot: 20-bin (pesato per lunghezza)
% ------------------------
figure('Color','w'); ax = axes; hold(ax,'on');
rectangle('Position',[-r -r 2*r 2*r], 'Curvature',[1 1], 'LineStyle',':', 'EdgeColor',[0 0 0], 'LineWidth',0.75);
plot([-r r], [-r r], 'k:', 'LineWidth', 0.75);
plot([-r r], [ r -r], 'k:', 'LineWidth', 0.75);

dxw = L_weighted .* cos(centers(:));
dyw = L_weighted .* sin(centers(:));
quiver(zeros(nbins,1), zeros(nbins,1), dxw, dyw, 0, 'LineWidth', 1.8);

axis equal; xlim([-r r]); ylim([-r r]); box on; ax.XTick=[]; ax.YTick=[];
title(sprintf('20-bin glyph (length-weighted, den=%s)', ternary(include_axes_in_denominator,'TOTAL','quad-defined')));

pct_move = 100 * sum(cw) / Wden; %#ok<NASGU>
annotation('textbox',[0.0 0.0 1.0 0.07], 'String', ...
    sprintf('On-axes: %d/%d  |  W_{den}=%.3g', sum(on_axes), N, Wden), ...
    'EdgeColor','none','HorizontalAlignment','center','FontSize',9);

%% 20-bin directional GLYPH (length-weighted), self-contained script
% Expects u, v in workspace (arrow components ΔA, ΔB). Produces a polar glyph
% with 20 sectors, ray lengths proportional to total movement per sector,
% colored by quadrant (NE/NO/SE/SO) with a larger visual scale.

% ------------------------
% PARAMETERS (tweak here)
% ------------------------
include_axes_in_denominator = true;  % true: normalize by ALL arrows (incl. on-axes)
scale_gain   = 1.8;                  % enlarge rays (increase for bigger arrows)
glyph_max    = 1.0;                  % nominal radius before scaling

% Quadrant colors (NE, NO, SE, SO)
C_NE = [0.95 0.60 0.15];  % orange
C_NO = [0.10 0.70 0.30];  % green
C_SE = [0.60 0.20 0.60];  % purple
C_SO = [0.20 0.60 0.90];  % blue

% ------------------------
% INPUT CHECKS
% ------------------------
assert(exist('u','var')==1 && exist('v','var')==1, 'u and v must exist in the workspace.');
u = u(:); v = v(:);
N = numel(u);  assert(N>0, 'u and v must be non-empty and same length.');
assert(numel(v)==N, 'u and v must have same length.');

% ------------------------
% ANGLES, LENGTHS, BINS
% ------------------------
nbins_per_quad = 5;
nbins = 4 * nbins_per_quad;                      % 20 bins
theta = atan2(v, u);                             % angle in (-pi, pi]
w     = hypot(u, v);                             % arrow lengths
on_axes = (u == 0) | (v == 0);                   % undefined direction

edges   = linspace(-pi, pi, nbins+1);            % 21 edges
centers = (edges(1:end-1) + edges(2:end)) / 2;   % 20 bin centers

theta_valid = theta(~on_axes);
w_valid     = w(~on_axes);

bin_idx = discretize(theta_valid, edges);        % 1..20
cw = accumarray(bin_idx(~isnan(bin_idx)), ...
                w_valid(~isnan(bin_idx)), [nbins 1], @sum, 0);  % sum of lengths per bin

% Denominator for normalization
if include_axes_in_denominator
    Wden = sum(w);
else
    Wden = sum(w_valid);
end
if Wden == 0, Wden = 1; end

% Ray lengths (fraction of total movement per bin), then scaled up for display
L_weighted = glyph_max * (cw / Wden);
Lw_scaled  = L_weighted * scale_gain;

% ------------------------
% PLOT
% ------------------------
figure('Color','w'); ax = axes; hold(ax,'on');

% Fit guide ring to longest ray
r = max(glyph_max, 1.05 * max([Lw_scaled; glyph_max]));

% Guide ring and diagonals
rectangle('Position',[-r -r 2*r 2*r], 'Curvature',[1 1], 'LineStyle',':', ...
          'EdgeColor',[0 0 0], 'LineWidth',0.75);
plot([-r r], [-r r], 'k:', 'LineWidth', 0.75);
plot([-r r], [ r -r], 'k:', 'LineWidth', 0.75);

% Rays for each sector, colored by quadrant (based on center angle)
dxw = Lw_scaled .* cos(centers(:));
dyw = Lw_scaled .* sin(centers(:));

for k = 1:nbins
    th = centers(k);
    % Quadrant mapping by angle
    if      (th > 0)        && (th <=  pi/2), col = C_NE;  % NE
    elseif  (th >  pi/2)    && (th <=  pi),   col = C_NO;  % NO
    elseif  (th > -pi/2)    && (th <=  0),    col = C_SE;  % SE
    else,                                      col = C_SO;  % SO (includes th in (-pi,-pi/2] and exactly -pi)
    end
    quiver(0,0,dxw(k),dyw(k),0,'LineWidth',2.2,'Color',col,'MaxHeadSize',0.8);
end

axis equal; xlim([-r r]); ylim([-r r]); box on; ax.XTick=[]; ax.YTick=[];
norm_str = ternary(include_axes_in_denominator, 'TOTAL', 'quad-defined');
title(sprintf('20-bin glyph (length-weighted, den=%s) — scaled ×%.1f', norm_str, scale_gain));

% Footer annotation
annotation('textbox',[0.0 0.0 1.0 0.07], 'String', ...
    sprintf('On-axes: %d/%d  |  W_{den}=%.3g', sum(on_axes), N, Wden), ...
    'EdgeColor','none','HorizontalAlignment','center','FontSize',9);



function y = ternary(cond, a, b)
    if cond, y = a; else, y = b; end
end




