function [WCO,Wxy,sWxy,periods,scales,coi,pvCO,WCoh] = ...
    AWCOh(x,y,dt,dj,lowPeriod,upPeriod,pad,mother,be,ga,...
          wT_type,wT_size,wS_type,wS_size, nSur,sur_type,p,q)
                      
%AWCO  Wavelet Coherency and Cross Wavelet Transform of two series
%
%   [WCO,WXY,SWXY,PERIODS,SCALES,COI,PVC0] = ...
%       AWCO(X,Y,DT,DJ,LOWPERIOD,UPPERIOD,PAD,MOTHER,BE,GA,...
%            WT_TYPE,WT_SIZE,WS_TYPE,WS_SIZE,NSUR,SUR_TYPE,P,Q)   
%   Computes the Wavelet Coherency matrix (WCO) of two series X and Y, 
%   using an Analytic Wavelet Transform with a Morlet or a Generalized 
%   Morse Wavelet.
%   Different windows for smoothing can be used.
%   Also computes the Cross-Wavelet Transform (smoothed (SXWY) or 
%   non-smoothed (WXY)).
%   It also gives the p-values for the wavelet coherency (PVCO); these
%   are computed with NSUR surrogate series. The surrogates can be 
%   constructed in two ways:  fitting an ARMA(P,Q) model and building new 
%   samples by bootstrap or fitting an ARMA(P,Q) model and construct new 
%   samples by drawing errors from a  Gaussian distribution.
%
%   It has two subfunctions: WAVELETCOHERENCY and WINDOWBARTLETT 
%
%   It uses function AWT from this toolbox
% 
%   Might need function WINDOW from the matlab Signal Processing Toolbox.
%
%   Might need functions GARCHFIT, GARCHSET and  GARCHSIM  from the
%   Econometrics Toolbox.
%  
%  To use the function with no need of any Matlab toolbox, you must select: 
%
%       WT_TYPE=0 or WT_TYPE='Bar' or WT_TYPE=[]
%       WS_TYPE=0 or WS_TYPE='Bar' or WS_TYPE=[]
%       SUR_TYPE=0 or SUR_TYPE='ARMABoot' or  SUR_TYPE=[]
%       Q=0 or Q=[] (or missing)
%   INPUTS:
%       X, Y - two vectors of the same length (series).
%
%   Optional inputs: 
%       DT - sampling rate (Default: 1)
%       DJ - frequency resolution (1/DJ = number of voices per octave);
%            (Default: 1./4)
%       LOWPERIOD - lower period of the decomposition (Default: 2*DT)
%       UPPERIOD - upper period of the decomposition 
%           (Default: length(X)*DT)
%       PAD - an integer (power of 2) defining the total length of the 
%             vectors X and Y after zero padding; if PAD is not a power of
%             2, pad with zeros to total length: 2^(next power of 2 + 2)
%       MOTHER - the mother wavelet function 
%              0 or 'Morlet' -> Morlet wavelet (default)
%              1 or  'GMW' ->  Generalized Morse Wavelet    
%       BE - the beta parameter for the GMW or the Morlet (w0) parameter 
%           Default values: 6.0 for 'Morlet'
%                           3.0 for 'GMW'                      
%       GA - the gamma parameter of the GMW
%            Default value: 3.0
%       WT_TYPE - type of window for smoothing in time direction                              
%       WT_SIZE - (half) size of window in time (default: 5).
%       WS_TYPE - type of window for smoothing in scale direction                              
%       WS_SIZE - (half) size of window in scale (default: 5)     
%       WT_TYPE,WS_TYPE -       
%          0 or 'Bar' -> Bartlett (default); in this case, there is no need
%                       of Signal Processing Toolbox
%          1 or 'Bla'-> Blackman
%          2 or 'Box' -> Rectangular
%          3 or 'Ham' -> Hamming 
%          4 or 'Han' -> Hanning                            
%          5 or 'Tri' -> Triangular
%          
%       NSUR - integer, number of surrogate series, if we want to compute
%              p-values for the Wavelet Coherency
%             (Default: 0 - no computation)           
%       SUR_TYPE - type of surrogate series
%           0 or 'ARMABoot'-> surrogates based on ARMA(P,Q) model+bootstrap
%                           (default)
%                           In this case, IF Q=0, there is no need 
%                           of Econometrics toolbox.
%           1 or 'Econ' -> surrogates based on ARMA(P,Q) model+Gauss;
%                          always needs Econometrics toolbox.       
%       P,Q -  non-negative integers, orders of the ARMA(P,Q) model 
%              (Default: P=0, Q=0).
%                
%   OUTPUTS:
%       WCO - Complex Wavelet Coherency of X and Y
%       WXY - Cross-Wavelet Transform of X and Y
%       SWXY - Smoothed Cross Wavelet Transform of X and Y
%       PERIODS - the vector of Fourier periods (in time units)
%                 that correspond to the the scales
%       SCALES - the vector of scales, given by 
%                s0*2^(j*dj);j=0,...,J1, where J1+1 is  number of scales 
%                and s0 is minimum scale
%       COI - the "cone-of-influence", which is a vector of n_x=length(X) 
%             points that contains the limit of the region where the 
%             wavelet transforms are influenced by edge effects
%       PVCO - p-values for Wavelet Coherency
%   Example:
%       t=1:1000; 
%       x=sin(t);
%       y=sin(t.*cos(t*0.1));
%       dt=1/6; dj=1/12;
%       lowPeriod=1; upPeriod=50;
%       pad=0;
%       mother='GMW';
%       be=3.0; ga=3.0;
%       wT_type='Box';
%       wT_size=2;
%       wS_type='Box';
%       wS_size='2;
%       nSur=5000;
%       sur_type='ARMABoot';
%       p=1;
%       q=1;
%      
%       [XCO,WXY,sWXY,periods,scales,coi,pvCO]=...
%        AWCO(x,y,dt,dj,lowPeriod,upPeriod,pad,mother,be,ga,...
%             wT_type,wT_size,wS_type,wS_size,nSur,sur_type,p,q);
%                            
%   See also: AWCOOUTPUT.


%   References:
%   [1] Aguiar-Conraria, L. and Soares, M.J. (2010)
%       "The Continuous Wavelet Transform: A Primer", NIPE Working paper.
%   [2] Cazelles, B., Chavez, M., de Magny, G. C., Guégan, J.-F. and 
%       Hales, S. (2007), "Time-Dependent Spectral Analysis of
%       Epidemiological Time-Series with Wavelets", Journal of the Royal 
%       Society Interface, 4, 625—36. 
%   [3] Lilly, J.M. and Olhede, S. C.(2009), "Higher-Order Properties of 
%       Analytic Wavelets", IEEE Trans. Signal Process. 57 (1), 146-160.
%   [4]  Lilly, J.M. and Olhede, S. C.(2009), "On the Analytic Wavelet 
%       Transform", eprint arXiv:0711.3834
%   [5] Torrence, C. and Compo, T.C., "A Prectical Guide to Wavelet 
%       Analysis" (1998), Bulletin of the American Meteorological Society,      
%       79, 605—618.
% 

%   Copyright 2010, L. Aguiar-Conraria and M.J. Soares
%
%
%   Luís AGUIAR-CONRARIA              Maria Joana SOARES                      
%   Dep. Economics                    Dep. Mathematics and Applications   
%   University of Minho               University of Minho
%   4710-057 Braga                    4710-057 Braga
%   PORTUGAL                          PORTUGAL
%                           
%   lfaguiar@eeg.uminho.pt            jsoares@math.uminho.pt 

%%%%%%%%%%%%%%%     Tests on inputs and default values    %%%%%%%%%%%%%%%%%
% Tests on inputs
if (nargin <2)
    error('You must input two series');
% Check that x and y are vectors
elseif (~(isvector(x)) || ~(isvector(y))) 
    error('x and y must be vectors');
end
% Check that vectors have the same length
n=length(x);
if n ~= length(y)
    error('Series are not of the same length');
end

% Default parameters
if ((nargin<3) || isempty(dt)), dt = 1; end 
if ((nargin<4) || isempty(dj)), dj = 1./4; end
if ((nargin<5) || isempty(lowPeriod)), lowPeriod = 2*dt; end
if ((nargin<6) || isempty(upPeriod)), upPeriod = n*dt; end
if ((nargin<7) || isempty(pad)), pad = 0; end
if ((nargin<8) || isempty(mother)), mother='Morlet'; end
if ((nargin<9) || isempty(be)), be = 0; end
if ((nargin<10)|| isempty(ga)), ga = 3.0; end
if ((nargin<11)|| isempty(wT_type)), wT_type = 'Bar'; end
if ((nargin<12) || isempty(wT_size)), wT_size = 5; end 
if ((nargin<13)|| isempty(wS_type)), wS_type='Bar'; end
if ((nargin<14) || isempty(wS_size)), wS_size = 5; end 
if ((nargin<15) || isempty(nSur)), nSur=0; end
if ((nargin<16) || isempty(sur_type)), sur_type = 0;end
if ((nargin<17) || isempty(p)), p = 0; end
if ((nargin<18) || isempty(q)), q = 0; end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if isnumeric(mother)
    if mother==0
        mother = 'Morlet';
    elseif mother==1
        mother = 'GMW';   
    else
        error('wrong type of mother wavelet')
    end
end

if isnumeric(wT_type)
    if wT_type==1
        wT_type = @blackman;
    elseif wT_type==2
        wT_type = @rectwin;
    elseif wT_type==3
        wT_type = @hamming;
    elseif wT_type==4
        wT_type = @hann;
    elseif wT_type==5
        wT_type = @triangle;
    elseif ~(wT_type==0)
        error('wrong window type')
    end
end

if ischar(wT_type)  
    if strcmpi(wT_type,'Bla')
        wT_type = @blackman;
    elseif strcmpi(wT_type,'Box')
        wT_type = @rectwin;
    elseif strcmpi(wT_type,'Ham')
        wT_type = @hamming;
    elseif strcmpi(wT_type,'Han')
        wT_type = @hann;
    elseif strcmpi(wT_type,'Tri')
        wT_type = @triang;    
    elseif ~(strcmpi(wT_type,'Bar'))
        error('wrong window type')
    end
end

if isnumeric(wS_type)
    if wS_type==1
        wS_type = @blackman;
    elseif wS_type==2
        wS_type = @rectwin;
    elseif wS_type==3
        wS_type = @hamming;
    elseif wS_type==4
        wS_type = @hann;
    elseif wS_type==5
        wS_type = @triangle;
    elseif ~(wS_type==0)
        error('wrong window type')
    end
end

if ischar(wS_type)  
    if strcmpi(wS_type,'Bla')
        wS_type = @blackman;
    elseif strcmpi(wS_type,'Box')
        wS_type = @rectwin;
    elseif strcmpi(wS_type,'Ham')
        wS_type = @hamming;
    elseif strcmpi(wS_type,'Han')
        wS_type = @hann;
    elseif strcmpi(wS_type,'Tri')
        wS_type = @triang;    
    elseif ~(strcmpi(wS_type,'Bar'))
        error('wrong window type')
    end
end

if isnumeric(sur_type)
    if sur_type==0
        sur_type = 'ARMABoot';
    elseif sur_type==1
        sur_type = 'ARMAEcon';   
    else
        error('wrong type of surrogates')
    end
end
%%%%%%%%%  Computation of WCO,WXY,sWXY, periods, scales and coi %%%%%%%%%%% 
%                 -- uses subfunction WaveletCoherency --
%                 
% Use column vectors 
x=x(:);
y=y(:);

% Normalize series
x=(x-mean(x))./std(x);
y=(y-mean(y))./std(y); 

 [WCO,Wxy,sWxy,periods,scales,coi,WCoh] = ...
    WaveletCoherency(x,y,dt,dj,lowPeriod,upPeriod,pad,mother,be,ga,...
               wT_type,wT_size,wS_type,wS_size);

%%%%%%%%%%%%%  Computation of p-values of Wavelet Coherency  %%%%%%%%%%%%%% 
%                      (only computed if nSur>0)                          
% -- uses function SurrogateARAMABoot or SurrogateARMAEcon of this toolbox
if nSur>0
    [nRows,nCols] = size(WCO);
    pvCO=zeros(size(WCO));  
    if strcmpi(sur_type,'ARMABoot')
        SeriesSurX = SurrogateARMABoot(x,p,q,nSur);
        SeriesSurY = SurrogateARMABoot(y,p,q,nSur);  
	elseif (strcmpi(sur_type,'ARMAEcon'))
        SeriesSurX = SurrogateARMAEcon15(x,p,q,nSur);
        SeriesSurY = SurrogateARMAEcon15(y,p,q,nSur); 	
    else
        error('wrong sur_type')
    end

    for iSur=1:nSur      
        xSur = SeriesSurX(:,iSur);
        ySur = SeriesSurY(:,iSur);
        WCOSur = WaveletCoherency(xSur,ySur,dt,dj,lowPeriod,upPeriod,...
                    pad,mother,be,ga,wT_type,wT_size,wS_type,wS_size);
        for jj=1:nRows
            for kk=1:nCols
                if abs( WCOSur(jj,kk) ) >=abs( WCO(jj,kk) )
                     pvCO(jj,kk) = pvCO(jj,kk)+1;
                end
            end
        end
    end
    pvCO=pvCO/nSur; 
else
    pvCO=[];
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                 SUBFUNCTION WAVELETCOHERENCY                           %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%           
function [WCO,Wxy,sWxy,periods,scales,coi,WCoh] = ...
    WaveletCoherency(x,y,dt,dj,lowPeriod,upPeriod,pad,mother,...
               be,ga,wT_type,wT_size,wS_type,wS_size)

%%%%%%%%%%%%%%      Wavelet Transform of x and y     %%%%%%%%%%%%%%%%%%%%
%                 (uses function AWT of this toolbox) 
[wav_x,periods,scales,coi] = AWT(x,dt,dj,lowPeriod,...
                                       upPeriod,pad,mother,be,ga);
wav_y = AWT(y,dt,dj,lowPeriod,upPeriod,pad,mother,be,ga);
                                         
%%%%%%%%%%%%%%   Computations of  Cross-Wavelet Transform    %%%%%%%%%%%%%
Wxy=wav_x .* conj(wav_y); %  Formula (23) in [1]
                          
%%%%%%%%%%%%%%%%%%%          Smoothing Process          %%%%%%%%%%%%%%%%%%%
sWxy=zeros(size(Wxy));
Wxx=abs(wav_x).^2; 
Wyy=abs(wav_y).^2;
[nScales,nTimes]= size(Wxy);

% Smoothing in scale direction
    wS_size=fix(wS_size/(2*dj));
    if (ischar(wS_type)&& strcmpi(wS_type,'Bar'))...
        ||(isnumeric(wS_type)&& wS_type==0)
        FIL=WindowBartlett(wS_size);
    else
        FIL = window(wS_type,wS_size); 
    end
    FIL = FIL./(sum(FIL));
    for iTime=1:nTimes
        Wxx(:,iTime) = conv(Wxx(:,iTime),FIL,'same');
        Wyy(:,iTime) = conv(Wyy(:,iTime),FIL,'same');
        sWxy(:,iTime) = conv(Wxy(:,iTime),FIL,'same');
    end
    
% Smoothing in time direction
    for iScale=1:nScales                                         
        wTSS = fix(periods(iScale)*wT_size/dt); % Size of window is 
                                                % adapted to scale
        if (ischar(wT_type)&& strcmpi(wT_type,'Bar'))...
            ||(isnumeric(wT_type)&& wT_type==0)
            FIL=WindowBartlett(wTSS);
        else
            FIL = window(wT_type,wTSS);
        end
        FIL = FIL./(sum(FIL));
        Wxx(iScale,:)=conv(Wxx(iScale,:),FIL,'same');
        Wyy(iScale,:)=conv(Wyy(iScale,:),FIL,'same');       
        sWxy(iScale,:)=conv(sWxy(iScale,:),FIL,'same');
    end
   
%%%%%%%%%%%     Computation of (Complex) Wavelet Coherency     %%%%%%%%%%%%
WCO = sWxy./ sqrt(Wxx .* Wyy); % Formula (25) of [1]
WCO(Wxx<1e-10)=0;
WCO(Wyy<1e-10)=0; % To avoid meaningless numbers when Wxx or Wyy are 
                  % nearly zero
%WCO = abs(WCO);  % Uncomment if absolute value of coherency is desired 
% sWxyR = real(sWxy); 
WxyR = real(Wxy); 
% WCoh = sWxyR ./ sqrt(Wxx .* Wyy);   
WCoh = WxyR ./ sqrt(Wxx .* Wyy);   
end
%%%%%%%%%%%%%%%% END OF SUBFUNCTION WAVELETCOHERENCY %%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                 SUBFUNCTION WINDOWBARTLETT                      %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
function w = WindowBartlett(n)
	if (~(n == round(n)) || n<=0)
        error('Size of window must be a positive integer')
	end
    w = 2*(0:(n-1)/2)/(n-1);
    if rem(n,2)
        w = [w w((n-1)/2:-1:1)]';
    else
        w = [w w(n/2:-1:1)]';
    end
end
%%%%%%%%%%%%%%%%%%%% END OF SUBFUNCTION WINDOWBARTLETT %%%%%%%%%%%%%%%%%%%% 
end
