#!/bin/bash

# computes the strength and location of the NH/SH Hadley cell as maximum/minimum of the meridional mass stream function between 20S:40N/40S:20N (reference?)

set -ex

source $(dirname $CONDA_EXE)/../etc/profile.d/conda.sh
conda activate ralfs_standard

infile=$(readlink -f ${1})
outfile=$(readlink -f ${2})
varname=mastrfu

cat << __EOF__ | sed -e "s/^   //" > compute_hadley_cell.py
   import xarray as xr
   import numpy as np
   
   indata=xr.open_dataset('${infile}')
   
   lat=indata.lat
   time=indata.time
   plev=indata.plev
   strf=indata.${varname}

   strfmaskedn=strf.where((strf.coords["lat"]<40)&(strf.coords["lat"]>-20))
   strfmaskeds=strf.where((strf.coords["lat"]<20)&(strf.coords["lat"]>-40))

   hadcn=np.empty(shape=(len(time),len(plev),1,1),dtype="f")
   hadcs=np.empty(shape=(len(time),len(plev),1,1),dtype="f")

   hadcn[:,:,:,:]=np.nan
   hadcs[:,:,:,:]=np.nan
   
   strfmins=np.empty(shape=(len(time),len(plev),1,1),dtype="f")
   strfmaxn=np.empty(shape=(len(time),len(plev),1,1),dtype="f")

   strfmins[:,:,:,:]=np.nan
   strfmaxn[:,:,:,:]=np.nan    

   for t in range(0,len(time)):
       for z in range(1,len(plev)):

           strftzn=strfmaskedn[t,z,:,0]
           strftzs=strfmaskeds[t,z,:,0]

           if np.count_nonzero(~np.isnan(strftzn)) != 0:
   
               strfmax=np.nanmax(strftzn)
               strfmaxn[t,z,0,0]=strfmax
               maxlatn=np.where(strftzn == strfmax)

               if len(maxlatn) == 1:
                   if strf[t,z,int(maxlatn[0][0]),0] > 0:
   
                       latindstart=int(maxlatn[0][0])
                       for y in range(latindstart,0,-1):
                           if strf[t,z,y,0] < 0:
                               hadcntmp=lat[y]-(lat[y+1]-lat[y])*(strf[t,z,y,0]/(strf[t,z,y+1,0]-strf[t,z,y,0]))
                               hadcn[t,z,0,0]=np.around(hadcntmp,decimals=5)
                               break
                      

           else:
               hadcn[t,z,0,0]=np.nan


           if np.count_nonzero(~np.isnan(strftzs)) != 0:

               strfmin=np.nanmin(strftzs)
               strfmins[t,z,0,0]=strfmin
               minlats=np.where(strftzs == strfmin)
      
               if len(minlats[0]) == 1:
                   if strf[t,z,int(minlats[0][0]),0] < 0:
                       latindstart=int(minlats[0][0])
                       for y in range(latindstart,len(lat)-1):
                           if strf[t,z,y,0] > 0:
                               hadcstmp=lat[y-1]-(lat[y]-lat[y-1])*(strf[t,z,y-1,0]/(strf[t,z,y,0]-strf[t,z,y-1,0]))
                               hadcs[t,z,0,0]=np.around(hadcstmp,decimals=5)
                               break

           else:
               hadcs[t,z,0,0]=np.nan
   
   outdata = xr.Dataset(
   {
   "hadc_nh_loc": (["time","plev","lat","lon"], hadcn),
   "hadc_sh_loc": (["time","plev","lat","lon"], hadcs),
   "hadc_nh_str": (["time","plev","lat","lon"], strfmaxn),
   "hadc_sh_str": (["time","plev","lat","lon"], strfmins),
   },
   coords={"time": (["time"], indata.time),"plev": (["plev"], indata.plev,{"long_name" : "Pressure level" , "units" : "Pa", "positive" : "Down"}),"lat": (["lat"], np.array([0],dtype="f8"),{"long_name" : "latitude", "units" : "degrees_north"}),"lon": (["lon"], np.array([0],dtype="f8"),{"long_name" : "longitude", "units" : "degrees_east"})
   },
   )
   
   outdata.to_netcdf("tmp.hadley.nc")
   
   quit()

__EOF__

python compute_hadley_cell.py

cdo selyear,1420/1781 tmp.hadley.nc ${outfile} && rm tmp.hadley.nc 

exit 0