# -*- 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/FeP_tend_integ_gcb2024_v2'
suptitle = r'Fe$_{\mathrm{P}}$ integrated tendency over the global ocean'

nyc=366+5
dtime=np.empty(nyc,dtype=np.int64)
for yr in range(0,nyc):
    dtime[yr] = 1653+yr

#print(dtime)

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

#path_spin='/NasB2021_2/ERF21S208/htsujino/G100-JRA55-do-Fe-GCB2024/20240501_JRA55do-RYF9091_spinup_001/linkdir/result/annual/'
#path_spin='/nas/ERF24S22QN/htsujino/G100-JRA55-do-Fe-GCB2024/20250228_JRA55do_IAF_sensitivity_x05/linkdir/result/annual/'
path_spin='../../data/G100-JRA55-do-Fe-GCB2024/20250228_JRA55do_IAF_sensitivity_x05/annual/'

ntm=nyc

#####
base_FeP = 'hs_FeP'
suffix = '.1653'
fep_file = path_spin + '/1653/' + base_FeP + suffix
ncg100 = netCDF4.Dataset(fep_file,'r')
dep_mod = ncg100.variables['depth'][:]
nz = len(ncg100.dimensions['depth'])
ncg100.close()
#####

fe_tmp=np.array(np.empty((nz)), dtype=np.float64)

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

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

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

for year in range(0,nyc):
    file_fep_remove = path_spin + '{yr:04}/horz_integ_Fe_tend_fepremv_glb.{yr:04}'.format(yr=year+1653)
    if (os.path.isfile(file_fep_remove)):
        f3 = open(file_fep_remove,'rb')
        fe_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fep_remove[year,:] = fe_tmp[:].copy()
        f3.close()

for year in range(nyc):
    file_fed_scav = path_spin + '{yr:04}/horz_integ_Fe_tend_fedscav_glb.{yr:04}'.format(yr=year+1653)
    if (os.path.isfile(file_fed_scav)):
        f3 = open(file_fed_scav,'rb')
        fe_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fed_scav[year,:] = fe_tmp[:]
        f3.close()

for year in range(nyc):
    file_fep_dsop = path_spin + '{yr:04}/horz_integ_Fe_tend_fepdsop_glb.{yr:04}'.format(yr=year+1653)
    if (os.path.isfile(file_fep_dsop)):
        f3 = open(file_fep_dsop,'rb')
        fe_tmp = np.fromfile(f3, dtype = '>f', count = nz)
        fep_dsop[year,:] = fe_tmp[:]
        f3.close()

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

ds_fep_remove = xr.Dataset( {'fep': (['time','depth'], fep_remove),},
                     coords = {'time' : dtime, 'depth': dep_mod,} )
#print(ds_fep_remove)

ds_fed_scav = xr.Dataset( {'fed': (['time','depth'], fed_scav),},
                     coords = {'time' : dtime, 'depth': dep_mod,} )
#print(ds_fed_scav)

ds_fep_dsop = xr.Dataset( {'fep': (['time','depth'], fep_dsop),},
                     coords = {'time' : dtime, 'depth': dep_mod,} )
#print(ds_fep_dsop)

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

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

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

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

da = ds_fep_remove.fep.sum(dim='depth',skipna=False)*365.25*86400*1e-6*1e-9
da.plot(ax=axes,ylim=[-650.0,700.0],color='blue',linewidth=2, label=r'Fe$_{\mathrm{P}}$ removal')

da = -ds_fed_scav.fed.sum(dim='depth',skipna=False)*365.25*86400*1e-6*1e-9
da.plot(ax=axes,ylim=[-650.0,700.0],color='darkgoldenrod',linewidth=2, label=r'Fe$_{\mathrm{D}}$ scavenging')

da = ds_fep_dsop.fep.sum(dim='depth',skipna=False)*365.25*86400*1e-6*1e-9
da.plot(ax=axes,ylim=[-650.0,700.0],color='darkgreen',linewidth=2, label=r'Fe$_{\mathrm{P}}$ desorption')

da = -ds_fed_scav.fed.sum(dim='depth',skipna=False)*365.25*86400*1e-6*1e-9 +ds_fep_dsop.fep.sum(dim='depth',skipna=False)*365.25*86400*1e-6*1e-9
da.plot(ax=axes,ylim=[-650.0,700.0],color='purple',linewidth=2, label='desorption + scavenging')

axes.grid()
axes.set_xlabel('')
axes.set_xticks(np.arange(1653,2024,50))
axes.set_xlabel('Year',fontsize=16)
axes.set_ylabel(r'10$^{9}$mol/yr',fontsize=16)
axes.set_yticks(np.linspace(-650.0,700.0,28))
axes.tick_params(labelsize=14)
axes.axhline(0.0,color='black',linestyle='dashed',linewidth=1,zorder=10)
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()