import xarray as xr
import numpy as np
from scipy import stats
from scipy.signal import butter, lfilter, filtfilt, freqz

#######TIME FILTERING FUNCTIONS#################

def butter_highpass(cutoff, fs, order=5):
    nyq = 0.5 * fs
    normal_cutoff = cutoff / nyq
    b, a = butter(order, normal_cutoff, btype='high', analog=False)
    return b, a

def butter_highpass_filter(data, cutoff, fs, order=5):
    b, a = butter_highpass(cutoff, fs, order=order)
    y = filtfilt(b, a, data, axis=0)
    #y = data.copy(data=y,deep=False) 
    data[:] = y
    return data

# Filter requirements.
order = 5
fs = 1.0      # sample rate, 1/d

# desired cutoff frequency of the filter, 1/d
highcut = 1./8.


def calculate_regression(selected_data, index_nao, j_dim, i_dim):
    regression = np.zeros((j_dim, i_dim))
    p_value = np.zeros((j_dim, i_dim))
    for j in range(j_dim):
        for i in range(i_dim):
            data = selected_data[:, j,i]
            reg = scipy.stats.linregress(index_nao,data) 
            regression[j, i] = reg[0] #slope
            p_value[j, i] = reg[3] #p_value
    return regression, p_value


def climatology_djf(path, start_year, end_year, var_name, a, b):
    
    var = xr.open_dataset(path + str(start_year) + '12' + '-' + str(start_year) + '12.nc')
    x = len(var.j.values)
    y = len(var.i.values)

    index_year = 0
    
    climatology_values = np.zeros(((end_year - start_year - 1), x, y))
    for i in np.arange(start_year, end_year-1, 1):
        var = np.empty((3, x, y))
        var[0] = xr.open_dataset(path + str(i) + '12' + '-' + str(i) + '12.nc')[var_name].values[0]
        for month in (1, 2):
            var[month] = xr.open_dataset(path + str(i + 1) + '0' + str(month) + '-' + str(i + 1) + '0' + str(month) + '.nc')[var_name].values[0]
        climatology_values[index_year, :, :] = np.mean(var, axis = 0) * a + b
        index_year += 1
        
    var_mean = np.mean(climatology_values, axis=0)

    return var_mean, climatology_values


def climatology_jja(path, start_year, end_year, var_name, a, b):
    
    var = xr.open_dataset(path + str(start_year) + '12' + '-' + str(start_year) + '12.nc')
    x = len(var.j.values)
    y = len(var.i.values)

    index_year = 0
    climatology_values = np.zeros(((end_year - start_year), x, y))
    for i in np.arange(start_year, end_year, 1):
        var = np.empty((3, x, y))
        for month in (6,7,8):
            var[month-6] = xr.open_dataset(path + str(i) + '0' + str(month) + '-' + str(i) + '0' + str(month) + '.nc')[var_name].values[0]
        climatology_values[index_year, :, :] = np.mean(var, axis = 0) * a + b
        index_year += 1
        print(i)
    
    var_mean = np.mean(var_climatology_values, axis=0)

    return var_mean, climatology_values


def climatology_vert(path, start_year, end_year, var_name, a, b, pr_lev):
    
    var = xr.open_dataset(path + str(start_year) + '12' + '-' + str(start_year) + '12.nc').sel(plev = pr_lev)
    x = len(var.j.values)
    y = len(var.i.values)

    index_year = 0
    climatology_values = np.zeros((end_year - start_year - 1, x, y))
    for i in np.arange(start_year, end_year-1, 1):
        var = np.zeros((3,x,y))
        var[0] = xr.open_dataset(path + str(i) + '12' + '-' + str(i) + '12.nc').sel(plev = pr_lev)[var_name].values[0]
        for month in (1, 2):
            var[month] = xr.open_dataset(path + str(i + 1) + '0' + str(month) + '-' + str(i + 1) + '0' + str(month) + '.nc').sel(plev = pr_lev)[var_name].values[0]
        climatology_values[index_year, :, :] = np.mean(var, axis = 0) * a + b 
        index_year += 1
        print(i)

    var_mean = np.mean(climatology_values, axis=0)

    return var_mean, climatology_values


def find_indices(lat_range, lon_range, lat, lon):
    lat_min=lat_range[0]
    lat_max=lat_range[1]
    lon_min=lon_range[0]
    lon_max=lon_range[1]
    indices = np.argwhere(
        (lat >= lat_min) & (lat <= lat_max) & (
            ((lon >= lon_min) & (lon <= 360)) |
            ((lon >= 0) & (lon <= lon_max))
        )
    )
    
    j_max = np.max(indices[(lon[indices[:, 0], indices[:, 1]] >= 0) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 0])
    j_min = np.min(indices[(lon[indices[:, 0], indices[:, 1]] >= 0) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 0])
    i_max_r = np.max(indices[(lon[indices[:, 0], indices[:, 1]] >= 0) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 1])
    i_min_r = np.min(indices[(lon[indices[:, 0], indices[:, 1]] >= 0) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 1])
    i_max_l = np.max(indices[(lon[indices[:, 0], indices[:, 1]] >= lon_min) & (lon[indices[:, 0], indices[:, 1]] <= 360)][:, 1])
    i_min_l = np.min(indices[(lon[indices[:, 0], indices[:, 1]] >= lon_min) & (lon[indices[:, 0], indices[:, 1]] <= 360)][:, 1])

    return j_max, j_min, i_max_r, i_max_l, i_min_r, i_min_l


def find_indices1(lat_range, lon_range, lat, lon):
    lat_min=lat_range[0]
    lat_max=lat_range[1]
    lon_min=lon_range[0]
    lon_max=lon_range[1]
    indices = np.argwhere(
        (lat >= lat_min) & (lat <= lat_max) & 
            (lon >= lon_min) & (lon <= lon_max)
    )
    
    j_max = np.max(indices[(lat[indices[:, 0], indices[:, 1]] >= lat_min) & (lat[indices[:, 0], indices[:, 1]] <= lat_max)][:, 0])
    j_min = np.min(indices[(lat[indices[:, 0], indices[:, 1]] >= lat_min) & (lat[indices[:, 0], indices[:, 1]] <= lat_max)][:, 0])
    i_max = np.max(indices[(lon[indices[:, 0], indices[:, 1]] >= lon_min) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 1])
    i_min = np.min(indices[(lon[indices[:, 0], indices[:, 1]] >= lon_min) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 1])
    
    return j_max, j_min, i_max, i_min


def meridional_gradient(var, lat):
    gradient = np.zeros(np.shape(var))
    j_max = len(lat[:,0]) - 1
    R = 6371
    for t in range(len(var[:,0,0])):
        for i in range(len(lat[0,:])):
            gradient[t,0,i] = (var[t,1,i] - var[t,0,i])/(R*(np.radians(lat[1,i]) - np.radians(lat[0,i])))
            gradient[t,j_max,i] = (var[t,j_max,i] - var[t,j_max - 1,i])/(R*(np.radians(lat[j_max,i]) - np.radians(lat[j_max - 1,i])))
            for j in range(len(lat[:,0])-2):
                gradient[t,j+1,i] = (var[t,j+2,i] - var[t,j,i])/(R*(np.radians(lat[j+2,i]) - np.radians(lat[j,i])))
        print(str(t + 1) + '/' + str(len(var[:,0,0])))
    gradient_mean = np.mean(gradient, axis = 0)
    return gradient_mean, gradient


def nao(path, start_year, end_year, j_max, j_min, i_max_r, i_max_l, i_min_r, i_min_l):
    
    SLP = []
    
    index_nao1 = []

    for i in np.arange(start_year,end_year - 1,1):
        season_var = []
        var = xr.open_dataset(path+str(i)+'12'+'-'+str(i)+'12.nc')
        var1 = var.sel(j=slice(j_min + 1, j_max + 1), i=slice(i_min_l + 1, i_max_l + 1))
        var2 = var.sel(j=slice(j_min + 1, j_max + 1), i=slice(i_min_r + 1, i_max_r + 1))
        var = xr.concat((var1,var2),dim = 'i')
        season_var.append(var)
        for month in (1,2):
            var = xr.open_dataset(path+str(i+1)+'0'+str(month)+'-'+str(i+1)+'0'+str(month)+'.nc')
            var1 = var.sel(j=slice(j_min + 1, j_max + 1), i=slice(i_min_l + 1, i_max_l + 1))
            var2 = var.sel(j=slice(j_min + 1, j_max + 1), i=slice(i_min_r + 1, i_max_r + 1))
            var = xr.concat((var1,var2),dim = 'i')
            season_var.append(var)
        season_combined_data = xr.concat(season_var, dim = 'time')
        season_mean = season_combined_data.mean(dim='time')
        index_nao1.append((season_mean.psl[j_Azo,i_Azo]/100) - (season_mean.psl[j_Rey,i_Rey]/100))
        SLP.append(season_mean/100)
        print(i)

    SLP = xr.concat(SLP, dim = 'time')

    index_nao1 = xr.concat(index_nao1, dim = 'time')
    index_nao1 = (index_nao1 - index_nao1.mean())/index_nao1.std()
    
    solver = Eof(SLP.psl)
    n=3
    
    eofs = solver.eofsAsCovariance(neofs=n)
    pcs = solver.pcs(pcscaling=1, npcs=n)
    
    index_nao2 = pcs.sel(mode=0)*-1

    if np.corrcoef(index_nao2,index_nao1)[0,1] > 0:
        index_nao2 = index_nao2
    else:
        index_nao2 = -index_nao2
        
    return index_nao1, index_nao2, eofs, solver


def precip_distr(path, start_year, end_year, season, var_name, j_max, j_min, i_max, i_min):
    
    dist = []
    
    if season == 'djf':
        for i in np.arange(start_year, end_year, 1):
            print(i)
            #december
            var = xr.open_dataset(path+str(i)+'12010000'+'-'+str(i)+'12311800.nc')[var_name]
            for day in range(31):
                for t in range(4):
                    var_area = var.sel(time = (var['time.day'] == day + 1)).values[t,j_min:j_max + 1,i_min:i_max + 1]
                    dist.append(np.mean(var_area, axis = (0,1)) * 3600 * 6)
            print('december')

            #january 
            var = xr.open_dataset(path+str(i)+'01010000'+'-'+str(i)+'01311800.nc')[var_name]
            for day in range(31):
                for t in range(4):
                    var_area = var.sel(time = (var['time.day'] == day + 1)).values[t,j_min:j_max + 1,i_min:i_max + 1]
                    dist.append(np.mean(var_area, axis = (0,1)) * 3600 * 6)
            print('january')
            #february
            try:
                var = xr.open_dataset(path+str(i)+'02010000'+'-'+str(i)+'02281800.nc')[var_name]
                days_number = 28
            except:
                var = xr.open_dataset(path+str(i)+'02010000'+'-'+str(i)+'02291800.nc')[var_name]
                days_number = 29
            for day in range(days_number):
                for t in range(4):
                    var_area = var.sel(time = (var['time.day'] == day + 1)).values[t,j_min:j_max + 1,i_min:i_max + 1]
                    dist.append(np.mean(var_area, axis = (0,1)) * 3600 * 6)
            print('february')
            
    elif season == 'jja':
        for i in np.arange(start_year, end_year, 1):
            print(i)
            #june
            var = xr.open_dataset(path+str(i)+'06010000'+'-'+str(i)+'06301800.nc')[var_name]
            for day in range(30):
                for t in range(4):
                    var_area = var.sel(time = (var['time.day'] == day + 1)).values[t,j_min:j_max + 1,i_min:i_max + 1]
                    dist.append(np.mean(var_area, axis = (0,1)) * 3600 * 6)
            print('june')
            #july
            var = xr.open_dataset(path+str(i)+'07010000'+'-'+str(i)+'07311800.nc')[var_name]
            for day in range(31):
                for t in range(4):
                    var_area = var.sel(time = (var['time.day'] == day + 1)).values[t,j_min:j_max + 1,i_min:i_max + 1]
                    dist.append(np.mean(var_area, axis = (0,1)) * 3600 * 6)
            print('july')
            #august
            var = xr.open_dataset(path+str(i)+'08010000'+'-'+str(i)+'08311800.nc')[var_name]
            for day in range(31):
                for t in range(4):
                    var_area = var.sel(time = (var['time.day'] == day + 1)).values[t,j_min:j_max + 1,i_min:i_max + 1]
                    dist.append(np.mean(var_area, axis = (0,1)) * 3600 * 6)
            print('august')
    else:
            print('wrong season')
            
    return dist


def select_area(lat_range, lon_range, lat, lon, var):
    lat_min=lat_range[0]
    lat_max=lat_range[1]
    lon_min=lon_range[0]
    lon_max=lon_range[1]
    indices = np.argwhere(
        (lat >= lat_min) & (lat <= lat_max) & (
            ((lon >= lon_min) & (lon <= 360)) |
            ((lon >= 0) & (lon <= lon_max))
        )
    )
    
    j_max = np.max(indices[(lon[indices[:, 0], indices[:, 1]] >= 0) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 0])
    j_min = np.min(indices[(lon[indices[:, 0], indices[:, 1]] >= 0) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 0])
    i_max_r = np.max(indices[(lon[indices[:, 0], indices[:, 1]] >= 0) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 1])
    i_min_r = np.min(indices[(lon[indices[:, 0], indices[:, 1]] >= 0) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 1])
    i_max_l = np.max(indices[(lon[indices[:, 0], indices[:, 1]] >= lon_min) & (lon[indices[:, 0], indices[:, 1]] <= 360)][:, 1])
    i_min_l = np.min(indices[(lon[indices[:, 0], indices[:, 1]] >= lon_min) & (lon[indices[:, 0], indices[:, 1]] <= 360)][:, 1])
    
    var = np.concatenate((var[j_min:j_max + 1,i_min_l:i_max_l + 1], var[j_min:j_max + 1,i_min_r:i_max_r + 1]), axis = 1)
    
    return var


def select_area1(lat_range, lon_range, lat, lon, var):
    lat_min=lat_range[0]
    lat_max=lat_range[1]
    lon_min=lon_range[0]
    lon_max=lon_range[1]
    indices = np.argwhere(
        (lat >= lat_min) & (lat <= lat_max) & 
            (lon >= lon_min) & (lon <= lon_max)
    )
    
    j_max = np.max(indices[(lat[indices[:, 0], indices[:, 1]] >= lat_min) & (lat[indices[:, 0], indices[:, 1]] <= lat_max)][:, 0])
    j_min = np.min(indices[(lat[indices[:, 0], indices[:, 1]] >= lat_min) & (lat[indices[:, 0], indices[:, 1]] <= lat_max)][:, 0])
    i_max = np.max(indices[(lon[indices[:, 0], indices[:, 1]] >= lon_min) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 1])
    i_min = np.min(indices[(lon[indices[:, 0], indices[:, 1]] >= lon_min) & (lon[indices[:, 0], indices[:, 1]] <= lon_max)][:, 1])

    var = var[j_min:j_max + 1,i_min:i_max + 1]
    
    return var
    

def temperature_distr(path, start_year, end_year, season, var_name, j_max, j_min, i_max, i_min, weights):
    
    dist = []
    
    if season == 'djf':
        for i in np.arange(start_year, end_year, 1):
            print(i)
            #december
            var = xr.open_dataset(path+str(i)+'1201'+'-'+str(i)+'1231.nc')[var_name]
            for day in range(31):
                var_area = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min:i_max + 1]
                var_area = var_area * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('december')

            #january 
            var = xr.open_dataset(path+str(i)+'0101'+'-'+str(i)+'0131.nc')[var_name]
            for day in range(31):
                var_area = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min:i_max + 1]
                var_area = var_area * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('january')
            #february
            try:
                var = xr.open_dataset(path+str(i)+'0201'+'-'+str(i)+'0228.nc')[var_name]
                days_number = 28
            except:
                var = xr.open_dataset(path+str(i)+'0201'+'-'+str(i)+'0229.nc')[var_name]
                days_number = 29
            for day in range(days_number):
                var_area = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min:i_max + 1]
                var_area = var_area * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('february')
            
    elif season == 'jja':
        for i in np.arange(start_year, end_year, 1):
            print(i)
            #june
            var = xr.open_dataset(path+str(i)+'0601'+'-'+str(i)+'0630.nc')[var_name]
            for day in range(30):
                var_area = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min:i_max + 1]
                var_area = var_area * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('june')
            #july
            var = xr.open_dataset(path+str(i)+'0701'+'-'+str(i)+'0731.nc')[var_name]
            for day in range(31):
                var_area = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min:i_max + 1]
                var_area = var_area * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('july')
            #august
            var = xr.open_dataset(path+str(i)+'0801'+'-'+str(i)+'0831.nc')[var_name]
            for day in range(31):
                var_area = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min:i_max + 1]
                var_area = var_area * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('august')
    else:
            print('wrong season')
            
    return dist


def temperature_distr_F(path, start_year, end_year, season, var_name, j_max, j_min, i_max_r, i_max_l, i_min_r, i_min_l, weights):
    
    dist = []
    
    if season == 'djf':
        for i in np.arange(start_year, end_year, 1):
            print(i)
            #december
            var = xr.open_dataset(path+str(i)+'1201'+'-'+str(i)+'1231.nc')[var_name]
            for day in range(31):
                var_area_l = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_l:i_max_l + 1]
                var_area_r = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_r:i_max_r + 1]
                var_area = np.concatenate((var_area_l, var_area_r), axis = 1) * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('december')

            #january 
            var = xr.open_dataset(path+str(i)+'0101'+'-'+str(i)+'0131.nc')[var_name]
            for day in range(31):
                var_area_l = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_l:i_max_l + 1]
                var_area_r = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_r:i_max_r + 1]
                var_area = np.concatenate((var_area_l, var_area_r), axis = 1) * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('january')
            #february
            try:
                var = xr.open_dataset(path+str(i)+'0201'+'-'+str(i)+'0228.nc')[var_name]
                days_number = 28
            except:
                var = xr.open_dataset(path+str(i)+'0201'+'-'+str(i)+'0229.nc')[var_name]
                days_number = 29
            for day in range(days_number):
                var_area_l = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_l:i_max_l + 1]
                var_area_r = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_r:i_max_r + 1]
                var_area = np.concatenate((var_area_l, var_area_r), axis = 1) * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('february')
            
    elif season == 'jja':
        for i in np.arange(start_year, end_year, 1):
            print(i)
            #june
            var = xr.open_dataset(path+str(i)+'0601'+'-'+str(i)+'0630.nc')[var_name]
            for day in range(30):
                var_area_l = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_l:i_max_l + 1]
                var_area_r = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_r:i_max_r + 1]
                var_area = np.concatenate((var_area_l, var_area_r), axis = 1) * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('june')
            #july
            var = xr.open_dataset(path+str(i)+'0701'+'-'+str(i)+'0731.nc')[var_name]
            for day in range(31):
                var_area_l = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_l:i_max_l + 1]
                var_area_r = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_r:i_max_r + 1]
                var_area = np.concatenate((var_area_l, var_area_r), axis = 1) * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('july')
            #august
            var = xr.open_dataset(path+str(i)+'0801'+'-'+str(i)+'0831.nc')[var_name]
            for day in range(31):
                var_area_l = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_l:i_max_l + 1]
                var_area_r = var.sel(time = (var['time.day'] == day + 1)).values[0,j_min:j_max + 1,i_min_r:i_max_r + 1]
                var_area = np.concatenate((var_area_l, var_area_r), axis = 1) * weights
                dist.append(np.sum(var_area)/np.sum(weights) - 273.15)
            print('august')
    else:
            print('wrong season')
            
    return dist


def variance(path, start_year, end_year, var_name, lat_range, lon_range, lat, lon, lat_NA, a, b, plev):
    
    var_values = []

    index_year = 0
    index_day = 0
    var_values = np.zeros((1805, len(lat_NA[:,0]), len(lat_NA[0,:])))
    
    if (var_name == 'ua' or var_name == 'va'):
        for i in np.arange(start_year,end_year,1):
            start = index_day
            #december
            var = xr.open_dataset(path+str(i)+'1201'+'-'+str(i)+'1231.nc')[var_name].sel(plev = plev)
            for day in range(31):
                var_range = select_area(lat_range, lon_range, lat, lon,var.sel(time = (var['time.day'] == day + 1)).values[0,:,:])
                var_values[index_day,:,:] = var_range
                index_day+=1

            #january
            var = xr.open_dataset(path+str(i)+'0101'+'-'+str(i)+'0131.nc')[var_name].sel(plev = plev)
            for day in range(31):
                var_range = select_area(lat_range, lon_range, lat, lon,var.sel(time = (var['time.day'] == day + 1)).values[0,:,:])
                var_values[index_day,:,:] = var_range
                index_day+=1

            #february
            try:
                var = xr.open_dataset(path+str(i)+'0201'+'-'+str(i)+'0228.nc')[var_name].sel(plev = plev)
                days_number = 28
            except:
                var = xr.open_dataset(path+str(i)+'0201'+'-'+str(i)+'0229.nc')[var_name].sel(plev = plev)
                days_number = 29
            for day in range(days_number):
                var_range = select_area(lat_range, lon_range, lat, lon,var.sel(time = (var['time.day'] == day + 1)).values[0,:,:])
                var_values[index_day,:,:] = var_range
                index_day+=1

            print(i)
            
            end = index_day
            
            var_values[start:end,:,:] = butter_highpass_filter(var_values[start:end,:,:], highcut, fs, order)
            print('start, end: ', start, end)
            
    else:
        for i in np.arange(start_year,end_year,1):
            start = index_day
            #december
            var = xr.open_dataset(path+str(i)+'1201'+'-'+str(i)+'1231.nc')[var_name]
            for day in range(31):
                var_range = select_area(lat_range, lon_range, lat, lon,var.sel(time = (var['time.day'] == day + 1)).values[0,:,:]*a +b)
                var_values[index_day,:,:] = var_range
                index_day += 1

            #january
            var = xr.open_dataset(path+str(i)+'0101'+'-'+str(i)+'0131.nc')[var_name]
            for day in range(31):
                var_range = select_area(lat_range, lon_range, lat, lon,var.sel(time = (var['time.day'] == day + 1)).values[0,:,:]*a +b)
                var_values[index_day,:,:] = var_range
                index_day+=1

            #february
            try:
                var = xr.open_dataset(path+str(i)+'0201'+'-'+str(i)+'0228.nc')[var_name]
                days_number = 28
            except:
                var = xr.open_dataset(path+str(i)+'0201'+'-'+str(i)+'0229.nc')[var_name]
                days_number = 29
            for day in range(days_number):
                var_range = select_area(lat_range, lon_range, lat, lon,var.sel(time = (var['time.day'] == day + 1)).values[0,:,:]*a +b)
                var_values[index_day,:,:] = var_range
                index_day+=1

            print(i)
            
            end = index_day
            
            var_values[start:end,:,:] = butter_highpass_filter(var_values[start:end,:,:], highcut, fs, order)
            print('start, end: ', start, end)

    var_mean = np.mean(var_values, axis = 0) 
    var_variance = np.zeros(np.shape(lat_NA))
    for day in range(len(var_values)):
        var_variance += (var_values[day,:,:] - var_mean) ** 2
    var_variance = var_variance/(len(var_values)-1)
    
    return var_variance


def zonal_mean_djf(path, start_year, end_year, var_name, a, b):
    
    var = xr.open_dataset(path + str(start_year) + '12' + '-' + str(start_year) + '12.nc')

    index_year = 0
    var_zonal = np.empty(((end_year - start_year - 1)*3, len(var.plev.values), len(var.j.values)))
    for i in np.arange(start_year, end_year-1, 1):
        var = xr.open_dataset(path+str(i)+'12'+'-'+str(i)+'12.nc')[var_name]
        var_zonal[index_year*3, :, :] = var.mean(dim='i').values[0,:,:] * a + b
        for month in (1, 2):
            var = xr.open_dataset(path+str(i+1)+'0'+str(month)+'-'+str(i+1)+'0'+str(month)+'.nc')[var_name]
            var_zonal[index_year*3, :, :] = var.mean(dim='i').values[0,:,:] * a + b
        index_year += 1
        print(i)

    var_mean = np.mean(var_zonal, axis=0)

    return var_mean, var_zonal


def zonal_mean_jja(path, start_year, end_year, var_name, a, b):
    
    var = xr.open_dataset(path + str(start_year) + '12' + '-' + str(start_year) + '12.nc')

    index_year = 0
    var_zonal = np.empty((end_year - start_year, len(var.plev.values), len(var.j.values)))
    for i in np.arange(start_year, end_year, 1):
        for month in (6,7,8):
            var = xr.open_dataset(path+str(i)+'0'+str(month)+'-'+str(i)+'0'+str(month)+'.nc')
            var_zonal[index_year*3 + month, :, :] = var.mean(dim='i').values[0,:,:] * a + b
        index_year += 1
        print(i)
    var_mean = np.mean(var_zonal, axis=0)

    return var_mean, var_zonal