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

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


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

    print("__________", opaths, "__________", months, rename_from, notnans)
    if not opaths:
        return None, None

    # read monthly dataframes and concatenate into one
    if rename_from != None:
        if rename_from == 'auto':
            if variable.lower() in feather.read_dataframe(opaths[0]).columns:
                rename_from = variable.lower()

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

    # 1d observations
    if 'station' in mon_dfs.columns:
        if notnans == False:
            final_df = mon_dfs
        else:
            if notnans == None:
                notnans = [st for st in mon_dfs['station'].unique()
                    if mon_dfs[mon_dfs['station']==st][variable].notnull().any()]
            mon_dfs_filter = mon_dfs['station']
            final_df = mon_dfs[mon_dfs_filter.isin(notnans)]
    # 2d observations
    else:
        final_df = mon_dfs

    return notnans, final_df


def retrieve_single_point(fname, tstep, lat, lon, variable):
    """ """
    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')
    else:
        da = ds[variable].sel(longitude=lon, latitude=lat, method='nearest')
    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 None, None, None, None

        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]

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

    curr = None
    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

    if curr is not None:
        return curr.strftime(fmt_date), curr.hour

    return cdate, edate.hour


def get_currdate_tstep(model_start, model_start_before, current_time_before, delay, selected_date, tstep=4):
    """ Returns date and timestep """

    # MODELS starting at 00 with 3 days of forecast
    print("Starting at 0 and NOT DELAYED: 3 days forecast!")
    cdo_tsteps = "1/25"
    delayed = (current_time_before and not delay) or (not current_time_before and delay)

    # MODELS not starting at 00
    if (model_start == 12 and not model_start_before) or \
            (current_time_before and model_start_before == 12):
        if delayed:
            print("Starting at 12 and DELAYED: 2 days forecast!")
            # DELAYED (2 days): models that starts at 12h considering end_date = current_date - 1
            if tstep < 4:
                selected_date = (datetime.strptime(selected_date, "%Y%m%d") -
                                 timedelta(days=1)).strftime("%Y%m%d")
                tstep = int(tstep) + 4
            else:
                tstep = int(tstep) - 4
            cdo_tsteps = "1/21"
        else:
            print("Starting at 12 and NOT DELAYED: 3 days forecast!")
            # NOT DELAYED (3 days): models that starts at 12h considering end_date = current_date
            selected_date = (datetime.strptime(selected_date, "%Y%m%d") -
                             timedelta(days=1)).strftime("%Y%m%d")
            tstep = int(tstep) + 4
            cdo_tsteps = "5/29"
    # MODELS starting at 00 with 2 days of forecast
    elif delayed:
        print("Starting at 0 and DELAYED: 2 days forecast!")
        cdo_tsteps = "5/25"

    return selected_date, tstep, cdo_tsteps
