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

mod_res_params = ['M07', 'M08', 'M10', 'M12', 'M13', 'M14', 'M15', 'M16']
img_res_params = ['M07_highres', 'M08_highres', 'M10_highres', 'M12_highres', 'M13_highres', 'M14_highres', 'M15_highres', 'M16_highres']


def run_all(directory):

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

        mod_tiles = []
        img_tiles = []

        for idx, mfile in enumerate(mod_files):
            if idx % 8 == 0:
                w_o_ext, ext = os.path.splitext(mfile)
                ifile = w_o_ext+'.highres'+ext
                if not os.path.exists(ifile):
                    continue
                print(mfile)
                run(mfile, ifile, mod_tiles, img_tiles)

        if len(mod_tiles) == 0:
            continue

        mod_nda = np.stack(mod_tiles)
        img_nda = np.stack(img_tiles)

        np.save('/data/Personal/rink/viirs/mod_res_'+str(cnt), mod_nda)
        np.save('/data/Personal/rink/viirs/img_res_'+str(cnt), img_nda)


def run(mod_res_filename, img_res_filename, mod_tiles, img_tiles):
    mod_h5f = h5py.File(mod_res_filename, 'r')
    img_h5f = h5py.File(img_res_filename, 'r')

    mod_tile_width = 64
    img_tile_width = mod_tile_width * 2

    mod_param = 'observation_data/M15'
    img_param = 'observation_data/M15_highres'

    mod_num_lines = mod_h5f[mod_param].shape[0]
    mod_num_pixels = mod_h5f[mod_param].shape[1]

    img_num_lines = img_h5f[img_param].shape[0]
    img_num_pixels = img_h5f[img_param].shape[1]

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

    mod_data = get_grid_values(mod_h5f, mod_param, 0, 0, None, mod_num_lines, mod_num_pixels, range_name=None)
    img_data = get_grid_values(img_h5f, img_param, 0, 0, None, img_num_lines, img_num_pixels, range_name=None)

    num_cntr_tiles = 2
    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 * 2

        i_m = i_c
        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)

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

    mod_h5f.close()
    img_h5f.close()


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 % 16 == 0:
                try:
                    h5f = h5py.File(mfile, 'r')
                except:
                    print('cant open file: ', mfile)
                    continue
                for param in mod_res_params:
                    name = 'observation_data/'+param
                    try:
                        gvals = get_grid_values_all(h5f, name, range_name=None, stride=10)
                        data_dct[param].append(gvals.flatten())
                    except:
                        print('problem reading file: ', mfile)
                        continue
                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)

    f = open('/home/rink/emis_rad_mean.pkl', 'wb')
    pickle.dump(mean_dct, f)
    f.close()

    f = open('/home/rink/emis_rad_std.pkl', 'wb')
    pickle.dump(std_dct, f)
    f.close()

    print(mean_dct)
    print(std_dct)


def run_test(directory):
    # files = glob.glob(directory + 'clavrx_snpp_viirs*.h5')
    files = Path(directory).rglob('clavrx_snpp_viirs*.h5')

    for file in files:
        h5f = h5py.File(file, 'r')
        try:
            opd_nl = get_grid_values_all(h5f, 'cld_opd_nlcomp')
            reff_nl = get_grid_values_all(h5f, 'cld_reff_nlcomp')
        except:
            continue

        if np.sum(np.isnan(opd_nl)) < opd_nl.size and np.sum(np.isnan(reff_nl)) < reff_nl.size:
            print(file)

        h5f.close()