from icing.pireps import pirep_icing
from deeplearning.amv_raob import get_images
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 = {}

pirep_file = '/home/rink/data/pireps/pireps_2019010000_2019063023.csv'

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())

hist = np.zeros(20)


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

    grd_vals = hfds[j_c-half_width:j_c+half_width+1, i_c-half_width:i_c+half_width+1]
    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():

    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

    for idx, time in enumerate(time_keys):
        if (idx % 4) != 0:
            continue
        print(100.0*(idx/len(time_keys)))
        reports = ice_dict[time]
        for tup in reports:
            lat, lon, fl, rpt_str = tup
            lat_s[0] = lat
            lon_s[0] = lon
            cc, ll = nav.earth_to_lc_s(lon_s, lat_s)

            try:
                clvr_ds = get_clavrx_datasource(time)
            except Exception:
                continue
            clvr_file = clvr_ds.get_file(time)[0]
            if clvr_file is None:
                continue
            if clvr_file != last_clvr_file:
                try:
                    h5f = h5py.File(clvr_file, 'r')
                    cnt += 1
                except Exception:
                    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

            gvals = get_grid_values(h5f, 'temp_10_4um_nom', ll[0], cc[0], 20)