function TF = time_freq_wavelet(data, frex, srate, timevec, normalization, baselinetime, mkplt)
% data: time x trial


% wavelet parameters
% min_freq = passband(1);
% max_freq = passband(2);

% other wavelet parameters
%frex = logspace(log10(min_freq),log10(max_freq),num_frex);
num_frex = length(frex);
time = -3:1/srate:3;
%s    = logspace(log10(3),log10(10),num_frex)./(2*pi*frex);
% s    =  3./(2*pi*frex); 
s    = 7./(2*pi*frex);  % 7 cycles

% Data parameters
n_trials = size(data,2);
n_timepnts = size(data,1);
 
 
% FFT parameters (use next-power-of-2)
n_wavelet     = length(time);
n_data        = n_trials*n_timepnts;
n_convolution = n_wavelet+n_data-1;
n_conv_pow2   = pow2(nextpow2(n_convolution));
 
half_of_wavelet_size = (length(time)-1)/2;

% FFT of data (note: this doesn't change on frequency iteration)


fft_data = fft(reshape(data,1,n_trials*n_timepnts),n_conv_pow2);

% initialize output time-frequency data
tf = zeros(num_frex, n_timepnts, n_trials);
tf_db = zeros(num_frex, n_timepnts, n_trials);
tf_avg_db = zeros(num_frex,n_timepnts);

phase = zeros(num_frex,n_timepnts, n_trials);

if normalization
    baseidx = dsearchn(timevec',baselinetime');
end
%% Wavelet-convolution in the frequency domain

for fi=1:length(frex)
    
    % create wavelet and get its FFT
    fft_wavelet = fft( sqrt(1/(s(fi)*sqrt(pi))) * exp(2*1i*pi*frex(fi).*time) .* exp(-time.^2./(2*(s(fi)^2))) , n_conv_pow2 );
    
    
    % convolution
    eegconv = ifft(fft_wavelet.*fft_data);
    eegconv = eegconv(1:n_convolution);
    eegconv = eegconv(half_of_wavelet_size+1:end-half_of_wavelet_size);
    eegconv = reshape(eegconv,n_timepnts, n_trials);
    
    % extract power
    power = abs(eegconv).^2;
    tf(fi,:,:) = power;
    
    % extract ITPC
    phase(fi,:,:) = angle(eegconv);
    
    if normalization
        % normalization for each trial
        tf_db(fi,:,:) = 10*log10(power./mean(power(baseidx(1):baseidx(2),:)));
        
        % Average power over trials and then normalization
        temppower = mean(power,2);
        tf_avg_db(fi,:) = 10*log10(temppower./mean(temppower(baseidx(1):baseidx(2))));
    end
    
end

TF.tf = tf;
TF.tf_db = tf_db;
TF.tf_avg_db = tf_avg_db;
TF.phase = phase;


%%
if mkplt
   
    contourf(timevec,frex,tf_avg_db,40,'linecolor','none')
    colormap jet
    colorbar
    ylabel('Frequency (Hz)')
    xlabel('Time(s)')
end