import xarray as xr
import numpy as np
from util.geos_nav import GEOSNavigation
from netCDF4 import Dataset
import pickle
from util.util import haversine_np


def romio_to_fgf(filename, goes_e_w='EAST'):
    ds = xr.open_dataset(filename, engine='cfgrib')

    lons = ds['longitude']
    lats = ds['latitude']

    lons = lons.values
    lats = lats.values

    lat_array, lon_array = np.meshgrid(lats, lons, indexing='ij')

    nav = GEOSNavigation()

    if goes_e_w == 'EAST':
        nav = GEOSNavigation()
    elif goes_e_w == 'WEST':
        nav = GEOSNavigation(sub_lon=-137.0)

    cc, ll = nav.earth_to_lc_s(lon_array.flatten(), lat_array.flatten())
    tlons, tlats = nav.lc_to_earth(cc, ll)
    dist = haversine_np(lon_array.flatten(), lat_array.flatten(), tlons, tlats)
    ok = np.invert(np.isnan(dist))
    print(np.average(dist[ok]))
    print(np.histogram(dist[ok]))

    cc = cc.reshape(lon_array.shape)
    ll = ll.reshape(lon_array.shape)

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

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

    ds.close()

    return cc, ll


def create_file(filename, fgf_elem, fgf_line):
    dim_x_len = fgf_elem.shape[1]
    dim_y_len = fgf_elem.shape[0]

    rootgrp = Dataset(filename, 'w', format='NETCDF4')
    dim_x = rootgrp.createDimension('romio_x', size=dim_x_len)
    dim_y = rootgrp.createDimension('romio_y', size=dim_y_len)

    romio_fgf_elems = rootgrp.createVariable('romio_fgf_elems', 'f4', ['romio_y', 'romio_x'])
    romio_fgf_lines = rootgrp.createVariable('romio_fgf_lines', 'f4', ['romio_y', 'romio_x'])

    romio_fgf_elems[:, :] = fgf_elem[:, :]
    romio_fgf_lines[:, :] = fgf_line[:, :]

    rootgrp.close()