from icing.pireps import pirep_icing
import numpy as np
import pickle
import os
from util.util import get_time_tuple_utc
from aeolus.datasource import CLAVRx
from util.geos_nav import GEOSNavigation
import h5py

clavrx_dir = '/apollo/cloud/scratch/ICING/'
dir_fmt = '%Y_%m_%d_%j'
# dir_list = [f.path for f in os.scandir('.') if f.is_dir()]
ds_dct = {}
ice_dict = None
no_ice_dict = None
time_keys = None
pirep_file = '/home/rink/data/pireps/pireps_2019010000_2019063023.csv'

ds_list = ['temp_8_5um_nom', 'temp_10_4um_nom', 'temp_11_0um_nom', 'temp_13_3um_nom', 'cld_height_acha',
           'supercooled_cloud_fraction', 'cld_temp_acha', 'cld_opd_acha', 'cloud_phase', 'solar_zenith_angle']
ds_grd_dct = {name: [] for name in ds_list}


def setup():
    ice_dict, no_ice_dict = pirep_icing(pirep_file)
    print('num obs: ice, no ice', len(ice_dict), len(no_ice_dict))
    time_keys = list(ice_dict.keys())
    return time_keys, ice_dict, no_ice_dict


def get_clavrx_datasource(timestamp):
    dt_obj, time_tup = get_time_tuple_utc(timestamp)
    date_dir_str = dt_obj.strftime(dir_fmt)
    ds = ds_dct.get(date_dir_str)
    if ds is None:
        ds = CLAVRx(clavrx_dir + date_dir_str + '/')
        ds_dct[date_dir_str] = ds
    return ds


def get_grid_values(h5f, grid_name, j_c, i_c, half_width, scale_factor_name='scale_factor', add_offset_name='add_offset'):
    hfds = h5f[grid_name]
    attrs = hfds.attrs
    ylen, xlen = hfds.shape
    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

    grd_vals = hfds[j_l:j_r, i_l:i_r]
    grd_vals = np.where(grd_vals == -999, np.nan, grd_vals)
    grd_vals = np.where(grd_vals == -32768, np.nan, grd_vals)

    if attrs is None:
        return grd_vals

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

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

    return grd_vals


def run(time_keys, pirep_dct, outfile=None):

    nav = GEOSNavigation(sub_lon=-75.0, CFAC=5.6E-05, COFF=-0.101332, LFAC=-5.6E-05, LOFF=0.128212, num_elems=2500, num_lines=1500)

    lon_s = np.zeros(1)
    lat_s = np.zeros(1)
    last_clvr_file = None
    last_h5f = None

    cnt = 0
    miss_a = 0
    miss_b = 0
    for idx, time in enumerate(time_keys):
        reports = pirep_dct[time]
        for tup in reports:
            lat, lon, fl, rpt_str = tup
            lat_s[0] = lat
            lon_s[0] = lon

            try:
                clvr_ds = get_clavrx_datasource(time)
            except Exception:
                miss_a += 1
                continue

            clvr_file = clvr_ds.get_file(time)[0]
            if clvr_file is None:
                miss_b += 1
                continue

            if clvr_file != last_clvr_file:
                try:
                    h5f = h5py.File(clvr_file, 'r')
                except Exception:
                    if h5f is not None:
                        h5f.close()
                    print('Problem with file: ', clvr_file)
                    continue
                if last_h5f is not None:
                    last_h5f.close()
                last_h5f = h5f
                last_clvr_file = clvr_file
            else:
                h5f = last_h5f

            cc, ll = nav.earth_to_lc_s(lon_s, lat_s)

            for didx, ds_name in enumerate(ds_list):
                gvals = get_grid_values(h5f, ds_name, ll[0], cc[0], 20)
                if gvals is not None:
                    ds_grd_dct[ds_name].append(gvals)

            cnt += 1

    print('num images: ', cnt, len(time_keys), miss_a, miss_b)


def create_file(filename, ds_list, ds_types):
    pass