#!/usr/bin/env python3
# -*- coding: utf-8 -*-
""" Utils module with utility functions """

import matplotlib as mpl
from matplotlib import cm
import xarray as xr
import numpy as np
import pandas as pd
import math
import feather
import os.path
from datetime import datetime
from datetime import timedelta


TIMES = {
    'animation': 900,
    'transition': 500,
    'slider_transition': 500
}


def concat_dataframes(fname_tpl, months, variable, rename_from=None, notnans=None):
    """ Concatenate monthly dataframes """

    # build feather files paths
    opaths = [fname_tpl.format(month, variable)
        for month in months if os.path.exists(fname_tpl.format(month, variable))]

    if not opaths:
        return None, None

    # read monthly dataframes and concatenate into one
    if rename_from:
        mon_dfs = pd.concat([feather.read_dataframe(opath)
            .rename(columns={rename_from: variable})
            for opath in opaths])
    # in case of models we don't rename the variable column
    else:
        mon_dfs = pd.concat([feather.read_dataframe(opath)
            for opath in opaths])

    if notnans is None:
        notnans = [st for st in mon_dfs['station'].unique()
            if not mon_dfs[mon_dfs['station']==st][variable].isnull().all()]

    return notnans, mon_dfs[mon_dfs['station'].isin(notnans)]


def retrieve_single_point(fname, tstep, lat, lon, variable, method='netcdf', forecast=True):
    """ """
    from data_handler import DEBUG
    if DEBUG: print(fname, tstep, lat, lon, variable)
    ds = xr.open_dataset(fname)
    if variable not in ds.variables:
        variable = variable.lower()
    # print('TIMESERIES', fname, variable, lon, lat)
    if 'lat' in ds.variables:
        da = ds[variable].sel(lon=lon, lat=lat, method='nearest')
        clat = 'lat'
        clon = 'lon'
    else:
        da = ds[variable].sel(longitude=lon, latitude=lat, method='nearest')
        clat = 'latitude'
        clon = 'longitude'
    if DEBUG: print(da)
    return da.values[tstep]


def retrieve_timeseries(fname, lat, lon, variable, method='netcdf', forecast=False):
    """ """
    if method == 'feather' and not forecast:
        df = feather.read_dataframe(fname)
        if 'lat' in df.columns:
            lat_col = 'lat'
            lon_col = 'lon'
        else:
            lat_col = 'latitude'
            lon_col = 'longitude'

        if variable not in df.columns:
            variable = variable.lower()

        if variable not in df.columns:
            return

        n_lon = find_nearest(df[lon_col].values, lon)
        n_lat = find_nearest(df[lat_col].values, lat)
        ts = df[(df[lat_col] == n_lat) & (df[lon_col] == n_lon)][['time',
            variable]].set_index('time')
        return n_lat, n_lon, ts.index, ts[variable]

    def preprocess(ds, n=8):
        return ds.isel(time=range(n))

    if forecast:
        ds = xr.open_dataset(fname)
    else:
        ds = xr.open_mfdataset(fname, concat_dim='time', combine='nested',
                preprocess=preprocess)
    if variable not in ds.variables:
        variable = variable.lower()
    # print('TIMESERIES', fname, variable, lon, lat)
    if 'lat' in ds.variables:
        da = ds[variable].sel(lon=lon, lat=lat, method='nearest')
        clat = 'lat'
        clon = 'lon'
    else:
        da = ds[variable].sel(longitude=lon, latitude=lat, method='nearest')
        clat = 'latitude'
        clon = 'longitude'
    return da[clat].values, da[clon].values, da.indexes['time'], da


def find_nearest(array, value):
    """ Find the nearest value of a couple of coordinates """
    return array[np.abs(array-value).argmin()]


def find_nearest2(array, value):
    """ Find the nearest value of a couple of coordinates """
    idx = np.searchsorted(array, value, side="left")
    if idx > 0 and (idx == len(array) or math.fabs(value - array[idx-1]) <
                    math.fabs(value - array[idx])):
        return array[idx-1]
    else:
        return array[idx]


def calc_matrix(n):
    """ Calculate the mosaic optimum matrix shape """
    sqrt_n = math.sqrt(n)
    ncols = sqrt_n == int(sqrt_n) and int(sqrt_n) or int(sqrt_n) + 1
    nrows = n%ncols > 0 and int(n/ncols)+1 or int(n/ncols)
    return ncols, nrows


def magnitude(num):
    """ Calculate magnitude """
    if num == 0:
        num = 1
    return int(math.floor(math.log10(num)))


def normalize_vals(vals, valsmin, valsmax, rnd=2):
    """ Normalize values to 0-1 scale """
    if len(vals) == 1:
        return np.array([0])
    vals = np.array(vals)
    if rnd < 2:
        rnd = 2
    return np.around((vals-valsmin)/(valsmax-valsmin), rnd)


def get_colorscale(bounds, colormap, discrete=True):
    """ Create colorscale """
    from data_handler import DEBUG
    if isinstance(colormap, str):
        colormap = cm.get_cmap(colormap)

    bounds = np.array(bounds).astype('float32')
    magn = magnitude(bounds[-1])
    n_bounds = normalize_vals(bounds, bounds[0], bounds[-1], magn)
    norm = mpl.colors.BoundaryNorm(bounds, colormap.N, clip=True)
    s_map = cm.ScalarMappable(norm=norm, cmap=colormap)

    norm_val = not (list(bounds) == list(n_bounds))

    colorscale = [[idx,
                   'rgba' + str(s_map.to_rgba(val,
                                              alpha=True,
                                              bytes=True,
                                              norm=norm_val))]
                  for idx, val in zip(n_bounds, bounds)]

    if discrete is True:
        for item in colorscale.copy():
            if colorscale.index(item) < len(colorscale)-2:
                colorscale.insert(colorscale.index(item)+1,
                                  [colorscale[colorscale.index(item)+1][0],
                                   colorscale[colorscale.index(item)][1]])
            elif len(colorscale) == 1:
                colorscale.insert(1, [1, colorscale[0][1]])

    return colorscale


def get_vis_edate(end_date, hour=None):
    """ Return default date and timestep for visibility. """
    from data_handler import DEBUG
    delay = timedelta(hours=8)
    half_day = timedelta(hours=12)
    fmt_full = "%Y%m%d %H:%M"
    fmt_date = "%Y%m%d"
    hours = (0, 6, 12, 18)

    now = datetime.now()
    now_hour = now.hour
    if hour is not None:
        now_hour = hour
    edate = datetime.strptime(end_date, fmt_date) + timedelta(hours=now_hour)
    cdate = edate.strftime(fmt_date)

    for idx, h in enumerate(hours[:-1]):
        curr = datetime.strptime("{} {:02d}:00".format(cdate, h), fmt_full)
        curr1 = datetime.strptime("{} {:02d}:00".format(cdate, hours[idx+1]), fmt_full)
        if edate < curr + delay:
            curr = curr - half_day
            curr1 = curr1 - half_day

        if (edate >= curr + delay) and (edate < curr1 + delay):
            if DEBUG: print("NOW", edate, "CURR", curr, "H", curr.hour)
            break

    return curr.strftime(fmt_date), curr.hour
