#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sun Apr 14 15:08:30 EDT 2019
#import os, os.path, sys
import xarray as xr, numpy as np

def geopotential_height_make(ds):
    '''construct geopotential height from pk,bk,ps,temp,z0'''
    Ra = 287 # J/K/kg
    g = 9.8 # m/s**-2
    ifile = '/tigress/wenchang/MODEL_OUT/FLOR_pkak.nc'
    _ds = xr.open_dataset(ifile)
    pk, bk = _ds.pk, _ds.bk
#     ifile = '/scratch/gpfs/wenchang/FLOR/work/CTL1860_noleap_tigercpu_intelmpi_18_576PE/POSTP/10010101.atmos_month.nc'
#     ds = xr.open_dataset(ifile)
    dim_t, dim_y, dim_x = 'time', 'grid_yt', 'grid_xt'
    phalf, pfull = ds.phalf, ds.pfull
    ps = ds.ps
    temp = ds.temp
    z0 = ds.zsurf

    p = ( pk + ps*bk ).transpose(dim_t, 'phalf', dim_y, dim_x)
    dlnp = np.log(p.shift(phalf=-1) / p).isel(phalf=slice(0,-1)) \
        .rename({'phalf': 'pfull'}).assign_coords(pfull=pfull)
    dz = Ra/g * temp * dlnp
    z = z0 + dz.isel(pfull=slice(-1, None, -1)).cumsum(dim='pfull') \
        .isel(pfull=slice(-1, None, -1)).rename({'pfull': 'phalf'}) \
        .assign_coords(phalf=phalf[0:-1])
    z = z.transpose(dim_t, 'phalf', dim_y, dim_x)

    z_half = xr.zeros_like(p)
    z_half.values[:, 0:-1, :, :] = z.values
    z_half.values[:, -1, :, :] = z0.values

    z_full = (z_half.shift(phalf=-1) + z_half ) / 2
    z_full = z_full.isel(phalf=slice(0,-1)).rename({'phalf': 'pfull'}) \
        .assign_coords(pfull=pfull)

    z_full.attrs['units'] = 'm'
    z_full.attrs['long_name'] = 'geopotential height'
    z_full.name = 'z_full'
    
    return z_full