import sys
import random
import numpy as np
import itertools
import pandas as pd

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.colors as matcolors

import climatools.muths as muths
import climatools.dates as dates
import climatools.misc as misc

def matplotlib_basic_colours():
    '''
    Returns a list of plot colours available in matplotlib.
    The colours in this list are easily distinguished from each
    other by eye.
    '''
    return ['b', 'g', 'r', 'c', 'm', 'y', 'k']

def matplotlib_nonnothing_linestyles(longname = False):
    '''
    Returns the list of plot linestyles available in matplotlib.
    Linestyles that are invisible are left out of this list.
    '''
    if longname:
        return [v for k, v in Line2D.markers.items() if v != '_draw_nothing']
    else:
        return [k for k, v in Line2D.lineStyles.items()
                if v != '_draw_nothing']

def matplotlib_nonnothing_markers(longname = False):
    '''
    Returns the list of plot markers available in matplotlib.
    Markers that are invisible are left out of this list.
    '''
    if longname:
        return [v for k, v in Line2D.markers.items() if v != 'nothing']
    else:
        return [k for k, v in Line2D.markers.items() if v != 'nothing']

def matplotlib_colour_linestyle_tuples(N = 10):
    '''
    Returns a shuffled list of unique tuples of plot colours and linestyles
    of length Npairs.
    INPUT:
    Npairs --- length of list returned,
    the number of unique tuples of plot colours and linestyles returned.
    '''
    colours = matplotlib_basic_colours()
    linestyles = matplotlib_nonnothing_linestyles()

    uniques = list(itertools.product(colours, linestyles))

    random.shuffle(uniques)
    return random.sample(uniques, N)

def matplotlib_colour_linestyle_marker_tuples(N = 10):
    colours = matplotlib_basic_colours()
    linestyles = matplotlib_nonnothing_linestyles()
    markers = matplotlib_nonnothing_markers()

    uniques = list(itertools.product(colours, linestyles, markers))

    random.shuffle(uniques)
    return random.sample(uniques, N)

def matplotlib_colormap_names(kind = 'divergent'):
    cmaps = dict([('Sequential', ['Blues', 'BuGn', 'BuPu',
                                  'GnBu', 'Greens', 'Greys', 'Oranges', 'OrRd',
                                  'PuBu', 'PuBuGn', 'PuRd', 'Purples', 'RdPu',
                                  'Reds', 'YlGn', 'YlGnBu', 'YlOrBr', 'YlOrRd']),
                  ('Sequential (2)', ['afmhot', 'autumn', 'bone', 'cool', 'copper',
                                      'gist_heat', 'gray', 'hot', 'pink',
                                      'spring', 'summer', 'winter']),
                  ('Diverging',      ['BrBG', 'bwr', 'coolwarm', 'PiYG', 'PRGn', 'PuOr',
                                      'RdBu', 'RdGy', 'RdYlBu', 'RdYlGn', 'Spectral',
                                      'seismic']),
                  ('Qualitative',    ['Accent', 'Dark2', 'Paired', 'Pastel1',
                                      'Pastel2', 'Set1', 'Set2', 'Set3']),
                  ('Miscellaneous',  ['gist_earth', 'terrain', 'ocean', 'gist_stern',
                                      'brg', 'CMRmap', 'cubehelix',
                                      'gnuplot', 'gnuplot2', 'gist_ncar',
                                      'nipy_spectral', 'jet', 'rainbow',
                                      'gist_rainbow', 'hsv', 'flag', 'prism'])])

    ccmaps = dict([('sequential', cmaps['Sequential'] + cmaps['Sequential (2)']),
                   ('diverging', cmaps['Diverging']),
                   ('qualitative', cmaps['Qualitative']),
                   ('misc', cmaps['Miscellaneous'])])
    return ccmaps[kind]

def dates_locators_by_timescale():
    locators = {'year': matplotlib.dates.YearLocator,
                'month': matplotlib.dates.MonthLocator,
                'day': matplotlib.dates.DayLocator,
                'hour': matplotlib.dates.HourLocator,
                'minute': matplotlib.dates.MinuteLocator}
    return locators

def symmetric_about_white_cmap_levels(rough_maxabs, Ncolours = 11):
    '''
    Returns min and max levels and step that will have 0 centred
    on the colour white on a diverging colourmap that diverges from
    the colour white
    INPUT:
    rough_maxabs -- rough maximum absolute value of the levels
                    (note the absolute value because max = -min
                    necessarily)
    Ncolours -- number of intervals on the colour bar
    '''
    try:
        assert Ncolours % 2 != 0
    except AssertionError:
        raise InputError('Number of colours needs to be an odd number.')

    int_dv = muths.round_to_1(2 * rough_maxabs / Ncolours)
    rough_dv = 2 * rough_maxabs / Ncolours
    ddv = int_dv - rough_dv
    maxabs = rough_maxabs + Ncolours / 2 * ddv
    return (- maxabs, maxabs, int_dv)

def round_levels_with_zero_centred_between_two(vmin0 = -1, vmax0 = 2, Nsteps0 = 10):
    '''
    Suppose you want to divide the range (vmin0, vmax0) into roughly Nsteps0 intervals, where vmin0 < 0 < vmax0.
    This function returns a new range(vmin, vmax) and new Nsteps that is
    as close to the desired one as possible, but are such that dv = (vmax - vmin) / Nsteps
    is a nice round value, and the value zero is centred on one of the Nsteps intervals.
    INPUT:
    vmin0 --- lower limit, must be negative.
    vmax0 --- upper limit, must be positive.
    Nsteps0 --- rough desired number of intervals between vmin0 and vmax0.
    OUTPUT:
    vmin --- lower limit
    vmax --- upper limit
    Nsteps --- number of intervals between vmin and vmax
    '''
    dv0 = (vmax0 - vmin0) / Nsteps0
    dv = muths.round_to_1(dv0)
    half_dv = .5 * dv
    vmin = 0 - dv * int(np.ceil((0 - vmin0 - half_dv) / dv)) - half_dv
    vmax = 0 + dv * int(np.ceil((vmax0 - 0 - half_dv) / dv)) + half_dv
    Nsteps = int((vmax - vmin) / dv)
    return (vmin, vmax, Nsteps)

def slice_divergent_colormap(cmap = None, vmin = -1, vmax = 2):
    '''
    A divergent colormap is one that is centred at white, while extending
    to the two sides of it are two distintive colour types.  For example,
    the divergent colormap bwr is centred on white; to one side are blues
    and to the other reds. The length on either side of white are the same
    by default, making it useful to map to intervals where the negative and
    positive limits are equal in maginitude, such as (-1, 1), (-241, 241), etc.

    This function slices the default colormap up such that it can be mapped
    to (keeping white mapped to zero)
    intervals that are not symmetrical about 0, such as (-100, 2), (-2, 349), etc.

    INPUT:
    cmap -- matplotlib colormap object (only really makes sense with a divergent one)
    vmin -- lower limit of interval (negative)
    vmax -- upper limit of interval (positive)
    OUTPUT:
    sliced_cmap -- sliced colormap with white mapped to zero of the interval (vmin, vmax)
    '''
    if not cmap:
        cmap = plt.get_cmap('bwr')

    real_range = vmax - vmin
    largest_abs_range = 2 * max(abs(vmin), abs(vmax))
    fractional_range = real_range / largest_abs_range

    if abs(vmin) > abs(vmax):
        cmap_min, cmap_max = 0, fractional_range
    elif abs(vmax) > abs(vmin):
        cmap_min, cmap_max = 1 - fractional_range, 1
    else:
        cmap_min, cmap_max = 0, 1

    sliced_cmap = matcolors.\
                  LinearSegmentedColormap.from_list('test_cmap',
                                                    cmap(np.linspace(cmap_min, cmap_max)))
    return sliced_cmap

def axes_beyond_ticks(ax, which = 'x'):
    '''
    Draws the axes with a range that is just wider than
    the existing range so that tick labels at the
    lowest and highest values do not overlap with neighbouring
    subplots. Useful when there is no spacing in between subplots.
    INPUT:
    which --- x for x axis, y for y axis
    '''
    if which == 'x':
        ticklocs = ax.xaxis.get_majorticklocs()
        interval = ticklocs[1] - ticklocs[0]
        ax.set_xlim((ticklocs[0] - .6 * interval,
                     ticklocs[-1] + .6 * interval))
    elif which == 'y':
        ticklocs = ax.yaxis.get_majorticklocs()
        interval = ticklocs[1] - ticklocs[0]
        ax.set_ylim((ticklocs[0] - .6 * interval,
                     ticklocs[-1] + .6 * interval))
    else:
        raise InputError('which has to be either x or y')
    return ax

def get_datetime_major_locator(xlist):
    '''
    Return a good dates locator for major ticks.
    INPUT:
    xlist --- list: [number of unique years amongst minor ticks,
                     number of unique months amongst minor ticks,
                     number of unique days amongst minor ticks,
                     number of unique hours amongst minor ticks,
                     number of unique minutes amongst minor ticks]
    OUTPUT:
    locator  -- one of [matplotlib.dates.YearLocator,
                        matplotlib.dates.MonthLocator,
                        matplotlib.dates.DayLocator,
                        matplotlib.dates.HourLocator,
                        matplotlib.dates.MinuteLocator]
    '''
    locators = [matplotlib.dates.YearLocator,
                matplotlib.dates.MonthLocator,
                matplotlib.dates.DayLocator,
                matplotlib.dates.HourLocator,
                matplotlib.dates.MinuteLocator]

    max_number_majorticks = 5

    for k, x in enumerate(xlist):
        if x <= 1:
            continue
        else:
            if x > max_number_majorticks:
                return locators[k - 1] if k > 0 else None
            else:
                return locators[k]
            return None

def get_N_unique_datetimeperiod_labels(ticklocs):
    '''
    Return  number of unique datetime labels for periods
    YEAR, MONTH, DAY, HOUR, MINUTE.
    INPUT:
    ticklocs --- tick locators in datetime in floating point numbers
                 (days since 0001-01-01 UTC, plus 1)
    OUTPUT --- list: [number of unique years amongst minor ticks,
                      number of unique months amongst minor ticks,
                      number of unique days amongst minor ticks,
                      number of unique hours amongst minor ticks,
                      number of unique minutes amongst minor ticks]
    '''
    minorticklocs = [matplotlib.dates.num2date(tickloc)
                    for tickloc in ticklocs] #ax.xaxis.get_minorticklocs()
    return [len(set([tickloc.year for tickloc in minorticklocs])),
            len(set([tickloc.month for tickloc in minorticklocs])),
            len(set([tickloc.day for tickloc in minorticklocs])),
            len(set([tickloc.hour for tickloc in minorticklocs])),
            len(set([tickloc.minute for tickloc in minorticklocs]))]

def split_to_major_minor_timescales(Nlabels):
    timescales = ('year', 'month', 'day', 'hour', 'minute',)

    maxN_major_ticks = 5

    for k, x in enumerate(Nlabels):
        if x <= 1:
            continue
        else:
            if x > maxN_major_ticks:
                return (timescales[: k], timescales[k:])
            else:
                return (timescales[: k + 1], timescales[k + 1:])
            return timescales[:]
    return major_timescales, minor_timescales

def get_datetime_tick_formats(timescales):
    directives = {'year':   '%Y',
                  'month':  '%m',
                  'day':    '%d',
                  'hour':   '%H',
                  'minute': '%M'}
    ymds = [directives[timescale] for timescale in timescales
            if timescale in ['year', 'month', 'day']]
    hms = [directives[timescale] for timescale in timescales
           if timescale in ['hour', 'minute']]
    return ' '.join(['/'.join(ymds), ':'.join(hms)])

def set_xaxis_datetime_ticklocs_ticklabels(xaxis, duration = pd.Timedelta(days = 1)):
    '''
    Set major and minor tick locations and labels
    INPUT:
    xaxis --- matplotlib.axis.xaxis/yaxis on which ticks and ticklabels are to be set
    duration --- duration which the axis spans in pandas.Timedelta
    '''
    # maximum allowed major and minor intervals
    maxN_major, maxN_minor = 5, 18

    timescales = ('year', 'month', 'day', 'hour', 'minute')

    locators = (
        matplotlib.dates.YearLocator,
        matplotlib.dates.MonthLocator,
        matplotlib.dates.DayLocator,
        matplotlib.dates.HourLocator,
        matplotlib.dates.MinuteLocator,
        )

    bylocators = (
        lambda x: matplotlib.dates.YearLocator(x),
        lambda x: matplotlib.dates.MonthLocator(bymonth = range(1, 13, x)),
        lambda x: matplotlib.dates.DayLocator(bymonthday = range(1, 32, x)),
        lambda x: matplotlib.dates.HourLocator(byhour = range(0, 24, x)),
        lambda x: matplotlib.dates.MinuteLocator(byminute = range(0, 60, x))
        )

    timedeltas = (pd.Timedelta(days = 365),
                  pd.Timedelta(days = 30),
                  pd.Timedelta(days = 1),
                  pd.Timedelta(hours = 1),
                  pd.Timedelta(minutes = 1))

    # get number of timescales in duration
    N_tscales = [duration / timedelta for timedelta in  timedeltas]

    # choose which timescales to be major and minor
    # get the index
    # essentially the rule is that the number of major intervals cannot
    # exceed maxN_major
    gr8r_than_1 = [n > 1 for n in N_tscales]
    if not any(gr8r_than_1):
        indx_major, idx_minor = -1, None
    elif all(gr8r_than_1):
        if N_tscales[0] > maxN_major:
            indx_major, indx_minor = None, 0
        else:
            index_major, indx_minor = 0, 1
    else:
        indx = gr8r_than_1.index(True)
        if N_tscales[indx] > maxN_major:
            indx_major, indx_minor = indx - 1, indx
        else:
            indx_major, indx_minor = indx, indx + 1

    # create major dates locator
    if indx_major:
        locator_major = locators[indx_major]()

    # create major tick label formatters
    if indx_major:
        fmt_major = get_datetime_tick_formats(timescales[: indx_major + 1])

    # create major ticks and labels
    if indx_major:
        xaxis.set_major_locator(locator_major)
        xaxis.set_major_formatter(matplotlib.dates.DateFormatter('\n' + fmt_major))

    # create minor dates locator
    if indx_minor:
        if N_tscales[indx_minor] <= maxN_minor:
            locator_minor = locators[indx_minor]()
        else:
            # scale down the number of minor ticks if there are too many
            N_to_group = int(np.ceil(N_tscales[indx_minor] / maxN_minor))
            locator_minor = bylocators[indx_minor](N_to_group)

    # create minor tick formatters
    if indx_minor:
        fmt_minor = get_datetime_tick_formats([timescales[indx_minor]])

    # create minor ticks and labels
    if indx_minor:
        xaxis.set_minor_locator(locator_minor)
        xaxis.set_minor_formatter(matplotlib.dates.DateFormatter(fmt_minor))

def plot_DataArray(ax, da,
                   datetime_label = True,
                   title = '', label = ':D',
                   colour = 'b', linestyle = '-', marker = '.',
                   ylabel = 'ylabel', ylim = None, yscale = 'linear'):
    '''
    Returns a line plot of time v.s some variable that is in an DataArray
    of dimensions (time, lon0, lat0), where lon0 and lat0 are fixed.
    INPUT:
    da --- xray DataArray
    datetime_label --- True to use Python datetime objects to represent time
    label --- string used to label the line, needed to the legend
    OUTPUT:
    ax --- matplotlib Axes object
    '''
    x = da.coords['time'].values
    y = da[{'lon': 0, 'lat': 0}].values

    ax.plot(x, y,
            label = label,
            color = colour, linestyle = linestyle, marker = marker)

    ax.set_title(title)

    # backround
    ax.set_axis_bgcolor((1., 1., 1.))
    ax.set_axisbelow(b = True)

    # y-axis
    ax.set_ylabel('{}'.format(ylabel) + \
                  ' [{}]'.format(da.units))
    ax.set_yscale(yscale)
    ax.set_ylim(ylim)
    ax.yaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(10))
    ax.yaxis.set_tick_params(length = 6., which = 'major')
    ax.yaxis.set_tick_params(length = 3., which = 'minor')
    ax.yaxis.grid(b = True, which = 'major', color = (0., 0., 0.))

    # x-axis
    ax.set_xlabel('{} [{}]'.format(da.coords['time'].attrs['long_name'],
                                   da.coords['time'].attrs['units']))
    ax.xaxis.set_major_locator(matplotlib.dates.DayLocator(interval = 1))
    ax.xaxis.set_minor_locator(matplotlib.dates.HourLocator(byhour = [6, 12, 18]))
    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%d\n%b'))
    ax.xaxis.grid(b = True, which = 'major', color = (0., 0., 0.))
    ax.xaxis.set_tick_params(length = 6., which = 'major')
    ax.xaxis.set_tick_params(length = 3., which = 'minor')
    return ax

def plot_vertical_profile(ax, da,
                          label = '', colour = 'b', linestyle = '-',
                          title = None,
                          xlabel = None, xlim = None, xscale = 'linear',
                          xlabels_rotate = 0.,
                          ylabel = None, ylim = None, yscale = 'linear'):
    '''
    Plots vertical profile (i.e. pressure against some variable) on a maplotlib Axes
    INPUT:
    xaxis_pow --- power of 10 to multiply xlabels by (if equals = 10, 150 because 1500 * 1e-1)
    xlabels_rotate --- angle to rotate xlabels by [degrees]
    '''
    if 'lev' in da:
        yname = 'lev'
    elif 'ilev' in da:
        yname = 'ilev'
    else:
        raise ValueError('vertical dimension must be either named lev or ilev')

    y = da.coords[yname].values
    x = da[{'lon': 0, 'lat': 0}].values

    ax.plot(x, y,
            label = label, color = colour, linestyle = linestyle, linewidth = 2.3)

    if title:
        ax.set_title(title)

    # background
    ax.set_axis_bgcolor((1., 1., 1.))
    ax.set_axisbelow(b = True)

    # y-axis
    if ylabel:
        ax.set_ylabel(ylabel)
    ax.set_ylim(ylim)
    ax.set_yscale(yscale)
    ax.yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(100))
    ax.yaxis.grid(b = True, which = 'major')

    # x-axis
    if xlabel:
        ax.set_xlabel(xlabel)
    ax.set_xscale(xscale)
    ax.xaxis.grid(b = True, which = 'major')

    if xlim:
        ax.set_xlim(xlim)
    else:
        xticklocs = ax.xaxis.get_majorticklocs()
        xtick_interval = xticklocs[1] - xticklocs[0]
        xaxis_pow = muths.pow_base10_for_decimal(xtick_interval, decimal = 1)
        ax.xaxis_pow = xaxis_pow
        ax.xaxis.set_major_formatter(
            matplotlib.ticker.FuncFormatter(lambda x, pos: '{:.1f}'.\
                                            format(10**xaxis_pow * x)))
        plt.setp(ax.xaxis.get_majorticklabels(), rotation = xlabels_rotate)

    return ax

def contourf_DataArray(ax, da,
                       datetime_label = True, mbar_label = True,
                       cmap_levels = None,
                       cmap = matplotlib.cm.jet):
    '''
    Returns a contour-fill plot of time vs level vs variable for variable
    of dimensions (time, level, lon0, lat0), where lon0 and lat0 are fixed.
    INPUT:
    da --- xray DataArray
    cmap_levels --- tuple (extend, vmin, vmax, vstep) of
                   whether to extend colormap,
                   minimum, maximum and step of contour levels to use
    cmap --- maplotlib colormap object
    '''
    if 'lev' in da:
        spatial_dim = 'lev'
    elif 'ilev' in da:
        spatial_dim = 'ilev'
    else:
        raise ValueError('spatial dimension to plot must be either lev or ilev')

    y = da[spatial_dim].values
    x = [pd.Timestamp(dtns).to_datetime() for dtns in da['time'].values]
    Z = da.transpose(spatial_dim, 'time')

    if cmap_levels:
        extend, lev_min, lev_max, lev_step = cmap_levels
    else:
        lev_min, lev_max, lev_step = Z.min(), Z.max(), (Z.max() - Z.min()) / 20
        extend = 'neither'

    if extend == 'both':
        cmap = slice_divergent_colormap(cmap = cmap, vmin = lev_min, vmax = lev_max)
        levels = np.linspace(lev_min, lev_max,
                             int((lev_max - lev_min) / lev_step) + 1)
    else:
        levels = np.arange(lev_min, lev_max + .1 * lev_step, lev_step)

    cs = ax.contourf(x, y, Z,
                     levels = levels,
                     cmap = cmap,
                     extend = extend)

    cs.cmap.set_over('yellow')
    cs.cmap.set_under('black')

    cbar = plt.colorbar(cs, ax = ax)
    cbar.ax.set_ylabel('[{}]'.format(da.units))
    cbar.set_ticks(levels)

    cbar.formatter.set_powerlimits((0, 0)) # this works as alternative to what's below

#    cbarlabel_decimal = 0
#    cbartick_pow = muths.pow_base10_for_decimal(lev_step,
#                                                decimal = cbarlabel_decimal)
#    print(lev_step, cbartick_pow)
#    cbar_format_string = '{}' #'{{:.{}f}}'.format(cbarlabel_decimal)
#    cbar.formatter = matplotlib.ticker.FuncFormatter(
#        lambda x, pos: cbar_format_string.format(10**int(cbartick_pow) * x))

#    cbar.ax.text(x = 1, y = 1.07, s = '1e{}'.format(- cbartick_pow))

    cbar.update_ticks()

    ax.set_title('{}\n{}'.format(da.attrs['case_name'],
                                 da.attrs['long_name']),
                 fontsize = 10)

    [spine.set_color((1., 1., 1)) for k, spine in ax.spines.items()]

    ax.invert_yaxis()
    ax.set_ylabel('{} [{}]'.format(spatial_dim, da[spatial_dim].units))
    ax.yaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(50))
    ax.yaxis.grid(b = True, which = 'major')
    ax.yaxis.set_tick_params(length = 6, which = 'major')
    ax.yaxis.set_tick_params(length = 3, which = 'minor')

    ax.set_xlabel('time [{}]'.format(da['time'].units))
    set_xaxis_datetime_ticklocs_ticklabels(ax.xaxis,
                                           duration = pd.Timedelta(x[-1] - x[0]))
    ax.xaxis.grid(b = True, which = 'minor')
    ax.xaxis.set_tick_params(length = 6, which = 'major')
    ax.xaxis.set_tick_params(length = 3, which = 'minor')
    return ax

def daytime_nighttime_shading(ax, dts,
                              hour_daystart = 8, hour_nightstart = 18):
    '''
    Shades in colours for day-time and night-time.
    INPUT:
    ax --- matplotlib Axes object
    dts --- an array/list of Pandas Timestamps (a Pandas DatetimeIndex)
    hour_daystart --- hour at which day-time starts
    hour_nightstart --- hour at which night-time starts
    '''
    # Change the Hour of each Timestamp to correspond to
    # that at which day-time or night-time starts
    dtsDN = dts.map(lambda x: dates.to_daytime_nighttime(\
        x, hour_daystart = hour_daystart,
        hour_nightstart = hour_nightstart))

    # Find and sort the boundries, datetimes at which each day-time
    # and night-time period starts.  This includes incomplete periods
    # at either ends
    DNbndries = sorted([dts[0]] + \
                       [ts.to_datetime() for ts in set(dtsDN)] + \
                       [dts[-1]])

    # Fill in colours between the boundaries
    for leftbnd, rightbnd in zip(DNbndries[: -1], DNbndries[1:]):
        if leftbnd.hour == hour_daystart or \
               rightbnd.hour == hour_nightstart:
            ax.axvspan(leftbnd, rightbnd,
                       color = 'yellow', alpha = .1, edgecolor = None)
        else:
            ax.axvspan(leftbnd, rightbnd,
                       color = 'black', alpha = .1, edgecolor = None)
    return ax

def get_common_cmap_levels(cmap_levels, Nsteps = 10):
    '''
    Given several sets of colormap levels, find a common set
    that covers all the sets, with (Nsteps + 1) levels
    INPUT:
    cmap_levels --- a list of sets of colormap levels
    Nsteps --- number of intervals in the return set of levels
    OUTPUT:
    (extend, colormap min, colormap max)
    '''
    cmap_minmin = sorted(cmap_levels, key = lambda x: x[1])[0]
    cmap_maxmax = sorted(cmap_levels, key = lambda x: x[2])[-1]

    if cmap_minmin[0] in ['min', 'both']:
        extend_min = True
    else:
        extend_min = False
    if cmap_maxmax[0] in ['max', 'both']:
        extend_max = True
    else:
        extend_max = False

    if extend_min and not extend_max:
        extend = 'min'
    elif extend_max and not extend_min:
        extend = 'max'
    elif extend_min and extend_max:
        extend = 'both'
    else:
        extend = 'neither'

    cmap_limits = (extend, cmap_minmin[1], cmap_maxmax[2])
    return get_nicenround_steps(cmap_limits, Nsteps = Nsteps)

def get_cmap_limits(da, quantile = .5):
    '''
    Returns colormap limits from an xray DataArray.
    INPUT:
    da --- DataArray
    quantile --- between 0 and 1
    OUTPUT:
    (extend, lev_min or abs(lev_max), lev_max)
    extend --- how to extend the colormap, can be \'min\', \'max\', or \'both\'
    lev_min --- value of minimum level
    lev_max --- value of maximum level
    abs(lev_max) --- for data array with positive and negative values,
                     the value of the level that is farthest from zero
    '''
    df = da.to_pandas()
    dfs = df.stack()

    if np.all(dfs == 0):
        return ('both', - 1e-10, 1e-10)
    elif np.all(dfs <= 0):
        lev_min = dfs.quantile(q = 1 - quantile)
        return ('min', lev_min, 0.)
    elif np.all(dfs >= 0):
        lev_max = dfs.quantile(q = quantile)
        return ('max', 0., lev_max)
    else:
        return ('both',
                dfs[dfs <= 0].quantile(q = 1 - quantile),
                dfs[dfs >= 0].quantile(q = quantile))

def get_nicenround_steps(cmap_limits, Nsteps = 10):
    '''
    Readjust given colormap limits in order to have
    nice and rounded intervals between levels.
    INPUT:
    cmap_limits --- tuple of extend, min level, max level
    Nsteps --- number of intervals in the colormap
    OUTPUT:
    tuple of extend, min level, max level, interval between level
    '''
    if cmap_limits[0] == 'both':
        cmap_min, cmap_max, cmap_Nsteps = round_levels_with_zero_centred_between_two(
            vmin0 = cmap_limits[1], vmax0 = cmap_limits[2], Nsteps0 = Nsteps)
        cmap_step = (cmap_max - cmap_min) / cmap_Nsteps
        return ('both', cmap_min, cmap_max, cmap_step)
    elif cmap_limits[0] in ['min', 'max']:
        extend, cmap_min, cmap_max = cmap_limits
        rough_cmap_step = (cmap_max - cmap_min) / Nsteps
        nice_cmap_step = muths.round_to_1(rough_cmap_step)
        cmap_Nsteps = int(np.ceil((cmap_max - cmap_min) / nice_cmap_step))
        if extend == 'min':
            # align the top
            nice_cmap_min = cmap_max - cmap_Nsteps * nice_cmap_step
            return (extend, nice_cmap_min, cmap_max, nice_cmap_step)
        elif extend == 'max':
            # align the bottom
            nice_cmap_max = cmap_min + cmap_Nsteps * nice_cmap_step
            return (extend, cmap_min, nice_cmap_max, nice_cmap_step)

def get_suitable_cmap(cmap_levels):
    '''
    Selects a divergent colormap if the colormap levels
    are positive and negative.
    INPUT:
    cmap_levels --- tuple (extend, cmap_min, cmap_max, cmap_step)
                    where extend is either \'min\', \'max\' or \'both\'.
    '''
    extend = cmap_levels[0]
    if extend == 'both':
        return plt.get_cmap('bwr')
    else:
        return plt.get_cmap('PuBuGn')

def contourf_interest_for_all_cases(d3sets, interest = 'CLOUD',
                                   cmap = plt.get_cmap('PuBuGn')):
    '''
    Function to contourf a set of DataArrays
    INPUT:
    d3sets --- a dictionary of xray.Datasets containing xray.DataArrays of
               dimensions (time, lev, lon, lat)
    interest --- which field to plot
    cmap --- matplotlib colormap object, i.e. which colormap to use for contourf
    '''
    cases = sorted(d3sets.keys())

    das = [d3sets[case][interest][{'lon': 0, 'lat': 0}] for case in cases]

    quantile, Nsteps = .9999, 17
    cmap_limitss = [get_cmap_limits(da, quantile = quantile) for da in das]
    cmap_levels = [get_nicenround_steps(cmap_limits, Nsteps = Nsteps)
                   for cmap_limits in cmap_limitss]

    common_cmap_levels = get_common_cmap_levels(cmap_levels, Nsteps = Nsteps)

    if cmap == 'auto-select':
        cmap = get_suitable_cmap(common_cmap_levels)

    Nplots = len(cases)

    fig, axes = plt.subplots(nrows = 1, ncols = Nplots,
                             figsize = (10. * Nplots, 5), dpi = 300)

    for ax, da in zip([axes] if Nplots == 1 else axes, das):
        ax = contourf_DataArray(ax, da,
                                cmap_levels = common_cmap_levels,
                                cmap = cmap)
    plt.tight_layout()
    return fig

def plotVS_timeaveraged_interest_for_all_cases(d3sets, diff_d3sets,
                                               interest = 'CLOUD',
                                               xscale = 'linear',
                                               bot_xlim = None, bot_xlabels_rotate = 0.,
                                               top_xlim = None, top_xlabels_rotate = 0.,
                                               yscale = 'linear', ylim = None,
                                               linestyles = None):

    vspairs = [[p.strip() for p in diff_case.split('-')]
               for diff_case in sorted(diff_d3sets.keys())]

    Nplots = len(vspairs)

    fig, axes = plt.subplots(figsize = (9, 9), dpi = 300,
                             nrows = 1, ncols = Nplots,
                             sharey = True)

    labels, handles = [], []

    axes.invert_yaxis() if Nplots ==1 else axes[0].invert_yaxis()

    for ax, vspair in zip([axes] if Nplots == 1 else axes, vspairs):
        x, y = vspair

        for model in vspair:
            if interest in ['AREI', 'AREL']:
                da = dates.average_over_time(d3sets[model][interest], key = lambda x: x > 0)
            else:
                da = dates.average_over_time(d3sets[model][interest])

            # plot each member in the comparison pair
            ax = plot_vertical_profile(ax, da,
                                       label = '{}'.format(model),
                                       colour = linestyles[model]['colour'],
                                       linestyle = linestyles[model]['linestyle'],
                                       xscale = xscale, xlabels_rotate = bot_xlabels_rotate,
                                       yscale = yscale,
                                       xlim = bot_xlim)

        ax = axes_beyond_ticks(ax, which = 'x')

        # plot difference on twiny
        diff_colour = (0.929, 0.329, 0.972)
        ax2 = ax.twiny()

        da = dates.average_over_time(diff_d3sets[' - '.join(vspair)][interest])

        ax2 = plot_vertical_profile(ax2, da,
                                    label = 'difference',
                                    colour = diff_colour, linestyle = '-',
                                    xscale = xscale, xlabels_rotate = top_xlabels_rotate,
                                    yscale = yscale,
                                    xlim = top_xlim)

        [ticklabel.set_color(diff_colour) for ticklabel in ax2.xaxis.get_ticklabels()]
        ax2 = axes_beyond_ticks(ax2, which = 'x')
        ax2.axvline(x = 0, color = 'r', alpha = .5)

        handles1, labels1 = ax.get_legend_handles_labels()
        handles2, labels2 = ax.get_legend_handles_labels()
        handles.extend(handles1 + handles2)
        labels.extend(labels1 + labels2)

    ## Make 1 legend for whole figure
    uhandles, ulabels = misc.any_unique_labels(handles, labels)
    uhandles, ulabels = zip(*sorted(zip(uhandles, ulabels), key = lambda x: x[1]))
    fig.legend(uhandles, ulabels,
               loc = 'center', ncol = 3,
               bbox_to_anchor = (.35, .91), prop = {'size': 14})

    fig.suptitle(da.attrs['long_name'])

    plt.figtext(x = 0.02, y = .5, s = 'lev [mbar]', rotation = 90.)
    plt.figtext(x = .45, y = 0.08, s = '{} [{}]'.format(interest, da.units))
    plt.figtext(x = .7, y = .91, s = 'difference')
    if hasattr(ax2, 'xaxis_pow'):
        plt.figtext(x = .9, y = 0.08, s = '1e{}'.format(- ax.xaxis_pow))
        plt.figtext(x = .9, y = .91, s = '1e{}'.format(- ax2.xaxis_pow))
    plt.subplots_adjust(wspace = 0., top = .84, bottom = .15)
    return fig

def plotVS_interest_for_all_cases(dsets, diff_dsets, interest = 'FLNT',
                                  left_ylim = None,
                                  right_ylim = None,
                                  linestyles = None):

    vspairs = [[p.strip() for p in diff_case.split('-')]
               for diff_case in sorted(diff_dsets.keys())]

    Nplots = len(vspairs)

    fig, axes = plt.subplots(figsize = (11, 3 * Nplots + 2.5), dpi = 300,
                             nrows = Nplots, ncols = 1)

    handles, labels = [], []

    for ax, vspair in zip([axes] if Nplots == 1 else axes, vspairs):
        x, y = vspair

        for model in vspair:
            ax = plot_DataArray(ax, dsets[model][interest],
                                label = '{}'.format(model),
                                colour = linestyles[model]['colour'],
                                linestyle = linestyles[model]['linestyle'],
                                marker = '',
                                ylabel = interest, ylim = left_ylim)

        da = diff_dsets[' - '.join(vspair)][interest]

        diff_colour = (0.929, 0.329, 0.972)
        ax2 = ax.twinx()
        ax2 = plot_DataArray(ax2, da, label = 'difference',
                             title = '',
                             colour = diff_colour, linestyle = '-', marker = '',
                             ylim = right_ylim)
        ax2.set_ylabel('')
        [label.set_color(diff_colour)
         for label in ax2.yaxis.get_ticklabels()]

        ### background shading for daytime and nighttime
        ax2 = daytime_nighttime_shading(ax2, da.coords['time'].to_pandas().index,
                                        hour_daystart = 6, hour_nightstart = 18)

        ## collect handles and labels for legend later
        handles1, labels1 = ax.get_legend_handles_labels()
        handles2, labels2 = ax2.get_legend_handles_labels()
        handles.extend(handles1 + handles2)
        labels.extend(labels1 + labels2)

    ## Make 1 legend for whole figure
    uhandles, ulabels = misc.any_unique_labels(handles, labels)
    uhandles, ulabels = zip(*sorted(zip(uhandles, ulabels),
                                    key = lambda x: x[1]))

    fig.suptitle(da.attrs['long_name'])

    fig.legend(uhandles, ulabels,
               loc = 'center', ncol = 3,
               bbox_to_anchor = (.45, .91), prop = {'size': 12})
    plt.subplots_adjust(wspace = 0., top = .84, bottom = .15)
    return fig

def plot_pdseries_indexVSvalues_linearlog(srss=None,
                                          names=None,
                                          colours=None,
                                          linestyles=None,
                                          markers=None,
                                          ylim=None, ylabel=None,
                                          xlim_linear=None, xlim_log=None,
                                          xlabel=None,
                                          title=None,
                                          figsize=(8, 5)):
    '''
    Plots index versus values for one or more Pandas.Series,
    for both linear and log y-scales.

    When y-scale is linear, x-axis limits are the minimum and maximum
    values for the range of y above 1.  When y-scale log, x-axis limits
    are the minumum and maximum values for the range of y below 1.

    Parameters
    ----------
    srss : list of Pandas series to plot
    names : corresponding list of strings to label the above;
            these are used in the legend
    colours : corresponding list of matplotlib colours
              (e.g. \'k\', \'r\', etc.)
    linestyles : corresponding list of matplotlib linestyles
                 (e.g. \'-\', \'--\', etc.)
    markers : corresponding list of matplotlib markers
              (e.g. \'o\', \'+\', etc.)
    ylim : a tuple of length two containing the lower and upper
           limits on the y-axis
    xlim_linear : a tuple of length two containing
                  the lower and upper limits on the linear x-axis
    xlim_log : a tuple of length two containing the lower and upper limits
               on the log x-axis
    title : string containing the title of the figure
    xlabel : string containing the x-axis label
    ylabel : string containing the y-axis label
    figsize : tuple of length two containing the width and height,
              in centimetres, of the figure
    fig : matplotlib figure object for the plot
    '''
    xys = list(itertools.chain(*[(srs.values, srs.index.values)
                                 for srs in srss]))

    if not markers:
        markers = [None for _ in range(len(xys))]

    if not linestyles:
        linestyles = ['-' for _ in range(len(xys))]

    if not colours:
        colour_cycle = itertools.cycle(matplotlib_basic_colours())
        colours = [next(colour_cycle) for _ in range(len(xys))]

    if not names:
        names = [k + 1 for k in range(len(xys))]

    fig, axs = plt.subplots(nrows=1, ncols=2, figsize=figsize)

    for yscale, ax in zip(('linear', 'log'), axs):

        lines = ax.plot(*xys)

        [plt.setp(line,
                  linestyle=style, color=colour, marker=marker,
                  linewidth=2.)
         for line, style, colour, marker
         in zip(lines, linestyles, colours, markers)]

        ax.set_title(title)
        ax.grid(b=True)
        ax.legend(names, loc='best')

        if ylim:
            ax.set_ylim(ylim)
        ax.set_yscale(yscale)
        ax.invert_yaxis()
        ax.set_ylabel(ylabel)

        if yscale == 'linear':
            if xlim_linear:
                ax.set_xlim(xlim_linear)
            else:
                if ylim:
                    ymin, ymax = ylim
                    xmin = min([srs[(srs.index > 1e0) & (srs.index < ymax)]
                                .min() for srs in srss])
                    xmax = max([srs[(srs.index > 1e0) & (srs.index < ymin)]
                                .max() for srs in srss])
                else:
                    xmin = min([srs[srs.index > 1e0].min() for srs in srss])
                    xmax = max([srs[srs.index > 1e0].max() for srs in srss])

                dx = xmax - xmin
                xmin -= .1 * dx
                xmax += .1 * dx
                ax.set_xlim((xmin, xmax))
        elif yscale == 'log':
            if xlim_log:
                ax.set_xlim(xlim_log)
            else:
                if ylim:
                    ymin, ymax = ylim
                    xmin = min([srs[(srs.index < 1e0) & (srs.index > ymin)]
                                .min() for srs in srss])
                    xmax = max([srs[(srs.index < 1e0) & (srs.index > ymin)]
                                .max() for srs in srss])
                else:
                    xmin = min([srs[srs.index < 1e0].min() for srs in srss])
                    xmax = max([srs[srs.index < 1e0].max() for srs in srss])

                dx = xmax - xmin
                xmin -= .1 * dx
                xmax += .1 * dx
                ax.set_xlim((xmin, xmax))

        ax.xaxis.get_major_formatter().set_powerlimits((0, 1))
        ax.set_xlabel(xlabel)
    return fig

def plot_pandas_series(ax=None,
                       srss=None, names=None,
                       colours=None, linestyles=None, markers=None,
                       values_vs_index=False,
                       logy=False, inverty=False, ylim=None, ylabel='ylabel',
                       logx=False, invertx=False, xlim=None, xlabel='xlabel',
                       title='title'):

    if values_vs_index:
        xys = list(itertools.chain(*[(srs.values, srs.index.values)
                                     for srs in srss]))
    else:
        xys = list(itertools.chain(*[(srs.index.values, srs.values)
                                     for srs in srss]))

    if not markers:
        markers = [None for _ in range(len(xys))]

    if not linestyles:
        linestyles = ['-' for _ in range(len(xys))]

    if not colours:
        colour_cycle = itertools.cycle(matplotlib_basic_colours())
        colours = [next(colour_cycle) for _ in range(len(xys))]

    if not names:
        names = [k + 1 for k in range(len(xys))]

    lines = ax.plot(*xys)

    [plt.setp(line,
              linestyle=style, color=colour, marker=marker,
              linewidth=1.7)
     for line, style, colour, marker
     in zip(lines, linestyles, colours, markers)]

    ax.set_title(title)
    ax.grid(b=True)
    ax.legend(names, loc='best')

    if ylim:
        ax.set_ylim(ylim)

    if logy:
        ax.set_yscale('log')

    if inverty:
        ax.invert_yaxis()

    ax.set_ylabel(ylabel)

    if xlim:
        ax.set_xlim(xlim)

    if logx:
        ax.set_xscale('log')

    ax.set_xlabel(xlabel)

    return ax


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from utils.sent_utils import inf_batch_iterator, iterate_minibatches_indices
from utils.common_utils import log

import sklearn.linear_model as linear_model
import tensorflow as tf
import numpy as np

def linear_mapping(x, y, n_jobs=8):
  # train least square
  projection = linear_model.LinearRegression(n_jobs=n_jobs)
  projection.fit(x, y)
  return projection.predict

def mlp_mapping(x, y, epochs=30, batch_size=256, lr=1e-3, wd=1e-4,
                activation=tf.nn.relu, norm=False):
  dim = y.shape[1]
  l2_reg = tf.keras.regularizers.l2(wd)
  model = tf.keras.models.Sequential()
  model.add(tf.keras.layers.Dense(dim, activation=activation, name='reg_fc',
                                  kernel_regularizer=l2_reg))
  model.add(tf.keras.layers.Dense(dim, name='reg_output',
                                  kernel_regularizer=l2_reg))
  if norm:
    model.add(tf.keras.layers.Lambda(lambda x: tf.nn.l2_normalize(x, axis=-1),
                                     name='reg_norm'))
  model.compile(tf.keras.optimizers.Adam(lr), loss='mse', metric='mse')
  model.fit(x, y, batch_size=batch_size, verbose=1, epochs=epochs)

  def mapping(z):
    mapped = model.predict(z, batch_size=1024)
    tf.keras.backend.clear_session()
    return mapped

  return mapping

class Generator(object):
  def __init__(self, output_dim, activation=tf.nn.relu):
    self.fc1 = tf.keras.layers.Dense(output_dim, activation=activation,
                                     name='gen_fc1')
    # self.fc2 = tf.keras.layers.Dense(output_dim, activation=activation,
    #                                  name='gen_fc2')
    self.output = tf.keras.layers.Dense(output_dim, name='gen_output')

  def forward(self, x):
    x = self.fc1(x)
    # x = self.fc2(x)
    x = self.output(x)
    return x

class Discriminator(object):
  def __init__(self, output_dim, activation=tf.nn.relu):
    self.fc1 = tf.keras.layers.Dense(output_dim, activation=activation,
                                     name='disc_fc1')
    # self.fc2 = tf.keras.layers.Dense(output_dim, activation=activation,
    #                                  name='disc_fc2')
    self.output = tf.keras.layers.Dense(1, name='disc_output')

  def forward(self, y):
    y = self.fc1(y)
    # y = self.fc2(y)
    y = self.output(y)
    return y

class WGANGP(object):
  """Learn a GAN that translates x to y"""
  def __init__(self, x_dim, y_dim, lr=5e-4, lmbda=10., gamma=0., beta1=0.5,
               activation=tf.nn.relu):
    self.generator = Generator(y_dim, activation)
    self.discriminator = Discriminator(y_dim, activation)
    self.input_x = tf.placeholder(tf.float32, shape=(None, x_dim), name='x')
    self.input_y = tf.placeholder(tf.float32, shape=(None, y_dim), name='y')

    self.real_score = self.discriminator.forward(self.input_y)
    self.fake_y = self.generator.forward(self.input_x)
    self.fake_score = self.discriminator.forward(self.fake_y)

    # calculate gradient penalty
    alpha = tf.random.uniform(shape=tf.shape(self.fake_y), minval=0.,
                              maxval=1.)
    self.interpolate_y = alpha * self.input_y + self.fake_y
    grad = tf.gradients(self.discriminator.forward(self.interpolate_y),
                        [self.interpolate_y])[0]
    grad_norm = tf.norm(grad, axis=1, ord='euclidean')
    self.grad_pen = tf.reduce_mean(tf.square(grad_norm - 1))

    t_vars = tf.trainable_variables()

    self.disc_params = [v for v in t_vars if v.name.startswith('disc')]
    self.disc_optimizer = tf.train.AdamOptimizer(lr, beta1=beta1)
    self.disc_loss = tf.reduce_mean(self.fake_score) - tf.reduce_mean(
      self.real_score)
    self.disc_train_ops = self.disc_optimizer.minimize(
        self.disc_loss + lmbda * self.grad_pen, var_list=self.disc_params)

    l2_loss = tf.reduce_sum(tf.square(self.fake_y - self.input_y), axis=-1)
    self.l2_loss = tf.reduce_mean(l2_loss)

    self.gen_params = [v for v in t_vars if v.name.startswith('gen')]
    self.gen_optimizer = tf.train.AdamOptimizer(lr, beta1=beta1)
    self.gen_loss = -tf.reduce_mean(self.fake_score)
    self.gen_train_ops = self.gen_optimizer.minimize(
        self.gen_loss + gamma * self.l2_loss, var_list=self.gen_params)

  def generate(self, sess, x):
    return sess.run(self.fake_y, {self.input_x: x})

  def train_disc_one_batch(self, sess, x, y):
    d_err, _ = sess.run([self.disc_loss, self.disc_train_ops],
                        {self.input_x: x, self.input_y: y})
    return d_err

  def train_gen_one_batch(self, sess, x, y):
    g_err, l2_err, _ = sess.run([-self.gen_loss, self.l2_loss,
                                 self.gen_train_ops],
                                {self.input_x: x, self.input_y: y})
    return g_err, l2_err

def gan_mapping(x, y, lr=1e-4, lmbda=10., gamma=0., beta1=0.5,
                activation=tf.nn.relu, epoch=30, disc_iters=10,
                batch_size=128):
  n_data, x_dim = x.shape
  y_dim = y.shape[1]
  model = WGANGP(x_dim, y_dim, lr=lr, lmbda=lmbda, gamma=gamma, beta1=beta1,
                 activation=activation)

  gen_sampler = inf_batch_iterator(n_data, batch_size)
  disc_sampler = inf_batch_iterator(n_data, batch_size)
  num_batch_per_epoch = n_data // batch_size + (n_data % batch_size) != 0

  sess = tf.Session()
  sess.run(tf.global_variables_initializer())
  for e in range(epoch):
    train_d_loss = []
    train_g_loss = []
    train_l2_loss = []
    for _ in range(num_batch_per_epoch):
      # train disc first
      for _ in range(disc_iters):
        disc_idx = next(disc_sampler)
        disc_x, disc_y = x[disc_idx], y[disc_idx]
        d_err = model.train_disc_one_batch(sess, disc_x, disc_y)
        train_d_loss.append(d_err)

      gen_idx = next(gen_sampler)
      gen_x, gen_y = x[gen_idx], y[gen_idx]
      g_err, l2_err = model.train_gen_one_batch(sess, gen_x, gen_y)
      train_g_loss.append(g_err)
      train_l2_loss.append(l2_err)

    train_d_loss = np.mean(train_d_loss)
    train_g_loss = np.mean(train_g_loss)
    train_l2_loss = np.mean(train_l2_loss)
    log('Epoch: {}, disc loss: {:.4f}, gen loss: {:.4f},'
        ' l2 loss: {:.4f}'.format(
          e + 1, train_d_loss, train_g_loss, train_l2_loss))

  def mapping(z):
    mapped = []
    for idx in iterate_minibatches_indices(len(z), batch_size=2048):
      batch_mapped = model.generate(sess, z[idx])
      mapped.append(batch_mapped)

    tf.keras.backend.clear_session()
    return np.vstack(mapped)

  return mapping


import webapp2
from models.snippet import Snippet
from models.user import User
from models.usersnippet import Rate
from webapp2_extras import jinja2, sessions
from handlers.register import RegisterHandler
from handlers.base import BaseHandler
from handlers.login import LoginHandler, LogoutHandler
from handlers.snippets \
    import CreateSnippetHandler, MySnippetsHandler, DeleteSnippetHandler, EditSnippetHandler,RateSnippetHandler

class MainHandler(BaseHandler):
    def get(self):
        snippets = Snippet.query().order(-Snippet.creation_date).fetch()
        session = self.session.get('user')
        is_logged = False
        logged_user = None
        if session is not None:
            is_logged = True
            logged_user = User.query(User.username == self.session.get('user')).get()

        for snippet in snippets:
            user = User.query(User.key == snippet.user).get()
            snippet.completeuser = user
            if logged_user:
                rate = Rate.query(Rate.snippet == snippet.key, Rate.user == logged_user.key).get()
                snippet.user_valoration = rate.positive_valoration if rate else None

        data = {
            "is_logged": is_logged,
            "snippets": snippets
        }

        jinja = jinja2.get_jinja2(app=self.app)
        self.response.write(jinja.render_template("index.html", **data))

config = {}
config['webapp2_extras.sessions'] = {
    'secret_key': 'prepre',
}

app = webapp2.WSGIApplication([
    ('/', MainHandler),
    ('/register', RegisterHandler),
    ('/logout', LogoutHandler),
    ('/login', LoginHandler),
    ('/newsnippet', CreateSnippetHandler),
    ('/mysnippets', MySnippetsHandler),
    ('/mysnippets/delete', DeleteSnippetHandler),
    ('/editsnippet', EditSnippetHandler),
    ('/rate', RateSnippetHandler)
], debug=True, config=config)


from __future__ import annotations

import torch.nn as nn
from typing import NamedTuple, List

class LightweightNode(NamedTuple):
    """Attach a symbolic name to an ``nn.Module``."""

    name:   str
    module: nn.Module

    @property
    def qualified_path(self) -> List[str]:
        return self.name.split('.')

    def __eq__(self, other: LightweightNode) -> bool:
        cond_type = isinstance(other, LightweightNode)
        cond_name = self.name == other.name
        cond_obj  = self.module is other.module
        return cond_type and cond_name and cond_obj

class LightweightNodeList(list):  # https://stackoverflow.com/a/24160909
    """A list sub-class accepting only ``LightweightNode`` items."""

    def append(self, item):
        if not isinstance(item, LightweightNode):
            raise TypeError
        super(LightweightNodeList, self).append(item)

    def insert(self, index, item):
        if not isinstance(item, LightweightNode):
            raise TypeError
        super(LightweightNodeList, self).insert(index, item)

    def __add__(self, item):
        if not isinstance(item, LightweightNode):
            raise TypeError
        super(LightweightNodeList, self).__add__(item)

    def __iadd__(self, item):
        if not isinstance(item, LightweightNode):
            raise TypeError
        super(LightweightNodeList, self).__iadd__(item)

    def __str__(self) -> str:
        """When printed, this string shows a nicely formatted list of named
        ``nn.Module``s.

        This functionality is useful, for instance, when users want to
        visually inspect a floating-point ``nn.Module`` to decide how to
        quantise it.

        This method overwrites the default behaviour of Python lists'
        ``__str__`` dunder method, which resolves to calling the ``__repr__``
        method of each item in the list (https://stackoverflow.com/a/727779).

        """

        # overwrite the default behaviour of Python lists' `__str__` method
        max_chars = max(map(lambda node: len(node.name), self))

        str_ = f""
        for node in self:
            str_ += f"\n"
            name = ".".join(reversed(node.qualified_path)).rjust(max_chars)
            str_ += f"{name}\t{node.module}"

        return str_

    def show(self) -> None:
        print(self.__str__())


"""
An example of using Python's built in `multiprocessing` or `discomp`
to test many hyperparameters.

To test 30 hyperparameter configurations on my MacBook Pro it took:
* 81 minutes and 30 seconds using `map`
* 34 minutes and 40 seconds using `multiprocessing.Pool().map`
* 8 minutes and 52 seconds using `discomp.Pool().map`

More background at https://dis.co/blog/machine-learning-with-disco/
"""

import os
import random
import multiprocessing

import numpy as np
import xgboost
from sklearn.model_selection import train_test_split
from sklearn import metrics

import pandas as pd
import discomp

os.environ["DISCO_LOGIN_USER"] = "your-email-or-use-environment-variables"
os.environ["DISCO_LOGIN_PASSWORD"] = "your-password-or-api-key"

print("loading data")
url = "https://discodemo.s3-us-west-2.amazonaws.com/yuval/creditcard.csv"
df = pd.read_csv(url, header=0)

def test_fit(params):
    # This function will be launched multiple times in other processes
    Y = df["Class"]
    X = df.drop(["Class"], axis=1)
    print("Dataset X shape: {}".format(X.shape))
    print("Dataset Y shape: {}".format(Y.shape))
    print("{}\n\n{}".format(X[:2], Y[:2]))

    # split data into train and test sets
    seed = 7
    test_size = 0.33
    X_train, X_test, y_train, y_test = train_test_split(
        X, Y, test_size=test_size, random_state=seed
    )

    print("fit model")
    model = xgboost.XGBRegressor(objective="reg:squarederror", **params)
    model.fit(X_train, y_train)

    # make predictions for test data
    print("predict")
    y_pred = model.predict(X_test)

    print("evaluating predictions")
    fpr, tpr, thresholds = metrics.roc_curve(y_test, y_pred)
    roc_auc = metrics.auc(fpr, tpr)

    print(f"roc_auc: {roc_auc:.2} params: {params}")
    return roc_auc

# Set up a random search through hyper parameters
param_options = []
param_test_count = 30
for i in range(param_test_count):
    params = dict(
        max_depth=random.randint(2, 100),
        num_boost_round=random.randint(2, 1000),
        learning_rate=random.uniform(0.0001, 1.0),
        gamma=random.randint(0, 20),
        subsample=random.uniform(0.00001, 1.0),
    )
    param_options.append(params)

print("distributing tests")
# For purely local `multiprocessing`, replace the `discomp` line with the following.
# with multiprocessing.Pool() as pool:
with discomp.Pool() as pool:
    # `discomp.Pool().map` will send the `test_fit` to multiple machines to run in parallel.
    # The amount of machines depends on the number of items in the `param_options` list.
    results = pool.map(test_fit, param_options)

    # For a serial example, replace the above two lines with the following.
    # results = list(map(test_fit, param_options))
    best_index = np.argmax(results)
    print(f"Ran {param_test_count} tests")
    print(f"Best result {results[best_index]:0.2f}")
    print(f"Best params: {param_options[best_index]}")


import glob2
import os
import cv2
import numpy as np
threshold = 35
input_dir = 'SCUT-EnsText/train_sets'
save_str = "stroke"

if not os.path.exists(os.path.join(input_dir, save_str)):
    os.mkdir(os.path.join(input_dir, save_str))
images = glob2.glob(os.path.join(input_dir, 'all_images', '*.jpg'))
print("#item: ", len(images))

for i, item in enumerate(images):
    label_img_path = os.path.join(input_dir, 'all_labels', os.path.basename(item))
    img = cv2.imread(item)
    label_img = cv2.imread(label_img_path)
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray_label_img = cv2.cvtColor(label_img, cv2.COLOR_BGR2GRAY)

    diff = np.abs(gray_img.astype(np.float32) - gray_label_img.astype(np.float32))
    diff_threshold = threshold
    diff[diff < diff_threshold] = 1
    diff[diff >= diff_threshold] = 0
    diff = (diff * 255).astype("uint8")
    cv2.imwrite(os.path.join(input_dir, save_str, os.path.basename(item)), diff)


# -*- coding:utf-8 -*-
import shutil

from flask_script import Command, Option, Server

import flask
import requests

from whyis.data_extensions import DATA_EXTENSIONS

import rdflib
from whyis.nanopub import Nanopublication
import tempfile

from whyis.namespace import np
from whyis.blueprint.nanopub.nanopub_utils import load_nanopub_graph
import sys

class LoadNanopub(Command):
    '''Add a nanopublication to the knowledge graph.'''

    _TEMP_STORE_DEFAULT = None

    def get_options(self):
        return [
            Option('-i', '--input', dest='input_file', required=True, help='Path to file containing nanopub', type=str),
            Option('-f', '--format', dest='file_format', default='trig', help='File format (default: trig; also turtle, json-ld, xml, nquads, nt, rdfa)', type=str),
            Option('-r', '--revises', dest='was_revision_of', help="URI of nanopublication that this is a revision of", type=str),
            Option("--temp-store", dest="temp_store", type=str, default=self._TEMP_STORE_DEFAULT, help="backing store type to use for temporary graphs; deprecated")
        ]

    def run(self, input_file, file_format, temp_store=_TEMP_STORE_DEFAULT, was_revision_of=None):
        flask.current_app.managed = True
        if was_revision_of is not None:
            wasRevisionOf = set(flask.current_app.db.objects(predicate=np.hasAssertion,
                                                             subject=rdflib.URIRef(was_revision_of)))
            if len(wasRevisionOf) == 0:
                print("Could not find active nanopublication to revise:", was_revision_of)
                return
            was_revision_of = wasRevisionOf
        g = rdflib.ConjunctiveGraph(identifier=rdflib.BNode().skolemize())
        if temp_store == "Oxigraph":
            g_store_tempdir = tempfile.mkdtemp()
            g.store.open(g_store_tempdir, True)
        else:
            g_store_tempdir = None
        # g = rdflib.ConjunctiveGraph(identifier=rdflib.BNode().skolemize())

        try:
            g1 = load_nanopub_graph(location=input_file, format=file_format, store=g.store)

            nanopubs = []
            for npub in flask.current_app.nanopub_manager.prepare(g):
                if was_revision_of is not None:
                    for r in was_revision_of:
                        print("Marking as revision of", r)
                        npub.pubinfo.add((npub.assertion.identifier, flask.current_app.NS.prov.wasRevisionOf, r))
                print('Prepared', npub.identifier)
                nanopubs.append(npub)
            flask.current_app.nanopub_manager.publish(*nanopubs)
            print("Published", npub.identifier)
        finally:
            if g_store_tempdir is not None:
                shutil.rmtree(g_store_tempdir)


from sklearn.cluster import KMeans, AgglomerativeClustering
import scipy.stats as ss
import seaborn as sns
import matplotlib.pyplot as plt
import math
import numpy as np
import json
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

def visualize_data(data_dict, train):
    fig, ax = plt.subplots(1, figsize=(15, 15))
    for i in range(len(data_dict)):
        if train:
            x = np.array(data_dict[f'a{i + 1}']['strains']).squeeze()
        else:
            x = np.array(data_dict[f'te_a{i + 1}']['strains']).squeeze()
        ax.plot(x, label=f'specimens a{i + 1}')

    ax.set_xlabel('time steps')
    ax.set_ylabel('strains')
    ax.legend(loc='upper left')
    if train:
        fig.suptitle('Training raw data')
    else:
        fig.suptitle('Testing raw data')
    plt.savefig('raw_data.png')

def plot_elbow(wcss):
    """
    Plot elbow method for k-means clustering
    :param wcss: List of 13 elements (from n_clusters=3 to n_clusters=15)
    :return: None
    """
    plt.figure()
    K = range(3, 16)
    plt.plot(K, wcss, 'bx-')
    plt.savefig('elbow.png')

def visualize_cluster_results(train_dict, n_clusters):
    """
    :param train_dict: dictionary with keys a1, a2, ..., a6 where the values are the clusters
    :param n_clusters: int (optional for using)
    :return:
    """

    """ TODO:
        Your code goes here
    """
    fig, ax = plt.subplots(1, figsize=(7, 7))
    for i in range(len(train_dict)):
        stepfunc = train_dict[f'a{i + 1}']
        stepfunc = np.array(stepfunc)
        stepfunc = stepfunc.flatten() + 1
        ax.plot(stepfunc, label=f'specimen a{i + 1}')

    ax.set_xlabel('time steps')
    ax.set_ylabel('clusters')
    ax.legend(loc='upper left')
    fig.suptitle('Sojourn time distributions')
    plt.savefig('Sojourn_distributions.png')

def plot_distributions(sjrn_t, n_clusters):
    plt.rcParams.update({'font.size': 14})
    exit = False
    if n_clusters < 4:
        cols = n_clusters
    else:
        cols = 4
    fig, ax = plt.subplots(
        nrows=int(math.ceil(n_clusters / cols)),
        ncols=cols, figsize=(25, 30),
        squeeze=False)
    indx = 0
    for i in range(int(math.ceil(n_clusters/cols))):
        for j in range(cols):
            sns.kdeplot(sjrn_t[:, indx], ax=ax[i, j])
            ax[i, j].set_title(f'Cluster {indx + 1}')
            ax[i, j].ticklabel_format(axis="y", style="sci", scilimits=(0, 0))
            indx += 1
            if indx == n_clusters:
                exit = True
                break
        if exit:
            break
    plt.subplots_adjust(hspace=0.4)
    plt.savefig('distributions.png')

def plot_fitted_and_initial_distributions(sjrn_t, n_clusters, dist_param_list):
    exit = False
    if n_clusters < 4:
        cols = n_clusters
    else:
        cols = 4
    fig, ax = plt.subplots(
        nrows=int(math.ceil(n_clusters / cols)),
        ncols=cols, figsize=(25, 20),
        squeeze=False)
    indx = 0
    for i in range(int(math.ceil(n_clusters/cols))):
        for j in range(cols):
            sns.kdeplot(sjrn_t[:, indx], ax=ax[i, j])

            c, loc, scale = dist_param_list[i]
            samples = ss.weibull_min.rvs(c, loc=loc, scale=scale, size=1000)
            sns.kdeplot(samples, ax=ax[i, j])
            ax[i, j].set_title(f'Cluster {indx + 1}')
            ax[i, j].ticklabel_format(axis="y", style="sci", scilimits=(0, 0))
            indx += 1
            if indx == n_clusters:
                exit = True
                break
        if exit:
            break
    plt.subplots_adjust(hspace=0.4)
    plt.savefig('fitted_initial_distributions.png')

def train(train_dict, data_dict, n_clusters, cluster_method, linkage):
    # use experiments a1 - a6 for training, a7-a8 for testing (named as te_a1, te_a2)

    x_train = []
    if cluster_method == 'kmeans':
        wcss = []
        for i in range(len(train_dict)):
            x_train.extend(data_dict[f'a{i + 1}']['strains'])
    else:  # in Agglomerative all samples should be used
        for i in range(len(data_dict)):
            if i < 6:
                x_train.extend(data_dict[f'a{i + 1}']['strains'])
            else:
                x_train.extend(data_dict[f'te_a{i + 1 - 6}']['strains'])
    x_train = np.array(x_train).squeeze().reshape(-1, 1)

    if cluster_method == 'kmeans':
        x_train = np.sort(x_train.reshape(-1))

        # Clustering using k-means, to check the optimal n_clusters with the elbow method (WCSS)
        # multiple k-means should be run in a loop, then you can plot the elbow (see plot_elbow() function defined above).
        # When you finish don't forget to remove the loop and comment again the plot_elbow(wcss) call...

        cluster_model = KMeans(n_clusters=n_clusters, n_init=10).fit(x_train.reshape(-1, 1))
        wcss.append(cluster_model.inertia_)
        # plot_elbow(wcss)

        _, indx = np.unique(cluster_model.labels_, return_index=True)
        lookup_labels = [cluster_model.labels_[i] for i in sorted(indx)]
    else:
        # Clustering using Agglomerative
        # Initialize and fit the clustering model using sklearn
        # The hyperparameters to use are the number of clusters and the linkage method
        cluster_model = AgglomerativeClustering(n_clusters=n_clusters, linkage=linkage)
        cluster_model = cluster_model.fit(x_train)
        lookup_labels = cluster_model.labels_

        """ TODO:
            Your code goes here
        """

    return cluster_model, lookup_labels, n_clusters

def predict(cluster_model, lookup_labels, data_dict, cluster_method, train):
    if train:
        clusters_dict = {f'a{i + 1}': [] for i in range(6)}
    else:
        clusters_dict = {f'te_a{i + 1}': [] for i in range(2)}

    for i in range(len(data_dict)):
        if cluster_method == 'kmeans':
            if train:
                pred = cluster_model.predict(
                    np.array(data_dict[f'a{i + 1}']['strains']).squeeze().reshape(-1, 1))
            else:
                pred = cluster_model.predict(
                    np.array(data_dict[f'te_a{i + 1}']['strains']).squeeze().reshape(-1, 1))
            labels = []
            for j in range(len(pred)):
                labels.append(np.where(lookup_labels == pred[j])[0])
            if train:
                clusters_dict[f'a{i + 1}'].extend(sorted(labels))
            else:
                clusters_dict[f'te_a{i + 1}'].extend(sorted(labels))
        else:
            if train:
                labels = cluster_model.fit_predict(
                    np.array(data_dict[f'a{i + 1}']['strains']).squeeze().reshape(-1, 1))
                clusters_dict[f'a{i + 1}'].extend(sorted(labels))
            else:
                labels = cluster_model.fit_predict(
                    np.array(data_dict[f'te_a{i + 1}']['strains']).squeeze().reshape(-1, 1))
                clusters_dict[f'te_a{i + 1}'].extend(sorted(labels))

    return clusters_dict

def sojourn_times(labels_dict, n_clusters, time_steps, train):
    """
    :param labels_dict: dictionary of cluster labels for each specimen
    For training, labels_dict has keys a1, a2, a3, a4, a5, a6, for testing it has keys te_a1, te_a2
    :param n_clusters: int --> number of clusters
    :param time_steps: int --> !!! don't forget the sojourn times should be measured in cycles not in steps !!!
    :param train: bool --> whether we are in training or testing phase (for naming the dictionary keys accordingly)
    :return: sojourn times array of shape [specimens, n_clusters] where specimens=train specimens for training, specimens=testing specimens for testing
    """

    # shape of sojourn times :  [specimens, n_clusters]
    sjrn_t = np.zeros((len(labels_dict), n_clusters))
    for i in range(len(labels_dict)):
        if train:
            labels = labels_dict[f'a{i + 1}']
        else:
            labels = labels_dict[f'te_a{i + 1}']
        for j in range(n_clusters):
            sjrn_t[i, j] = labels.count(j)
        sjrn_t[i, :] *= time_steps  # rescale to cycles
    sjrn_t[np.where(sjrn_t == 0)] = time_steps
    return sjrn_t

def fit_distribution_per_cluster(sojourn_times):
    dist = ss.weibull_min
    dist_param = []
    for i in range(sojourn_times.shape[1]):
        # store a tuple of parameters in order (if applicable): (shape, loc, scale)
        dist_param.append(dist.fit(sojourn_times[:, i], floc=0))
    return dist_param

def create_pdf_samples(dist_param_list, n_clusters):
    samples = []
    for i in range(n_clusters):
        c, loc, scale = dist_param_list[i]
        samples.append(np.sort(ss.weibull_min.rvs(c, loc=loc, scale=scale, size=10)))
    return samples

def rul_estimation(samples, data_clusters, sojourn_times, time_steps, n_clusters):
    """
    :param samples: List of np.array, the samples generated from the fitted distributions for each cluster
    :param data_clusters: dictionary of cluster labels for each specimen for testing phase (has keys te_a1, te_a2)
    :param sojourn_times: np.array of shape [test specimens, n_clusters]
    :param time_steps: int
    :param n_clusters: int
    :return: dictionary of rul estimations for each specimen
    """

    ruls = {f'te_a{i + 1}': {'mean_rul': []} for i in range(len(data_clusters))}

    for i in range(len(ruls)):
        prev = 0
        for j in range(n_clusters):
            time = np.arange(0, sojourn_times[i, j], time_steps)
            tot_rul = np.zeros((samples[j].shape[0]))
            rul = np.zeros((time.shape[0]))
            for k in range(j, n_clusters):
                tot_rul += samples[k]
            tot_mean = tot_rul.mean()
            for k in range(time.shape[0]):
                rul[k] = tot_mean - time[k]

            ruls[f'te_a{i + 1}']['mean_rul'].extend(rul)
            prev += sojourn_times[i, j]
    return ruls

def visualize_rul(test_dict, pred_rul_dict, time_steps):
    """
    Function to visualize the test RUL estimations for the specimens te_a1 and te_a2.
    Feel free to use any variable you want from the main function.
    :return: None
    """
    fig, ax = plt.subplots(1, 2, figsize=(23, 35), squeeze=False)
    indx = 0
    for i in range(1):
        for j in range(2):
            true_rul = np.array(test_dict[f'te_a{indx + 1}']['ruls']).squeeze()
            pred_rul = np.array(pred_rul_dict[f'te_a{indx + 1}']['mean_rul']).squeeze()
            x = np.arange(0, true_rul[0] + time_steps, time_steps)
            if pred_rul.shape[0] != true_rul.shape[0]:
                diff = pred_rul.shape[0] - true_rul.shape[0]
                pred_rul = np.delete(pred_rul, obj=-np.arange(diff))
            ax[i, j].plot(x, true_rul, linestyle='dashed', color='black', label='True RUL')
            ax[i, j].plot(x, pred_rul, color='blue')
            ax[i, j].ticklabel_format(axis="y", style="sci", scilimits=(0, 0))
            ax[i, j].set_xlabel('time steps')
            ax[i, j].set_ylabel('RUL')
            ax[i, j].set_title(f'Test specimen te_{indx + 1}')
            indx += 1
            ax[i, j].legend()
    fig.suptitle('Mean RUL predictions for train-test specimens')
    plt.savefig('rul_predictions.png')
    return None

def rul_estimation_with_uncertainty(samples, data_clusters, sojourn_times, time_steps, n_clusters):
    """

    :param samples: List of np.array, the samples generated from the fitted distributions for each cluster
    :param data_clusters: dictionary of cluster labels for each specimen for testing phase (has keys te_a1, te_a2)
    :param sojourn_times: np.array of shape [test specimens, n_clusters]
    :param time_steps: int
    :param n_clusters: int
    :return: dictionary of rul estimations for each specimen for mean, min and max of RUL corresponding to 95% confidence interval
    """

    # Hint: use scipy.stats.bootstrap (read the docs!) to estimate the confidence intervals by feeding the samples as input

    ruls = {f'te_a{i + 1}': {'mean_rul': [], 'min_rul': [], 'max_rul': []}
            for i in range(len(data_clusters))}

    for i in range(len(ruls)):
        """ TODO:
            Your code goes here
        """

    return ruls

def visualize_rul_with_uncertainty(test_dict, pred_rul_dict, time_steps):
    """
    Function to visualize the test RUL estimations for the specimens te_a1 and te_a2 with uncertainty.
    Feel free to use any variable you want from the main function.
    """

    """ TODO: OPTIONAL
        Your code goes here:
    """

    return None

def mse_loss(test_dict, pred_rul_dict):
    """
    :param test_dict: Nested dictionary with true RUL data
    :param pred_rul_dict: Nested dictionary with predicted RUL data
    :return: MSE loss
    """
    # Hint: Check the format of the dictionaries (nested dicts) and convert it to numpy array)
    actual_1 = np.array(test_dict['te_a1']['ruls'])
    mean_1 = np.array(pred_rul_dict['te_a1']['mean_rul'])
    actual_2 = np.array(test_dict['te_a2']['ruls'])
    mean_2 = np.array(pred_rul_dict['te_a2']['mean_rul'])
    # se1 = 0.0
    # se2 = 0.0
    # for i in range(len(mean_1)):
    #     se1_i = (actual_1[i] - mean_1[i])**2
    #     se1 = se1 + se1_i
    # for i in range(len(mean_2)):
    #     se2_i = (actual_2[i] - mean_2[i])**2
    #     se2 = se2 + se2_i
    se1 = np.sum((actual_1 - mean_1)**2)
    se2 = np.sum((actual_2 - mean_2)**2)
    mse = ((se1 / len(mean_1))**0.5 + (se2 / len(mean_2))**0.5) / 2

    """ TODO:
        Your code goes here
    """
    return mse

if __name__ == '__main__':
    print("Code running...\n")

    time_steps = 50  # every 50 seconds is one step, it's easier to use time steps, then convert to seconds for the results phase

    cluster_algorithm = ['kmeans', 'agglomerative']

    # Part 1: Clustering
    print('\n Part 1: Clustering')
    print('\n 1. Load two json files, the train.json and test.json for the training and testing data.'
          '\n 6 specimens for training and 2 for testing. Then, visualize the dataset.')
    input("\n Proceed ? Press any button to continue... ")

    # load from json
    # Train dictionary with keys: 'a1', 'a2', 'a3', 'a4', 'a5', 'a6' for the first 6 specimens
    with open('data/train.json', 'r') as fp:
        train_dict = json.load(fp)
    # Test dictionary with keys: 'te_a1', 'te_a2' for the remaining 2 specimens
    with open('data/test.json', 'r') as fp:
        test_dict = json.load(fp)

    # Stack together train, test dictionaries for utilization in the next steps
    data_dict = {}
    data_dict.update(train_dict)
    data_dict.update(test_dict)

    visualize_data(train_dict, train=True)
    visualize_data(test_dict, train=False)

    # 2. Create clusters using a training set
    print('\n 2. Create clusters using a training set')
    input("\n Proceed ? Press any button to continue... ")
    while True:
        choose_algo = input(
            '\n 2. Please choose a clustering algorithm from [K-means, Agglomerative]. Type "0" for Kmeans, "1" for Agglomerative : ')
        if choose_algo in ['0', '1']:
            break
        else:
            print('\n Wrong value given. Please choose one of the available clustering algorithms ( type 0 or 1)')

    cluster_method = cluster_algorithm[int(choose_algo)]

    if int(choose_algo) == 0:
        while True:
            n_clusters = input(
                '\n 3. Please enter the number of clusters for K-means in range [3, 15] : ')
            if 3 <= int(n_clusters) <= 15:
                break
            else:
                print(
                    '\n Number of clusters is out of the given range. Please make sure the number you typed belongs to the desired range')
        n_clusters = int(n_clusters)
        linkage = None
    else:
        while True:
            n_clusters = input(
                '\n 3. Please enter the number of clusters for Agglomerative in range [3, 15] : ')
            if 3 <= int(n_clusters) <= 15:
                break
            else:
                print(
                    '\n Number of clusters is out of the given range. Please make sure the number you typed belongs to the desired range')
        n_clusters = int(n_clusters)
        while True:
            linkage = input(
                "\n Please define the hyperparameter linkage for Agllomerative algorithm. "
                "Choose between {'ward', 'complete', 'average'} : ")
            if linkage in ['ward', 'complete', 'average']:
                break
            else:
                print(
                    '\n The given linkage method is not available. Please make sure you typed the method correctly.')

    # Build the cluster model
    cluster_model, lookup_labels, n_clusters = train(
        train_dict, data_dict, n_clusters, cluster_method, linkage)

    # Part 2: Sojourn time distributions
    print('\n Part 2: Sojourn time distributions')
    input("\n Proceed ? Press any button to continue... ")

    # Store on dictionaries the corresponding labels where the raw data were clustered
    train_labels_dict = predict(cluster_model, lookup_labels,
                                train_dict, cluster_method, train=True)
    test_labels_dict = predict(cluster_model, lookup_labels, test_dict, cluster_method, train=False)

    # 1. Degradation history for each training specimen in terms of label (cluster) sequence
    print('\n 1. Degradation history for each training specimen in terms of label (cluster) sequence')

    visualize_cluster_results(train_labels_dict, n_clusters=n_clusters)

    # 2. Identifying sojourn times at each cluster for each training specimen
    print('\n 2. Identifying sojourn times at each cluster for each training specimen')
    train_sojourn_time_arr = sojourn_times(train_labels_dict, n_clusters, time_steps, train=True)

    # 3.Plot probability distributions on sojourn times
    print('\n 3. Fitting probability distributions on sojourn times')
    plot_distributions(train_sojourn_time_arr, n_clusters)

    # 5. Output
    print('\n 5. Output')
    # fit the chosen family of distributions to the sojourn times
    train_dist_param_list = fit_distribution_per_cluster(train_sojourn_time_arr)
    # generate samples of the fitted distributions
    samples = create_pdf_samples(train_dist_param_list, n_clusters)

    # 6. Plot at the same graph the fitted distributions and the initial pdf of sojourn times
    print('\n 6. Plot at the same graph the fitted distributions and the initial pdf of sojourn times')
    plot_fitted_and_initial_distributions(train_sojourn_time_arr, n_clusters, train_dist_param_list)

    # Part 3: Prognostics
    print('\n Part 3: Prognostics')
    input("\n Proceed ? Press any button to continue... ")

    # 1. Use the tesing data to create the test sojourn times
    print('\n 1. Use the tesing data to create the test sojourn times')

    test_sojourn_time_arr = sojourn_times(test_labels_dict, n_clusters, time_steps, train=False)

    # 2. Calculate the RUL for each testing specimen
    print('\n 2. Calculate the RUL for each testing specimen')
    test_pred_rul_dict = rul_estimation(
        samples, test_labels_dict, test_sojourn_time_arr, time_steps, n_clusters)

    # 3. Visualize the RUL for each testing specimen
    print('\n 3. Visualize the RUL for each testing specimen')
    visualize_rul(test_dict, test_pred_rul_dict, time_steps)

    # Part 4: Error calculation
    print('\n 4. Error calculation ')
    input("\n Proceed ? Press any button to continue... ")

    test_mse = mse_loss(test_dict, test_pred_rul_dict)
    print(f"\n MSE loss for testing is {round(test_mse, 4)}")

    # Part 5. Optional task: Write a script able to predict not only the mean RUL but also the related X% confidence intervals (CIs) of a testing specimen
    # and plot the estimated mean RUL values versus the actual RUL (y-axis: mean RUL, 95% CIs, actual RUL/ x-axis: timestep).

    print('\n 5. Optional task: Write a script able to predict not only the mean RUL but also the related X% confidence intervals (CIs) of a testing specimen '
          'and plot the estimated mean RUL values versus the actual RUL (y-axis: mean RUL, 95% CIs, actual RUL/ x-axis: timestep).')

    test_pred_rul_dict = rul_estimation_with_uncertainty(
        samples, test_labels_dict, test_sojourn_time_arr, time_steps, n_clusters)

    visualize_rul_with_uncertainty(test_dict, test_pred_rul_dict, time_steps)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import telebot
import telebot.types as types

import re
import cash
import botan
import config
import chisla_fun

#import configTest as config
bot = telebot.TeleBot(config.token)

# Инлайн-режим с непустым запросом
@bot.inline_handler(lambda query: len(query.query) > 0)
def query_text(query):
    #botan.track(config.botan_key, query.from_user.id, {}, 'inline запрос')
    kb = types.InlineKeyboardMarkup()
    # Добавляем колбэк-кнопку с содержимым "test"
    kb.add(types.InlineKeyboardButton(text="Нажми меня", callback_data="test"))
    results = []
    single_msg = types.InlineQueryResultArticle(
        id="1", title="Press me",
        input_message_content=types.InputTextMessageContent(message_text="Решил проверить, что я умею?)"),
        reply_markup=kb
    )
    results.append(single_msg)
    bot.answer_inline_query(query.id, results)

#Если сообщение на которое тыкнули из инлайн-режима
@bot.callback_query_handler(lambda call: call.inline_message_id)
def callback_inline(call):
    if call.inline_message_id:
        if call.data == "test":
            bot.edit_message_text(inline_message_id=call.inline_message_id, text="Бдыщь")

# Если сообщение из чата с ботом
@bot.callback_query_handler(lambda call: call.message)
def callback_inline(call):
    if call.data == "test":
        bot.edit_message_text(chat_id=call.message.chat.id, message_id=call.message.message_id, text=call.message.text)
        text = 'Добро пожаловать!\nВыберите интересующий вас раздел:'
        keyboardStart = telebot.types.InlineKeyboardMarkup()
        callback_button0 = telebot.types.InlineKeyboardButton(text="Клиент", callback_data="cl")
        callback_button1 = telebot.types.InlineKeyboardButton(text="Гость", callback_data="g")
        keyboardStart.add(callback_button0,callback_button1)
        bot.send_message(call.message.chat.id, text,reply_markup=keyboardStart)
    if call.data == "cl":
        text = "Для записи на консультацию *позвоните* по номеру:\n\n[+79165127377](tel:+79165127377)"
        # bot.edit_message_text(chat_id=call.message.chat.id, message_id=call.message.message_id, text=call.message.text)
        bot.send_message(call.message.chat.id, text, parse_mode='Markdown')
    if call.data == "g":
        text = "Для посчёта вашего числа введите полностью *ФИО* или *дату рождения*.\nДата рождения в формате:\n*ГГГГ ММ ДД* или *ДД ММ ГГГГ*"
        # bot.edit_message_text(chat_id=call.message.chat.id, message_id=call.message.message_id, text=call.message.text)
        bot.send_message(call.message.chat.id, text, parse_mode='Markdown')
    if call.data == "sll":
        keyboard = telebot.types.InlineKeyboardMarkup()
        for text in cash.keyb["students"]:
            callback_button = telebot.types.InlineKeyboardButton(text=text, callback_data=cash.keyb["students"][text])
            keyboard.add(callback_button)
        bot.edit_message_text(chat_id=call.message.chat.id, message_id=call.message.message_id, text=call.message.text)
        bot.send_message(call.message.chat.id, text='Вы Выбрали раздел Студент.', reply_markup=keyboard)

@bot.message_handler(commands=['start'])
def send_welcome(message):
    botan.track(config.botan_key, message.chat.id, message, '/start')
    text = 'Добро пожаловать!\nЯ помогу тебе разобраться с числами!\nВыберите интересующий вас раздел:'
    keyboardStart = telebot.types.InlineKeyboardMarkup()
    callback_button0 = telebot.types.InlineKeyboardButton(text="Посчитать моё число", callback_data="g")
    callback_button1 = telebot.types.InlineKeyboardButton(text="Записаться на консультацию", callback_data="cl")
    keyboardStart.add(callback_button0)
    keyboardStart.add(callback_button1)
    bot.send_message(message.chat.id, text, reply_markup=keyboardStart)

@bot.message_handler(commands=['help'])
def send_help(message):
    botan.track(config.botan_key, message.chat.id, message, '/help')
    text = 'Помошь:\n/help - это сообещние\n/about - о боте\n\nДля содсчёт "собственного числа имени" введите ФИО полностью.\nДля подсчёта "собственного числа даты вождения" введите дату рождения в формате ДД ММ ГГГГ или ГГГГ ММ ДД.'
    bot.send_message(message.chat.id, text)

@bot.message_handler(commands=['about'])
def send_about(message):
    botan.track(config.botan_key, message.chat.id, message, '/about')
    text = 'Бот предназначен для помощи в посчёте "своих чисел".\nОбратная связь/разработчик: @pa_antya'
    bot.send_message(message.chat.id, text)

@bot.message_handler(commands=['test'])
def send_help(message):
    #botan.track(config.botan_key, message.chat.id, message, '/test')
    text = "Бот работает"
    keyboardTest = telebot.types.InlineKeyboardMarkup()
    callback_button_test = telebot.types.InlineKeyboardButton(text="ok!", callback_data="ok")
    keyboardTest.add(callback_button_test)
    bot.send_message(message.chat.id, text, reply_markup=keyboardTest)

# Обычный режим
@bot.message_handler(content_types=["text"])
def any_msg(message):
    print(message.text)
    pdate = re.compile(r'(\d{1,4})[\-_:;\\\|\/ ,.]{0,}(\d{1,2})[\-_:;\\\|\/ ,.]{0,}(\d{1,4})')
    pfio = re.compile(r'([a-zA-ZА-Яа-яёЁ\-\s]{2,})')

    mfio = pfio.match(message.text)
    mdate = pdate.match(message.text)

    #print(mfio)
    #print(mdate)
    if (mfio or mdate):
        if mdate:
            #print('Match found: "{}"'.format(mdate.group()))
            chislo = chisla_fun.date(message)
        if mfio:
            #print('Match found: "{}"'.format(mfio.group()))
            chislo = chisla_fun.fio(message)
        text = 'Вы ввели: *{}*\nВыше число: *{}*'.format(message.text, chislo)
        bot.send_message(message.chat.id, text, parse_mode='Markdown')
    else:
        text = 'Введены *не верные* данные.\nПерепроверьте, пожалуйста:\n{}\n*Если* данные *верные* - свяжитесь с разработчиком в раздете /about.'.format(message.text)
        bot.send_message(message.chat.id, text,parse_mode='Markdown')

    text = 'Выберите интересующий вас раздел:'
'''keyboardText = telebot.types.InlineKeyboardMarkup()
    callback_button0 = telebot.types.InlineKeyboardButton(text="Клиент", callback_data="cl")
    callback_button1 = telebot.types.InlineKeyboardButton(text="Гость", callback_data="g")
    keyboardText.add(callback_button0, callback_button1)
    bot.send_message(message.chat.id, text, reply_markup=keyboardText)
'''
if __name__ == '__main__':
    try:
        bot.polling(none_stop=True)
    except Exception:
        print('error Exception')

#YandexMetrica.setCustomAppVersion("1.13.2");


#!/usr/bin/env python
# coding: utf-8

# In[2]:

import pandas as pd
import matplotlib.pyplot as plt

# In[3]:

a = pd.read_excel('Sentiment_Analysis_1.xlsx')
data = pd.read_csv("disinformation_full_1.csv")
types=[]
for i in a['Sources of News']:
    b=i.split(',')
    types.append(len(b))
a['Number of News Sources']=types

# In[116]:

## STRONG RESULTS##

#a.groupby('Age')[['Frequency of sharing news']].mean().loc[['Less than 18 years','18 to 30 years','31 to 45 years','45 to 60 years', 'More than 60 years']].plot.barh()

#a.groupby('Age')['Verify authenticity of news'].mean().loc[['Less than 18 years','18 to 30 years','31 to 45 years','45 to 60 years', 'More than 60 years']].plot.barh()
#plt.legend(loc=(1.04,0))
#plt.show()

# In[4]:

print(data.columns)

# In[19]:

dark=[]
for i in (data['Money1 or 2']):
    if i =='B':
        dark.append(-1)
    elif i == 'A':
        dark.append(1)
    else:
        dark.append(0)
data['Image Formatted/Plain']=dark

# In[21]:

data.groupby('Age')[['Image Formatted/Plain']].mean().loc[['Less than 18 years','18 to 30 years','31 to 45 years','45 to 60 years', 'More than 60 years']].plot.barh()
plt.xlabel('Likelihood Scale')
plt.legend(loc=(1.04,0))
plt.show()

# In[139]:

bel=[]
fai=[]

for i in a.columns:
    if 'Faith' in i:
        fai.append(i)
    elif 'Believability' in i:
        bel.append(i)

# In[140]:

for i in range(len(bel)):
    (a.groupby('Age')[[bel[i],fai[i]]].mean().loc[['Less than 18 years','18 to 30 years','31 to 45 years','45 to 60 years', 'More than 60 years']].plot.barh())
    plt.xlabel('Likelihood Scale')
    plt.legend(loc=(1.04,0))
    plt.show()

# In[126]:

# In[128]:

for i in range(4):
    (a.groupby('Age')[[pm[i],who[i],doc[i]]].mean().loc[['Less than 18 years','18 to 30 years','31 to 45 years','45 to 60 years', 'More than 60 years']].plot.barh())
    plt.xlabel('Likelihood Scale')
    plt.legend(loc=(1.04,0))
    plt.show()

# In[63]:

for i in range(3):
    (a.groupby('Age')[[fact[i],fact1[i]]].mean().loc[['Less than 18 years','18 to 30 years','31 to 45 years','45 to 60 years', 'More than 60 years']].plot.barh())
    plt.xlabel('Likelihood Scale')
    plt.legend(loc=(1.04,0))
    plt.show()

# In[69]:

fake1=[]
fake2=[]
for i in fake:
    if '1' in i:

        fake1.append(i)
    else:
        fake2.append(i)

for i in range(3):
    (a.groupby('Age')[[fake1[i],fake2[i]]].mean().loc[['Less than 18 years','18 to 30 years','31 to 45 years','45 to 60 years', 'More than 60 years']].plot.barh())
    plt.xlabel('Likelihood Scale')
    plt.legend(loc=(1.04,0))
    plt.show()

# In[71]:

print(trust)
trust1=[]
trust2=[]
for i in trust:
    if '1' in i:

        trust1.append(i)
    else:
        trust2.append(i)

for i in range(3):
    (a.groupby('Age')[[trust1[i],trust2[i]]].mean().loc[['Less than 18 years','18 to 30 years','31 to 45 years','45 to 60 years', 'More than 60 years']].plot.barh())
    plt.xlabel('Likelihood Scale')
    plt.legend(loc=(1.04,0))
    plt.show()


from enum import Enum

class Platform(Enum):
    win32 = 0
    elf32 = 1
    elf64 = 2

    def __str__(self):
        return self.name

    def output_file_extension(self):
        if self == Platform.win32:
            return "exe"
        else:
            return "out"

    @staticmethod
    def by_name(type_name):
        if type_name == Platform.win32.name:
            return Platform.win32
        elif type_name == Platform.elf32.name:
            return Platform.elf32
        elif type_name == Platform.elf64.name:
            return Platform.elf64
        else:
            raise AttributeError(
                "platform %s is not supported, user %s, %s or %s instead" % (
                type_name, Platform.win32, Platform.elf32, Platform.elf64))


import boto3
from dotenv import load_dotenv

load_dotenv()
dynamo = boto3.client("dynamodb", region_name="sa-east-1")
tabela = "NPCs"

def pk(nome: str, grupo: str):
    return {"Nome": {"S": nome}, "Grupo": {"S": grupo}}

def getGrupos() -> list:
    res = dynamo.scan(
        TableName=tabela,
        IndexName="Grupo-index",
        Select="SPECIFIC_ATTRIBUTES",
        AttributesToGet=["Grupo"],
    )
    return list(set([r["Grupo"]["S"] for r in res["Items"]]))

def getNPCs(grupo: str) -> list:
    res = dynamo.query(
        TableName=tabela,
        IndexName="Grupo-index",
        KeyConditionExpression="Grupo = :grupo",
        ExpressionAttributeValues={":grupo": {"S": grupo}},
    )
    return [r["Nome"]["S"] for r in res["Items"]]

def getNPCInfo(nome: str, grupo: str):
    res = dynamo.get_item(TableName=tabela, Key=pk(nome, grupo))
    r = {}
    for k in res["Item"]:
        r[k] = res["Item"][k]["S"]
    return r

def addNPC(nome: str, grupo: str) -> bool:
    params = pk(nome, grupo)
    try:
        dynamo.put_item(TableName=tabela, Item=params)
        return True
    except Exception as e:
        print(e)
        return False

def deleteNPC(nome: str, grupo: str) -> bool:
    params = pk(nome, grupo)
    try:
        dynamo.delete_item(TableName=tabela, Key=params)
        return True
    except Exception as e:
        print(e)
        return False

def editItemNPC(params: dict, atributo: str, valor: str) -> bool:
    for k in params:
        params[k] = {"S": params[k]}
    params[atributo] = {"S": valor}
    try:
        dynamo.put_item(TableName=tabela, Item=params)
        return True
    except Exception as e:
        print(e)
        return False

def deleteItemNPC(nome: str, grupo: str, atributo: str) -> bool:
    params = pk(nome, grupo)
    dynamo.update_item(
        TableName=tabela, Key=params, UpdateExpression=f"REMOVE {atributo}"
    )
    return True

if __name__ == "__main__":
    t = getGrupos()
    print(t)
    p = getNPCs(t[0])
    print(p)
    print(getNPCInfo(p[0], t[0]))