% Forensic Environmental Cochlear Modelling
% Toy implementation of:
% - 1D transmission-line style BM displacement
% - "3D-like" BM displacement maps (pos vs freq)
% - Perceptual ERB-like filterbank, IHC, Poisson spikes
% - Environmental Auditory Risk Index (EARI) components
%
% Thiago Mathias – based on the methodological description in:
% Verhulst et al., 2018; Saremi et al., 2022; Hanekom et al., 2019;
% Hrncirik et al., 2023; Manis et al., 2018 (conceptual only).

clear; close all; clc;

%% 1. Stimulus definition (chirp in quiet vs noisy conditions)
Fs        = 44100;        % sampling rate (Hz)
T         = 3;            % duration (s)
t         = (0:1/Fs:T-1/Fs).';

f_start   = 500;          % chirp start frequency (Hz)
f_end     = 8000;         % chirp end frequency (Hz)

% Linear chirp in log-frequency domain
chirp_sig = chirp(t,f_start,T,f_end,'logarithmic');

% SPL scaling (relative; we just use dB differences)
SPL_ctrl  = 30;           % dB SPL (control)
SPL_noise = 85;           % dB SPL (noisy)
gain_ctrl = 10^(SPL_ctrl/20);
gain_nois = 10^(SPL_noise/20);

x_ctrl    = chirp_sig * gain_ctrl;
x_nois    = chirp_sig * gain_nois;

%% 2. Cochlear / BM model (Greenwood + Gaussian tuning)
% This section produces:
% - BM displacement maps (freq vs position) for CTRL & NOISY
% - Peak BM displacement vs position
%
% 1D transmission-line concept: we mimic local admittance with a
% Gaussian place–frequency map, level-dependent gain and bandwidth.

L_mm      = 35;                   % cochlear length (mm)
Npos      = 200;                  % positions along BM
Nfreq     = 200;                  % frequency samples

pos_mm    = linspace(0,L_mm,Npos);                 % from base (0) to apex (L)
freq_vec  = logspace(log10(300),log10(8000),Nfreq);

% Greenwood function parameters (human-like)
A_g       = 165.4;
a_g       = 2.1;
k_g       = 0.88;

% Characteristic frequency at each position
CF        = A_g * (10.^(a_g*(1 - pos_mm/L_mm)) - k_g);  % Hz

% Gaussian tuning width (in log10 frequency)
sigma_ctrl = 0.15;   % narrow filter – sharp tuning
sigma_nois = 0.30;   % broader filter – degraded tuning

% Preallocate matrices: pos x freq
BM_ctrl_lin = zeros(Npos,Nfreq);
BM_nois_lin = zeros(Npos,Nfreq);

for ip = 1:Npos
    fcf = CF(ip);
    log_ratio = log10(freq_vec/fcf);
    
    % Level-dependent gain (arbitrary units)
    gain_pos_ctrl = 1;            % baseline
    gain_pos_nois = 10^(40/20);   % ~40 dB stronger
    
    BM_ctrl_lin(ip,:) = gain_pos_ctrl * ...
        exp(- (log_ratio.^2)/(2*sigma_ctrl^2));
    
    BM_nois_lin(ip,:) = gain_pos_nois * ...
        exp(- (log_ratio.^2)/(2*sigma_nois^2));
end

% Convert to dB and normalise so that global max = 0 dB, min ~ -60 dB
BM_ctrl_dB = 20*log10(BM_ctrl_lin + eps);
BM_nois_dB = 20*log10(BM_nois_lin + eps);

max_all = max([BM_ctrl_dB(:); BM_nois_dB(:)]);
BM_ctrl_dB = BM_ctrl_dB - max_all;
BM_nois_dB = BM_nois_dB - max_all;
BM_ctrl_dB = max(BM_ctrl_dB,-60);    % floor at -60 dB
BM_nois_dB = max(BM_nois_dB,-60);

% Peak displacement vs position (for line plot)
peak_ctrl = max(BM_ctrl_dB,[],2);
peak_nois = max(BM_nois_dB,[],2);

%% 3. Plot BM displacement maps (Figure ~ BM map CTRL vs NOISY)
figure('Position',[100 100 900 380]);
subplot(1,2,1);
imagesc(freq_vec,pos_mm,BM_ctrl_dB);
set(gca,'YDir','normal','XScale','log');
xlabel('Frequency (Hz)');
ylabel('Cochlear position (mm)');
title('Control (30 dB SPL) – BM Displacement Map');
c = colorbar;
c.Label.String = 'Normalised BM displacement (dB)';

subplot(1,2,2);
imagesc(freq_vec,pos_mm,BM_nois_dB);
set(gca,'YDir','normal','XScale','log');
xlabel('Frequency (Hz)');
ylabel('Cochlear position (mm)');
title('Noise (85 dB SPL) – BM Displacement Map');
c = colorbar;
c.Label.String = 'Normalised BM displacement (dB)';

%% 4. Peak BM displacement vs position (line plot)
figure('Position',[200 200 650 420]);
plot(pos_mm,peak_ctrl,'b','LineWidth',2); hold on;
plot(pos_mm,peak_nois,'r','LineWidth',2);
xlabel('Position from base (mm)');
ylabel('Peak BM magnitude (dB, arb.)');
title('Peak Basilar Membrane Displacement — CTRL vs NOISY');
legend({'Controlled (30 dB)','Noisy (85 dB)'},'Location','SouthEast');
grid on;

%% 5. Perceptual pipeline: ERB-style filterbank + IHC + Poisson spikes
% We implement:
%   - bank of bandpass filters centred at ERB-spaced frequencies
%   - nonlinear compression
%   - rectification + low-pass integration
%   - Poisson spike generation

% Downsampled rate for efficiency
Fs_ds   = 10000;
x_ctrl_ds = resample(x_ctrl,Fs_ds,Fs);
x_nois_ds = resample(x_nois,Fs_ds,Fs);
t_ds    = (0:length(x_ctrl_ds)-1)'/Fs_ds;

% Centre frequencies for channels
Nch     = 40;
f_min   = 300;
f_max   = 8000;
cf_vec  = logspace(log10(f_min),log10(f_max),Nch);

% ERB filter bandwidth (simple approx)
erb = @(f) 24.7*(4.37*f/1000 + 1);

% Allocate matrices channel x time
env_ctrl = zeros(Nch,length(x_ctrl_ds));
env_nois = zeros(Nch,length(x_nois_ds));

gamma = 0.3;   % compression exponent
tau_lpf = 0.005;  % IHC low-pass time constant (~5 ms)
fc_lpf  = 1/(2*pi*tau_lpf);  % approx low-pass cutoff

for k = 1:Nch
    bw = erb(cf_vec(k));
    f1 = max(50, cf_vec(k) - bw);
    f2 = min(Fs_ds/2-100, cf_vec(k) + bw);
    
    % Simple bandpass (requires Signal Processing Toolbox)
    y_ctrl = bandpass(x_ctrl_ds,[f1 f2],Fs_ds);
    y_nois = bandpass(x_nois_ds,[f1 f2],Fs_ds);
    
    % Nonlinear compression (OHC-like)
    u_ctrl = abs(y_ctrl).^gamma;
    u_nois = abs(y_nois).^gamma;
    
    % Half-wave rectification + low-pass filtering (IHC)
    u_ctrl = max(u_ctrl,0);
    u_nois = max(u_nois,0);
    
    env_ctrl(k,:) = lowpass(u_ctrl,fc_lpf,Fs_ds);
    env_nois(k,:) = lowpass(u_nois,fc_lpf,Fs_ds);
end

% Instantaneous firing rate (arbitrary scaling)
alpha_rate = 300;        % spikes/s per unit envelope
lambda_ctrl = alpha_rate * env_ctrl;
lambda_nois = alpha_rate * env_nois;

% Generate Poisson spikes at 1 kHz sampling for raster display
Fs_spk   = 1000;
dt_spk   = 1/Fs_spk;
t_spk    = 0:dt_spk:T-dt_spk;
idx_step = round(Fs_ds/Fs_spk);

lambda_ctrl_ds = lambda_ctrl(:,1:idx_step:end);
lambda_nois_ds = lambda_nois(:,1:idx_step:end);

% Probability per bin
P_ctrl = lambda_ctrl_ds * dt_spk;
P_nois = lambda_nois_ds * dt_spk;

rng(1); % reproducible
spk_ctrl = rand(size(P_ctrl)) < P_ctrl;
spk_nois = rand(size(P_nois)) < P_nois;

%% 6. Neural activation heatmaps (rate in dB)
rate_ctrl_dB = 10*log10(lambda_ctrl + eps);
rate_nois_dB = 10*log10(lambda_nois + eps);

% Normalise for plotting
max_rate = max([rate_ctrl_dB(:); rate_nois_dB(:)]);
rate_ctrl_dB = rate_ctrl_dB - max_rate;
rate_nois_dB = rate_nois_dB - max_rate;
rate_ctrl_dB = max(rate_ctrl_dB,-55);
rate_nois_dB = max(rate_nois_dB,-55);

figure('Position',[150 150 900 380]);
subplot(1,2,1);
imagesc(t_ds,cf_vec,rate_ctrl_dB);
set(gca,'YDir','normal');
xlabel('Time (s)');
ylabel('Frequency (Hz)');
title('Neural Activation Heatmap — Controlled');
c = colorbar;
c.Label.String = 'Relative activation (dB)';

subplot(1,2,2);
imagesc(t_ds,cf_vec,rate_nois_dB);
set(gca,'YDir','normal');
xlabel('Time (s)');
ylabel('Frequency (Hz)');
title('Neural Activation Heatmap — Noisy');
c = colorbar;
c.Label.String = 'Relative activation (dB)';

%% 7. Neuromorphic spike rasters
figure('Position',[200 200 900 380]);
subplot(1,2,1);
hold on;
for k = 1:Nch
    t_k = t_spk(spk_ctrl(k,:));
    y_k = ones(size(t_k))*cf_vec(k);
    plot(t_k,y_k,'.','MarkerSize',4);
end
xlabel('Time (s)');
ylabel('Frequency (Hz)');
title('Neuromorphic Spikes — Controlled');
ylim([f_min f_max]);
set(gca,'YDir','normal');

subplot(1,2,2);
hold on;
for k = 1:Nch
    t_k = t_spk(spk_nois(k,:));
    y_k = ones(size(t_k))*cf_vec(k);
    plot(t_k,y_k,'.','MarkerSize',4);
end
xlabel('Time (s)');
ylabel('Frequency (Hz)');
title('Neuromorphic Spikes — Noisy');
ylim([f_min f_max]);
set(gca,'YDir','normal');

%% 8. Environmental Auditory Risk Index (EARI)
% Toy example: component scores normalised in [0,1]
BM_displacement_score = 0.9;    % stronger mechanical load
selectivity_loss_score = 0.65;  % moderate–high loss
neural_variability_score = 1.0; % very high variability

scores = [BM_displacement_score, ...
          selectivity_loss_score, ...
          neural_variability_score];

labels = {'BM displacement','Selectivity loss','Neural variability'};

figure('Position',[200 200 900 420]);

% Left: component bar chart
subplot(1,2,1);
barh(scores);
xlim([0 1]);
set(gca,'YTick',1:3,'YTickLabel',labels);
xlabel('Normalised component score (0–1)');
title('EARI Components');
grid on;

% Right: risk gauge (LOW–MODERATE–HIGH)
subplot(1,2,2);
hold on;
x0 = 0; x1 = 1;
y0 = 0; h = 1;

% coloured blocks
rectangle('Position',[x0,y0,(x1-x0)/3,h],'FaceColor',[0.7 0.9 0.7],'EdgeColor','none'); % low
rectangle('Position',[x0+(x1-x0)/3,y0,(x1-x0)/3,h],'FaceColor',[0.98 0.93 0.65],'EdgeColor','none'); % moderate
rectangle('Position',[x0+2*(x1-x0)/3,y0,(x1-x0)/3,h],'FaceColor',[0.95 0.7 0.7],'EdgeColor','none'); % high

% labels
text(0.17,1.05,'LOW','HorizontalAlignment','center');
text(0.5,1.05,'MODERATE','HorizontalAlignment','center');
text(0.83,1.05,'HIGH','HorizontalAlignment','center');

% overall risk score (here we use mean of components)
EARI = mean(scores);  % 0–1
plot([EARI EARI],[y0-0.1 y0+h+0.1],'k','LineWidth',2);
plot(EARI,y0+h+0.12,'kv','MarkerFaceColor','k','MarkerSize',8);

xlim([0 1]);
ylim([-0.1 1.3]);
axis off;
title('Environmental Auditory Risk Index (EARI)');