function [frex,tf]=mwavelet(time,signal,srate,nfrex)
% Time-frequency analysis funciton, modified from "A better way to define
% and describe Morlet wavelets for time-frequency analysis" by Michael X
% Cohen. bioRxiv preprint doi: https://doi.org/10.1101/397182. August 2018.
% time=time (x) axis; signal=amplitude (y axis); srtate=sample rate;
% nfrex=number of frequencies calculated; tf=time-frequency matrix (size
% pnts x frex)
pnts = length(time);
frex = linspace(2,25,nfrex);
fwhm = linspace(.8,.7,nfrex);
% setup wavelet and convolution parameters
wavet = -2:1/srate:2;
halfw = floor(length(wavet)/2)+1;
nConv = pnts + length(wavet) - 1;
% initialize time-frequency matrix
tf = zeros(length(frex),pnts);
% spectrum of data
dataX = fft(signal,nConv);

% loop over frequencies
for fi=1:length(frex)
    
    % create wavelet
    waveX = fft( exp(2*1i*pi*frex(fi)*wavet).*exp(-(4*log(2)*wavet).^2/fwhm(fi).^2),nConv );
    waveX = waveX./max(waveX); % normalize
    
    % convolve
    as = ifft( waveX.*dataX );
    % trim and reshape
    as = as(halfw:end-halfw+1);
    
    % power
    tf(fi,:) = abs(as).^2;
end
