# -*- coding: utf-8 -*-
import sys
import os
import json
import netCDF4
from netCDF4 import Dataset, num2date
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import datetime
import xarray as xr

######

outfile = './fig/Fe_integ_spinup'
suptitle = r'Fe$_\mathrm{D}$ and Fe$_\mathrm{P}$ integrated over the global ocean'

nyc=366
dtime=np.arange(2018-nyc*5+1,2025,1)
print(dtime)

####################################################################

#path_gcb2024='/NasB2021_2/ERF21S208/htsujino/G100-JRA55-do-Fe-GCB2024/'
#path_gcb2024_2='/nas/ERF24S22QN/htsujino/G100-JRA55-do-Fe-GCB2024/'
path_gcb2024='../../data/G100-JRA55-do-Fe-GCB2024/'
path_gcb2024_2='../../data/G100-JRA55-do-Fe-GCB2024/'

cycle_1st='20240501_JRA55do-RYF9091_spinup_001'
cycle_2nd='20240515_JRA55do-RYF9091_spinup_002'
cycle_3rd='20240524_JRA55do-RYF9091_spinup_003'
cycle_4th='20240601_JRA55do-RYF9091_spinup_004'
#cycle_5th='20240609_JRA55do_IAF_product_005'
#cycle_5th='20240612_JRA55do-RYF9091_product_005'
cycle_5th='20250228_JRA55do_IAF_sensitivity_x05'

base_FeD = 'hs_FeD'

suffix = '.1001'
#fed_file = path_gcb2024 + cycle_1st + '/linkdir/result/annual/1001/' + base_FeD + suffix
fed_file = path_gcb2024 + cycle_1st + '/annual/1001/' + base_FeD + suffix
ncg100 = netCDF4.Dataset(fed_file,'r')
dep_mod = ncg100.variables['depth'][:]
nz = len(ncg100.dimensions['depth'])
ncg100.close()

ntm=nyc*5+6
fed_spin=np.array(np.empty((ntm,nz)), dtype=np.float64)
fed_spin[0:ntm,0:nz]=np.nan
fed_tmp=np.array(np.empty((nz)), dtype=np.float64)

fep_spin=np.array(np.empty((ntm,nz)), dtype=np.float64)
fep_spin[0:ntm,0:nz]=np.nan
fep_tmp=np.array(np.empty((nz)), dtype=np.float64)

for year in range(0,nyc):

    # 1st
    #file_fed_int = path_gcb2024 + cycle_1st + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1001)
    file_fed_int = path_gcb2024 + cycle_1st + '/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1001)
    if (os.path.isfile(file_fed_int)):
        f3 = open(file_fed_int,'rb')
        fed_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fed_spin[year+nyc*0,:] = fed_tmp[:].copy()
        f3.close()

    #file_fep_int = path_gcb2024 + cycle_1st + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1001)
    file_fep_int = path_gcb2024 + cycle_1st + '/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1001)
    if (os.path.isfile(file_fep_int)):
        f3 = open(file_fep_int,'rb')
        fep_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fep_spin[year+nyc*0,:] = fep_tmp[:].copy()
        f3.close()

    # 2nd
    #file_fed_int = path_gcb2024 + cycle_2nd + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1001)
    file_fed_int = path_gcb2024 + cycle_2nd + '/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1001)
    if (os.path.isfile(file_fed_int)):
        f3 = open(file_fed_int,'rb')
        fed_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fed_spin[year+nyc*1,:] = fed_tmp[:].copy()
        f3.close()

    #file_fep_int = path_gcb2024 + cycle_2nd + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1001)
    file_fep_int = path_gcb2024 + cycle_2nd + '/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1001)
    if (os.path.isfile(file_fep_int)):
        f3 = open(file_fep_int,'rb')
        fep_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fep_spin[year+nyc*1,:] = fep_tmp[:].copy()
        f3.close()

    # 3rd
    #file_fed_int = path_gcb2024 + cycle_3rd + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1001)
    file_fed_int = path_gcb2024 + cycle_3rd + '/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1001)
    if (os.path.isfile(file_fed_int)):
        f3 = open(file_fed_int,'rb')
        fed_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fed_spin[year+nyc*2,:] = fed_tmp[:].copy()
        f3.close()

    #file_fep_int = path_gcb2024 + cycle_3rd + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1001)
    file_fep_int = path_gcb2024 + cycle_3rd + '/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1001)
    if (os.path.isfile(file_fep_int)):
        f3 = open(file_fep_int,'rb')
        fep_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fep_spin[year+nyc*2,:] = fep_tmp[:].copy()
        f3.close()

    # 4th
    #file_fed_int = path_gcb2024 + cycle_4th + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1001)
    file_fed_int = path_gcb2024 + cycle_4th + '/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1001)
    if (os.path.isfile(file_fed_int)):
        f3 = open(file_fed_int,'rb')
        fed_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fed_spin[year+nyc*3,:] = fed_tmp[:].copy()
        f3.close()

    #file_fep_int = path_gcb2024 + cycle_4th + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1001)
    file_fep_int = path_gcb2024 + cycle_4th + '/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1001)
    if (os.path.isfile(file_fep_int)):
        f3 = open(file_fep_int,'rb')
        fep_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fep_spin[year+nyc*3,:] = fep_tmp[:].copy()
        f3.close()


for year in range(nyc+6):

    # 5th
    #file_fed_int = path_gcb2024_2 + cycle_5th + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1653)
    file_fed_int = path_gcb2024_2 + cycle_5th + '/annual/' + '{yr:04}/horz_integ_FeD_glb.{yr:04}'.format(yr=year+1653)
    if (os.path.isfile(file_fed_int)):
        f3 = open(file_fed_int,'rb')
        fed_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fed_spin[year+nyc*4,:] = fed_tmp[:].copy()
        f3.close()

    #file_fep_int = path_gcb2024_2 + cycle_5th + '/linkdir/result/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1653)
    file_fep_int = path_gcb2024_2 + cycle_5th + '/annual/' + '{yr:04}/horz_integ_FeP_glb.{yr:04}'.format(yr=year+1653)
    if (os.path.isfile(file_fep_int)):
        f3 = open(file_fep_int,'rb')
        fep_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fep_spin[year+nyc*4,:] = fep_tmp[:].copy()
        f3.close()

#################

ds_fed_spin = xr.Dataset( {'fed': (['time','depth'], fed_spin),},
                     coords = {'time' : dtime, 'depth': dep_mod} )

ds_fep_spin = xr.Dataset( {'fep': (['time','depth'], fep_spin),},
                     coords = {'time' : dtime, 'depth': dep_mod,} )


#################

fig = plt.figure(figsize = (11,8))
fig.suptitle( suptitle, fontsize=20 )

#################

axes = fig.add_subplot(1,1,1)

da = ds_fed_spin.fed.sum(dim='depth',skipna=False)*1e-17
da.plot(ax=axes,ylim=[-1.0,10.0],color='blue',linewidth=2, label=r"Fe$_\mathrm{D}$")
da = ds_fep_spin.fep.sum(dim='depth',skipna=False)*1e-17
da.plot(ax=axes,ylim=[-1.0,10.0],color='tomato',linewidth=2, label=r"Fe$_\mathrm{P}$")

axes.grid()
axes.set_xlabel('')
axes.set_xticks(np.arange(2018-nyc*5+1,2024,nyc))
axes.set_xlabel('Year',fontsize=16)
axes.set_ylabel(r'10$^{11}$mol',fontsize=16)
axes.set_yticks(np.linspace(-1.0,12.0,14))
axes.tick_params(labelsize=14)
leg = axes.legend(bbox_to_anchor=(0.95,0.05),loc='lower right')

################

plt.subplots_adjust(left=0.15,right=0.9,bottom=0.08,top=0.92,hspace=0.22)
#
outpdf = outfile+'.pdf'
outpng = outfile+'.png'

plt.savefig(outpng, bbox_inches='tight', pad_inches=0.05)
plt.savefig(outpdf, bbox_inches='tight', pad_inches=0.05)
print("figure is saved to " + outpng + " and " + outpdf)
plt.show()