# -*- coding: utf-8 -*-
from netCDF4 import Dataset
from netCDF4 import MFDataset
from flux import comp_flux
from ncfile import create_file
import numpy as np

######### Namelist, here you can modify things ######
minjk=0
maxjk=75

nb_station=1

year_start=2015
year_end  =2015
ntimestep=12

temp_field='thetao'
salt_field='so'


input_dir='/cluster/shared/imr/robinson/OUTPUT_NORESM_RCP8.5_VRUNOFF/'
#diag_dir='/cluster/work/users/robinson/TEST/'
diag_dir='/cluster/work/users/annebs/NEMO/Fluxes/'
mesh_file='/cluster/projects/nn5001k/robinson/NAA10KM_DATA/mesh_mask.nc'




### Here you code a section with its name.
### Example bso_aw is Atlantic water at Barents Sea opening, i0/j0 are coordinates of the first point of your section
### and i1/j1 second point on the model grid. VERY IMPORTANT : i1=>i0, ALWAYS, but you are free to have j0>j1 or j0<j1 or j0=j1
### fltm is the minimum temperature, flsm is the minimum salinity, fltp is max temperature, flsp is max salinity


sec =["fao_shet"]
i0  =[ 527      ]
j0  =[ 256      ]
i1  =[ 564      ]
j1  =[ 232      ]
fltm=[ -10.     ]
flsm=[   0.     ]
fltp=[  30.     ]
flsp=[  50.     ]
flfp=[   0      ]
flfn=[   0      ]
flag=[   1      ]

### Take into account sea ice horizontal fluxes of mass (converted to liquid water volume) and salt, yes or no (1/0)

iflag=0

####### End of namelist #########

ft=Dataset(mesh_file)
gphit=ft.variables['gphit'][:,:,:]
glamt=ft.variables['glamt'][:,:,:]
e1t=ft.variables['e1t'][:,:,:]
e2t=ft.variables['e2t'][:,:,:]
e3t=ft.variables['e3t_0'][:,:,:,:]
tmask=ft.variables['tmask'][:,:,:,:]
gdept=ft.variables['gdept_1d'][:,:]
gdep=gdept[0,:]
ft.close()
gphit=gphit[0,:,:]
glamt=glamt[0,:,:]
e1t=e1t[0,:,:]
e2t=e2t[0,:,:]
e3t=e3t[0,:,:,:]*tmask[0,:,:,:]
volmask=e1t*e2t*e3t*tmask





for ny in range(year_start,year_end+1):

    year=str(ny)
    data_file_u=input_dir+'NAA10KM_1m_'+year+'0101_'+year+'1231_grid_U.nc'
    data_file_v=input_dir+'NAA10KM_1m_'+year+'0101_'+year+'1231_grid_V.nc'
    data_file_t=input_dir+'NAA10KM_1m_'+year+'0101_'+year+'1231_grid_T.nc'
    data_file_i=input_dir+'NAA10KM_1m_'+year+'0101_'+year+'1231_icemod.nc'
    
    salt=np.zeros(ntimestep)
    freh=np.zeros(ntimestep)
    
    time_axis=np.arange(ntimestep)
    nadd=(ny-1958)*ntimestep
    time_axis=[x+nadd for x in time_axis]
        
    glam=[[0.]]
    gphi=np.copy(glam)
   
  


    for ns in range(0,nb_station):
        if flag[ns]==1: 
       
           bsoflux, mtemp, msalt=comp_flux(temp_field,salt_field,volmask,iflag,data_file_i,data_file_u,data_file_v,data_file_t,fltm[ns],flsm[ns],fltp[ns],flsp[ns],i0[ns],j0[ns],i1[ns],j1[ns],minjk,'vfl')
           if flfp[ns]==1:
              nflux=np.where(bsoflux < 0.)
              bsoflux[nflux]=0.
           if flfn[ns]==1:
              pflux=np.where(bsoflux > 0.)
              bsoflux[pflux]=0.
              
           f2=np.sum(bsoflux,axis=3)
           f3=np.sum(f2,axis=2)
           output_file_name=diag_dir+sec[ns]+'_vfl_'+year+'.nc'
           field_name='Volume_Flow'
           create_file(output_file_name,field_name,time_axis,glam,gphi,gdep,f3)
           
           
           bsoflux, mtemp, msalt=comp_flux(temp_field,salt_field,volmask,iflag,data_file_i,data_file_u,data_file_v,data_file_t,fltm[ns],flsm[ns],fltp[ns],flsp[ns],i0[ns],j0[ns],i1[ns],j1[ns],minjk,'hfl')
           if flfp[ns]==1:
              nflux=np.where(bsoflux < 0.)
              bsoflux[nflux]=0.
           if flfn[ns]==1:
              pflux=np.where(bsoflux > 0.)
              bsoflux[pflux]=0.
           
           f2=np.sum(bsoflux,axis=3)
           f3=np.sum(f2,axis=2)
           output_file_name=diag_dir+sec[ns]+'_hfl_'+year+'.nc'          
           field_name='Heat_Flow'
           create_file(output_file_name,field_name,time_axis,glam,gphi,gdep,f3)


           
           bsoflux, mtemp, msalt=comp_flux(temp_field,salt_field,volmask,iflag,data_file_i,data_file_u,data_file_v,data_file_t,fltm[ns],flsm[ns],fltp[ns],flsp[ns],i0[ns],j0[ns],i1[ns],j1[ns],minjk,'sfl')
           if flfp[ns]==1:
              nflux=np.where(bsoflux < 0.)
              bsoflux[nflux]=0.
           if flfn[ns]==1:
              pflux=np.where(bsoflux > 0.)
              bsoflux[pflux]=0.

           f2=np.sum(bsoflux,axis=3)
           f3=np.sum(f2,axis=2)
           output_file_name=diag_dir+sec[ns]+'_sfl_'+year+'.nc'
           field_name='Salt_Flow'
           create_file(output_file_name,field_name,time_axis,glam,gphi,gdep,f3)