import os
import tempfile
import glob
from datetime import datetime, timedelta
import logging
import numpy as np
from scipy import interpolate
import netCDF4 as ncdf
from pyproj import Transformer

from buoy_tracking import geo_tools, model_fetch


#DEFAULT_TOPAZ_PATH = (os.path.join(os.path.dirname(__file__), "topaz"))
DEFAULT_TOPAZ_PATH = tempfile.gettempdir()
TOPAZ_DATETIME_FORMAT = "%Y-%m-%d-%H"

#logging.basicConfig(filename='tasking.log', level=logging.INFO)
logger = logging.getLogger(__name__)


def simple_forecast(target_time, drift_track, full_forecast=False):
    """
    Forecasts the position of the buoy at the given target time using
    a simple extrapolation based on previous drift speed and rotation.
    TO ADD: filtering is done to prevent wild forecasts due to inaccurate GPS coordinates.
    :param target_time: Time in UTC for the forecasted position
    :param drift_track: Drift track to use for extrapolation, array of (datetime, lat, lon) tuples.
    :param full_forecast: True; Return the position at every forecast timestep
                          False; Return only the target_time position
    :return: (latitude, longitude) forecast
    """

    # Initialize the record of motion ([speed, angular_vel])
    motion_record = np.zeros((len(drift_track)-2, 2))

    for i in range(1, len(drift_track)-1):
        pos_next = drift_track[i+1]
        pos_curr = drift_track[i]
        pos_last = drift_track[i-1]

        # dt[i-1] = pos_next[1:]  # ONLY FOR PLOTTING

        # Calculate d_pos/dt, which requires 2 positions
        distance = geo_tools.calc_distance(pos_curr[1:], pos_next[1:])

        time_delta = pos_next[0] - pos_curr[0]
        speed = distance / time_delta.total_seconds()  # in km/s

        # Calculate d_theta/dt, which requires 2 vectors and 3 positions
        bearing_curr = geo_tools.calc_bearing(pos_curr[1:], pos_next[1:])
        bearing_last = geo_tools.calc_bearing(pos_last[1:], pos_curr[1:])
        time_delta = pos_next[0] - pos_last[0]
        angular_velocity = (geo_tools.bearing_diff(bearing_curr, bearing_last)
                            / time_delta.total_seconds())  # in degrees/s

        # Record the speed and direction data
        motion_record[i-1] = [speed, angular_velocity]

    initial_time = drift_track[-1][0]
    lead_time = target_time - initial_time

    # We will forecast the position iteratively
    #   (this makes it easier to ignore the difference between arc length and a straight line)
    step_length = timedelta(minutes=15).total_seconds()
    # This rounding means the forecast lead time is only accurate to the step size
    num_steps = int(lead_time.total_seconds() / step_length)

    forecast_drift = [(0, 0, 0) for _ in range(int(num_steps/4))]

    # Determine the distance and direction traveled for each step
    d_per_step = np.average(motion_record[:, 0]) * step_length
    w_per_step = np.average(motion_record[:, 1]) * step_length

    # Initialize the forecast position to the last known location
    forecast_position = [drift_track[-1][1], drift_track[-1][2]]
    # Initialize the bearing to the last known bearing
    # bearing_curr will hold the last position data from the for loop
    forecast_bearing = bearing_curr

    # Iteratively calculate the final forecast
    forecast_time = initial_time
    for i in range(num_steps):
        # Add the average angular velocity to the current trajectory (bearing)
        forecast_bearing += w_per_step
        forecast_position = geo_tools.update_position(forecast_position[0], forecast_position[1],
                                                      forecast_bearing, d_per_step)
        forecast_time += timedelta(seconds=step_length)
        if i % 4 == 3:
            forecast_drift[int(i/4)] = (forecast_time, forecast_position[0], forecast_position[1])

    if full_forecast:
        return forecast_drift
    else:
        return forecast_drift[-1]


def advanced_forecast(buoy_position, target_time, full_forecast=False, topaz_path=None,
                      retry=True, force_update=False, return_topaz_stats=False):

    # To track the time since last topaz download
    #    return true if the data was refreshed, false otherwise
    topaz_update = False
    valid_file = None

    if topaz_path is None:
        topaz_path = DEFAULT_TOPAZ_PATH
    elif topaz_path == 'azure':
        a = 1

    topaz_files = glob.glob(os.path.join(topaz_path, '*.nc'))
    
    # Search for a valid existing forecast
    for topaz_filename in topaz_files:
        # Reconstruct the topaz start/end times from the stored filename
        topaz_start = datetime.strptime(os.path.basename(topaz_filename).split('_')[0],
                                        TOPAZ_DATETIME_FORMAT)

        topaz_end = datetime.strptime(os.path.basename(topaz_filename).split('_')[1],
                                        TOPAZ_DATETIME_FORMAT)

        # Returns True if the needed time window falls entirely within the existing topaz data range
        existing_valid = check_existing_topaz(topaz_start, topaz_end,
                                              buoy_position[0], target_time,
                                              topaz_filename)

        if not existing_valid or force_update:
            # Delete the existing data if it exists
            try:
                os.remove(topaz_filename)
            except FileNotFoundError:
                pass
            except PermissionError:
                pass
            except OSError:
                pass
        else:
            valid_file = topaz_filename
            valid_start = topaz_start
            valid_end = topaz_end

    # If no valid file was found, download a new one.
    #   If we found a valid file, recall stored variables
    if valid_file is None:
        dl_success, topaz_filename, topaz_start, topaz_end = fetch_topaz_forecast(buoy_position, target_time, topaz_path)
        # return just the last known point if the model download failed
        if not dl_success:
            logger.error("Error using advanced forecast, model download failed")
            forecast_drift = [buoy_position]
            if return_topaz_stats:
                return forecast_drift, topaz_update, topaz_start, topaz_end
            else:
                return forecast_drift
        else:
            topaz_update = True
    else:
        topaz_filename = valid_file
        topaz_start = valid_start
        topaz_end = valid_end

    # Load the variables from the file
    topaz_vars = load_topaz_vars(topaz_filename)

    # Run the actual forecast
    try:
        forecast_drift = _advanced_forecast(buoy_position, target_time, topaz_vars, full_forecast=full_forecast)
    except ValueError:
        # This happens when the buoy position falls out of the model ROI bounds
        # Retry with a forced update
        if retry and not topaz_update:
            # output warning to log?
            return advanced_forecast(buoy_position, target_time, topaz_path=topaz_path,
                                     full_forecast=full_forecast, retry=False, force_update=True)
        else:
            logger.error("Error using advanced forecast")
            forecast_drift = [buoy_position]
            pass
            # Do something to warn the user

    # Return the start and end time of the topaz file that was (maybe) downloaded, and also the forecast result
    if return_topaz_stats:
        return forecast_drift, topaz_update, topaz_start, topaz_end
    else:
        return forecast_drift


def _advanced_forecast(buoy_position, target_time, topaz_vars, full_forecast=False):
    """
    Forecasts the position of the buoy at the given target time using
    TOPAZ Sea ice velocity field forecasts.
    :param buoy_id: The name of the buoy to forecast
    :param target_time: Time in UTC for the forecasted position
    :param topaz_vars: Data from the topaz forecast as [t, lats, lons, u, v]
        These variables are returned from load_topaz_vars()
    :param buoy_position: (datetime, lat, lon) of the buoy position
    :return: (latitude, longitude) forecast
    """
    t, lats, lons, u, v = topaz_vars
    # Topaz time axis is in units of hours since 01/01/1950
    time_units_origin = datetime(1950, 1, 1)

    buoy_time = buoy_position[0]
    buoy_lat = buoy_position[1]
    buoy_lon = buoy_position[2]

    # Find the hours since 1950 for the buoys latest position
    #   Round to the nearest whole hour
    time_diff = buoy_time - time_units_origin
    buoy_time_topazfmt = np.around(time_diff.total_seconds() / 3600)

    lead_time = target_time - buoy_time
    lead_time_hours = int(np.around(lead_time.total_seconds() / 3600))
    # remainder_time = lead_time - timedelta(hours=lead_time_hours)

    # Find the index for the buoy time in the topaz array
    t0 = (np.abs(t - buoy_time_topazfmt)).argmin()

    # Find the index of the grid cell with the smallest difference from the buoy lat/lon
    ll_idx = np.unravel_index((np.abs(lons - buoy_lon) + np.abs(lats - buoy_lat)).argmin(),
                              np.shape(lons))

    # If the first value is masked, they all are as mask is based on geography
    #  This covers when buoy drifts out of forecast realm (e.g. washes ashore)
    if u[t0, ll_idx[0], ll_idx[1]] is np.ma.masked:
        logger.warning("Buoy likely outside region of interest")
        logger.error("Advanced forecast requested masked location")
        return([buoy_position])

    # Trim the u/v grids to a local window around the buoy location
    # Local grid size (in each direction from center), 12km gridcell = 1 cells per 12hr at 1km/hr
    lgs = round(lead_time_hours / 24) + 2

    # Sanitize the local grid size based on lat/lon grid size
    a, b, i, j = local_grid_bounds_check(lats, ll_idx, lgs)

    # trim u and v
    u = u[t0:t0 + lead_time_hours + 4, a:b, i:j]
    v = v[t0:t0 + lead_time_hours + 4, a:b, i:j]
    lats = lats[a:b, i:j]
    lons = lons[a:b, i:j]

    # u = u[t0:t0 + lead_time_hours + 4, ll_idx[0] - lgs:ll_idx[0] + lgs, ll_idx[1] - lgs:ll_idx[1] + lgs]
    # v = v[t0:t0 + lead_time_hours + 4, ll_idx[0] - lgs:ll_idx[0] + lgs, ll_idx[1] - lgs:ll_idx[1] + lgs]
    # lats = lats[ll_idx[0] - lgs:ll_idx[0] + lgs, ll_idx[1] - lgs:ll_idx[1] + lgs]
    # lons = lons[ll_idx[0] - lgs:ll_idx[0] + lgs, ll_idx[1] - lgs:ll_idx[1] + lgs]

    # Reproject the lat/lon grid to one with units of meters
    # Create the projection transformations for wgs84 (4326) and polar stereo (3413)
    transformer_forward = Transformer.from_crs('epsg:4326', 'epsg:3413', always_xy=True)
    transformer_reverse = Transformer.from_crs('epsg:3413', 'epsg:4326', always_xy=True)

    #  x/y grid
    grid_x = np.zeros(np.shape(lons))
    grid_y = np.zeros(np.shape(lats))
    grid_t = np.linspace(0, lead_time_hours+3, lead_time_hours+4)     # 1 hourly temporal grid

    # For each grid cell find x/y values from lon/lat
    for i in range(np.shape(grid_y)[0]):
        for j in range(np.shape(grid_y)[1]):
            grid_x[i, j],  grid_y[i, j] = transformer_forward.transform(lons[i, j], lats[i, j])

    # Trim and round the reprojected units to create a regular grid
    grid_x = np.around(np.average(grid_x, axis=0), 0)
    grid_y = np.around(np.average(grid_y, axis=1), 0)

    # Project the buoy coordinates onto the xy-grid
    x, y = transformer_forward.transform(buoy_lon, buoy_lat)
    
    # U and V are in units of m/s
    # Convert to m/hr
    scale_factor = 3600
    u = np.multiply(u, scale_factor)
    v = np.multiply(v, scale_factor)

    # Run the forecast using a rk4 solver over the velocity fields
    forecast_drift_xy = rk4(u, v, 0, x, y, .5, lead_time_hours, grid_x, grid_y, grid_t, full_forecast=full_forecast)

    # Project the solution back to lon/lat
    # len(forecast_drift)==1 if full_forecast==False
    forecast_drift = []
    for tn, x, y in forecast_drift_xy:
        new_lon, new_lat = transformer_reverse.transform(x, y)
        forecast_drift.append((buoy_time+timedelta(hours=tn), new_lat, new_lon))

    return forecast_drift


def rk4(u, v, t0, x0, y0, step_size, final_step, grid_x, grid_y, grid_t, full_forecast=False):
    """
    Runge-Kutta 4th order solver to determine final position from known velocity fields
    :param u: Velocity field grid (in x direction)
    :param v: Velocity field grid (in y direction)
    :param t0: Index of initial time
    :param x0: Initial x position (in raw meters, not an index)
    :param y0: Initial y position (in raw meters, not an index)
    :param step_size: Step size to take for each iteration (units of t0)
    :param final_step: Time where solver stops
    :param grid_x: 2d grid with x coordinates
    :param grid_y: 2d grid with y coordinates
    :param grid_t: 1d grid with time coordinates
    :return: (latitude, longitude) forecast
    """
    # Initialize variables with initial time/position
    h = step_size
    x, y = x0, y0
    tn = t0

    if full_forecast:
        forecast_drift = []

    # Interpolate u and v values over the x/y/t grid so that we can
    #   determine velocities at fractional grid cells (x and y are true meters, not indices)
    u_interp = interpolate.RegularGridInterpolator((grid_t, grid_x, grid_y), u, bounds_error=True)
    v_interp = interpolate.RegularGridInterpolator((grid_t, grid_x, grid_y), v, bounds_error=True)

    # Find final position iteratively based on step size
    while tn <= final_step:

        # Evaluate the function at f(tn, yn)
        k1 = [u_interp([tn, x, y])[0],
              v_interp([tn, x, y])[0]]
        k2 = [u_interp([tn + (h/2), x + ((h*k1[0])/2), y + ((h*k1[0])/2)])[0],
              v_interp([tn + (h/2), x + ((h*k1[1])/2), y + ((h*k1[1])/2)])[0]]
        k3 = [u_interp([tn + (h/2), x + ((h*k2[0])/2), y + ((h*k2[0])/2)])[0],
              v_interp([tn + (h/2), x + ((h*k2[1])/2), y + ((h*k2[1])/2)])[0]]
        k4 = [u_interp([tn + h, x + h*k3[0], y + h*k3[0]])[0],
              v_interp([tn + h, x + h*k3[1], y + h*k3[1]])[0]]

        # Update for the next iteration
        tn += h
        x = x + (1/6)*h*(k1[0] + 2*k2[0] + 2*k3[0] + k4[0])
        y = y + (1/6)*h*(k1[1] + 2*k2[1] + 2*k3[1] + k4[1])

        if full_forecast:
            forecast_drift.append((tn, x, y))

    if full_forecast:
        return forecast_drift
    else:
        forecast_drift = [(tn, x, y)]
        return forecast_drift


def local_grid_bounds_check(big_grid, idx, lgs):
    """
    Checks whether an index and a window of fits within a bigger grid,
        and returns the proper indicies if not. 
    big_grid: full grid to draw window from
    idx: index for window center point [x, y]
    lgs: local grid size; added in each diminsion to idx
    """
    # Add 1 to upper index to get a grid centered on idx
    #   (because thats how arrays are indexed)
    xmax, ymax = np.shape(big_grid)
    a = idx[0] - lgs
    b = idx[0] + lgs + 1
    i = idx[1] - lgs
    j = idx[1] + lgs + 1
    if a < 0:
        a = 0
    if b >= xmax:
        b = xmax - 1
    if i < 0:
        i = 0
    if j >= ymax:
        j = ymax - 1

    return a, b, i, j


def load_topaz_vars(topaz_filename):
    # The advanced forecast needs to download sea ice velocity field forecasts
    topaz = ncdf.Dataset(topaz_filename, mode='r')

    t = topaz.variables['time'][:]
    lats = topaz.variables['latitude'][:]
    lons = topaz.variables['longitude'][:]
    u = topaz.variables['uice'][:, :, :]
    v = topaz.variables['vice'][:, :, :]
    u, v = np.nan_to_num(u), np.nan_to_num(v)

    return t, lats, lons, u, v


def fetch_topaz_forecast(buoy_position, target_time, output_dir):
    """
    :param buoy_position: (datetime, lat, lon) of the buoy position
    :param target_time: Time in UTC for the forecasted position
    """

    # Determine the range of data we need based on target forecast time
    # End date is the target time:
    #       Add a small buffer of 4 hours
    date_end = target_time + timedelta(hours=4)
    # Start time is the buoys last know position:
    date_start = buoy_position[0]

    forecast_hours = (date_end - date_start).total_seconds() / 3600

    # Calculates an appoximate region of interest based on the length of forecast and 1km/h drift
    lat_min, lat_max, lon_min, lon_max = calc_latlon_window(buoy_position[1], buoy_position[2], forecast_hours)

    # If no existing file was found, download a new one.
    # Select the date range as the needed window, plus a buffer.
    #   This makes it more likely that future forecasts will be able to reuse this data file
    date_min = date_start - timedelta(hours=24)
    # Don't want to add too much time here, because the forecast skill drops off with increased
    #   lead times. I.e. we want to be redownloading this file every day or two.
    date_max = date_end + timedelta(hours=24)
    output_filename = '{}_{}_velocityfield.nc'.format(datetime.strftime(date_min, TOPAZ_DATETIME_FORMAT), 
                                                      datetime.strftime(date_max, TOPAZ_DATETIME_FORMAT))
    
    url = 'https://nrt.cmems-du.eu/motu-web/Motu'
    username = os.getenv('TOPAZ_USERNAME') #credentials.LOGIN['topaz_username']
    password = os.getenv('TOPAZ_PASSWORD') #credentials.LOGIN['topaz_password']
    
    # Authenticate with the model data server
    conn_session = model_fetch.copernicus_marine_login(url, username, password)
    # Define the model parameters we want to download
    payload = model_fetch.order_payload(date_min, date_max)
    # Request the data
    dl_success = model_fetch.download_data(url, payload, os.path.join(output_dir, output_filename), conn_session)

    # cmd = ("python -m motuclient --motu https://nrt.cmems-du.eu/motu-web/Motu "
    #        "--service-id ARCTIC_ANALYSIS_FORECAST_PHYS_002_001_a-TDS "
    #        "--product-id dataset-topaz4-arc-1hr-myoceanv2-be "
    #        "--longitude-min -180 --longitude-max 180 --latitude-min 68 --latitude-max 90 "
    #        '--date-min "{}" --date-max "{}" '
    #        "--variable latitude --variable longitude --variable uice --variable vice "
    #        "--out-dir {} --out-name {} "
    #        "--user {} --pwd {}").format(#lon_min, lon_max, lat_min, lat_max,
    #                                         date_min, date_max, output_dir,
    #                                         output_filename, username, password)

    #os.system(cmd)
    return dl_success, os.path.join(output_dir, output_filename), date_min, date_max


def check_existing_topaz(topaz_start, topaz_end, date_start, date_end, topaz_filename):
    """
    Check if the existing forecast file matchs the time window needs
    :param topaz_start: datetime; beginning of existing forecast window
    :param topaz_end: datetime; end of existing forecast window
    :param date_start: datetime; beginning of forecast need
    :param date_end: datetime; end of forecast need
    return true if existing topaz is usable, false if a new one must be acquired
    """

    # Try to open the existing dataset
    try:
        topaz = ncdf.Dataset(topaz_filename, mode='r')
    except OSError:
        return False

    # If our needed window falls within the range of this file,
    #   we do not need to download a new set.
    if topaz_start <= date_start and topaz_end >= date_end:
        return True
    else:
        return False


def calc_latlon_window(lat, lon, forecast_hours):

    # Assume maximum average drift of 4km/hr
    max_distance_km = float(forecast_hours * 4)
    max_distance_deglat = max_distance_km / 111.0     # 1 degree lat is about 111km at 70N
    max_distance_deglon = max_distance_km / 36.0      # 1 degree lon is about 36km at 70.7N

    # Set a minimum download size
    if max_distance_deglat < 0.5:
        max_distance_deglat = 0.5
    if max_distance_deglon < 0.5:
        max_distance_deglon = 0.5

    lat_min = round(lat - max_distance_deglat, 3)
    lat_max = round(lat + max_distance_deglat, 3)
    lon_min = round(lon - max_distance_deglat, 3)
    lon_max = round(lon + max_distance_deglat, 3)

    # Lat can't go below the approx coast position
    if lat_min < 70:
        lat_min = 70

    print(lat, lon, lat_min, lat_max, lon_min, lon_max)

    return lat_min, lat_max, lon_min, lon_max