% Piano Sustain Effect Script, method 1
% Based on Lehtonen et al. 2007 paper

MIDI = 60; % Directory has 60, 67
METHOD = 1; % This is for save filename for analysis only

% Mixing coefficient and energy detection parameters
% Table V in the paper specifies mixing coefficients based on key index:
% Keys 1-16:   0.015
% Keys 17-28:  0.01
% Keys 29-40:  0.008
% Keys 41-54:  0.006
% Keys 55-88:  0.005

% For MIDI note 60 (middle C)
mixing_coeff = 0.005;  

noPedal_path = sprintf('../data/sostenuto_%d.wav', MIDI);
pedal_path = sprintf('../data/long_pedal_%d.wav', MIDI);

fs = 44100;  

% Input and convert to mono
[input_signal, fs_file] = audioread(noPedal_path);
[pedal, Fs2] = audioread(pedal_path);

if fs_file ~= fs
    input_signal = resample(input_signal, fs, fs_file);
end

input_signal = input_signal(:, 1); 
pedal = pedal(:,1);

% Initialize parameters for 12 string models
fundamental_freqs = 27.5 * 2.^((0:11)/12);  % Starting from A0
delay_lengths = floor(fs ./ fundamental_freqs + 0.5);

% Damping coefficients (table IV g)
damping_coeffs = [0.9918, 0.9903, 0.9942, 0.9928, 0.9929, 0.9941,...
                 0.9941, 0.9954, 0.9947, 0.9958, 0.9938, 0.9929];

% Dispersion filter coefficients
B = [3.0e-4, 2.9e-4, 2.7e-4, 2.6e-4, 2.5e-4, 2.4e-4,...
     2.3e-4, 2.2e-4, 2.1e-4, 2.0e-4, 1.9e-4, 1.8e-4];
Dk = sqrt(B) * 2;
dispersion_coeffs = (1 - Dk) ./ (1 + Dk);

% Initialize single delay line for each string model
delay_lines = cell(1, 12);
for i = 1:12
    delay_lines{i} = zeros(1, delay_lengths(i));
end

% Initialize pointers
read_ptrs = ones(1, 12);
write_ptrs = ones(1, 12);

% Initialize filter states
dispersion_states = zeros(1, 12);
lowpass_states = zeros(1, 12);

% Tone corrector coefficient from eq (6)
b_tone = [0.1164, -0.1164]; 
a_tone = [1, -1.880, 0.8836]; 

% Initialize output
output_signal = zeros(size(input_signal));

% Main loop
for n = 1:length(input_signal)
    
    % Apply tone correction
    in_sample = filter(b_tone, a_tone, input_signal(n));
    
    % Process through all string models
    resonance = zeros(1, 12);
    
    for string = 1:12

        % Get current pointers and delay line
        read_ptr = read_ptrs(string);
        write_ptr = write_ptrs(string);
        delay_line = delay_lines{string};
        
        % Read from delay line
        delay_out = delay_line(read_ptr);
        
        % Apply dispersion filter to input
        dispersed = dispersion_coeffs(string) * in_sample + dispersion_states(string);
        dispersed = dispersed / (1 + dispersion_coeffs(string)^2);
        dispersion_states(string) = in_sample;
        
        % Apply lowpass filter
        c = -0.197; % Frequency-dependent decay coefficient
        b_k = damping_coeffs(string) * (1 + c);
        filtered = b_k * delay_out / (1 + c * lowpass_states(string));
        lowpass_states(string) = delay_out;
        
        % Write to delay line
        delay_line(write_ptr) = dispersed + 0.2 * filtered;
        
        % Update resonance
        resonance(string) = filtered;
        
        % Update pointers
        read_ptrs(string) = mod(read_ptr, length(delay_line)) + 1;
        write_ptrs(string) = mod(write_ptr, length(delay_line)) + 1;
        
    end
    
    % Mix original signal with resonance
    output_signal(n) = input_signal(n) + mixing_coeff * sum(resonance);
end


% Normalize output to prevent clipping
output_signal = output_signal / max(abs(output_signal));

% Plot 
figure('Name', 'Piano Sustain Effect Analysis, Lehtonen');

% Time domain comparison
subplot(4,1,1);
t = (0:length(input_signal)-1)/fs;
plot(t, input_signal, 'b', t, output_signal, 'r');
title('Time Domain Comparison');
legend('Original', 'With Sustain');
xlabel('Time (s)');
ylabel('Amplitude');

% Spectrograms
subplot(4,1,2);
spectrogram(input_signal, hamming(2048), 1024, 2048, fs, 'yaxis');
title('Spectrogram - Original Signal');

subplot(4,1,3);
spectrogram(output_signal, hamming(2048), 1024, 2048, fs, 'yaxis');
title('Spectrogram - With Sustain Effect');

% Frequency response comparison
subplot(4,1,4);
[pxx_in, f] = pwelch(input_signal, hamming(2048), 1024, [], fs);
[pxx_out, ~] = pwelch(output_signal, hamming(2048), 1024, [], fs);
plot(f, 10*log10(pxx_in), 'b', f, 10*log10(pxx_out), 'r');
title('Frequency Response Comparison');
xlabel('Frequency (Hz)');
ylabel('Power/Frequency (dB/Hz)');
legend('Original', 'With Sustain');

% Play the results
disp('Playing with sustain effect...');
soundsc(output_signal, fs);


% Save the result
% synth_path = sprintf('../result/%d_%d.wav', METHOD, MIDI);
% audiowrite(synth_path, output_signal, fs);
% disp('Processed audio saved');