# Compute Nino 3.4 index
#
# Use OISST and NorCPM output
# Calculate Nino 3.4 SST anomaly
# Plot Nino 3.4 index
# Period for climatology: 1993-2016 (Copernicus)
#
#
# Script version 2.1
#
# Script works with:
#   Python version 3.10
#
#   Package version
#     numpy:  1.26.4
#     pandas: 2.2.0
#
#
#
# Ver1.0.0: Created by Mariko Koseki, 15.05.2024
# Ver2.0.0: Updated by Mariko, 12.06.2024
# ver2.1: updated by Mariko, 08.12.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('')



##--- Input -----------------------------------------##
args = sys.argv
if len(args) == 2:
    input_date = args[1]
    if len(input_date) == 6:
        yyyy = int(input_date[0:4])
        mm = int(input_date[4:6])
        dd = 1

    else:
        print('')
        print('---How to use this script---')
        print('python plot_enso_oisst.py <yyyymm>')
        print('<yyyymm> = Forecast start date')
        print('Example: python plot_enso_oisst.py 202312')
        print('')
        sys.exit()

else:
    print('')
    print('---How to use this script---')
    print('python plot_enso_oisst.py <yyyymm>')
    print('<yyyymm> = Forecast start date')
    print('Example: python plot_enso_oisst.py 202312')
    print('')
    sys.exit()



if mm >= 8:
    #print(mm)
    start_month_mdl = mm
    start_year_mdl = yyyy
    end_month_mdl = mm - 7
    end_year_mdl = yyyy + 1

    start_month_obs = mm - 6
    start_year_obs = yyyy
    end_month_obs = mm - 1
    end_year_obs = yyyy
else:
    start_month_mdl = mm
    start_year_mdl = yyyy
    end_month_mdl = mm + 5
    end_year_mdl = yyyy

    if start_month_mdl == 1:
        start_month_obs = mm + 6
        start_year_obs = yyyy - 1
        end_month_obs = 12
        end_year_obs = yyyy - 1
    elif start_month_mdl == 7:
        start_month_obs = mm - 6
        start_year_obs = yyyy
        end_month_obs = mm - 1
        end_year_obs = yyyy

    else:
        start_month_obs = mm + 6
        start_year_obs = yyyy - 1
        end_month_obs = mm - 1
        end_year_obs = yyyy


mon = mm -1

start_year = '{:0>4d}'.format(yyyy)
start_month = '{:0>2d}'.format(mm)
start_date = '{:0>2d}'.format(dd)

issue_date = start_year + start_month + start_date #ver2.0.0
print('')
print('Issued: ', issue_date)
print('')


forecast_start = '{:0>2d}/{:0>2d}/{:0>4d}'.format(dd,mm,yyyy)
print('Forecast start: ', forecast_start)


forecast_start_date = '{:0>4d}{:0>2d}'.format(start_year_mdl,start_month_mdl) #string
forecast_end_date = '{:0>4d}{:0>2d}'.format(end_year_mdl,end_month_mdl) #string

obs_start_date = '{:0>4d}{:0>2d}'.format(start_year_obs,start_month_obs) #string
obs_end_date = '{:0>4d}{:0>2d}'.format(end_year_obs,end_month_obs) #string

print('')
print('Forecast start date: ', forecast_start_date) #Forecast start date
print('Forecast end date: ', forecast_end_date) #Forecast end date

print('Observation start date: ', obs_start_date) #Observation start date
print('Observation end date: ', obs_end_date) #Observation end date
print('')



## Create list of year_month and month ##
start_date = '{:0>4d}-{:0>2d}'.format(start_year_obs,start_month_obs)
end_date = datetime.datetime(end_year_obs,end_month_obs,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()

start_date_long = datetime.datetime(start_year_obs,start_month_obs,1)
start_date_long = start_date_long - relativedelta(months=1)
start_date_long = start_date_long.strftime('%Y-%m')
end_date_long = datetime.datetime(end_year_mdl,end_month_mdl,1)
end_date_long = end_date_long + relativedelta(months=2)
end_date_long = end_date_long.strftime('%Y-%m')
ym_str_list_long = pd.date_range(start_date_long, end_date_long, freq='ME').strftime("%Y_%m").tolist()

#print(ym_str_list, ym_str_list_all)
m_str_list = pd.date_range(start_date, end_date, freq='ME').strftime("%m").tolist()
#print(m_str_list)
m_int_list = [int(s)-1 for s in m_str_list]
#print(m_int_list)


## Select NorCPM system ##
'''
Set system number 'system_num'
'''
system_num = '1'
#system_num = '2'

print('')
print('system: ', system_num)
print('')



##--- Functions -----------------------------------------##
def nino34_norcpm_grid(sst_norcpm, lats_norcpm, lons_norcpm):
    '''
    # Compute area average of SST in the Nino 3.4 region
    # Nino 3.4 region (NorCPM grid): 5~-5 (5N-5S), -170~-120 (120W-170W)
    '''
    ## Check index of Nino3.4 area ##
    lat_nino34_start_ind_nor = np.where((lats_norcpm >= -5) & (lats_norcpm <= 5))[0][0]
    lat_nino34_end_ind_nor = np.where((lats_norcpm >= -5) & (lats_norcpm <= 5))[0][-1]
    #print(lat_nino34_start_ind_nor, lat_nino34_end_ind_nor)

    lon_nino34_start_ind_nor = np.where((lons_norcpm >= -170) & (lons_norcpm <= -120))[0][0]
    lon_nino34_end_ind_nor = np.where((lons_norcpm >= -170) & (lons_norcpm <= -120))[0][-1]
    #print(lon_nino34_start_ind_nor, lon_nino34_end_ind_nor) 

    ## Extract Nino3.4 area ##
    nino_sst_norcpm = sst_norcpm[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_norcpm_clim.shape)


    ## Compute Nino 3.4 area average of SST ##
    nino34_norcpm = np.nanmean(np.nanmean(nino_sst_norcpm, 0), 0)
    return nino34_norcpm



def nino34_oisst_grid(oisst, lats_oi, lons_oi):
    '''
    # Compute area average of SST in the Nino 3.4 region
    # Nino 3.4 region (IOSST grid): 5~-5 (5N-5S), 190~240 (120W-170W)
    '''
    ## Check index of Nino3.4 area ##
    lat_nino34_start_ind_oi = np.where((lats_oi >= -5) & (lats_oi <= 5))[0][0]
    lat_nino34_end_ind_oi = np.where((lats_oi >= -5) & (lats_oi <= 5))[0][-1]
    #print(lat_nino34_start_ind_oi, lat_nino34_end_ind_oi)

    lon_nino34_start_ind_oi = np.where((lons_oi >= 190) & (lons_oi <= 240))[0][0]
    lon_nino34_end_ind_oi = np.where((lons_oi >= 190) & (lons_oi <= 240))[0][-1]
    #print(lon_nino34_start_ind_oi, lon_nino34_end_ind_oi)


    ## Extract Nino3.4 area ##
    nino_oisst = oisst[lat_nino34_start_ind_oi:lat_nino34_end_ind_oi + 1, lon_nino34_start_ind_oi:lon_nino34_end_ind_oi +1]
    #print(nino_oisst.shape)


    ## Compute Nino 3.4 area average of SST ##
    nino34_oisst = np.nanmean(np.nanmean(nino_oisst, 0), 0)
    return nino34_oisst




##--- 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 ------------------------------------------## # ver2.1
## Set path to input files ##
'''
Set 'norcpm_monthly; oisst_monthly' (path to input NetCDF file)
'''
norcpm_monthly = '/nird/datapeak/NS9873K/www/norcpm/forecast/monthly/'
norcpm_clim  = '/nird/datapeak/NS9873K/norcpm/validation/sst/clim/bccr_1_sea_surface_temperature_clim_' + str(start_year_clm) + '_' + str(end_year_clm) + '.nc'
oisst_monthly = '/nird/datapeak/NS9873K/norcpm/validation/observation/NOAA/OISST/monthly/'
oisst_clim = '/nird/datapeak/NS9873K/norcpm/validation/observation/NOAA/OISST/clim/OISST_clim_' + str(start_year_clm) + '_' + str(end_year_clm) + '.nc'



## Set path to output files ##
'''
Set 'out_fig_dir' (path to output file)
'''
#out_fig_dir  = '/nird/datapeak/NS9873K/norcpm/validation/enso/fig/clm_' + str(start_year_clm) + '_' + str(end_year_clm) + '/' # <--- tmp forlder
out_fig_dir  = '/nird/datapeak/NS9873K/www/norcpm/forecast/plots/version2/' #ver2.0.0                            <--- folder for archiving figures
if not os.path.exists(out_fig_dir):
    os.makedirs(out_fig_dir)




##--- Read netCDF files ----------------------------------##
#################
## Climatology ##
#################
print('climatology of NorCPM: ', norcpm_clim)
print('climatology of OISST: ', oisst_clim)
print('')

nc_norcpm_clim = netCDF4.Dataset(norcpm_clim, 'r')
nc_oisst_clim = netCDF4.Dataset(oisst_clim, 'r')


## Check length of variables ##
nlon_norclm = len(nc_norcpm_clim.dimensions['longitude']) # Longitude: 360
nlat_norclm = len(nc_norcpm_clim.dimensions['latitude']) # Latitude: 180
nlead_norclm = len(nc_norcpm_clim.dimensions['lead']) # Lead time: 6
nmonth_norclm = len(nc_norcpm_clim.dimensions['month']) # month: 12

nlon_oiclm = len(nc_oisst_clim.dimensions['lon']) # Longitude: 360
nlat_oiclm = len(nc_oisst_clim.dimensions['lat']) # Latitude: 180
nmonth_oiclm = len(nc_oisst_clim.dimensions['month']) # month: 12


## Read variables ##
### lead time ###
leads_norclm = nc_norcpm_clim.variables['lead'][:] # lead time: 1-6

### month ###
months_norclm = nc_norcpm_clim.variables['month'][:] # month: 1-12

months_oiclm = nc_oisst_clim.variables['month'][:] # month: 12


### latitude, longitude ###
lons_norclm = nc_norcpm_clim.variables['longitude'][:] # Longitude: -179.5 - 179.5
lats_norclm = nc_norcpm_clim.variables['latitude'][:] # Latitude: 89.5 - -89.5

lons_oiclm = nc_oisst_clim.variables['lon'][:] # Lon: 0.5 - 359.5
lats_oiclm = nc_oisst_clim.variables['lat'][:] # Lat: 89.5 - -89.5


#print(nlon_norclm, nlat_norclm, nlead_norclm, nmonth_norclm)
#print(nlon_oiclm, nlat_oiclm, nmonth_oiclm)
#print(leads_norclm, months_norclm, lons_norclm, lats_norclm)
#print(months_oiclm, lons_oiclm, lats_oiclm)
#print('')


### SST ###
sst_norcpm_clim = nc_norcpm_clim.variables['clim_sst'][:][:][:][:]
oisst_clim = nc_oisst_clim.variables['clim_oisst'][:][:][:]


## Convert variables into Numpy ndarray ##
sst_norcpm_clim_np = np.array(sst_norcpm_clim)
sst_norcpm_clim_np = sst_norcpm_clim_np[:, mon, :, :]
oisst_clim_np = np.array(oisst_clim)

#print(sst_norcpm_clim_np.shape)
#print(oisst_clim_np.shape)





##############
## Forecast ##
##############
nc_file_forecast = norcpm_monthly + str(start_year_mdl) + str(start_month_mdl).zfill(2) + '/sea_surface_temperature_bccr_' \
        + system_num + '_' + str(start_year_mdl) + '_' + str(start_month_mdl).zfill(2) + '.nc'
print('')
print('forecast file name: ', nc_file_forecast)
nc_sst = netCDF4.Dataset(nc_file_forecast, 'r')


## Check length of variables ##
nlon_sst = len(nc_sst.dimensions['longitude']) # Longitude: 360
nlat_sst = len(nc_sst.dimensions['latitude']) # Latitude: 180
nmem_sst = len(nc_sst.dimensions['number']) # Ensemble member: 60
nlead_sst = len(nc_sst.dimensions['time']) # Lead time: 6


## Read variables ##
### lead time ###
time_sst = nc_sst.variables['time'][:] # lead time: 1-6
timeunits = nc_sst.variables['time'].getncattr('units')
timecalendar = nc_sst.variables['time'].getncattr('calendar')
dtime = netCDF4.num2date(time_sst,timeunits,timecalendar)
dtime = [ f'{i.year}-{i.month:02}-{i.day:02}' for i in dtime]


### latitude, longitude ###
lons_sst = nc_sst.variables['longitude'][:] # Longitude: -179.5 - 179.5
lats_sst = nc_sst.variables['latitude'][:] # Latitude: 89.5 - -89.5

#print('')
#print(nlon_sst, nlat_sst, nmem_sst, nlead_sst)
#print(time_sst, lons_sst, lats_sst)
#print('')

#sst_forecast = nc_sst.variables['sst'][lead][mem][:][:]
sst_forecast = nc_sst.variables['sst'][:][:][:][:]

## Convert variables into Numpy ndarray ##
sst_forecast_np = np.array(sst_forecast)
#print(sst_forecast.shape)


###########
## OISST ##
###########
nc_file_oisst1 = oisst_monthly + ym_str_list[0] + '.nc'
nc_file_oisst2 = oisst_monthly + ym_str_list[1] + '.nc'
nc_file_oisst3 = oisst_monthly + ym_str_list[2] + '.nc'
nc_file_oisst4 = oisst_monthly + ym_str_list[3] + '.nc'
nc_file_oisst5 = oisst_monthly + ym_str_list[4] + '.nc'
nc_file_oisst6 = oisst_monthly + ym_str_list[5] + '.nc'

print('')
print('oisst: ', nc_file_oisst1)
print('oisst: ', nc_file_oisst2)
print('oisst: ', nc_file_oisst3)
print('oisst: ', nc_file_oisst4)
print('oisst: ', nc_file_oisst5)
print('oisst: ', nc_file_oisst6)
print('')

nc_oisst1 = netCDF4.Dataset(nc_file_oisst1, 'r')
nc_oisst2 = netCDF4.Dataset(nc_file_oisst2, 'r')
nc_oisst3 = netCDF4.Dataset(nc_file_oisst3, 'r')
nc_oisst4 = netCDF4.Dataset(nc_file_oisst4, 'r')
nc_oisst5 = netCDF4.Dataset(nc_file_oisst5, 'r')
nc_oisst6 = netCDF4.Dataset(nc_file_oisst6, 'r')


## Check length of variables ##
nlon_oisst = len(nc_oisst1.dimensions['lon']) # Longitude: 360
nlat_oisst = len(nc_oisst1.dimensions['lat']) # Latitude: 180


# Read variables ##
### latitude, longitude ###
lons_oisst = nc_oisst1.variables['lon'][:] # Lon: 0.5 - 359.5
lats_oisst = nc_oisst1.variables['lat'][:] # Lat: 89.5 - -89.5

#print('')
#print(nlon_oisst, nlat_oisst)
#print(lons_oisst, lats_oisst)
#print('')



oisst1 = nc_oisst1.variables['sst'][:][:][:].squeeze() #(time, lat, lon)=(1, 180, 360) -> (180, 360)
add_offset1 = nc_oisst1.variables['sst'].add_offset; scale_factor1 = nc_oisst1.variables['sst'].scale_factor
oisst_real1 = (oisst1 * scale_factor1) + add_offset1
oisst2 = nc_oisst2.variables['sst'][:][:][:].squeeze() #(time, lat, lon)=(1, 180, 360) -> (180, 360)
add_offset2 = nc_oisst2.variables['sst'].add_offset; scale_factor2 = nc_oisst2.variables['sst'].scale_factor
oisst_real2 = (oisst2 * scale_factor2) + add_offset2
oisst3 = nc_oisst3.variables['sst'][:][:][:].squeeze() #(time, lat, lon)=(1, 180, 360) -> (180, 360)
add_offset3 = nc_oisst3.variables['sst'].add_offset; scale_factor3 = nc_oisst3.variables['sst'].scale_factor
oisst_real3 = (oisst3 * scale_factor3) + add_offset3
oisst4 = nc_oisst4.variables['sst'][:][:][:].squeeze() #(time, lat, lon)=(1, 180, 360) -> (180, 360)
add_offset4 = nc_oisst4.variables['sst'].add_offset; scale_factor4 = nc_oisst4.variables['sst'].scale_factor
oisst_real4 = (oisst4 * scale_factor4) + add_offset4
oisst5 = nc_oisst5.variables['sst'][:][:][:].squeeze() #(time, lat, lon)=(1, 180, 360) -> (180, 360)
add_offset5 = nc_oisst5.variables['sst'].add_offset; scale_factor5 = nc_oisst5.variables['sst'].scale_factor
oisst_real5 = (oisst5 * scale_factor5) + add_offset5
oisst6 = nc_oisst6.variables['sst'][:][:][:].squeeze() #(time, lat, lon)=(1, 180, 360) -> (180, 360)
add_offset6 = nc_oisst6.variables['sst'].add_offset; scale_factor6 = nc_oisst6.variables['sst'].scale_factor
oisst_real6 = (oisst6 * scale_factor6) + add_offset6

# Convert variables into Numpy ndarray ##
oisst_np1 = np.array(oisst_real1)
oisst_np2 = np.array(oisst_real2)
oisst_np3 = np.array(oisst_real3)
oisst_np4 = np.array(oisst_real4)
oisst_np5 = np.array(oisst_real5)
oisst_np6 = np.array(oisst_real6)
#print(oisst_np1.shape)




##--- Call functions --------------------------------------##
### Compute area average of SST in the Nino 3.4 region ###

#################
## Climatology ##
#################
nino34_norcpm_clm_all = np.zeros((nlead_norclm))
for lead in range(0, nlead_norclm):
    clm_sst_norcpm = sst_norcpm_clim_np[lead, :, :]
    nino34_nor_clm = nino34_norcpm_grid(clm_sst_norcpm, lats_norclm, lons_norclm)
    #print('climatology of NorCPM: ', nino34_nor_clm)
    nino34_norcpm_clm_all[lead] = nino34_nor_clm

print('')
print('all climatologies of NorCPM: ', nino34_norcpm_clm_all)

#print(oisst_clim_np.shape)
nino34_oisst_clm_list = []
for mon in m_int_list:
    nino34_oisst_clm = nino34_oisst_grid(oisst_clim_np[mon, :, :], lats_oiclm, lons_oiclm)
    #print('climatology of OISST: ', nino34_oisst_clm)
    nino34_oisst_clm_list.append(nino34_oisst_clm)
nino34_oisst_clm_all = np.array(nino34_oisst_clm_list)

print('')
print('all climatologies of OISST: ', nino34_oisst_clm_all)


##############
## Forecast ##
##############
nino34_forecast_all = np.zeros((nlead_sst, nmem_sst)) #(6, 60)
for lead in range(0, nlead_sst):
    for mem in range(0, nmem_sst):
        sst_norcpm = sst_forecast_np[lead, mem, :, :]
        nino34_forecast = nino34_norcpm_grid(sst_norcpm, lats_sst, lons_sst)
        nino34_forecast_all[lead, mem] = nino34_forecast
        #print(nino34_forecast)
        #print(nino34_forecast_all.shape)

#print(nino34_forecast_all)

nino34_forecast_ens_mean = np.mean(nino34_forecast_all, axis=1)
print('')
print('ensemble mean: ', nino34_forecast_ens_mean)
#print('shape of nino34_forecast_all: ', nino34_forecast_all.shape)



###########
## OISST ##
###########
nino34_oisst1 = nino34_oisst_grid(oisst_np1, lats_oisst, lons_oisst)
nino34_oisst2 = nino34_oisst_grid(oisst_np2, lats_oisst, lons_oisst)
nino34_oisst3 = nino34_oisst_grid(oisst_np3, lats_oisst, lons_oisst)
nino34_oisst4 = nino34_oisst_grid(oisst_np4, lats_oisst, lons_oisst)
nino34_oisst5 = nino34_oisst_grid(oisst_np5, lats_oisst, lons_oisst)
nino34_oisst6 = nino34_oisst_grid(oisst_np6, lats_oisst, lons_oisst)

nino34_oisst_list = [nino34_oisst1, nino34_oisst2, nino34_oisst3, nino34_oisst4, nino34_oisst5, nino34_oisst6]

nino34_oisst_all = np.array(nino34_oisst_list)
print('')
print('OISST: ', nino34_oisst_all)



## --- Calculate percentiles ------------------------------------------------##
percentile_10_all = []
percentile_90_all = []
for lead in range(0, nlead_sst):
    percentile_10 = np.percentile(nino34_forecast_all[lead, :], 10)
    percentile_90 = np.percentile(nino34_forecast_all[lead, :], 90)
    percentile_10_all.append(percentile_10)
    percentile_90_all.append(percentile_90)

#print(percentile_10_all)
#print(percentile_90_all)




##--- Calculate Nino 3.4 SST anomaly --------------------------------------##
nino34_index_norcpm = nino34_forecast_ens_mean - nino34_norcpm_clm_all
nino34_index_oisst = nino34_oisst_all - nino34_oisst_clm_all
nino34_index_norcpm_per_10 = percentile_10_all - nino34_norcpm_clm_all
nino34_index_norcpm_per_90 = percentile_90_all - nino34_norcpm_clm_all

nino34_ens_mem_list = []
for mem in range(0, nmem_sst):
    nino34_each_mem = nino34_forecast_all[:, mem] - nino34_norcpm_clm_all
    nino34_ens_mem_list.append(nino34_each_mem)
nino34_ens_mem = np.array(nino34_ens_mem_list)


nino34_boundary = np.zeros((2))
nino34_boundary[0] = nino34_index_oisst[-1]
nino34_boundary[1] = nino34_index_norcpm[0]

#print(nino34_index_oisst)
#print(nino34_index_norcpm)
#print(nino34_index_norcpm_per_10, nino34_index_norcpm_per_90)
#print(nino34_boundary)
#print(nino34_ens_mem.shape)




##--- Plot Nino 3.4 index -----------------------------------------##
### Create a new figure and a set of subplots ###
fig1 = plt.figure(figsize=(11,7))
ax1 = fig1.add_subplot(111)


### Create x ###
norcpm_time = []
for n in range(6):  
    s_format1 = '%Y-%m-%d'
    dt1 = datetime.datetime.strptime(dtime[n], s_format1)
    norcpm_time.append(dt1)

oisst_time = []
for o in range(6):
    s_format2 = '%Y_%m'
    dt2 = datetime.datetime.strptime(ym_str_list[o], s_format2)
    oisst_time.append(dt2)

time_all = []
for a in range(14):
    s_format3 = '%Y_%m'
    dt3 = datetime.datetime.strptime(ym_str_list_long[a], s_format3)
    time_all.append(dt3)

time_boundary = []
time_boundary.append(oisst_time[-1])
time_boundary.append(norcpm_time[0])


### Plot x, y ###
ax1.fill_between(norcpm_time, nino34_index_norcpm_per_10, nino34_index_norcpm_per_90, facecolor='cyan', alpha=0.5, label = '10-90 percentile range')

ax1.plot(norcpm_time, nino34_ens_mem[0,:], color='r', alpha=0.5, linewidth = 0.5, label = 'Ensemble member')
for mem in range(1, nmem_sst):
    ax1.plot(norcpm_time, nino34_ens_mem[mem, :], color='r', alpha=0.5, linewidth = 0.5)

ax1.plot(time_all, [0]*len(time_all), '--', color='grey')
ax1.plot(time_boundary, nino34_boundary, '--' ,color='green', linewidth = 1)
ax1.plot(oisst_time, nino34_index_oisst, color='black', linewidth = 2.5, marker = 'o', label= 'Observation (' + obs_start_date + '-' + obs_end_date + ')')
ax1.plot(norcpm_time, nino34_index_norcpm, color='b', linewidth = 2.5, marker = 'o', label = 'Ensemble mean (' + forecast_start_date + '-' + forecast_end_date + ')')


### Get legend, handles and labels ###
hans, labs = ax1.get_legend_handles_labels()


### Set xlim, ylim ###
span = [(datetime.datetime(start_year_obs, start_month_obs, 1) - datetime.timedelta(days=15)), (datetime.datetime(end_year_mdl, end_month_mdl, calendar.monthrange(end_year_mdl, end_month_mdl)[1]) - datetime.timedelta(days=15))]
ax1.set_xlim(span)
ax1.set_ylim([-3, 3])


### Set X axis ###
#ax1.xaxis.set_major_locator(mdates.MonthLocator(interval=2)) # <--- interval: 2 month
ax1.xaxis.set_major_locator(mdates.MonthLocator(interval=1)) # <--- interval: 1 month
#datefmt = mdates.DateFormatter("%Y-%m") # <--- format: yyyy-mm
datefmt = mdates.DateFormatter("%b") # <--- format: Jan
ax1.xaxis.set_major_formatter(datefmt)


### Set Y axis ###
plt.gca().yaxis.set_major_formatter(plt.FormatStrFormatter('%.1f'))


### Set label size ###
ax1.tick_params(labelsize = 13, pad = 10)


### Add legend ###
l1 = ax1.legend(handles=hans[:5], labels=labs[:5], frameon=False, loc='lower right', fontsize = 'small')
#l2 = ax1.legend(handles=hans[5:], labels=labs[5:], frameon=False, loc='lower left', fontsize = 'small')
#ax1.add_artist(l1)


### Add labels ###
#ax1.set_xlabel('Date', fontsize = 12.5)
ax1.set_ylabel('Temperature anomaly [K]', fontsize = 14)



### Add title ###
fig_title = 'NorCPM seasonal Nino3.4 forecast'
ax_title = '(forecast start: ' + forecast_start +')'

fig1.suptitle(fig_title, fontsize = 18)
ax1.set_title(ax_title, loc="right", fontsize = 16)



'''
### Add text ###
fig_text = 'Issued: ' + issue_date
ax_pos = ax1.get_position()
fig1.text(ax_pos.x1 - 0.03, ax_pos.y1 - 0.85, fig_text, fontsize = 10)
'''


### Save figures ###
figs_dir = out_fig_dir + '/' + start_year + start_month + '/'
if not os.path.exists(figs_dir):
    os.makedirs(figs_dir)

### File name ###
'''
<bccr system number>_<variable name>_<aggregation>_<lead time>
'''
#fig_name1 = figs_dir + 'bccr_' + system_num + '_enso_oisst.png' # <--- temp forlder
fig_name1 = figs_dir + 'bccr_' + system_num + '_enso_s' + issue_date +'.png' # #ver2.0.0 <--- archiving figure
fig1.savefig(fig_name1, dpi=300, format='png')



plt.show()



##--- Close netCDF file ---------------------------------------##
nc_norcpm_clim.close()
nc_oisst_clim.close()
nc_sst.close()
nc_oisst1.close()
nc_oisst2.close()
nc_oisst3.close()
nc_oisst4.close()
nc_oisst5.close()
nc_oisst6.close()

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