from icing.pireps import pirep_icing
import numpy as np
import pickle
import os
from util.util import get_time_tuple_utc, GenericException
from aeolus.datasource import CLAVRx, GOESL1B
from util.geos_nav import GEOSNavigation
import h5py
import re
import datetime
from datetime import timezone

goes_date_format = '%Y%j%H'
goes16_directory = '/arcdata/goes/grb/goes16'  # /year/date/abi/L1b/RadC
#clavrx_dir = '/apollo/cloud/scratch/ICING/'
clavrx_dir = '/ships19/cloud/scratch/ICING/'
dir_fmt = '%Y_%m_%d_%j'
# dir_list = [f.path for f in os.scandir('.') if f.is_dir()]
ds_dct = {}
goes_ds_dct = {}
ice_dict = None
no_ice_dict = None
time_keys = None
#pirep_file = '/home/rink/data/pireps/pireps_2019010000_2019063023.csv'
pirep_file = '/home/rink/data/pireps/pireps_20180101_20200331.csv'

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_55um_nom', 'refl_0_65um_nom', 'refl_0_86um_nom', 'refl_1_38um_nom',
               'refl_1_60um_nom']
l1b_grd_dct = {name: [] for name in l1b_ds_list}

ds_list = ['cld_height_acha', 'cld_geo_thick', 'cld_press_acha', 'sensor_zenith_angle', 'cloud_type', 'supercooled_prob_acha',
           'supercooled_cloud_fraction', 'cld_temp_acha', 'cld_opd_acha', 'cloud_phase', 'solar_zenith_angle',
           'cloud_mask', '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']
ds_grd_dct = {name: [] for name in ds_list}

a_clvr_file = '/home/rink/data/clavrx/clavrx_OR_ABI-L1b-RadC-M3C01_G16_s20190020002186.level2.nc'


def setup():
    ice_dict, no_ice_dict = pirep_icing(pirep_file)
    return 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_goes_datasource(timestamp):
    dt_obj, time_tup = get_time_tuple_utc(timestamp)

    yr_dir = str(dt_obj.timetuple().tm_year)
    date_dir = dt_obj.strftime(dir_fmt)

    files_path = goes16_directory + '/' + yr_dir + '/' + date_dir + '/abi' + '/L1b' + '/RadC/'
    ds = goes_ds_dct.get(date_dir)
    if ds is None:
        ds = GOESL1B(files_path)
        goes_ds_dct[date_dir] = 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 create_file(filename, data_dct, ds_list, lon_c, lat_c, time_s, fl_alt_s):
    h5f_expl = h5py.File(a_clvr_file, 'r')
    h5f = h5py.File(filename, 'w')

    for ds_name in ds_list:
        data = data_dct[ds_name]
        h5f.create_dataset(ds_name, data=data)

    h5f.create_dataset('longitude', data=lon_c)
    h5f.create_dataset('latitude', data=lat_c)
    h5f.create_dataset('time', data=time_s)
    h5f.create_dataset('icing_altitude', data=fl_alt_s)

    # copy relevant attributes
    for ds_name in ds_list:
        h5f[ds_name].attrs.create('standard_name', data=h5f_expl[ds_name].attrs.get('standard_name'))
        h5f[ds_name].attrs.create('long_name', data=h5f_expl[ds_name].attrs.get('long_name'))
        h5f[ds_name].attrs.create('units', data=h5f_expl[ds_name].attrs.get('units'))

    h5f.close()


def run(pirep_dct, outfile=None, outfile_l1b=None):
    time_keys = list(pirep_dct.keys())

    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

    lon_c = []
    lat_c = []
    time_s = []
    fl_alt_s = []
    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:
                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')
                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)
            if cc[0] < 0:
                continue

            cnt_a = 0
            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_a += 1

            cnt_b = 0
            for didx, ds_name in enumerate(l1b_ds_list):
                gvals = get_grid_values(h5f, ds_name, ll[0], cc[0], 20)
                if gvals is not None:
                    l1b_grd_dct[ds_name].append(gvals)
                    cnt_b += 1

            if cnt_a > 0 and cnt_a != len(ds_list):
                raise GenericException('weirdness')
            if cnt_b > 0 and cnt_b != len(l1b_ds_list):
                raise GenericException('weirdness')

            if cnt_a == len(ds_list) and cnt_b == len(l1b_ds_list):
                lon_c.append(lon_s[0])
                lat_c.append(lat_s[0])
                time_s.append(time)
                fl_alt_s.append(fl)

    data_dct = {}
    for ds_name in ds_list:
        data_dct[ds_name] = np.array(ds_grd_dct[ds_name])
    lon_c = np.array(lon_c)
    lat_c = np.array(lat_c)
    time_s = np.array(time_s)
    fl_alt_s = np.array(fl_alt_s)

    if outfile is not None:
        create_file(outfile, data_dct, ds_list, lon_c, lat_c, time_s)

    data_dct = {}
    for ds_name in l1b_ds_list:
        data_dct[ds_name] = np.array(l1b_grd_dct[ds_name])

    if outfile_l1b is not None:
        create_file(outfile_l1b, data_dct, l1b_ds_list, lon_c, lat_c, time_s, fl_alt_s)


def analyze(ice_dct, no_ice_dct):

    last_file = None
    ice_files = []
    ice_times = []
    for ts in list(ice_dct.keys()):
        try:
            ds = get_goes_datasource(ts)
            goes_file, t_0, _ = ds.get_file(ts)
            if goes_file is not None and goes_file != last_file:
                ice_files.append(goes_file)
                ice_times.append(t_0)
                last_file = goes_file
        except Exception:
            continue

    last_file = None
    no_ice_files = []
    no_ice_times = []
    for ts in list(no_ice_dct.keys()):
        try:
            ds = get_goes_datasource(ts)
            goes_file, t_0, _ = ds.get_file(ts)
            if goes_file is not None and goes_file != last_file:
                no_ice_files.append(goes_file)
                no_ice_times.append(t_0)
                last_file = goes_file
        except Exception:
            continue

    ice_times = np.array(ice_times)
    no_ice_times = np.array(no_ice_times)

    itrsct_vals, comm1, comm2 = np.intersect1d(no_ice_times, ice_times, return_indices=True)

    ice_indexes = np.arange(len(ice_times))

    ucomm2 = np.setxor1d(comm2, ice_indexes)
    np.random.shuffle(ucomm2)
    ucomm2 = ucomm2[0:8000]

    files_comm = []
    for i in comm2:
        files_comm.append(ice_files[i])

    files_extra = []
    times_extra = []
    for i in ucomm2:
        files_extra.append(ice_files[i])
        times_extra.append(ice_times[i])

    files = files_comm + files_extra
    times = itrsct_vals.tolist() + times_extra
    times = np.array(times)

    sidxs = np.argsort(times)
    for i in sidxs:
        filename = os.path.split(files[i])[1]
        so = re.search('_s\\d{11}', filename)
        dt_str = so.group()
        print(dt_str[2:])


def apply_qc_icing_pireps(icing_alt, cld_top_hgt, cld_phase, cld_opd, bt_11um, cld_mask):
    opd_threshold = 6
    closeness = 100.0  # meters
    num_obs = len(icing_alt)
    cld_mask = cld_mask.reshape((num_obs, -1))
    cld_top_hgt = cld_top_hgt.reshape((num_obs, -1))
    cld_phase = cld_phase.reshape((num_obs, -1))
    cld_opd = cld_opd.reshape((num_obs, -1))
    bt_11um = bt_11um.reshape((num_obs, -1))

    keep_0 = np.logical_or(cld_mask == 2, cld_mask == 3)  # cloudy
    keep_1 = np.invert(np.isnan(cld_top_hgt))
    keep_2 = np.invert(np.isnan(bt_11um))
    keep_3 = np.invert(np.isnan(cld_opd))
    keep = keep_0 & keep_1 & keep_2 & keep_3

    keep = np.where(keep, cld_top_hgt[:, ] > icing_alt[:], False)
    keep = np.where(keep,
        np.invert((cld_phase[:, ] == 3) &
            np.logical_and(cld_top_hgt[:, ]+closeness > icing_alt[:], cld_top_hgt[:, ]-closeness < icing_alt[:])),
                 False)

    keep = np.where(keep, (cld_opd[:, ] > opd_threshold) & (cld_phase[:, ] == 3) & (cld_top_hgt[:, ] > icing_alt[:]), False)
    keep = np.where(keep, np.invert((cld_phase[:, ] == 3) & (cld_opd[:,] < 0.1) & (cld_top_hgt[:, ] > icing_alt[:]), False))
    keep = np.where(keep, np.invert(bt_11um[:, ] > 270.0), False)
    keep = np.where(keep, np.invert(bt_11um[:, ] < 228.0), False)

    return keep