import re
import datetime
from datetime import timezone

import numpy as np
import xarray as xr
import rasterio
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import h5py
from util.util import get_grid_values_all
from util.gfs_reader import *
from util.geos_nav import GEOSNavigation
from aeolus.datasource import GFSfiles

gfs_files = GFSfiles('/Users/tomrink/data/contrail/gfs/')
# GEOSNavigation needs to be updated to support GOES-18
# geos_nav = GEOSNavigation()


def load_image(image_path):
    # Extract date time string from image path
    datetime_regex = '_\\d{8}_\\d{6}'
    datetime_string = re.search(datetime_regex, image_path)
    if datetime_string:
        datetime_string = datetime_string.group()
    dto = datetime.datetime.strptime(datetime_string, '_%Y%m%d_%H%M%S').replace(tzinfo=timezone.utc)
    ts = dto.timestamp()

    img = mpimg.imread(image_path)

    return img, ts


def get_contrail_mask_image(image, thresh=0.157):
    image = np.where(image > thresh,1, 0)
    return image


def extract(mask_image, image_ts, clavrx_path):
    gfs_file, _, _ = gfs_files.get_file(image_ts)
    gfs_h5f = h5py.File(gfs_file, 'r')
    xr_dataset = xr.open_dataset(gfs_file)

    clvrx_h5f = h5py.File(clavrx_path, 'r')
    cloud_top_press = get_grid_values_all(clvrx_h5f, 'cld_press_acha').flatten()
    clvrx_lons = get_grid_values_all(clvrx_h5f, 'longitude').flatten()
    clvrx_lats = get_grid_values_all(clvrx_h5f, 'latitude').flatten()

    contrail_idxs = (mask_image == 1).flatten()
    print('number of contrail pixels: ', np.sum(contrail_idxs))

    # Assuming GOES FD for now -------------------
    # elems, lines = np.meshgrid(np.arange(5424), np.arange(5424))
    # lines, elems = lines.flatten(), elems.flatten()
    # See note above regarding support for GOES-18
    # contrail_lines, contrail_elems = lines[contrail_idxs], elems[contrail_idxs]
    # contrail_lons, contrail_lats = geos_nav.lc_to_earth(contrail_elems, contrail_lines)

    contrail_press = cloud_top_press[contrail_idxs]
    contrail_lons, contrail_lats = clvrx_lons[contrail_idxs], clvrx_lats[contrail_idxs]

    keep = np.invert(np.isnan(contrail_press))
    contrail_press = contrail_press[keep]
    contrail_lons = contrail_lons[keep]
    contrail_lats = contrail_lats[keep]

    wind = get_point_s(xr_dataset, ['u-wind','v-wind'], contrail_lons, contrail_lats, contrail_press)