function idx = enso(sst,t,varargin)
% enso calculates the El Nino Southern Oscillation index from sea surface
% temperatures based on the ENSO definition proposed by Trenberth 1997. 
% 
%% Syntax 
% 
%  idx = enso(sst,t) 
%  idx = enso(sst,t,lat,lon)
%  idx = enso(sst,t,lat,lon,'region',NinoRegion)
%  idx = enso(sst,t,mask)
%  idx = enso(...,'smoothing',months)   
% 
%% Description 
% 
% idx = enso(sst,t) calculates El Nino Southern Oscillation index from 
% a time series of sea surface temperatures sst and their corresponding 
% times t. sst can be a vector of sea surface temperatures that have been
% averaged over a region of interest, or sst can be a 3D matrix whose third
% dimension correponds to times t. If sst is a 3D matrix, a time series is
% automatically generated by averaging all the grid cells in sst for each
% time step. 
% 
% idx = enso(sst,t,lat,lon) calculates the Nino 3.4 index for 3D sst time
% series and corresponding grid coordinates lat,lon. Using this syntax,
% grid cells within the Nino 3.4 region are automatically determined and
% the Nino index is calculated from the area-averaged time series of ssts
% within that region. 
% 
% idx = enso(sst,t,lat,lon,'region',NinoRegion) allows any of the following
% Nino regions, entered as string (in 'single quotes') 
%  * '1+2'
%  * '3'
%  * '3.4' 
%  * '4'
%  * 'ONI' 
% 
% idx = enso(sst,t,mask) calculates the sst index using the unweighted mean
% of sst grid cells corresponding to true values in a 2D logical mask. 
% 
% idx = enso(...,'smoothing',months) defines the moving average window
% in months. Default value is 5, following Trenberth 1997. The only exception
% is for the ONI region, which has a smoothing window of 3 months. To turn off
% averaging, set 'smoothing',false. 
% 
%% Examples
% For examples and a description of methods, type 
% 
%   cdt enso
%
%% Reference
% 
% Trenberth, Kevin E. "The Definition of El Nino." Bulletin of the  
% American Meteorological Society 78.12 (1997): 2771-2778. 
% https://doi.org/10.1175/1520-0477(1997)078%3C2771:TDOENO%3E2.0.CO;2
% 
%% Author Info
% This function was written by Kaustubh Thirumalai of the University of 
% Arizona, January 2019.
% http://www.kaustubh.info
% 
% See also: 

%% Initial error checks: 

narginchk(2,Inf) 
assert(ismember(length(t),size(sst))==1,'Error: length of t must match dimensions of sst.') 

% Try and check for Nino3.4 input
if mean(sst,'omitnan')<20 
    warning('Error: Nino3.4 SSTs cooler than 20 C? Something seems awry.');
end

%% Set defaults: 

window_months = 5; % moving window months
AutoMask = false; % doesn't try to determine which grid cells to analyze by default. 
NinoRegion = '3.4'; 
mask = true(size(sst,1),size(sst,2)); 

%% Parse inputs: 

% Convert time from datetime, datestr, or datevec:
t = datenum(t); % (If it's already datenum, nothing changes.)

% Define the Nino Region:
tmp = strncmpi(varargin,'region',3); 
if any(tmp)
   NinoRegion = varargin{find(tmp)+1}; 
   AutoMask = true; 
   assert(~isnumeric(NinoRegion),'Error: Nino Region must be defined by string, not a number. (This is to allow such regions as ''1+2''.)')
   assert(ismember(lower(NinoRegion),{'1+2','3','3.4','4','oni'}),'Error: Unrecognized Nino region ',NinoRegion'.')
end

% Note: Smoothing preference is checked after defining Nino regions. 

% Check for user-defined grids and masks: 
if nargin>2
   if islogical(varargin{1})
      mask = varargin{1}; 
      assert(isequal([size(sst,1) size(sst,2)],size(mask)),'Error: dimensions of sst and the mask do not agree.') 
   end
   
   if nargin>3
      if islatlon(varargin{1},varargin{2})
         Lat = varargin{1}; 
         Lon = varargin{2}; 
         AutoMask = true; 
         assert(~isvector(Lat),'Error: lat,lon grids must be matrices as if generated by meshgrid or cdtgrid.')
         assert(isequal(size(Lat),size(Lon),[size(sst,1) size(sst,2)]),'Error: Dimensions of Lat, Lon, and the sst grid must all agree.') 
      end
   end
   
end

%% Define the mask 

switch lower(NinoRegion)
   case '1+2'
      latrange = [-10 0]; 
      lonrange = [-90 -80]; 
      
   case '3'
      latrange = [-5 5]; 
      lonrange = [-150 -90]; 
      
   case '3.4'
      latrange = [-5 5]; 
      lonrange = [-170 -120]; 
      
   case '4'
      latrange = [-5 5]; 
      lonrange = [-160 -150]; 
      
   case 'oni'
      latrange = [-5 5]; 
      lonrange = [-170 -120]; 
      window_months = 3; % The only difference from standard Nino 3.4
      
   otherwise
      error('This should have been caught earlier, so I''m really not sure how we got here, but the Nino region is not recognized.') 
end

% Check smoothing doing this now in case oni (3 month) is chosen but not explicitly defined: 
tmp = strncmpi(varargin,'smoothing',3);
if any(tmp) 
   window_months = varargin{find(tmp)+1}; 
   assert(isscalar(window_months),'Error: Moving average window must be a scalar value in months.') 
end

if AutoMask   
   assert(exist('Lat','var')==1,'Automatic region masking requires defining Lat and Lon as the third and fourth inputs to the enso function.') 
   mask = geomask(Lat,Lon,latrange,lonrange);
end

%% Get time series

if ndims(sst)>2
   if AutoMask
      % Make it area-weighted: 
      A = cdtarea(Lat,Lon); 
      sst = local(sst,mask,'weight',A); 
   else
      sst = local(sst,mask); 
   end
end

%% Calculate ENSO index: 

% The pre-moving-averaged index is just the sst anomaly: 
idx = deseason(sst,t,'detrend','linear') - mean(sst,'omitnan');

% Moving average: 
if window_months % window_months might be a logical "false" or it might be a number of months.
   
   % Define the moving average radius in days, from user-defined moving window width (diameter):
   rad_days = window_months*365.25/24; 
   
   % Calculate the moving average using scatstat1:
   idx = scatstat1(t,idx,rad_days,@mean,'omitnan'); 
end

end

