util.py

import metpy
import numpy as np
import xarray as xr
import datetime
from datetime import timezone
from metpy.units import units
from metpy.calc import thickness_hydrostatic
from collections import namedtuple
import os
import h5py
import pickle
from netCDF4 import Dataset
from util.setup import ancillary_path
from scipy.interpolate import RectBivariateSpline, interp2d
from scipy.ndimage import gaussian_filter

LatLonTuple = namedtuple('LatLonTuple', ['lat', 'lon'])

homedir = os.path.expanduser('~') + '/'

# --- CLAVRx Radiometric parameters and metadata ------------------------------------------------
l1b_ds_list = ['temp_10_4um_nom', 'temp_11_0um_nom', 'temp_12_0um_nom', 'temp_13_3um_nom', 'temp_3_75um_nom',
               'temp_6_2um_nom', 'temp_6_7um_nom', 'temp_7_3um_nom', 'temp_8_5um_nom', 'temp_9_7um_nom',
               'refl_0_47um_nom', 'refl_0_65um_nom', 'refl_0_86um_nom', 'refl_1_38um_nom', 'refl_1_60um_nom']

l1b_ds_types = {ds: 'f4' for ds in l1b_ds_list}
l1b_ds_fill = {l1b_ds_list[i]: -32767 for i in range(10)}
l1b_ds_fill.update({l1b_ds_list[i+10]: -32768 for i in range(5)})
l1b_ds_range = {ds: 'actual_range' for ds in l1b_ds_list}

# --- CLAVRx L2 parameters and metadata
ds_list = ['cld_height_acha', 'cld_geo_thick', 'cld_press_acha', 'sensor_zenith_angle', 'supercooled_prob_acha',
           'supercooled_cloud_fraction', 'cld_temp_acha', 'cld_opd_acha', 'solar_zenith_angle',
           'cld_reff_acha', 'cld_reff_dcomp', 'cld_reff_dcomp_1', 'cld_reff_dcomp_2', 'cld_reff_dcomp_3',
           'cld_opd_dcomp', 'cld_opd_dcomp_1', 'cld_opd_dcomp_2', 'cld_opd_dcomp_3', 'cld_cwp_dcomp', 'iwc_dcomp',
           'lwc_dcomp', 'cld_emiss_acha', 'conv_cloud_fraction', 'cloud_type', 'cloud_phase', 'cloud_mask']

ds_types = {ds_list[i]: 'f4' for i in range(23)}
ds_types.update({ds_list[i+23]: 'i1' for i in range(3)})
ds_fill = {ds_list[i]: -32768 for i in range(23)}
ds_fill.update({ds_list[i+23]: -128 for i in range(3)})
ds_range = {ds_list[i]: 'actual_range' for i in range(23)}
ds_range.update({ds_list[i]: None for i in range(3)})

ds_types.update(l1b_ds_types)
ds_fill.update(l1b_ds_fill)
ds_range.update(l1b_ds_range)

ds_types.update({'temp_3_9um_nom': 'f4'})
ds_types.update({'cloud_fraction': 'f4'})
ds_fill.update({'temp_3_9um_nom': -32767})
ds_fill.update({'cloud_fraction': -32768})
ds_range.update({'temp_3_9um_nom': 'actual_range'})
ds_range.update({'cloud_fraction': 'actual_range'})


def get_fill_attrs(name):
    if name in ds_fill:
        v = ds_fill[name]
        if v is None:
            return None, '_FillValue'
        else:
            return v, None
    else:
        return None, '_FillValue'


class GenericException(Exception):
    def __init__(self, message):
        self.message = message


class EarlyStop:
    def __init__(self, window_length=3, patience=5):
        self.patience = patience
        self.min = np.finfo(np.single).max
        self.cnt = 0
        self.cnt_wait = 0
        self.window = np.zeros(window_length, dtype=np.single)
        self.window.fill(np.nan)

    def check_stop(self, value):
        self.window[:-1] = self.window[1:]
        self.window[-1] = value

        if np.any(np.isnan(self.window)):
            return False

        ave = np.mean(self.window)
        if ave < self.min:
            self.min = ave
            self.cnt_wait = 0
            return False
        else:
            self.cnt_wait += 1

        if self.cnt_wait > self.patience:
            return True
        else:
            return False


def get_time_tuple_utc(timestamp):
    dt_obj = datetime.datetime.fromtimestamp(timestamp, timezone.utc)
    return dt_obj, dt_obj.timetuple()


def get_datetime_obj(dt_str, format_code='%Y-%m-%d_%H:%M'):
    dto = datetime.datetime.strptime(dt_str, format_code).replace(tzinfo=timezone.utc)
    return dto


def get_timestamp(dt_str, format_code='%Y-%m-%d_%H:%M'):
    dto = datetime.datetime.strptime(dt_str, format_code).replace(tzinfo=timezone.utc)
    ts = dto.timestamp()
    return ts


def add_time_range_to_filename(pathname, tstart, tend=None):
    filename = os.path.split(pathname)[1]
    w_o_ext, ext = os.path.splitext(filename)

    dt_obj, _ = get_time_tuple_utc(tstart)
    str_start = dt_obj.strftime('%Y%m%d%H')
    filename = w_o_ext+'_'+str_start

    if tend is not None:
        dt_obj, _ = get_time_tuple_utc(tend)
        str_end = dt_obj.strftime('%Y%m%d%H')
        filename = filename+'_'+str_end
    filename = filename+ext

    path = os.path.split(pathname)[0]
    path = path+'/'+filename
    return path


def haversine_np(lon1, lat1, lon2, lat2):
    """
    Calculate the great circle distance between two points
    on the earth (specified in decimal degrees)

    (lon1, lat1) must be broadcastable with (lon2, lat2).

    """
    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])

    dlon = lon2 - lon1
    dlat = lat2 - lat1

    a = np.sin(dlat/2.0)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2.0)**2

    c = 2 * np.arcsin(np.sqrt(a))
    km = 6367 * c
    return km


def bin_data_by(a, b, bin_ranges):
    nbins = len(bin_ranges)
    binned_data = []

    for i in range(nbins):
        rng = bin_ranges[i]
        idxs = (b >= rng[0]) & (b < rng[1])
        binned_data.append(a[idxs])

    return binned_data


def bin_data_by_edges(a, b, edges):
    nbins = len(edges) - 1
    binned_data = []

    for i in range(nbins):
        idxs = (b >= edges[i]) & (b < edges[i+1])
        binned_data.append(a[idxs])

    return binned_data


def get_bin_ranges(lop, hip, bin_size=100):
    bin_ranges = []
    delp = hip - lop
    nbins = int(delp/bin_size)

    for i in range(nbins):
        rng = [lop + i*bin_size, lop + i*bin_size + bin_size]
        bin_ranges.append(rng)

    return bin_ranges


# t must be monotonic increasing
def get_breaks(t, threshold):
    t_0 = t[0:t.shape[0]-1]
    t_1 = t[1:t.shape[0]]
    d = t_1 - t_0
    idxs = np.nonzero(d > threshold)
    return idxs


# return indexes of ts where value is within ts[i] - threshold < value < ts[i] + threshold
# eventually, if necessary, fully vectorize (numpy) this is possible
# threshold units: seconds
def get_indexes_within_threshold(ts, value, threshold):
    idx_s = []
    t_s = []
    for k, v in enumerate(ts):
        if (ts[k] - threshold) <= value <= (ts[k] + threshold):
            idx_s.append(k)
            t_s.append(v)
    return idx_s, t_s


def pressure_to_altitude(pres, temp, prof_pres, prof_temp, sfc_pres=None, sfc_temp=None, sfc_elev=0):
    if not np.all(np.diff(prof_pres) > 0):
        raise GenericException("target pressure profile must be monotonic increasing")

    if pres < prof_pres[0]:
        raise GenericException("target pressure less than top of pressure profile")

    if temp is None:
        temp = np.interp(pres, prof_pres, prof_temp)

    i_top = np.argmax(np.extract(prof_pres <= pres, prof_pres)) + 1

    pres_s = prof_pres.tolist()
    temp_s = prof_temp.tolist()

    pres_s = [pres] + pres_s[i_top:]
    temp_s = [temp] + temp_s[i_top:]

    if sfc_pres is not None:
        if pres > sfc_pres:  # incoming pressure below surface
            return -999.0

        prof_pres = np.array(pres_s)
        prof_temp = np.array(temp_s)
        i_bot = prof_pres.shape[0] - 1

        if sfc_pres > prof_pres[i_bot]:  # surface below profile bottom
            pres_s = pres_s + [sfc_pres]
            temp_s = temp_s + [sfc_temp]
        else:
            idx = np.argmax(np.extract(prof_pres < sfc_pres, prof_pres))
            if sfc_temp is None:
                sfc_temp = np.interp(sfc_pres, prof_pres, prof_temp)
            pres_s = prof_pres.tolist()
            temp_s = prof_temp.tolist()
            pres_s = pres_s[0:idx+1] + [sfc_pres]
            temp_s = temp_s[0:idx+1] + [sfc_temp]

    prof_pres = np.array(pres_s)
    prof_temp = np.array(temp_s)

    prof_pres = prof_pres[::-1]
    prof_temp = prof_temp[::-1]

    prof_pres = prof_pres * units.hectopascal
    prof_temp = prof_temp * units.kelvin
    sfc_elev = sfc_elev * units.meter

    z = thickness_hydrostatic(prof_pres, prof_temp) + sfc_elev

    return z.magnitude


# http://fourier.eng.hmc.edu/e176/lectures/NM/node25.html
def minimize_quadratic(xa, xb, xc, ya, yb, yc):
    x_m = xb + 0.5*(((ya-yb)*(xc-xb)*(xc-xb) - (yc-yb)*(xb-xa)*(xb-xa)) / ((ya-yb)*(xc-xb) + (yc-yb)*(xb-xa)))
    return x_m


# Return index of nda closest to value. nda must be 1d
def value_to_index(nda, value):
    diff = np.abs(nda - value)
    idx = np.argmin(diff)
    return idx


def find_bin_index(nda, value_s):
    idxs = np.arange(nda.shape[0])
    iL_s = np.zeros(value_s.shape[0])
    iL_s[:,] = -1

    for k, v in enumerate(value_s):
        above = v >= nda
        if not above.any():
            continue

        below = v < nda
        if not below.any():
            continue

        iL = idxs[above].max()
        iL_s[k] = iL

    return iL_s.astype(np.int32)


# array solzen must be degrees, missing values must NaN. For small roughly 50x50km regions only
def is_day(solzen, test_angle=80.0):
    solzen = solzen.flatten()
    solzen = solzen[np.invert(np.isnan(solzen))]
    if len(solzen) == 0 or np.sum(solzen <= test_angle) < len(solzen):
        return False
    else:
        return True


# array solzen must be degrees, missing values must NaN. For small roughly 50x50km regions only
def is_night(solzen, test_angle=100.0):
    solzen = solzen.flatten()
    solzen = solzen[np.invert(np.isnan(solzen))]
    if len(solzen) == 0 or np.sum(solzen >= test_angle) < len(solzen):
        return False
    else:
        return True


def check_oblique(satzen, test_angle=70.0):
    satzen = satzen.flatten()
    satzen = satzen[np.invert(np.isnan(satzen))]
    if len(satzen) == 0 or np.sum(satzen <= test_angle) < len(satzen):
        return False
    else:
        return True


def get_median(tile_2d):
    tile = tile_2d.flatten()
    return np.nanmedian(tile)


def get_grid_values_all(h5f, grid_name, scale_factor_name='scale_factor', add_offset_name='add_offset',
                        fill_value_name='_FillValue', valid_range_name='valid_range', actual_range_name='actual_range', fill_value=None, stride=None):
    hfds = h5f[grid_name]
    attrs = hfds.attrs

    if attrs is None:
        raise GenericException('No attributes object for: '+grid_name)

    if stride is None:
        grd_vals = hfds[:,]
    else:
        grd_vals = hfds[::stride, ::stride]

    if fill_value_name is not None:
        attr = attrs.get(fill_value_name)
        if attr is not None:
            if np.isscalar(attr):
                fill_value = attr
            else:
                fill_value = attr[0]
            grd_vals = np.where(grd_vals == fill_value, np.nan, grd_vals)
    elif fill_value is not None:
        grd_vals = np.where(grd_vals == fill_value, np.nan, grd_vals)

    if valid_range_name is not None:
        attr = attrs.get(valid_range_name)
        if attr is not None:
            low = attr[0]
            high = attr[1]
            grd_vals = np.where(grd_vals < low, np.nan, grd_vals)
            grd_vals = np.where(grd_vals > high, np.nan, grd_vals)

    if scale_factor_name is not None:
        attr = attrs.get(scale_factor_name)
        if attr is not None:
            if np.isscalar(attr):
                scale_factor = attr
            else:
                scale_factor = attr[0]
            grd_vals = grd_vals * scale_factor

    if add_offset_name is not None:
        attr = attrs.get(add_offset_name)
        if attr is not None:
            if np.isscalar(attr):
                add_offset = attr
            else:
                add_offset = attr[0]
            grd_vals = grd_vals + add_offset

    if actual_range_name is not None:
        attr = attrs.get(actual_range_name)
        if attr is not None:
            low = attr[0]
            high = attr[1]
            grd_vals = np.where(grd_vals < low, np.nan, grd_vals)
            grd_vals = np.where(grd_vals > high, np.nan, grd_vals)

    return grd_vals


# dt_str_0: start datetime string in format YYYY-MM-DD_HH:MM
# dt_str_1: stop datetime string, if not None num_steps is computed
# format_code: default '%Y-%m-%d_%H:%M'
# num_steps with increment of days, hours, minutes or seconds
# dt_str_1 and num_steps cannot both be None
# return num_steps+1 lists of datetime strings and timestamps (edges of a numpy histogram)
def make_times(dt_str_0, dt_str_1=None, format_code='%Y-%m-%d_%H:%M', num_steps=None, days=None, hours=None, minutes=None, seconds=None):
    if days is not None:
        inc = 86400*days
    elif hours is not None:
        inc = 3600*hours
    elif minutes is not None:
        inc = 60*minutes
    else:
        inc = seconds

    dt_obj_s = []
    ts_s = []
    dto_0 = datetime.datetime.strptime(dt_str_0, format_code).replace(tzinfo=timezone.utc)
    ts_0 = dto_0.timestamp()

    if dt_str_1 is not None:
        dto_1 = datetime.datetime.strptime(dt_str_1, format_code).replace(tzinfo=timezone.utc)
        ts_1 = dto_1.timestamp()
        num_steps = int((ts_1 - ts_0)/inc)

    dt_obj_s.append(dto_0)
    ts_s.append(ts_0)
    dto_last = dto_0
    for k in range(num_steps):
        dt_obj = dto_last + datetime.timedelta(seconds=inc)
        dt_obj_s.append(dt_obj)
        ts_s.append(dt_obj.timestamp())
        dto_last = dt_obj

    return dt_obj_s, ts_s


def make_histogram(values, edges):
    h = np.histogram(values, bins=edges)
    return h


def normalize(data, param, mean_std_dict, copy=True):

    if mean_std_dict.get(param) is None:
        return data

    if copy:
        data = data.copy()

    shape = data.shape
    data = data.flatten()

    mean, std, lo, hi = mean_std_dict.get(param)
    data -= mean
    data /= std

    not_valid = np.isnan(data)
    data[not_valid] = 0

    data = np.reshape(data, shape)

    return data


def add_noise(data, noise_scale=0.01, seed=None, copy=True):
    if copy:
        data = data.copy()

    shape = data.shape
    data = data.flatten()

    if seed is not None:
        np.random.seed(seed)
    rnd = np.random.normal(loc=0, scale=noise_scale, size=data.size)
    data += rnd

    data = np.reshape(data, shape)

    return data


def denormalize(data, param, mean_std_dict, copy=True):
    if copy:
        data = data.copy()

    if mean_std_dict.get(param) is None:
        return data

    shape = data.shape
    data = data.flatten()

    mean, std, lo, hi = mean_std_dict.get(param)
    data *= std
    data += mean

    data = np.reshape(data, shape)

    return data


def scale(data, param, mean_std_dict, copy=True):
    if copy:
        data = data.copy()

    if mean_std_dict.get(param) is None:
        return data

    shape = data.shape
    data = data.flatten()

    _, _, lo, hi = mean_std_dict.get(param)

    data -= lo
    data /= (hi - lo)

    not_valid = np.isnan(data)
    data[not_valid] = 0

    data = np.reshape(data, shape)

    return data


f = open(ancillary_path+'geos_crs_goes16_FD.pkl', 'rb')
geos_goes16_fd = pickle.load(f)
f.close()

f = open(ancillary_path+'geos_crs_goes16_CONUS.pkl', 'rb')
geos_goes16_conus = pickle.load(f)
f.close()

f = open(ancillary_path+'geos_crs_H08_FD.pkl', 'rb')
geos_h08_fd = pickle.load(f)
f.close()


def get_cartopy_crs(satellite, domain):
    if satellite == 'GOES16':
        if domain == 'FD':
            geos = geos_goes16_fd
            xlen = 5424
            xmin = -5433893.0
            xmax = 5433893.0
            ylen = 5424
            ymin = -5433893.0
            ymax = 5433893.0
        elif domain == 'CONUS':
            geos = geos_goes16_conus
            xlen = 2500
            xmin = -3626269.5
            xmax = 1381770.0
            ylen = 1500
            ymin = 1584175.9
            ymax = 4588198.0
    elif satellite == 'H08':
        geos = geos_h08_fd
        xlen = 5500
        xmin = -5498.99990119
        xmax = 5498.99990119
        ylen = 5500
        ymin = -5498.99990119
        ymax = 5498.99990119

    return geos, xlen, xmin, xmax, ylen, ymin, ymax


def get_grid_values(h5f, grid_name, j_c, i_c, half_width, num_j=None, num_i=None, scale_factor_name='scale_factor', add_offset_name='add_offset',
                    fill_value_name='_FillValue', valid_range_name='valid_range', actual_range_name='actual_range', fill_value=None):
    hfds = h5f[grid_name]
    attrs = hfds.attrs
    if attrs is None:
        raise GenericException('No attributes object for: '+grid_name)

    ylen, xlen = hfds.shape

    if half_width is not None:
        j_l = j_c-half_width
        i_l = i_c-half_width
        if j_l < 0 or i_l < 0:
            return None

        j_r = j_c+half_width+1
        i_r = i_c+half_width+1
        if j_r >= ylen or i_r >= xlen:
            return None
    else:
        j_l = j_c
        j_r = j_c + num_j + 1
        i_l = i_c
        i_r = i_c + num_i + 1

    grd_vals = hfds[j_l:j_r, i_l:i_r]

    if fill_value_name is not None:
        attr = attrs.get(fill_value_name)
        if attr is not None:
            if np.isscalar(attr):
                fill_value = attr
            else:
                fill_value = attr[0]
            grd_vals = np.where(grd_vals == fill_value, np.nan, grd_vals)
    elif fill_value is not None:
        grd_vals = np.where(grd_vals == fill_value, np.nan, grd_vals)

    if valid_range_name is not None:
        attr = attrs.get(valid_range_name)
        if attr is not None:
            low = attr[0]
            high = attr[1]
            grd_vals = np.where(grd_vals < low, np.nan, grd_vals)
            grd_vals = np.where(grd_vals > high, np.nan, grd_vals)

    if scale_factor_name is not None:
        attr = attrs.get(scale_factor_name)
        if attr is not None:
            if np.isscalar(attr):
                scale_factor = attr
            else:
                scale_factor = attr[0]
            grd_vals = grd_vals * scale_factor

    if add_offset_name is not None:
        attr = attrs.get(add_offset_name)
        if attr is not None:
            if np.isscalar(attr):
                add_offset = attr
            else:
                add_offset = attr[0]
            grd_vals = grd_vals + add_offset

    if actual_range_name is not None:
        attr = attrs.get(actual_range_name)
        if attr is not None:
            low = attr[0]
            high = attr[1]
            grd_vals = np.where(grd_vals < low, np.nan, grd_vals)
            grd_vals = np.where(grd_vals > high, np.nan, grd_vals)

    return grd_vals


def concat_dict_s(t_dct_0, t_dct_1):
    keys_0 = list(t_dct_0.keys())
    nda_0 = np.array(keys_0)

    keys_1 = list(t_dct_1.keys())
    nda_1 = np.array(keys_1)

    comm_keys, comm0, comm1 = np.intersect1d(nda_0, nda_1, return_indices=True)

    comm_keys = comm_keys.tolist()

    for key in comm_keys:
        t_dct_1.pop(key)
    t_dct_0.update(t_dct_1)

    return t_dct_0


# Example GOES file to retrieve GEOS parameters in MetPy form (CONUS)
exmp_file_conus = '/Users/tomrink/data/OR_ABI-L1b-RadC-M6C14_G16_s20193140811215_e20193140813588_c20193140814070.nc'
# Full Disk
exmp_file_fd = '/Users/tomrink/data/OR_ABI-L1b-RadF-M6C16_G16_s20212521800223_e20212521809542_c20212521809596.nc'

# keep for reference
# if domain == 'CONUS':
#     exmpl_ds = xr.open_dataset(exmp_file_conus)
# elif domain == 'FD':
#     exmpl_ds = xr.open_dataset(exmp_file_fd)
# mdat = exmpl_ds.metpy.parse_cf('Rad')
# geos = mdat.metpy.cartopy_crs
# xlen = mdat.x.values.size
# ylen = mdat.y.values.size
# exmpl_ds.close()

# Taiwan domain:
# lon, lat = 120.955098, 23.834310
# elem, line = (1789, 1505)
# # UR from Taiwan
# lon, lat = 135.0, 35.0
# elem_ur, line_ur = (2499, 995)
taiwan_i0 = 1079
taiwan_j0 = 995
taiwan_lenx = 1420
taiwan_leny = 1020
# geos.transform_point(135.0, 35.0, ccrs.PlateCarree(), False)
# geos.transform_point(106.61, 13.97, ccrs.PlateCarree(), False)
taiwain_extent = [-3342, -502, 1470, 3510]  # GEOS coordinates, not line, elem


# ------------ This code will not be needed when we implement a Fully Convolutional CNN -----------------------------------
# Generate and return tiles of name_list parameters
def make_for_full_domain_predict(h5f, name_list=None, satellite='GOES16', domain='FD', res_fac=1):
    w_x = 16
    w_y = 16
    i_0 = 0
    j_0 = 0
    s_x = int(w_x / res_fac)
    s_y = int(w_y / res_fac)

    geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
    if satellite == 'H08':
        xlen = taiwan_lenx
        ylen = taiwan_leny
        i_0 = taiwan_i0
        j_0 = taiwan_j0

    grd_dct = {name: None for name in name_list}

    cnt_a = 0
    for ds_name in name_list:
        fill_value, fill_value_name = get_fill_attrs(ds_name)
        gvals = get_grid_values(h5f, ds_name, j_0, i_0, None, num_j=ylen, num_i=xlen, fill_value_name=fill_value_name, fill_value=fill_value)
        if gvals is not None:
            grd_dct[ds_name] = gvals
            cnt_a += 1

    if cnt_a > 0 and cnt_a != len(name_list):
        raise GenericException('weirdness')

    grd_dct_n = {name: [] for name in name_list}

    n_x = int(xlen/s_x) - 1
    n_y = int(ylen/s_y) - 1

    r_x = xlen - (n_x * s_x)
    x_d = 0 if r_x >= w_x else int((w_x - r_x)/s_x)
    n_x -= x_d

    r_y = ylen - (n_y * s_y)
    y_d = 0 if r_y >= w_y else int((w_y - r_y)/s_y)
    n_y -= y_d

    ll = [j_0 + j*s_y for j in range(n_y)]
    cc = [i_0 + i*s_x for i in range(n_x)]

    for ds_name in name_list:
        for j in range(n_y):
            j_ul = j * s_y
            j_ul_b = j_ul + w_y
            for i in range(n_x):
                i_ul = i * s_x
                i_ul_b = i_ul + w_x
                grd_dct_n[ds_name].append(grd_dct[ds_name][j_ul:j_ul_b, i_ul:i_ul_b])

    grd_dct = {name: None for name in name_list}
    for ds_name in name_list:
        grd_dct[ds_name] = np.stack(grd_dct_n[ds_name])

    return grd_dct, ll, cc


# rho_water = 1.
# rho_ice = 0.917

rho_water = 1000000.0  # g m^-3
rho_ice = 917000.0  # g m^-3

# real(kind=real4), parameter:: Rho_Water = 1.0    !g / m ^ 3
# real(kind=real4), parameter:: Rho_Ice = 0.917    !g / m ^ 3
#
# !--- compute
# cloud
# water
# path
# if (Iphase == 0) then
# Cwp_Dcomp(Elem_Idx, Line_Idx) = 0.55 * Tau * Reff * Rho_Water
# Lwp_Dcomp(Elem_Idx, Line_Idx) = 0.55 * Tau * Reff * Rho_Water
# else
# Cwp_Dcomp(Elem_Idx, Line_Idx) = 0.667 * Tau * Reff * Rho_Ice
# Iwp_Dcomp(Elem_Idx, Line_Idx) = 0.667 * Tau * Reff * Rho_Ice
# endif


def compute_lwc_iwc(iphase, reff, opd, geo_dz):
    xy_shape = iphase.shape

    iphase = iphase.flatten()
    keep_0 = np.invert(np.isnan(iphase))

    reff = reff.flatten()
    keep_1 = np.invert(np.isnan(reff))

    opd = opd.flatten()
    keep_2 = np.invert(np.isnan(opd))

    geo_dz = geo_dz.flatten()
    keep_3 = np.logical_and(np.invert(np.isnan(geo_dz)), geo_dz > 1.0)

    keep = keep_0 & keep_1 & keep_2 & keep_3

    lwp_dcomp = np.zeros(reff.shape[0])
    iwp_dcomp = np.zeros(reff.shape[0])
    lwp_dcomp[:] = np.nan
    iwp_dcomp[:] = np.nan

    ice = iphase == 1 & keep
    no_ice = iphase != 1 & keep

    # compute ice/liquid water path, g m-2
    reff *= 1.0e-06  # convert microns to meters

    iwp_dcomp[ice] = 0.667 * opd[ice] * rho_ice * reff[ice]
    lwp_dcomp[no_ice] = 0.55 * opd[no_ice] * rho_water * reff[no_ice]

    iwp_dcomp /= geo_dz
    lwp_dcomp /= geo_dz

    lwp_dcomp = lwp_dcomp.reshape(xy_shape)
    iwp_dcomp = iwp_dcomp.reshape(xy_shape)

    return lwp_dcomp, iwp_dcomp


def make_for_full_domain_predict_viirs_clavrx(h5f, name_list=None, res_fac=1, day_night='DAY', use_dnb=False):
    w_x = 16
    w_y = 16
    i_0 = 0
    j_0 = 0
    s_x = int(w_x / res_fac)
    s_y = int(w_y / res_fac)

    ylen = h5f['scan_lines_along_track_direction'].shape[0]
    xlen = h5f['pixel_elements_along_scan_direction'].shape[0]

    use_nl_comp = False
    if (day_night == 'NIGHT' or day_night == 'AUTO') and use_dnb:
        use_nl_comp = True

    grd_dct = {name: None for name in name_list}

    cnt_a = 0
    for ds_name in name_list:
        name = ds_name

        if use_nl_comp:
            if ds_name == 'cld_reff_dcomp':
                name = 'cld_reff_nlcomp'
            elif ds_name == 'cld_opd_dcomp':
                name = 'cld_opd_nlcomp'

        fill_value, fill_value_name = get_fill_attrs(name)
        gvals = get_grid_values(h5f, name, j_0, i_0, None, num_j=ylen, num_i=xlen, fill_value_name=fill_value_name, fill_value=fill_value)
        if gvals is not None:
            grd_dct[ds_name] = gvals
            cnt_a += 1

    if cnt_a > 0 and cnt_a != len(name_list):
        raise GenericException('weirdness')

    # TODO: need to investigate discrepencies with compute_lwc_iwc
    # if use_nl_comp:
    #     cld_phase = get_grid_values(h5f, 'cloud_phase', j_0, i_0, None, num_j=ylen, num_i=xlen)
    #     cld_dz = get_grid_values(h5f, 'cld_geo_thick', j_0, i_0, None, num_j=ylen, num_i=xlen)
    #     reff = grd_dct['cld_reff_dcomp']
    #     opd = grd_dct['cld_opd_dcomp']
    #
    #     lwc_nlcomp, iwc_nlcomp = compute_lwc_iwc(cld_phase, reff, opd, cld_dz)
    #     grd_dct['iwc_dcomp'] = iwc_nlcomp
    #     grd_dct['lwc_dcomp'] = lwc_nlcomp

    grd_dct_n = {name: [] for name in name_list}

    n_x = int(xlen/s_x) - 1
    n_y = int(ylen/s_y) - 1

    r_x = xlen - (n_x * s_x)
    x_d = 0 if r_x >= w_x else int((w_x - r_x)/s_x)
    n_x -= x_d

    r_y = ylen - (n_y * s_y)
    y_d = 0 if r_y >= w_y else int((w_y - r_y)/s_y)
    n_y -= y_d

    ll = [j_0 + j*s_y for j in range(n_y)]
    cc = [i_0 + i*s_x for i in range(n_x)]

    for ds_name in name_list:
        for j in range(n_y):
            j_ul = j * s_y
            j_ul_b = j_ul + w_y
            for i in range(n_x):
                i_ul = i * s_x
                i_ul_b = i_ul + w_x
                grd_dct_n[ds_name].append(grd_dct[ds_name][j_ul:j_ul_b, i_ul:i_ul_b])

    grd_dct = {name: None for name in name_list}
    for ds_name in name_list:
        grd_dct[ds_name] = np.stack(grd_dct_n[ds_name])

    lats = get_grid_values(h5f, 'latitude', j_0, i_0, None, num_j=ylen, num_i=xlen)
    lons = get_grid_values(h5f, 'longitude', j_0, i_0, None, num_j=ylen, num_i=xlen)

    ll_2d, cc_2d = np.meshgrid(ll, cc, indexing='ij')

    lats = lats[ll_2d, cc_2d]
    lons = lons[ll_2d, cc_2d]

    return grd_dct, ll, cc, lats, lons


def make_for_full_domain_predict2(h5f, satellite='GOES16', domain='FD', res_fac=1):
    w_x = 16
    w_y = 16
    i_0 = 0
    j_0 = 0
    s_x = int(w_x / res_fac)
    s_y = int(w_y / res_fac)

    geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
    if satellite == 'H08':
        xlen = taiwan_lenx
        ylen = taiwan_leny
        i_0 = taiwan_i0
        j_0 = taiwan_j0

    n_x = int(xlen/s_x)
    n_y = int(ylen/s_y)

    solzen = get_grid_values(h5f, 'solar_zenith_angle', j_0, i_0, None, num_j=ylen, num_i=xlen)
    satzen = get_grid_values(h5f, 'sensor_zenith_angle', j_0, i_0, None, num_j=ylen, num_i=xlen)
    solzen = solzen[0:(n_y-1)*s_y:s_y, 0:(n_x-1)*s_x:s_x]
    satzen = satzen[0:(n_y-1)*s_y:s_y, 0:(n_x-1)*s_x:s_x]

    return solzen, satzen
# -------------------------------------------------------------------------------------------


def prepare_evaluate(h5f, name_list, satellite='GOES16', domain='FD', res_fac=1, offset=0):
    w_x = 16
    w_y = 16
    i_0 = 0
    j_0 = 0
    s_x = int(w_x / res_fac)
    s_y = int(w_y / res_fac)

    geos, xlen, xmin, xmax, ylen, ymin, ymax = get_cartopy_crs(satellite, domain)
    if satellite == 'H08':
        xlen = taiwan_lenx
        ylen = taiwan_leny
        i_0 = taiwan_i0
        j_0 = taiwan_j0

    n_x = int(xlen/s_x) - 1
    n_y = int(ylen/s_y) - 1

    r_x = xlen - (n_x * s_x)
    x_d = 0 if r_x >= w_x else int((w_x - r_x)/s_x)
    n_x -= x_d

    r_y = ylen - (n_y * s_y)
    y_d = 0 if r_y >= w_y else int((w_y - r_y)/s_y)
    n_y -= y_d

    ll = [(offset+j_0) + j*s_y for j in range(n_y)]
    cc = [(offset+i_0) + i*s_x for i in range(n_x)]

    grd_dct_n = {name: [] for name in name_list}

    cnt_a = 0
    for ds_name in name_list:
        fill_value, fill_value_name = get_fill_attrs(ds_name)
        gvals = get_grid_values(h5f, ds_name, j_0, i_0, None, num_j=ylen, num_i=xlen, fill_value_name=fill_value_name, fill_value=fill_value)
        if gvals is not None:
            grd_dct_n[ds_name] = gvals
            cnt_a += 1

    if cnt_a > 0 and cnt_a != len(name_list):
        raise GenericException('weirdness')

    solzen = get_grid_values(h5f, 'solar_zenith_angle', j_0, i_0, None, num_j=ylen, num_i=xlen)
    satzen = get_grid_values(h5f, 'sensor_zenith_angle', j_0, i_0, None, num_j=ylen, num_i=xlen)
    solzen = solzen[0:(n_y-1)*s_y:s_y, 0:(n_x-1)*s_x:s_x]
    satzen = satzen[0:(n_y-1)*s_y:s_y, 0:(n_x-1)*s_x:s_x]

    grd_dct = {name: None for name in name_list}
    for ds_name in name_list:
        grd_dct[ds_name] = np.stack(grd_dct_n[ds_name])

    return grd_dct, solzen, satzen, ll, cc


flt_level_ranges_str = {k: None for k in range(5)}
flt_level_ranges_str[0] = '0_2000'
flt_level_ranges_str[1] = '2000_4000'
flt_level_ranges_str[2] = '4000_6000'
flt_level_ranges_str[3] = '6000_8000'
flt_level_ranges_str[4] = '8000_15000'

# flt_level_ranges_str = {k: None for k in range(1)}
# flt_level_ranges_str[0] = 'column'


def get_cf_nav_parameters(satellite='GOES16', domain='FD'):
    param_dct = None

    if satellite == 'H08':  # We presently only have FD
        param_dct = {'semi_major_axis': 6378.137,
                     'semi_minor_axis': 6356.7523,
                     'perspective_point_height': 35785.863,
                     'latitude_of_projection_origin': 0.0,
                     'longitude_of_projection_origin': 140.7,
                     'inverse_flattening': 298.257,