# Read sea ice area/extent in the Arctic
#  - NERSC (SSMI)
#  - OISST (ver2, ver2.1)
#
# Plot sea ice area/extent comparing two observational data
#
#
# Script version 1.3.0
#
# Python version: 3.10.13
#
# Package version
#   numpy:  1.26.4
#   pandas:  2.2.0
#   matplotlib:  3.7.3
#
# Ver1.0.0: Created by Mariko Koseki, 24.05.2024 
# Ver1.0.1: Updated by Mariko, 24.05.2024
# Ver1.1.0: Updated by Mariko, 03.06.2024
# Ver1.2.0: Updated by Mariko, 04.06.2024
# Ver1.3.0: Updated by Mariko, 12.06.2024


import numpy as np
import pandas as pd
import datetime
from datetime import timedelta
from dateutil.relativedelta import relativedelta
import statistics
import math
import collections
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('matplotlib: ', matplotlib.__version__)



##--- Input -----------------------------------------##
args = sys.argv
if len(args) == 3:
    date_start = args[1]
    date_end = args[2]
    if (len(date_start) == 6) & (len(date_end) == 6):
        yyyy1 = int(date_start[0:4])
        mm1 = int(date_start[4:6])
        yyyy2 = int(date_end[0:4])
        mm2 = int(date_end[4:6])
        if yyyy1 == yyyy2:
            pass

        else:
            print('')
            print('---How to use this script---')
            print('python compare_seaice_nersc_oisst.py <yyyymm1> <yyyymm2>')
            print('<yyyy1mm1> = start month; <yyyy1mm2> = end month')
            print('yyyy1 = yyyy2')
            print('Example: python compare_seaice_nersc_oisst.py 202301 202312')
            print('')
            sys.exit()


    else:
        print('')
        print('---How to use this script---')
        print('python compare_seaice_nersc_oisst.py <yyyymm1> <yyyymm2>')
        print('<yyyy1mm1> = start month; <yyyy1mm2> = end month')
        print('yyyy1 = yyyy2')
        print('Example: python compare_seaice_nersc_oisst.py 202301 202312')
        print('')
        sys.exit()


else:
    print('')
    print('---How to use this script---')
    print('python compare_seaice_nersc_oisst.py <yyyymm1> <yyyymm2>')
    print('<yyyymm1> = start month; <yyyymm2> = end month')
    print('yyyy1 = yyyy2')
    print('Example: python compare_seaice_nersc_oisst.py 202301 202312')
    print('')
    sys.exit()



if mm2 == 12:
    start_month = mm1
    start_year = yyyy1
    end_month = 1
    end_year = yyyy2 + 1

else:
    start_month = mm1
    start_year = yyyy1
    end_month = mm2 + 1
    end_year = yyyy2



##--- Set version ---------------------------------------## ver1.1.0
ver21 = 'v2.1'
ver2 = 'v2'


##--- Set path ------------------------------------------##
## Set path to input files ##
'''
Set path to input directory ###
'''
### NERSC ###
seaice_dir = '/nird/projects/NS9873K/norcpm/validation/seaicepre/'
out_obs_post = seaice_dir + 'obs_post/' # Observation (post-processed)


### OISST ###
out_ice_oisst21 = '/nird/projects/NS9873K/norcpm/validation/seaicepre/oisst/' + ver21 + '/' #ver1.1.0, ver1.2.0
out_ice_oisst2 = '/nird/projects/NS9873K/norcpm/validation/seaicepre/oisst/' + ver2 + '/' #ver1.2.0



## Output: Set path to output file ##
'''
Set path to output file
'''
out_dir = seaice_dir + 'fig/'

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 ------------------------------------##
strdt = datetime.date(start_year, start_month, 1)
enddt = datetime.date(end_year, end_month, calendar.monthrange(end_year, end_month)[1])
days_num = (enddt - strdt).days + 1

ymd_str_list = []
for i in range(days_num):
    ymd_str_list.append(strdt + timedelta(days=i))

#print(ymd_str_list)
#print('start date: ', ymd_str_list[0])
#print('end date: ', ymd_str_list[-1])


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

#print(ym_str_list)




##--- Read CSV files ------------------------------##
## Paths to CSV files ##
### NERSC ###
file_nersc = out_obs_post + 'sea_ice_obs_post_' + str(start_year) + '.csv'

### OISST ### #ver1.2.0
file_oisst_list21 = []
file_oisst_list2 = []
for ym in ym_str_list:
    YEAR = ym[0:4]
    MONTH = ym[4:]

    file_oisst21 = out_ice_oisst21 + 'seaice_oisst_' + YEAR + MONTH + '.csv'
    file_oisst2 = out_ice_oisst2 + 'seaice_oisst_' + YEAR + MONTH + '.csv'
    file_oisst_list21.append(file_oisst21)
    file_oisst_list2.append(file_oisst2)

print('')
print('file of nersc: ', file_nersc)
print('')
print('file of OISST ver2.1: ', file_oisst_list21)
print('')
print('file of OISST ver2: ', file_oisst_list2)
print('')



print('')
print('check if sea ice data from oisst (ver2.1) exists...')
print('')
file_oisst_list21_ext = [] #ver1.2.0
for oisst_file21 in file_oisst_list21:
    if not os.path.exists(oisst_file21):
        print('')
        print('OISST (ver2.1) file name: ',oisst_file21)
        print('no such a directory!!!')
        print('-----------------')
        #sys.exit()

    else:
        file_oisst_list21_ext.append(oisst_file21)
        print('')
        print('OISST (ver2.1) file name: ',oisst_file21)
        print('Data exists!')
        print('-----------------')

print(file_oisst_list21_ext)



print('')
print('check if sea ice data from oisst (ver2) exists...')
print('')
file_oisst_list2_ext = [] #ver1.2.0
for oisst_file2 in file_oisst_list2:
    if not os.path.exists(oisst_file2):
        print('')
        print('OISST (ver2) file name: ',oisst_file2)
        print('no such a directory!!!')
        print('-----------------')
        #sys.exit()

    else:
        file_oisst_list2_ext.append(oisst_file2)
        print('')
        print('OISST (ver2) file name: ',oisst_file2)
        print('Data exists!')
        print('-----------------')

print(file_oisst_list2_ext)



print('')
print('')
print('')
print('check if sea ice data from NERSC exists...')
print('')
if not os.path.exists(file_nersc):
    print('')
    print('NERSC file name: ',file_nersc)
    print('no such a directory!!!')
    #sys.exit()

else:
    print('')
    print('NERSC file name: ',file_nersc)
    print('Data exists!')
    print('-----------------')


#sys.exit()

## CSV files -> pandas.DataFrame
### NERSC ###
if os.path.exists(file_nersc): #ver1.2.0
    ice_nersc = pd.read_csv(file_nersc)
    #print(ice_nersc)

### OISST ###
if len(file_oisst_list21_ext) > 0: #ver1.2.0
    ice_oisst_all21 = [] #ver1.2.0
    for file_oisst21 in file_oisst_list21_ext:
        ice_oisst21 = pd.read_csv(file_oisst21)
        #print(ice_oisst21)
        ice_oisst_all21.append(ice_oisst21)

if len(file_oisst_list2_ext) > 0: #ver1.2.0
    ice_oisst_all2 = [] #ver1.2.0
    for file_oisst2 in file_oisst_list2_ext:
        ice_oisst2 = pd.read_csv(file_oisst2)
        #print(ice_oisst2)
        ice_oisst_all2.append(ice_oisst2)


## Combine DataFrame objects ##
if len(file_oisst_list21_ext) > 0: #ver1.2.0
    ice_oisst21_mer= pd.concat(ice_oisst_all21, ignore_index=True, axis=0) #ver1.2.0
if len(file_oisst_list2_ext) > 0: #ver1.2.0
    ice_oisst2_mer= pd.concat(ice_oisst_all2, ignore_index=True, axis=0) #ver1.2.0
#print(ice_oisst21_mer)


## Change 'object' to 'datetime' ##
### NERSC ###
if os.path.exists(file_nersc): #ver1.2.0
    ice_nersc['DATE'] = pd.to_datetime(ice_nersc['DATE'])

### OISST ###
if len(file_oisst_list21_ext) > 0: #ver1.2.0
    ice_oisst21_mer['dtime'] = pd.to_datetime(ice_oisst21_mer['dtime']) #ver1.2.0
if len(file_oisst_list2_ext) > 0: #ver1.2.0
    ice_oisst2_mer['dtime'] = pd.to_datetime(ice_oisst2_mer['dtime']) #ver1.2.0



## Extract simulation period ##
### NERSC ###
if os.path.exists(file_nersc): #ver1.2.0
    ice_nersc = ice_nersc[(ice_nersc['DATE'] >= datetime.datetime(int(start_year), int(start_month), 1)) \
        & (ice_nersc['DATE'] <= datetime.datetime(int(end_year), int(end_month), calendar.monthrange(int(end_year), int(end_month))[1]))]



## Reset index ##
### NERSC ###
if os.path.exists(file_nersc): #ver1.2.0
    ice_nersc_res = ice_nersc.reset_index()
    #print(ice_nersc_res)



## Extract variables ##
### NERSC ###
if os.path.exists(file_nersc): #ver1.2.0
    time_nersc =ice_nersc_res['DATE']
    seaice_area_nersc = ice_nersc_res['AREA_Interpolated'] #ver1.3.0
    seaice_extent_nersc = ice_nersc_res['EXTENT_Interpolated'] #ver1.3.0


### OISST ###
if len(file_oisst_list21_ext) > 0: #ver1.2.0
    time_oisst21 =ice_oisst21_mer['dtime']
    seaice_area_oisst21 = ice_oisst21_mer['sea_ice_area'] #ver1.3.0
    seaice_extent_oisst21 = ice_oisst21_mer['sea_ice_extent'] #ver1.3.0
if len(file_oisst_list2_ext) > 0: #ver1.2.0
    time_oisst2 =ice_oisst2_mer['dtime']
    seaice_area_oisst2 = ice_oisst2_mer['sea_ice_area'] #ver1.3.0
    seaice_extent_oisst2 = ice_oisst2_mer['sea_ice_extent'] #ver1.3.0


#print(seaice_nersc)




##--- Plot sea ice area -----------------------------------------------##
## Create a new figure and a set of subplots ##
fig1 = plt.figure(figsize=(10,6))
ax1 = fig1.add_subplot(111)



## Plot x, y ##
### NERSC ###
if os.path.exists(file_nersc): #ver1.2.0
    ax1.plot(time_nersc, seaice_area_nersc, color='black', linewidth = 2.5, label= 'NERSC (sea ice area)')
    ax1.plot(time_nersc, seaice_extent_nersc, color='blue', linewidth = 2.5, label= 'NERSC (sea ice extent)')



### OISST ###
'''
if len(file_oisst_list2_ext) > 0: #ver1.2.0
    ax1.plot(time_oisst2, seaice_area_oisst2, color='green', linewidth = 5, label= 'OISST (sea ice area; ' + ver2 + ')', alpha=0.7) #ver1.3.0
    ax1.plot(time_oisst2, seaice_extent_oisst2, color='cyan', linewidth = 5, label= 'OISST (sea ice extent; ' + ver2 + ')', alpha=0.7) #ver1.3.0
'''
if len(file_oisst_list21_ext) > 0: #ver1.2.0
    ax1.plot(time_oisst21, seaice_area_oisst21, color='red', linewidth = 2.5, label= 'OISST (sea ice area; ' + ver21 + ')') #ver1.3.0
    ax1.plot(time_oisst21, seaice_extent_oisst21, color='magenta', linewidth = 2.5, label= 'OISST (sea ice extent; ' + ver21 + ')') #ver1.3.0


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


## Set xlim, ylim ## #Ver1.0.1
span = [datetime.datetime(start_year, start_month, 1), datetime.datetime(yyyy2, mm2, calendar.monthrange(yyyy2, mm2)[1])]
ax1.set_xlim(span)
ax1.set_ylim([0, 16000000])



## 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 ##
ax1.yaxis.set_major_formatter(ptick.ScalarFormatter(useMathText=True))
ax1.ticklabel_format(style="sci", axis="y", scilimits=(6,6)) # 1*10^6


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


## Add grid ##
ax1.minorticks_on()
ax1.grid(color = "gray", linestyle = "--")


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


## Add labels, title ##
ax1.set_ylabel('Sea Ice Area/Extent [$\mathrm{km^2}$]', fontsize = 12)

ax1.set_title('Daily Arctic sea ice area/extent (' + str(start_year) + ')', fontsize = 15)



## Save figures ##
fig_dir = out_dir + '/compare/'
if not os.path.exists(fig_dir):
    os.makedirs(fig_dir)

fig_name1 = fig_dir + 'compare_seaice_area_extent_' + str(start_year) + '_nersc_oisst.png' #ver1.1.0, ver1.2.0, ver1.3.0

fig1.savefig(fig_name1, dpi=300, format='png')

plt.show()

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