import numpy as np
import h5py
from util.util import get_grid_values
import glob
import os
from pathlib import Path

emis_params = ['temp_10_4um_nom', 'temp_11_0um_nom', 'temp_12_0um_nom', 'temp_13_3um_nom', 'temp_3_9um_nom',
               'temp_6_7um_nom']
l2_params = ['cloud_fraction', 'cld_temp_acha', 'cld_press_acha']


def run_all(directory):

    cnt = 0
    for p in os.scandir(directory):
        if not p.is_dir():
            continue
        cnt += 1
        l1b_files = glob.glob(directory+p.name+'/'+'clavrx_snpp_viirs*.uwssec*.nc')

        l1b_tiles = []
        l2_tiles = []

        for idx, l1b_f in enumerate(l1b_files):
            if idx % 8 == 0:
                # w_o_ext, ext = os.path.splitext(l1b_f)
                # l2_f = w_o_ext+'.highres'+ext
                # if not os.path.exists(l2_f):
                #     continue
                print(l1b_f)

                try:
                    l1b_h5f = h5py.File(l1b_f, 'r')
                except:
                    print('cant open file: ', l1b_f)
                    continue

                # try:
                #     l2_h5f = h5py.File(l2_f, 'r')
                # except:
                #     print('cant open file: ', l2_f)
                #     l1b_h5f.close()
                #     continue

                try:
                    run(l1b_h5f, None, l1b_tiles, l2_tiles, factor=1)
                except Exception as e:
                    print(e)
                    l1b_h5f.close()
                    # l2_h5f.close()
                    continue

                l1b_h5f.close()
                # l2_h5f.close()

        if len(l1b_tiles) == 0 or len(l2_tiles) == 0:
            continue

        l1b_nda = np.stack(l1b_tiles)
        l2_nda = np.stack(l2_tiles)

        np.save('/data/Personal/rink/viirs_clavrx/l1b_'+str(cnt), l1b_nda)
        np.save('/data/Personal/rink/viirs_clavrx/l2_'+str(cnt), l2_nda)


def run(l1b_h5f, l2_h5f, l1b_tiles, l2_tiles, factor=2):
    if l2_h5f is None:
        l2_h5f = l1b_h5f

    l1b_param_name = emis_params[0]
    l2_param_name = l2_params[0]

    mod_tile_width = 64
    img_tile_width = mod_tile_width * factor

    mod_num_lines = l1b_h5f[l1b_param_name].shape[0]
    mod_num_pixels = l1b_h5f[l1b_param_name].shape[1]

    img_num_lines = l2_h5f[l2_param_name].shape[0]
    img_num_pixels = l2_h5f[l2_param_name].shape[1]

    mod_num_y_tiles = int(mod_num_lines / mod_tile_width)
    mod_num_x_tiles = int(mod_num_pixels / mod_tile_width)

    l1b_grd_s = []
    l2_grd_s = []

    for param in emis_params:
        l1b_grd_s.append(get_grid_values(l1b_h5f, param, 0, 0, None, mod_num_lines, mod_num_pixels, range_name=None))

    for param in l2_params:
        l2_grd_s.append(get_grid_values(l2_h5f, param, 0, 0, None, img_num_lines, img_num_pixels, range_name=None))

    mod_data = np.stack(l1b_grd_s)
    img_data = np.stack(l2_grd_s)

    num_cntr_tiles = 10
    i_c = int(mod_num_pixels / num_cntr_tiles)  # center
    j_skip = int(mod_num_y_tiles / num_cntr_tiles) * mod_tile_width
    for k in range(num_cntr_tiles):
        j_c = k * j_skip
        j_m = j_c
        j_i = j_m * factor

        i_m = i_c
        i_i = i_m * factor

        nda = mod_data[:, j_m:j_m + mod_tile_width, i_m:i_m + mod_tile_width]
        l1b_tiles.append(nda)

        nda = img_data[:, j_i:j_i + img_tile_width, i_i:i_i + img_tile_width]
        l2_tiles.append(nda)

    # for j in range(mod_num_y_tiles):
    #     j_m = j * mod_tile_width
    #     j_i = j_m * 2
    #     for i in range(mod_num_x_tiles):
    #         i_m = i * mod_tile_width
    #         i_i = i_m * 2
    #
    #         nda = mod_data[j_m:j_m+mod_tile_width, i_m:i_m+mod_tile_width]
    #         mod_tiles.append(nda)
    #         nda = img_data[j_i:j_i+img_tile_width, i_i:i_i+img_tile_width]
    #         img_tiles.append(nda)


# def hack():
#     l1b_nda = np.load('/Users/tomrink/data/viirs_clavrx/l1b_all.npy')
#     l2_nda = np.load('/Users/tomrink/data/viirs_clavrx/l2_all.npy')
#
#     num_samples = l1b_nda.shape[0]
#     idxs = np.arange(num_samples)
#     vld_idxs = idxs[::10]
#     trn_idxs = idxs[np.in1d(idxs, vld_idxs, invert=True)]
#
#     l1b_trn = l1b_nda[trn_idxs]
#     l2_trn = l2_nda[trn_idxs]
#     l1b_vld = l1b_nda[vld_idxs]
#     l2_vld = l2_nda[vld_idxs]
#
#     print(l1b_trn.shape, l2_trn.shape, l1b_vld.shape, l2_vld.shape)
#
#     np.save('/Users/tomrink/data/viirs_clavrx/l1b_all_trn.npy', l1b_trn)
#     np.save('/Users/tomrink/data/viirs_clavrx/l1b_all_vld.npy', l1b_vld)
#     np.save('/Users/tomrink/data/viirs_clavrx/l2_all_trn.npy', l2_trn)
#     np.save('/Users/tomrink/data/viirs_clavrx/l2_all_vld.npy', l2_vld)

# def run_mean_std(directory):
#
#     data_dct = {name: [] for name in mod_res_params}
#     mean_dct = {name: 0 for name in mod_res_params}
#     std_dct = {name: 0 for name in mod_res_params}
#
#     for p in os.scandir(directory):
#         if not p.is_dir():
#             continue
#         mod_files = glob.glob(directory+p.name+'/'+'VNP02MOD*.uwssec.nc')
#
#         for idx, mfile in enumerate(mod_files):
#             if idx % 8 == 0:
#                 h5f = h5py.File(mfile, 'r')
#                 for param in mod_res_params:
#                     name = 'observation_data/'+param
#                     gvals = get_grid_values_all(h5f, name, range_name=None, stride=10)
#                     data_dct[param].append(gvals.flatten())
#                 print(mfile)
#                 h5f.close()
#
#     for param in mod_res_params:
#         data = data_dct[param]
#         data = np.concatenate(data)
#
#         mean_dct[param] = np.nanmean(data)
#         std_dct[param] = np.nanstd(data)