# Compute Nino 3.4 index (Nino 3.4 SST anomaly) from OSTIA
#
# Period for climatology: 1993-2016 (Copernicus)
#
#
# Script version 1.0
#
# Script works with:
#   Python version 3.10
#
#   Package version
#     numpy:  
#     pandas: 
#
#
#
# Ver1.0: Created by Mariko Koseki, 03.07.2025




##--- Import modules ------------------------------##
import netCDF4
import numpy as np
import pandas as pd
import math
import datetime
from datetime import timedelta
from dateutil.relativedelta import relativedelta
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.ticker as ptick
import os
import glob
import sys
import platform
import calendar



print('--- Python version ---')
print(platform.python_version())

print('--- Package version ---')
print('numpy: ', np.__version__)
print('pandas: ', pd.__version__)
print('-------------------------')
print('')



##--- Functions -----------------------------------------##
def nino34_sst(sst, lats, lons):
    '''
    # Compute area average of SST in the Nino 3.4 region
    # Nino 3.4 region: 5~-5 (5N-5S), 190~240 (120W-170W)
    '''
    ## Check index of Nino3.4 area ##
    lat_nino34_start_ind_nor = np.where((lats >= -5) & (lats <= 5))[0][0]
    lat_nino34_end_ind_nor = np.where((lats >= -5) & (lats <= 5))[0][-1]
    #print(lat_nino34_start_ind_nor, lat_nino34_end_ind_nor)

    lon_nino34_start_ind_nor = np.where((lons >= 190) & (lons <= 240))[0][0]
    lon_nino34_end_ind_nor = np.where((lons >= 190) & (lons <= 240))[0][-1]
    #print(lon_nino34_start_ind_nor, lon_nino34_end_ind_nor) 

    ## Extract Nino3.4 area ##
    nino_sst = sst[lat_nino34_start_ind_nor:lat_nino34_end_ind_nor + 1, lon_nino34_start_ind_nor:lon_nino34_end_ind_nor +1]
    #print(nino_sst.shape)


    ## Compute Nino 3.4 area average of SST ##
    nino34_sst = np.nanmean(np.nanmean(nino_sst, 0), 0)
    #print('')
    #print('nino3.4 SST: ', nino34_sst)
    #print('')
    return nino34_sst




##--- Set year, month for climatology ------------------------------------##
'''
Set 'start_year_clm' and 'end_year_clm'
Copernicus: 1993-2016
'''
start_year_clm = 1993 # 1993
end_year_clm = 2016 # 2016




##--- Set path ------------------------------------------##
## Set path to input files ##
'''
Set path to input NetCDF files
'''
ostia_monthly = '/nird/projects/NS9873K/norcpm/validation/observation/MET/OSTIA/monthly/'
ostia_clim = '/nird/projects/NS9873K/norcpm/validation/observation/MET/OSTIA/clim/OSTIA_sst_clim_' + str(start_year_clm) + '_' + str(end_year_clm) + '.nc'



## Set path to output files ##
'''
Set path to output file
'''
out_dir  = '/nird/projects/NS9873K/norcpm/validation/enso/OSTIA/'
if not os.path.exists(out_dir):
    os.makedirs(out_dir)





##--- Read netCDF files ----------------------------------##
#################
## Climatology ##
#################
print('climatology of OSTIA: ', ostia_clim)
print('')

nc_ostia_clim = netCDF4.Dataset(ostia_clim, 'r')


## Check length of variables ##
nlon_clm = len(nc_ostia_clim.dimensions['lon']) # Longitude: 360
nlat_clm = len(nc_ostia_clim.dimensions['lat']) # Latitude: 180
nmonth_clm = len(nc_ostia_clim.dimensions['month']) # month: 12
#print(nlon_clm, nlat_clm, nmonth_clm)


## Read variables ##
### month ###
months_clm = nc_ostia_clim.variables['month'][:] # month: 1-12 
#print(months_clm)


### latitude, longitude ###
lons_clm = nc_ostia_clim.variables['lon'][:] # Lon: 0.5 - 359.5
lats_clm = nc_ostia_clim.variables['lat'][:] # Lat: 89.5 - -89.5
#print(lons_clm, lats_clm)


### SST ###
ostia_clim = nc_ostia_clim.variables['clim_ostia'][:][:][:]


## Convert variables into Numpy ndarray ##
ostia_clim_np = np.array(ostia_clim)



### Call functions ###
'''
# Compute area average of SST in the Nino 3.4 region
'''
mm1 = 0; mm2 = 12
mm_int_list = []
for mm in range(mm1, mm2):
    mm_int_list.append(mm)

#print(mm_int_list)

nino34_clm_list = []
for mon in range(0, nmonths_clim):
    nino34_ostia_clm = nino34_sst(ostia_clim_np[mon, :, :], lats_clm, lons_clm)
    
    print('climatology of OSTIA: ', nino34_ostia_clm)
    nino34_clm_list.append(nino34_ostia_clm)
nino34_clm_all = np.array(nino34_clm_list).T

print('')
print('all climatologies of OSTIA: ', nino34_clm_all)
print('')


### Close netCDF file ###
nc_ostia_clim.close()



###########
## OSTIA ##
###########
### Create list of year-month ###
start_year = 1982
end_year = 2025


y_str_list = ["{}".format(y) for y in range(start_year,end_year+1)]
m_str_list = ["{:02}".format(m) for m in range(1,13)]

'''
start_date = '{:0>4d}-{:0>2d}'.format(start_year,start_month)
end_date = datetime.datetime(end_year,end_month,1)
end_date = end_date + relativedelta(months=1)
end_date = end_date.strftime('%Y-%m')
ym_str_list = pd.date_range(start_date, end_date, freq='ME').strftime('%Y-%m').tolist()
'''

print(y_str_list, m_str_list)



## Set length of variables ##
nyear = len(y_str_list)
nmonth = len(m_str_list)



##--- Set year, month ------------------------------------##
#for ym in ym_str_list:
    #year_int = int(ym[0:4])
    #month_int = int(ym[5:])

nino34idx_all = []
for year_num in range(0, nyear):
    yyyy = y_str_list[year_num]
    print('year: ', yyyy)


    nino34_yearly = []
    for mon_num in range(0, nmonth):
        mm = m_str_list[mon_num]
        print('month: ', mm)



        file_ostia = ostia_monthly + yyyy + '_' +  mm + '.nc'
        print('file name: ', file_ostia)



        if not os.path.exists(file_ostia):
            print('no such a directory!!!')

            nino34_ostia = np.nan
            nino34_yearly.append(nino34_ostia)

            continue

        else:
            print('Data exists!')
            print('')

            ## Read netCDF files ##
            nc = netCDF4.Dataset(file_ostia, 'r')


            ## Read variables ##
            ### latitude, longitude ###
            lons = nc.variables['lon'][:] # Longitude
            lats = nc.variables['lat'][:] # Latitude


            ### SST ###
            ostia_sst = nc.variables['sst'][:][:][:].squeeze() #(time, lat, lon)=(1, 180, 360) -> (180, 360)

        
            ## Convert variables into Numpy ndarray ##
            ostia_sst_np = np.array(ostia_sst)


            ### Call functions ###
            '''
            # Compute area average of SST in the Nino 3.4 region
            '''
            nino34_ostia = nino34_sst(ostia_sst_np, lats, lons)
            #print('SST anomaly (OSTIA): ', nino34_ostia)


            nino34_yearly.append(nino34_ostia)

            # Close netCDF files
            nc.close()

    print(nino34_yearly)

    ## Compute Nino3.4 index ##
    nino34_idx = nino34_yearly - nino34_clm_all
    nino34idx_all.append(nino34_idx)

## List -> Numpy array ##
nino34idx_all_np = np.array(nino34idx_all).T
print('')
print(nino34idx_all_np)
print(nino34idx_all_np.shape)




## Convert np array into Pandas DataFrame ##
month = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
all_np = np.column_stack([month, nino34idx_all_np])

df = pd.DataFrame(all_np, columns=['month','1982', '1983', '1984', '1985', '1986', '1987', '1988', '1989', '1990', '1991', '1992', '1993', '1994', '1995',\
        '1996', '1997', '1998', '1999', '2000', '2001', '2002', '2003', '2004', '2005', '2006', '2007', '2008', '2009', '2010',\
        '2011', '2012', '2013', '2014', '2015', '2016', '2017', '2018', '2019', '2020', '2021', '2022', '2023', '2024', '2025'])

print(df)




##--- Save results as CSV --------------------------------------##
### Save DataFrame as CSV ###
df.to_csv(out_dir + 'nino34_idx_ostia_all.csv', index=False)
print('Saved Nino3.4 index!')
print('')


print('')
print('Completed!!')
print('')