# Compute climatologies of SST from monthly mean of the IO-SST-V2.1
# Period: 1993-2020 -> Copernicus: 1993-2016 #ver1.1.0
#
#
# Script version 1.1.0
#
# Script works with:
#   Python version 3.10
#
#   Package version
#     numpy: 1.26.3
#     pandas: 2.2.0
#
#
#
# Ver1.0.0: Created by Mariko Koseki, 19.04.2024
# Ver1.1.0: updated by Mariko, 14.05.2024


##--- Import modules ------------------------------##
import netCDF4
import numpy as np
from numpy import dtype
import pandas as pd
import math
import datetime
from datetime import timedelta
#from dateutil.relativedelta import relativedelta
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('')




##--- Set path ------------------------------------------##
## Set path to input files ##
'''
Set 'norcpm_monthly' (path to input NetCDF file)
'''
oisst_monthly = '/nird/projects/NS9873K/norcpm/validation/observation/NOAA/OISST/monthly/'


## Set path to output files ##
'''
Set 'out_clm_dir' (path to output file)
'''
out_dir  = '/nird/projects/NS9873K/norcpm/validation/observation/NOAA/OISST/clim/'
if not os.path.exists(out_dir):
    os.makedirs(out_dir)


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

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




##--- Set year, month ------------------------------------##
'''
Set 'start_year' and 'end_year'
Copernicus: 1993-2016 #ver1.1.0
'''
start_year = 1993 # 1993
end_year = 2017 # 2016
y_str_list = ["{}".format(y) for y in range(start_year,end_year)]
m_str_list = ["{:02}".format(m) for m in range(1,13)]



'''
##--- Test ------------------------------------##
yyyy = '2023'
mm = '12'
nc_oisst = oisst_monthly + yyyy + '_' + mm + '.nc'
print('file name: ', nc_oisst)
print('')

nc = netCDF4.Dataset(nc_oisst, 'r')

## Check length of variables ##
nlon = len(nc.dimensions['lon']) # Longitude:
nlat = len(nc.dimensions['lat']) # Latitude:
print(nlon, nlat)

## Read variables ##
oisst = nc.variables['sst'][:][:][:] #(time, lat, lon)
print(oisst.shape) #(1, 180, 360)
'''
'''
float sst(time, lat, lon) ;
                sst:long_name = "Daily sea surface temperature" ;
                sst:units = "Celsius" ;
                sst:_FillValue = -999.f ;
                sst:missing_value = -999.f ;
                sst:cell_methods = "time: mean zlev: mean" ;
                sst:add_offset = 0.f ;
                sst:scale_factor = 1.f ;

<<If both scale_factor and add_offset attributes are present, the data are first scaled before the offset is added.>>
# real_value = (display_value * scale_factor) + add_offset
'''
'''
add_offset = nc.variables['sst'].add_offset
scale_factor = nc.variables['sst'].scale_factor
oisst_real = (oisst * scale_factor) + add_offset
print(add_offset, scale_factor)
print('{:.20g}'.format(oisst[0,20,20]))
print('{:.20g}'.format(oisst_real[0,20,20]))

time = nc.variables['time'][:] # "days since 1978-01-01 12:00:00"
timeunits = nc.variables['time'].getncattr('units')
timecalendar = nc.variables['time'].getncattr('calendar')
date_sst = netCDF4.num2date(time,timeunits,timecalendar)
date_sst = [ f'{i.year}-{i.month:02}' for i in date_sst]
print(date_sst)

nc.close()

sys.exit()
###--- Test end --------------------------------###
'''




## Set length of variables ##
nyear = len(y_str_list)
nmonth = len(m_str_list)
nlat = 180; nlon = 360

months = np.arange(1, nmonth + 1)


##--- Read netCDF files ----------------------------------##
## Calculate climatology ##
'''
## Example: 2023_12.nc
dimensions:
	time = UNLIMITED ; // (1 currently)
	lat = 180 ;
	lon = 360 ;
variables:
    	float time(time) ;
		time:standard_name = "time" ;
		time:long_name = "Center time of the day" ;
		time:units = "days since 1978-01-01 12:00:00" ;
		time:calendar = "standard" ;
		time:axis = "T" ;
	double lon(lon) ;
		lon:standard_name = "longitude" ;
		lon:long_name = "longitude" ;
		lon:units = "degrees_east" ;
		lon:axis = "X" ;
	double lat(lat) ;
		lat:standard_name = "latitude" ;
		lat:long_name = "latitude" ;
		lat:units = "degrees_north" ;
		lat:axis = "Y" ;
	float sst(time, lat, lon) ;
		sst:long_name = "Daily sea surface temperature" ;
		sst:units = "Celsius" ;
		sst:_FillValue = -999.f ;
		sst:missing_value = -999.f ;
		sst:cell_methods = "time: mean zlev: mean" ;
		sst:add_offset = 0.f ;
		sst:scale_factor = 1.f ;
'''


oisst_clim_3d = np.zeros((nmonth, nlat, nlon)) #(12, 180, 360)
for mon_num in range(0, nmonth):
    mm = m_str_list[mon_num]
    print('month: ', mm)




    ##--- Calculate average for nyear ------------------------------------------## #ver1.1.0
    oisst_clim_mean = np.zeros((nyear, nlat, nlon)) #(nyear, 180, 360)
    for year_num in range(0, nyear):
        yyyy = y_str_list[year_num]

        print('year: ', yyyy)
        nc_oisst = oisst_monthly + yyyy + '_' + mm + '.nc'
        print('file name: ', nc_oisst)
        print('')


        nc = netCDF4.Dataset(nc_oisst, 'r')
        '''
        ## Check length of variables ##
        #nlon = len(nc.dimensions['lon']) # Longitude:
        #nlat = len(nc.dimensions['lat']) # Latitude:
        
        #print(nlon, nlat)
        '''


        ## Read variables ##
        '''
        time = nc.variables['time'][:] # "days since 1978-01-01 12:00:00"
        timeunits = nc.variables['time'].getncattr('units')
        timecalendar = nc.variables['time'].getncattr('calendar')
        date_sst = netCDF4.num2date(time,timeunits,timecalendar)
        date_sst = [ f'{i.year}-{i.month:02}' for i in date_sst]
        '''

        lons = nc.variables['lon'][:] # Longitude
        lats = nc.variables['lat'][:] # Latitude


        oisst = nc.variables['sst'][:][:][:] #(time, lat, lon)=(1, 180, 360)
        add_offset = nc.variables['sst'].add_offset
        scale_factor = nc.variables['sst'].scale_factor
        oisst_real = (oisst * scale_factor) + add_offset

        ## Convert variables into Numpy ndarray ##
        oisst_np = np.array(oisst_real)


        ## Add results into a zeros numpy ndarray ##
        oisst_clim_mean[year_num,:,:] = oisst_np


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


    mean_2d = np.nanmean(oisst_clim_mean, axis=0) #(180, 360)
    oisst_clim_3d[mon_num,:,:] = mean_2d
print('file shape of 3D', oisst_clim_3d.shape) #(12, 180, 360)



print('')
print('Calculation completed')
print('')
print('------------------------')
print('')



##--- Save results as a new NetCDF file ------------------------------##
print('Save climatology')

## File name ##
outfile = out_dir + 'OISST_clim_' + str(start_year) + '_' + str(end_year-1) + '.nc'
nc2 = netCDF4.Dataset(outfile ,'w', format="NETCDF4")

## Define dimensions ##
nc2.createDimension('lat', nlat)
nc2.createDimension('lon', nlon)
nc2.createDimension('month', nmonth)

## Define variables ##
### latitude ###
lat = nc2.createVariable('lat', dtype('float32').char, ('lat',))
lat.units = 'degrees_north'
lat.long_name = 'latitude'
lat[:] = lats

### longitude ###
lon = nc2.createVariable('lon', dtype('float32').char, ('lon',))
lon.units = 'degrees_east'
lon.long_name = 'longitude'
lon[:] = lons


### month ###
month = nc2.createVariable('month', dtype('int32').char, ('month',))
month.units = 'month'
month.long_name = 'month'
month[:] = months

### climatology ###
clim_sst = nc2.createVariable('clim_oisst', dtype('float32').char, ('month', 'lat', 'lon',))
clim_sst.units = 'Celsius'
clim_sst.long_name = 'climatology of OI-SST'
clim_sst[:,:,:] = oisst_clim_3d


### close NetCDF file ###
nc2.close()


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