% Piano Sustain Effect Script
% Method 3

MIDI = 67; % Data folder has 60, 67
METHOD = 3; % This is for save filename for analysis only
START = 0.1; % Start time
PRESS = 0.0; % Press time


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

% Load the recordings
[no_decay, fs] = audioread(noPedal_path); 
[pedaled_decay, ~] = audioread(pedal_path);


if size(no_decay, 2) > 160
    no_decay = mean(no_decay, 2);
end
if size(pedaled_decay, 2) > 1
    pedaled_decay = mean(pedaled_decay, 2);
end

% Check the lengths of the two signals
len1 = length(no_decay);
len2 = length(pedaled_decay);

% Pad zeros to the shorter signal
if len1 > len2
    pedaled_decay = [pedaled_decay; zeros(len1 - len2, size(pedaled_decay, 2))];
elseif len2 > len1
    no_decay = [no_decay; zeros(len2 - len1, size(no_decay, 2))];
end

% Parameters
window_size = 4096;
hop_size = window_size/4;
nfft = window_size*2;
t = (0:length(no_decay)-1)'/fs;

% Attack portion
start_time = round(START * fs);
attack_duration = round(PRESS * fs);
attack = no_decay(1:attack_duration+start_time);


% STFT
window = hann(window_size);
[S_held, f, t_stft] = stft(no_decay, fs, 'Window', window, ...
    'OverlapLength', window_size-hop_size, 'FFTLength', nfft);
[S_pedaled, ~, ~] = stft(pedaled_decay, fs, 'Window', window, ...
    'OverlapLength', window_size-hop_size, 'FFTLength', nfft);

% Frequency band analysis
num_bands = 20;  % Analyze in frequency bands
band_edges = logspace(log10(50), log10(fs/2), num_bands+1);
modes_per_band = 15;  % Number of modes to find in each band

frequencies = [];
amplitudes = [];
decay_rates = [];
phases = [];

% Analyze each frequency band
for band = 1:num_bands
    % Find frequency indices for this band
    band_idx = f >= band_edges(band) & f < band_edges(band+1);
    band_spec = mean(abs(S_pedaled(band_idx,:)), 2);
    
    % Find peaks in this band
    [pks, locs] = findpeaks(band_spec, 'NPeaks', modes_per_band, 'SortStr', 'descend');
    
    if ~isempty(pks)
        % Convert local indices to frequency values
        band_freqs = f(band_idx);
        band_frequencies = band_freqs(locs);
        
        % Store parameters
        frequencies = [frequencies; band_frequencies];
        amplitudes = [amplitudes; pks];
        
        % Calculate phases and decay rates for each mode in band
        for i = 1:length(band_frequencies)
            [~, freq_idx] = min(abs(f - band_frequencies(i)));
            mode_evolution = abs(S_pedaled(freq_idx,:));
            phases = [phases; angle(S_pedaled(freq_idx,1))];
            
            % Estimate decay rate
            env = log(mode_evolution + eps);
            valid_points = env > max(env) - 6;
            if sum(valid_points) > 2
                time_valid = t_stft(valid_points);
                env_valid = env(valid_points);
                p = polyfit(time_valid, env_valid, 1);
                decay_rates = [decay_rates; max(-p(1), 0.05)];
            else
                decay_rates = [decay_rates; 1];
            end
        end
    end
end

% Normalize amplitudes
amplitudes = amplitudes / max(amplitudes);

% Frequency-dependent modifications
for i = 1:length(frequencies)
    % Slower decay for lower frequencies
    freq_factor = exp(-frequencies(i) / 1000);
    decay_rates(i) = decay_rates(i) * (0.5 + 0.5 * freq_factor);
    
    % Adjust amplitudes based on frequency
    freq_weight = exp(-frequencies(i) / 2000);
    amplitudes(i) = amplitudes(i) * freq_weight;
end

% Synthesize with mode coupling
y_synth = zeros(size(pedaled_decay));
y_synth(1:attack_duration+start_time) = attack;

% Generate coupled modes for decay part
t_decay = t(attack_duration+start_time+1:end);
decay_sum = zeros(size(t_decay));

for i = 1:length(frequencies)
  
    mode_signal = amplitudes(i) * exp(-decay_rates(i) * t_decay) .* ...
    cos(2*pi*frequencies(i)*t_decay + phases(i) + randn(size(t_decay)) * 0.0001);

    
    % Add coupling with nearby modes
    for j = 1:length(frequencies)
        if i ~= j
            freq_diff = abs(frequencies(i) - frequencies(j));
            if freq_diff < 50  % Couple modes within 50 Hz
                coupling = 0.2 * exp(-freq_diff / 10);
                mode_signal = mode_signal + coupling * amplitudes(j) * ...
                    exp(-decay_rates(j) * t_decay) .* ...
                    cos(2 * pi * frequencies(j) * t_decay + phases(j));
            end
        end
    end
    
    decay_sum = decay_sum + mode_signal;
end

% Add synthesized decay
y_synth(attack_duration+start_time+1:end) = decay_sum;


% Multi-band envelope matching with fixed butterworth filter
num_env_bands = 20;
band_edges = logspace(log10(20), log10(fs/2-100), num_env_bands+1);
window_len = round(0.1 * fs);

for band = 1:num_env_bands
    % Normalize and bound frequencies for butterworth filter
    f1 = max(min(band_edges(band)/(fs/2), 0.99), 0.01);
    f2 = max(min(band_edges(band+1)/(fs/2), 0.99), 0.01);
    
    % Create bandpass filter
    if band == 1
        % Lowpass for first band
        [b, a] = butter(2, f2, 'low');
    elseif band == num_env_bands
        % Highpass for last band
        [b, a] = butter(2, f1, 'high');
    else
        % Bandpass for middle bands
        [b, a] = butter(2, [f1, f2], 'bandpass');
    end
    
    % Apply filtering
    pedaled_band = filtfilt(b, a, pedaled_decay);
    synth_band = filtfilt(b, a, y_synth);
    
    % Calculate and apply envelope
    pedaled_env = movmean(abs(pedaled_band).^2, window_len).^0.5;
    synth_env = movmean(abs(synth_band).^2, window_len).^0.5;
    
    env_scale = pedaled_env ./ (synth_env + eps);
    env_scale = smoothdata(env_scale, 'gaussian', round(0.05 * fs));
    
    y_synth = y_synth + (synth_band .* env_scale - synth_band);
end

% Final normalization
y_synth = y_synth / max(abs(y_synth)) * max(abs(pedaled_decay));

% Plot results
figure('Position', [100, 100, 1200, 800]);

subplot(2,2,1);
plot(t, pedaled_decay);
title('Pedaled Note (Time Domain)');
xlabel('Time (s)');
ylabel('Amplitude');

subplot(2,2,2);
plot(t, y_synth);
title('Synthesized Note (Time Domain)');
xlabel('Time (s)');
ylabel('Amplitude');

% Plot spectrograms
window = hamming(2048);
noverlap = 1024;
nfft = 4096;

subplot(2,2,3);
spectrogram(pedaled_decay, window, noverlap, nfft, fs, 'yaxis');
title('Pedaled Note Spectrogram');
colorbar;

subplot(2,2,4);
spectrogram(y_synth, window, noverlap, nfft, fs, 'yaxis');
title('Synthesized Note Spectrogram');
colorbar;

% Make colorbars the same scale
ax3 = subplot(2,2,3);
ax4 = subplot(2,2,4);
c3 = caxis(ax3);
c4 = caxis(ax4);
clim = [min(c3(1), c4(1)), max(c3(2), c4(2))];
caxis(ax3, clim);
caxis(ax4, clim);

% Play sounds
% disp('Playing pedaled decay...');
% sound(pedaled_decay, fs);
% pause(length(pedaled_decay)/fs + 0.5);

disp('Playing synthesized decay...');
sound(y_synth, fs);

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