# Author: <NAME> <<EMAIL>>
"""Plot topographic maps of sensor space data."""
from __future__ import annotations

import dataclasses
from itertools import repeat
from math import floor, sqrt
from typing import Any, Dict, Literal, Sequence, Tuple, Union

import matplotlib as mpl
import matplotlib.axes
import numpy as np
from scipy import interpolate, linalg
from scipy.spatial import ConvexHull

from .._colorspaces import UNAMBIGUOUS_COLORS
from .._data_obj import NDVarArg, CategorialArg, IndexArg, Dataset
from .._text import ms
from ._base import (
    CMapArg, ColorArg,
    PlotType, EelFigure, PlotData, AxisData, DataLayer,
    Layout, ImLayout, VariableAspectLayout,
    ColorMapMixin, TimeSlicerEF, TopoMapKey, XAxisMixin, YLimMixin,
)
from ._utsnd import _ax_butterfly, _ax_im_array, _plt_im
from ._sensors import SENSORMAP_FRAME, SensorMapMixin, _plt_map2d

InterpolationArg = Literal[None, 'nearest', 'linear', 'spline']
SensorLabelsArg = Literal['', 'none', 'index', 'name', 'fullname']

class Topomap(SensorMapMixin, ColorMapMixin, TopoMapKey, EelFigure):
    """Plot individual topogeraphies

    Parameters
    ----------
    y
        Data to plot.
    xax
        Create a separate plot for each cell in this model.
    ds
        If a Dataset is provided, data can be specified as strings.
    sub
        Specify a subset of the data.
    vmax
        Upper limits for the colormap (default is determined from data).
    vmin
        Lower limit for the colormap (default ``-vmax``).
    cmap
        Colormap (default depends on the data).
    contours
        Contours to draw on topomaps. Can be an int (number of contours,
        including ``vmin``/``vmax``), a sequence (values at which to draw
        contours), or a ``**kwargs`` dict (must contain at least the "levels"
        key). Default is no contours.
    proj
        The sensor projection to use for topomaps (or one projection per plot).
    res
        Resolution of the topomaps (width = height = ``res``).
    interpolation
        Method for interpolating topo-map between sensors (default is based on
        mne-python).
    clip : bool | 'even' | 'circle'
        Outline for clipping topomaps: 'even' to clip at a constant distance
        (default), 'circle' to clip using a circle.
    clip_distance
        How far from sensor locations to clip (1 is the axes height/width).
    head_radius
        Radius of the head outline drawn over sensors (on sensor plots with
        normalized positions, 0.45 is the outline of the topomap); 0 to plot no
        outline; tuple for separate (right, anterior) radius.
        The default is determined automatically.
    head_pos
        Head outline position along the anterior axis (0 is the center, 0.5 is
        the top end of the plot).
    im_interpolation
        Topomap image interpolation (see Matplotlib's
        :meth:`~matplotlib.axes.Axes.imshow`). Matplotlib 1.5.3's SVG output
        can't handle uneven aspect with ``interpolation='none'``, use
        ``interpolation='nearest'`` instead.
    sensorlabels
        Show sensor labels. For 'name', any prefix common to all names
        is removed; with 'fullname', the full name is shown.
    mark
        Sensors which to mark.
    mcolor : matplotlib color
        Color for marked sensors.
    axtitle
        Title for the individual axes. The default is to show the names of the
        epochs, but only if multiple axes are plotted.
    xlabel
        Label below the topomaps (default is no label; ``True`` to use ``y``
        names).
    margins
        Layout parameter.
    ...
        Also accepts :ref:`general-layout-parameters`.

    Notes
    -----
    Keys:
     - ``t``: open a ``Topomap`` plot for the region under the mouse pointer.
     - ``T``: open a larger ``Topomap`` plot with visible sensor names for the
       map under the mouse pointer.
    """
    def __init__(
            self,
            y: Union[NDVarArg, Sequence[NDVarArg]],
            xax: CategorialArg = None,
            ds: Dataset = None,
            sub: IndexArg = None,
            vmax: float = None,
            vmin: float = None,
            cmap: CMapArg = None,
            contours: Union[int, Sequence, Dict] = None,
            # topomap args
            proj: str = 'default',
            res: int = None,
            interpolation: InterpolationArg = None,
            clip: Union[bool, str] = 'even',
            clip_distance: float = 0.05,
            head_radius: Union[float, Tuple[float, float]] = None,
            head_pos: Union[float, Sequence[float]] = 0,
            im_interpolation: str = None,
            # sensor-map args
            sensorlabels: SensorLabelsArg = None,
            mark: IndexArg = None,
            mcolor: ColorArg = None,
            # layout
            axtitle: Union[bool, Sequence[str]] = True,
            xlabel: Union[bool, str] = None,
            margins: Dict[str, float] = None,
            **kwargs,
    ):
        data = PlotData.from_args(y, ('sensor',), xax, ds, sub)
        self.plots = []
        ColorMapMixin.__init__(self, data.data, cmap, vmax, vmin, contours, self.plots)
        if isinstance(proj, str):
            proj = repeat(proj, data.n_plots)
        elif not isinstance(proj, Sequence):
            raise TypeError(f"proj={proj!r}")
        elif len(proj) != data.n_plots:
            raise ValueError(f"proj={proj!r}: need as many proj as axes ({data.n_plots})")

        layout = ImLayout(data.plot_used, 1.1, 2, margins, axtitle=axtitle, **kwargs)
        EelFigure.__init__(self, data.frame_title, layout)
        self._set_axtitle(axtitle, data, verticalalignment='top', pad=-1)

        # plots
        axes_data = data.for_plot(PlotType.IMAGE)
        for ax, layers, proj_ in zip(self.axes, axes_data, proj):
            h = _ax_topomap(ax, layers, clip, clip_distance, sensorlabels, mark, mcolor, None, proj_, res, im_interpolation, xlabel, self._vlims, self._cmaps, self._contours, interpolation, head_radius, head_pos)
            self.plots.append(h)

        TopoMapKey.__init__(self, self._topo_data)
        SensorMapMixin.__init__(self, [h.sensors for h in self.plots])
        self._show()

    def _fill_toolbar(self, tb):
        ColorMapMixin._fill_toolbar(self, tb)
        SensorMapMixin._fill_toolbar(self, tb)

    def _topo_data(self, event):
        if event.inaxes:
            ax_i = self.axes.index(event.inaxes)
            p = self.plots[ax_i]
            return p.data, p.title, p.proj

class TopomapBins(SensorMapMixin, ColorMapMixin, TopoMapKey, EelFigure):
    """Topomaps in time-bins

    Parameters
    ----------
    y
        Data to plot.
    xax
        Create a separate plot for each cell in this model.
    ds
        If a Dataset is provided, data can be specified as strings.
    sub
        Specify a subset of the data.
    bin_length
        Length ofthe time bins for topo-plots.
    tstart
        Beginning of the first time bin (default is the beginning of ``y``).
    tstop
        End of the last time bin (default is the end of ``y``).
    vmax
        Upper limits for the colormap (default is determined from data).
    vmin
        Lower limit for the colormap (default ``-vmax``).
    cmap
        Colormap (default depends on the data).
    contours
        Contours to draw on topomaps. Can be an int (number of contours,
        including ``vmin``/``vmax``), a sequence (values at which to draw
        contours), or a ``**kwargs`` dict (must contain at least the "levels"
        key). Default is no contours.
    proj
        The sensor projection to use for topomaps.
    res
        Resolution of the topomaps (width = height = ``res``).
    interpolation
        Method for interpolating topo-map between sensors (default is based on
        mne-python).
    clip : bool | 'even' | 'circle'
        Outline for clipping topomaps: 'even' to clip at a constant distance
        (default), 'circle' to clip using a circle.
    clip_distance
        How far from sensor locations to clip (1 is the axes height/width).
    head_radius
        Radius of the head outline drawn over sensors (on sensor plots with
        normalized positions, 0.45 is the outline of the topomap); 0 to plot no
        outline; tuple for separate (right, anterior) radius.
        The default is determined automatically.
    head_pos
        Head outline position along the anterior axis (0 is the center, 0.5 is
        the top end of the plot).
    im_interpolation
        Topomap image interpolation (see Matplotlib's
        :meth:`~matplotlib.axes.Axes.imshow`). Matplotlib 1.5.3's SVG output
        can't handle uneven aspect with ``interpolation='none'``, use
        ``interpolation='nearest'`` instead.
    sensorlabels
        Show sensor labels. For 'name', any prefix common to all names
        is removed; with 'fullname', the full name is shown.
    mark : Sensor index
        Sensors which to mark.
    mcolor : matplotlib color
        Color for marked sensors.
    ...
        Also accepts :ref:`general-layout-parameters`.

    Notes
    -----
    Keys:
     - ``t``: open a ``Topomap`` plot for the map under the mouse pointer.
     - ``T``: open a larger ``Topomap`` plot with visible sensor names for the
       map under the mouse pointer.
    """
    def __init__(
            self,
            y: Union[NDVarArg, Sequence[NDVarArg]],
            xax: CategorialArg = None,
            ds: Dataset = None,
            sub: IndexArg = None,
            bin_length: float = 0.050,
            tstart: float = None,
            tstop: float = None,
            vmax: float = None,
            vmin: float = None,
            cmap: CMapArg = None,
            contours: Union[int, Sequence, Dict] = None,
            # topomap args
            proj: str = 'default',
            res: int = None,
            interpolation: InterpolationArg = None,
            clip: Union[bool, str] = 'even',
            clip_distance: float = 0.05,
            head_radius: Union[float, Tuple[float, float]] = None,
            head_pos: Union[float, Sequence[float]] = 0,
            im_interpolation: str = None,
            # sensor-map args
            sensorlabels: SensorLabelsArg = None,
            mark: IndexArg = None,
            mcolor: ColorArg = None,
            **kwargs,
    ):
        data = PlotData.from_args(y, ('sensor', 'time'), xax, ds, sub)
        self._plots = []
        data._cannot_skip_axes(self)
        bin_data = data.for_plot(PlotType.IMAGE).bin(bin_length, tstart, tstop)
        ColorMapMixin.__init__(self, data.data, cmap, vmax, vmin, contours, self._plots)

        # create figure
        time = bin_data.y0.get_dim('time')
        n_bins = len(time)
        n_rows = bin_data.n_plots
        layout = Layout(n_bins * n_rows, 1, 1.5, tight=False, nrow=n_rows, ncol=n_bins, **kwargs)
        EelFigure.__init__(self, data.frame_title, layout)
        self._plots.extend(repeat(None, n_bins * n_rows))

        for column, t in enumerate(time):
            t_data = bin_data.sub_time(t)
            for row, layers in enumerate(t_data):
                i = row * n_bins + column
                ax = self.axes[i]
                self._plots[i] = _ax_topomap(ax, layers, clip, clip_distance, sensorlabels, mark, mcolor, None, proj, res, im_interpolation, None, self._vlims, self._cmaps, self._contours, interpolation, head_radius, head_pos)

        self._set_axtitle((str(t) for t in time), axes=self.axes[:len(time)])
        TopoMapKey.__init__(self, self._topo_data)
        SensorMapMixin.__init__(self, [h.sensors for h in self._plots])
        self._show()

    def _fill_toolbar(self, tb):
        ColorMapMixin._fill_toolbar(self, tb)
        SensorMapMixin._fill_toolbar(self, tb)

    def _topo_data(self, event):
        if event.inaxes:
            ax_i = self.axes.index(event.inaxes)
            p = self._plots[ax_i]
            return p.data, p.title, p.proj

class TopoButterfly(ColorMapMixin, TimeSlicerEF, TopoMapKey, YLimMixin, XAxisMixin, EelFigure):
    """Butterfly plot with corresponding topomaps

    Parameters
    ----------
    y : (list of) NDVar
        Data to plot.
    xax : None | categorial
        Create a separate plot for each cell in this model.
    ds : Dataset
        If a Dataset is provided, data can be specified as strings.
    sub : str | array
        Specify a subset of the data.
    vmax : scalar
        Upper limits for the colormap (default is determined from data).
    vmin : scalar
        Lower limit for the colormap (default ``-vmax``).
    cmap : str
        Colormap (default depends on the data).
    contours : int | sequence | dict
        Contours to draw on topomaps. Can be an int (number of contours,
        including ``vmin``/``vmax``), a sequence (values at which to draw
        contours), or a ``**kwargs`` dict (must contain at least the "levels"
        key). Default is no contours.
    color : matplotlib color
        Color of the butterfly plots.
    linewidth : scalar
        Linewidth for plots (defult is to use ``matplotlib.rcParams``).
    t
        Time to display in the topomap.
    proj : str
        The sensor projection to use for topomaps.
    res : int
        Resolution of the topomaps (width = height = ``res``).
    interpolation
        Method for interpolating topo-map between sensors (default is based on
        mne-python).
    clip : bool | 'even' | 'circle'
        Outline for clipping topomaps: 'even' to clip at a constant distance
        (default), 'circle' to clip using a circle.
    clip_distance : scalar
        How far from sensor locations to clip (1 is the axes height/width).
    head_radius : scalar | tuple
        Radius of the head outline drawn over sensors (on sensor plots with
        normalized positions, 0.45 is the outline of the topomap); 0 to plot no
        outline; tuple for separate (right, anterior) radius.
        The default is determined automatically.
    head_pos : scalar
        Head outline position along the anterior axis (0 is the center, 0.5 is
        the top end of the plot).
    im_interpolation : str
        Topomap image interpolation (see Matplotlib's
        :meth:`~matplotlib.axes.Axes.imshow`). Matplotlib 1.5.3's SVG output
        can't handle uneven aspect with ``interpolation='none'``, use
        ``interpolation='nearest'`` instead.
    sensorlabels
        Show sensor labels. For 'name', any prefix common to all names
        is removed; with 'fullname', the full name is shown.
    mark : Sensor index
        Sensors to mark in the topo-map. To highlight sensors in the butterfly
        plot, consider using :meth:`NDVar.mask` on ``y``.
    mcolor : matplotlib color
        Color for marked sensors.
    xlabel
        X-axis label. By default the label is inferred from the data.
    ylabel
        Y-axis label. By default the label is inferred from the data.
    xticklabels
        Specify which axes should be annotated with x-axis tick labels.
        Use ``int`` for a single axis, a sequence of ``int`` for multiple
        specific axes, or one of ``'left' | 'bottom' | 'all' | 'none'``.
    yticklabels
        Specify which axes should be annotated with y-axis tick labels.
        Use ``int`` for a single axis, a sequence of ``int`` for multiple
        specific axes, or one of ``'left' | 'bottom' | 'all' | 'none'``.
    axtitle : bool | sequence of str
        Title for the individual axes. The default is to show the names of the
        epochs, but only if multiple axes are plotted.
    xlim : scalar | (scalar, scalar)
        Initial x-axis view limits as ``(left, right)`` tuple or as ``length``
        scalar (default is the full x-axis in the data).
    ...
        Also accepts :ref:`general-layout-parameters`.

    Notes
    -----
    Topomap control:
     - LMB click in a butterfly plot fixates the topomap time
     - RMB click in a butterfly plot removes the time point, the topomaps
       follow the mouse pointer
     - ``.``: Increment the current topomap time (got right)
     - ``,``: Decrement the current topomap time (go left)
     - ``t``: open a ``Topomap`` plot for the time point under the mouse
       pointer
     - ``T``: open a larger ``Topomap`` plot with visible sensor names for the
       time point under the mouse pointer

    Navigation:
     - ``↑``: scroll up
     - ``↓``: scroll down
     - ``←``: scroll left
     - ``→``: scroll right
     - ``home``: scroll to beginning
     - ``end``: scroll to end
     - ``f``: x-axis zoom in (reduce x axis range)
     - ``d``: x-axis zoom out (increase x axis range)
     - ``r``: y-axis zoom in (reduce y-axis range)
     - ``c``: y-axis zoom out (increase y-axis range)
    """
    _default_xlabel_ax = -2

    def __init__(
            self,
            y: Union[NDVarArg, Sequence[NDVarArg]],
            xax: CategorialArg = None,
            ds: Dataset = None,
            sub: IndexArg = None,
            vmax: float = None,
            vmin: float = None,
            cmap: CMapArg = None,
            contours: Union[int, Sequence, Dict] = None,
            color: Any = None,
            linewidth: float = None,
            # topomap args
            t: float = None,
            proj: str = 'default',
            res: int = None,
            interpolation: InterpolationArg = None,
            clip: Union[bool, str] = 'even',
            clip_distance: float = 0.05,
            head_radius: Union[float, Tuple[float, float]] = None,
            head_pos: Union[float, Sequence[float]] = 0,
            im_interpolation: str = None,
            # sensor-map args
            sensorlabels: SensorLabelsArg = None,
            mark: IndexArg = None,
            mcolor: ColorArg = None,
            # layout
            xlabel: Union[bool, str] = True,
            ylabel: Union[bool, str] = True,
            xticklabels: Union[str, int, Sequence[int]] = 'bottom',
            yticklabels: Union[str, int, Sequence[int]] = 'left',
            axtitle: Union[bool, Sequence[str]] = True,
            frame: bool = True,
            xlim: Union[float, Tuple[float, float]] = None,
            **kwargs,
    ):
        data = PlotData.from_args(y, ('sensor', None), xax, ds, sub)
        data._cannot_skip_axes(self)
        xdim = data.dims[1]
        self._topomap_data = data.for_plot(PlotType.IMAGE)

        # create figure
        row_titles = self._set_axtitle(axtitle, data, data.n_plots)
        layout = VariableAspectLayout(data.n_plots, 3, 10, aspect=(None, 1), ax_frames=(frame, False), row_titles=row_titles, **kwargs)
        EelFigure.__init__(self, data.frame_title, layout)

        self.bfly_axes = self.axes[0::2]
        self.topo_axes = self.axes[1::2]
        self.bfly_plots = []
        self.topo_plots = []
        self.t_markers = []  # vertical lines on butterfly plots

        ColorMapMixin.__init__(self, data.data, cmap, vmax, vmin, contours,
                               self.topo_plots)

        self._topo_kwargs = {
            'clip': clip,
            'clip_distance': clip_distance,
            'head_radius': head_radius,
            'head_pos': head_pos,
            'proj': proj,
            'contours': self._contours,
            'res': res,
            'interpolation': interpolation,
            'im_interpolation': im_interpolation,
            'sensorlabels': sensorlabels,
            'mark': mark,
            'mcolor': mcolor,
        }

        # plot epochs (x/y are in figure coordinates)
        for ax, layers in zip(self.bfly_axes, data.for_plot(PlotType.LINE)):
            h = _ax_butterfly(ax, layers, 'time', 'sensor', mark, color, linewidth, self._vlims, clip)
            self.bfly_plots.append(h)

        # decorate axes
        self._configure_axis_dim('x', data.time_dim, xlabel, xticklabels, self.bfly_axes)
        self._configure_axis_data('y', data, ylabel, yticklabels, self.bfly_axes)

        # setup callback
        XAxisMixin._init_with_data(self, data.data, xdim, xlim, self.bfly_axes)
        YLimMixin.__init__(self, self.bfly_plots + self.topo_plots)
        TimeSlicerEF.__init__(self, xdim, data.time_dim, self.bfly_axes, False, initial_time=t)
        TopoMapKey.__init__(self, self._topo_data)
        self._t_label = None  # time label under lowest topo-map
        self._frame .store_canvas()

        self._show(crosshair_axes=self.bfly_axes)
        self._init_controller()

    def _fill_toolbar(self, tb):
        ColorMapMixin._fill_toolbar(self, tb)

    def _update_topo(self, t):
        if not self.topo_plots:
            data = self._topomap_data.sub_time(t)
            for ax, layers in zip(self.topo_axes, data):
                p = _ax_topomap(ax, layers, cmaps=self._cmaps, vlims=self._vlims, **self._topo_kwargs)
                self.topo_plots.append(p)
        else:
            data = self._topomap_data.sub_time(t, data_only=True)
            for p, layers in zip(self.topo_plots, data):
                p.set_data(layers)

    def _topo_data(self, event):
        ax = event.inaxes
        if ax is None:
            return
        p = self.bfly_plots[ax.id // 2]
        if ax in self.bfly_axes:
            t = event.xdata
        elif ax in self.topo_axes:
            t = self._current_time
        else:
            return
        seg = [l.sub(time=t) for l in p.data]
        return seg, f"{ms(t)} ms", self._topo_kwargs['proj']

    def _on_leave_axes_status_text(self, event):
        return "Topomap: t = %.3f" % self._current_time

    def _update_time(self, t, fixate):
        TimeSlicerEF._update_time(self, t, fixate)
        self._update_topo(t)
        if fixate:
            # add time label
            text = "t = %i ms" % round(t * 1e3)
            if self._t_label:
                self._t_label.set_text(text)
            else:
                ax = self.topo_axes[-1]
                self._t_label = ax.text(.5, -0.1, text, ha='center', va='top')
            self.canvas.draw()  # otherwise time label does not get redrawn
        elif self._time_fixed:
            self._t_label.remove()
            self._t_label = None
            self.canvas.draw()  # otherwise time label does not get redrawn
        elif hasattr(self.canvas, 'redraw'):
            self.canvas.redraw(self.topo_axes)

class _plt_topomap(_plt_im):
    """Topomap plot

    Parameters
    ----------
    ...
    im_frame : scalar
        Empty space beyond outmost sensors in the im plot.
    vmax : scalar
        Override the colorspace vmax.
    interpolation
        Method for interpolating topo-map between sensors.
    """
    _aspect = 'equal'

    def __init__(
            self,
            ax: mpl.axes.Axes,
            layer: DataLayer,
            proj: str,
            res: int,
            im_interpolation: str,
            vlims,
            cmaps,
            contours,
            interpolation: InterpolationArg,
            clip: str,
            clip_distance: float,
    ):
        # store attributes
        self._proj = proj
        self._visible_data = layer.y.sensor._visible_sensors(proj)
        self._grid = np.linspace(0, 1, res)
        self._mgrid = tuple(np.meshgrid(self._grid, self._grid))
        if interpolation is None and layer.y.x.dtype.kind in 'biu':
            interpolation = 'nearest'
        self._method = interpolation

        # clip mask
        if clip:
            locs = layer.y.sensor.get_locs_2d(self._proj, frame=SENSORMAP_FRAME)
            if clip == 'even':
                hull = ConvexHull(locs)
                points = locs[hull.vertices]
                default_head_radius = sqrt(np.min(np.sum((points - [0.5, 0.5]) ** 2, 1)))
                # find offset due to clip_distance
                tangents = np.roll(points, -1, axis=0) - np.roll(points, 1, axis=0)
                verticals = np.dot(tangents, [[0, -1], [1, 0]])
                verticals /= np.sqrt(np.sum(verticals ** 2, 1)[:, None])
                verticals *= clip_distance
                # apply offset
                points += verticals
                mask = mpl.patches.Polygon(points, transform=ax.transData)
            elif clip == 'circle':
                clip_radius = sqrt(np.max(np.sum((locs - [0.5, 0.5]) ** 2, 1)))
                mask = mpl.patches.Circle((0.5, 0.5), clip_radius, transform=ax.transData)
                default_head_radius = clip_radius
            else:
                raise ValueError('clip=%r' % (clip,))
        else:
            mask = None
            default_head_radius = None

        self._default_head_radius = default_head_radius
        _plt_im.__init__(self, ax, layer, cmaps, vlims, contours, (0, 1, 0, 1), im_interpolation, mask)

    def _data_from_ndvar(self, ndvar):
        v = ndvar.get_data(('sensor',))
        locs = ndvar.sensor.get_locs_2d(self._proj, frame=SENSORMAP_FRAME)
        if self._visible_data is not None:
            v = v[self._visible_data]
            locs = locs[self._visible_data]

        if self._method is None:
            # interpolate data
            xi, yi = self._mgrid

            # code adapted from mne-python topmap _griddata()
            xy = locs[:, 0] + locs[:, 1] * -1j
            d = np.abs(xy - xy[:, None])
            diagonal_step = len(locs) + 1
            d.flat[::diagonal_step] = 1.

            g = (d * d) * (np.log(d) - 1.)
            g.flat[::diagonal_step] = 0.
            try:
                weights = linalg.solve(g, v.ravel())
            except ValueError:
                if np.isnan(v).any():
                    raise NotImplementedError("Can't interpolate sensor data with NaN")
                unique_locs = np.unique(locs, axis=0)
                if len(unique_locs) < len(locs):
                    raise NotImplementedError("Error determining sensor map projection due to more than one sensor in a single location; try using a different projection.")
                raise

            m, n = xi.shape
            out = np.empty_like(xi)

            g = np.empty(xy.shape)
            for i in range(m):
                for j in range(n):
                    d = np.abs(xi[i, j] + -1j * yi[i, j] - xy)
                    mask = np.where(d == 0)[0]
                    if len(mask):
                        d[mask] = 1.
                    np.log(d, out=g)
                    g -= 1.
                    g *= d * d
                    if len(mask):
                        g[mask] = 0.
                    out[i, j] = g.dot(weights)
            return out
        elif self._method == 'spline':
            k = int(floor(sqrt(len(locs)))) - 1
            tck = interpolate.bisplrep(locs[:, 1], locs[:, 0], v, kx=k, ky=k)
            return interpolate.bisplev(self._grid, self._grid, tck)
        else:
            isnan = np.isnan(v)
            if np.any(isnan):
                nanmap = interpolate.griddata(locs, isnan, self._mgrid, self._method)
                mask = nanmap > 0.5
                v = np.where(isnan, 0, v)
                vmap = interpolate.griddata(locs, v, self._mgrid, self._method)
                np.place(vmap, mask, np.NaN)
                return vmap
            return interpolate.griddata(locs, v, self._mgrid, self._method)

class _ax_topomap(_ax_im_array):
    """Axes with a topomap

    Parameters
    ----------
    mark : list of IDs
        highlight a subset of the sensors
    """
    def __init__(
            self,
            ax: matplotlib.axes.Axes,
            layers: AxisData,
            clip: str = 'even',  # even or circle (only applies if interpolation is None)
            clip_distance: float=0.05,  # distance from outermost sensor for clip=='even'
            sensorlabels: SensorLabelsArg = None,
            mark: IndexArg = None,
            mcolor: ColorArg = None,
            mmarker=None,
            proj: str = 'default',  # topomap projection method
            res: int = None,  # topomap image resolution
            im_interpolation: str = None,  # matplotlib imshow interpolation method
            xlabel: Union[bool, str] = None,
            vlims: dict = {},
            cmaps: dict = {},
            contours: dict = {},
            interpolation: InterpolationArg = None,  # topomap interpolation method
            head_radius: Union[float, Sequence[float]] = None,
            head_pos: Union[float, Tuple[float, float]] = 0.,
            head_linewidth: float = None,
    ):
        self.ax = ax
        self.data = layers  # will not update from .set_data()
        self.proj = proj
        sensor_dim = layers.y0.sensor

        if xlabel is True:
            xlabel = layers.y0.name
        if im_interpolation is None:
            im_interpolation = 'bilinear'
        if res is None:
            res = 64 if interpolation is None else 100

        ax.set_axis_off()
        self.plots = [_plt_topomap(ax, layer, proj, res, im_interpolation, vlims, cmaps, contours, interpolation, clip, clip_distance) for layer in layers]

        # head outline
        if head_radius is None and clip == 'circle' and interpolation is None and sensor_dim._topomap_outlines(proj) == 'top':
            head_radius = self.plots[0]._default_head_radius

        # plot sensors
        self.sensors = _plt_map2d(ax, sensor_dim, proj, 1, '.', 1, 'k', mark, mcolor, mmarker, sensorlabels, False, head_radius, head_pos, head_linewidth)

        ax.set_aspect('equal')
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)

        if isinstance(xlabel, str):
            x, y = ax.transData.inverted().transform(ax.transAxes.transform((0.5, 0)))
            ax.text(x, y, xlabel, ha='center', va='top')

    def set_ylim(self, bottom, top):  # Alias for YLimMixin
        self.set_vlim(bottom, top)

@dataclasses.dataclass
class _TopoWindow:
    """Helper class for TopoArray.

    Maintains a topomap corresponding to one segment with flexible time point.
    """
    ax: matplotlib.axes.Axes  # topomap-axes
    parent: _ax_im_array  # array-plot
    topomap_args: dict
    connectionstyle: str = "angle3,angleA=90,angleB=0"
    label_position: Literal['above', 'below', 'none'] = 'above'
    color: ColorArg = UNAMBIGUOUS_COLORS['bluish green']
    annotation_xy: Tuple[float, float] = (0.5, 1.05)
    # internal plot handles
    t_line = None
    pointer = None
    text_pointer = None
    plot = None
    t = None

    def update(self, t):
        if t is not None:
            if self.t_line:
                self.t_line.remove()
            self.t_line = self.parent.ax.axvline(t, c=self.color)

            t_str = f"{ms(t)} ms"
            if self.pointer:
                self.pointer.axes = self.parent.ax
                self.pointer.xy = (t, 1)
                if self.text_pointer:
                    self.text_pointer.set_text(t_str)
                self.pointer.set_visible(True)
            else:
                text = t_str if self.label_position == 'above' else ''
                arrowprops = {'arrowstyle': '-', 'shrinkB': 0, 'color': self.color}
                if self.connectionstyle:
                    arrowprops['connectionstyle'] = self.connectionstyle
                self.pointer = self.parent.ax.annotate(text, (t, 0), xycoords='data', xytext=self.annotation_xy, textcoords=self.ax.transData, horizontalalignment='center', verticalalignment='center', arrowprops=arrowprops, zorder=4)
                if self.label_position == 'above':
                    self.text_pointer = self.pointer
                elif self.label_position == 'below':
                    self.text_pointer = self.ax.text(0.5, 0, t_str, va='top', ha='center', transform=self.ax.transAxes)

            if self.plot is None:
                layers = self.parent.data.sub_time(t)
                self.plot = _ax_topomap(self.ax, layers, **self.topomap_args)
            else:
                layers = self.parent.data.sub_time(t, data_only=True)
                self.plot.set_data(layers)
            self.t = t

    def clear(self):
        self.ax.cla()
        self.ax.set_axis_off()
        self.plot = None
        self.t = None
        if self.t_line:
            self.t_line.remove()
            self.t_line = None
        if self.pointer:
            self.pointer.remove()
            self.pointer = None

    def add_contour(self, meas, level, color):
        if self.plot:
            self.plot.add_contour(meas, level, color)

    def set_cmap(self, cmap, meas):
        if self.plot:
            self.plot.set_cmap(cmap, meas)

    def set_vlim(self, v, vmax=None, meas=None):
        if self.plot:
            self.plot.set_vlim(v, vmax, meas)

class TopoArray(ColorMapMixin, TopoMapKey, XAxisMixin, EelFigure):
    """Channel by sample plots with topomaps for individual time points

    Parameters
    ----------
    y
        Data to plot.
    xax
        Create a separate plot for each cell in this model.
    ds
        If a Dataset is provided, data can be specified as strings.
    sub
        Specify a subset of the data.
    vmax
        Upper limits for the colormap (default is determined from data).
    vmin
        Lower limit for the colormap (default ``-vmax``).
    cmap
        Colormap (default depends on the data).
    contours
        Contours to draw on topomaps. Can be an int (number of contours,
        including ``vmin``/``vmax``), a sequence (values at which to draw
        contours), or a ``**kwargs`` dict (must contain at least the "levels"
        key). Default is no contours.
    ntopo
        number of topomaps per array-plot.
    t
        Time points for topomaps.
    xlim : scalar | (scalar, scalar)
        Initial x-axis view limits as ``(left, right)`` tuple or as ``length``
        scalar (default is the full x-axis in the data).
    proj
        The sensor projection to use for topomaps.
    res
        Resolution of the topomaps (width = height = ``res``).
    interpolation
        Method for interpolating topo-map between sensors (default is based on
        mne-python).
    clip : bool | 'even' | 'circle'
        Outline for clipping topomaps: 'even' to clip at a constant distance
        (default), 'circle' to clip using a circle.
    clip_distance
        How far from sensor locations to clip (1 is the axes height/width).
    head_radius
        Radius of the head outline drawn over sensors (on sensor plots with
        normalized positions, 0.45 is the outline of the topomap); 0 to plot no
        outline; tuple for separate (right, anterior) radius.
        The default is determined automatically.
    head_pos
        Head outline position along the anterior axis (0 is the center, 0.5 is
        the top end of the plot).
    im_interpolation
        Topomap image interpolation (see Matplotlib's
        :meth:`~matplotlib.axes.Axes.imshow`). Matplotlib 1.5.3's SVG output
        can't handle uneven aspect with ``interpolation='none'``, use
        ``interpolation='nearest'`` instead.
    sensorlabels
        Show sensor labels. For 'name', any prefix common to all names
        is removed; with 'fullname', the full name is shown. Set to ``''`` to
        hide sensor position markers completely.
    mark : Sensor index
        Sensors which to mark.
    mcolor : matplotlib color
        Color for marked sensors.
    axtitle
        Title for the individual axes. The default is to show the names of the
        epochs, but only if multiple axes are plotted.
    xlabel
        X-axis label. By default the label is inferred from the data.
    ylabel
        Y-axis label. By default the label is inferred from the data.
    xticklabels
        Specify which axes should be annotated with x-axis tick labels.
        Use ``int`` for a single axis, a sequence of ``int`` for multiple
        specific axes, or one of ``'left' | 'bottom' | 'all' | 'none'``.
    yticklabels
        Specify which axes should be annotated with y-axis tick labels.
        Use ``int`` for a single axis, a sequence of ``int`` for multiple
        specific axes, or one of ``'left' | 'bottom' | 'all' | 'none'``.
    connectionstyle
        Style for the connections between the image array-plot and the
        topo-maps. Set to ``''`` for straight connections. See
        `Matplotlib demo <https://matplotlib.org/stable/gallery/userdemo/connectionstyle_demo.html>`_.
    connection_color
        Color for connection line.
    topo_labels
        Where to label time on topo-maps.
    ...
        Also accepts :ref:`general-layout-parameters`.

    Notes
    -----
     - LMB click on a topomap selects it for tracking the mouse pointer
         - LMB on the array plot fixates the topomap time point
     - RMB on a topomap removes the topomap

    """
    _make_axes = False

    def __init__(
            self,
            y: Union[NDVarArg, Sequence[NDVarArg]],
            xax: CategorialArg = None,
            ds: Dataset = None,
            sub: IndexArg = None,
            vmax: float = None,
            vmin: float = None,
            cmap: CMapArg = None,
            contours: Union[int, Sequence, Dict] = None,
            ntopo: int = None,
            t: Sequence[float] = (),
            xlim: Union[float, Tuple[float, float]] = None,
            # topomap args
            proj: str = 'default',
            res: int = None,
            interpolation: InterpolationArg = None,
            clip: Union[bool, str] = 'even',
            clip_distance: float = 0.05,
            head_radius: Union[float, Tuple[float, float]] = None,
            head_pos: Union[float, Sequence[float]] = 0,
            im_interpolation: str = None,
            # sensor-map args
            sensorlabels: SensorLabelsArg = None,
            mark: IndexArg = None,
            mcolor: ColorArg = None,
            # layout
            axtitle: Union[bool, Sequence[str]] = True,
            xlabel: Union[bool, str] = True,
            ylabel: Union[bool, str] = True,
            xticklabels: Union[str, int, Sequence[int]] = 'bottom',
            yticklabels: Union[str, int, Sequence[int]] = 'left',
            connectionstyle: str = "angle3,angleA=90,angleB=0",
            connection_color: ColorArg = UNAMBIGUOUS_COLORS['bluish green'],
            topo_labels: Literal['above', 'below', 'none'] = 'above',
            **kwargs,
    ):
        if ntopo is None:
            ntopo = len(t) if t else 3

        data = PlotData.from_args(y, ('time', 'sensor'), xax, ds, sub).for_plot(PlotType.IMAGE)

        # create figure
        if 'ncol' not in kwargs and 'nrow' not in kwargs:
            kwargs['nrow'] = 1
        layout = Layout(data.plot_used, 1.5, 3, tight=False, **kwargs)
        EelFigure.__init__(self, data.frame_title, layout)
        all_plots = []
        ColorMapMixin.__init__(self, data.data, cmap, vmax, vmin, contours, all_plots)
        TopoMapKey.__init__(self, self._topo_data)

        # save important properties
        self._data = data
        self._ntopo = ntopo
        self._default_xlabel_ax = -1 - ntopo
        self._proj = proj

        # prepare axes
        if layout.user_axes:
            self.axes = layout.user_axes
        else:
            x_frame_l = .6 / layout.axw / data.n_plots
            x_frame_r = .025 / data.n_plots
            gs = self.figure.add_gridspec(layout.nrow * 2, layout.ncol * ntopo, left=x_frame_l, right=1 - x_frame_r, bottom=.05, top=.9, wspace=.1, hspace=.3)
            if layout.nrow == 1:
                for col, used in enumerate(data.plot_used):
                    if not used:
                        continue
                    self.figure.add_subplot(gs[0, col*ntopo:(col+1)*ntopo], picker=True)
                    for j in range(ntopo):
                        self.figure.add_subplot(gs[1, col*ntopo+j], picker=True, xticks=[], yticks=[])
            elif layout.ncol == 1:
                for row, used in enumerate(data.plot_used):
                    if not used:
                        continue
                    self.figure.add_subplot(gs[row*2, 0:ntopo], picker=True)
                    for j in range(ntopo):
                        self.figure.add_subplot(gs[row*2+1, j], picker=True, xticks=[], yticks=[])
            else:
                raise ValueError("Layout with multiple columns and rows; set either ncol=1 or nrow=1")
            self.axes = self.figure.axes

        # im_array plots
        self._array_axes = []
        self._array_plots = []
        self._topo_windows = []
        topomap_args = dict(clip=clip, clip_distance=clip_distance, sensorlabels=sensorlabels, mark=mark, mcolor=mcolor, proj=proj, res=res, im_interpolation=im_interpolation, vlims=self._vlims, cmaps=self._cmaps, contours=self._contours, interpolation=interpolation, head_radius=head_radius, head_pos=head_pos)
        for i, layers in enumerate(data):
            ax_i = i * (ntopo + 1)
            ax = self.axes[ax_i]
            ax.ID = i
            ax.type = 'main'
            im_plot = _ax_im_array(ax, layers, 'time', im_interpolation, self._vlims, self._cmaps, self._contours)
            self._array_axes.append(ax)
            self._array_plots.append(im_plot)
            if i > 0:
                ax.yaxis.set_visible(False)

            # topo plots
            for j in range(ntopo):
                ax = self.axes[ax_i + 1 + j]
                ax.ID = i * ntopo + j
                ax.type = 'window'
                win = _TopoWindow(ax, im_plot, topomap_args, connectionstyle, topo_labels, connection_color)
                self.axes.append(ax)
                self._topo_windows.append(win)
        all_plots.extend(self._array_plots)
        all_plots.extend(self._topo_windows)

        # if t argument is provided, set topo-map time points
        if t:
            if np.isscalar(t):
                t = [t]
            self.set_topo_ts(*t)

        self._set_axtitle(axtitle, data, self._array_axes)
        self._configure_axis_dim('x', data.y0.time, xlabel, xticklabels, self._array_axes)
        self._configure_axis_dim('y', 'sensor', ylabel, yticklabels, self._array_axes, False, data.data)

        # setup callback
        XAxisMixin._init_with_data(self, data.data, 'time', xlim, self._array_axes)
        self._selected_window = None
        self.canvas.mpl_connect('pick_event', self._pick_handler)
        self._frame .store_canvas()
        self._show(crosshair_axes=self._array_axes)

    def _fill_toolbar(self, tb):
        ColorMapMixin._fill_toolbar(self, tb)

    def _topo_data(self, event):
        ax = event.inaxes
        if ax in self._array_axes:
            t = event.xdata
            data = self._array_plots[ax.ID].data.sub_time(t)
        else:
            topo_window = self._topo_windows[ax.ID]
            t = topo_window.t
            if t is None:
                return
            data = topo_window.plot.data
        return data, f"{ms} ms", self._proj

    def _iter_plots(self):
        "Iterate through non-empty plots"
        yield from self._array_plots
        for w in self._topo_windows:
            if w.plot is not None:
                yield w.plot

    def set_cmap(self, cmap, meas=None):
        """Change the colormap

        Parameters
        ----------
        cmap : str | colormap
            New colormap.
        meas : None | str
            Measurement to which to apply the colormap. With None, it is
            applied to all.
        """
        self._cmaps[meas] = cmap
        for p in self._iter_plots():
            p.set_cmap(cmap, meas)
        self.draw()

    def set_topo_t_single(self, topo_id, t):
        """
        Set the time for a single topomap.

        Parameters
        ----------
        topo_id : int
            Index of the topomap (numbered throughout the figure).
        t : scalar or ``None``
            time point; ``None`` clears the topomap
        """
        # get window ax
        w = self._topo_windows[topo_id]
        w.clear()

        if t is not None:
            w.update(t)

        self.canvas.draw()

    def set_topo_t(self, topo_id, t):
        """
        Set the time point for a topo-map (same for all array plots)

        Parameters
        ----------
        topo_id : int
            Index of the topomap (numberd for each array-plot).
        t : scalar or ``None``
            time point; ``None`` clears the topomap

        See Also
        --------
        .set_topo_ts : set several topomap time points at once
        .set_topo_t_single : set the time point of a single topomap
        """
        for i in range(len(self._array_plots)):
            _topo = self._ntopo * i + topo_id
            self.set_topo_t_single(_topo, t)

    def set_topo_ts(self, *t_list):
        """Set the time points displayed in topo-maps across all array-plots"""
        for i, t in enumerate(t_list):
            self.set_topo_t(i, t)

    def _pick_handler(self, pickevent):
        mouseevent = pickevent.mouseevent
        ax = pickevent.artist
        if ax.type == 'window':
            button = mouseevent.button  # 1: Left
            window = self._topo_windows[ax.ID]
            if button == 1:
                self._selected_window = window
            elif button in (2, 3):
                Id = window.ax.ID % self._ntopo
                self.set_topo_t(Id, None)
            else:
                pass
        elif (ax.type == 'main') and (self._selected_window is not None):
            self._selected_window.clear()  # to side track pdf export transparency issue
            # update corresponding topo_windows
            t = mouseevent.xdata
            Id = self._selected_window.ax.ID % self._ntopo
            self.set_topo_t(Id, t)

            self._selected_window = None
            self.canvas.draw()

    def _on_motion_sub(self, event):
        if (self._selected_window is not None and event.inaxes and
                event.inaxes.type == 'main' and
                event.xdata in self._data.plot_data[event.inaxes.ID].y0.time):
            self._selected_window.update(event.xdata)
            return {self._selected_window.ax}
        return set()


#!/usr/bin/env python

import gym
import mj_transfer
from mj_transfer.robots import NewBrightEnv

from time import time, sleep
from scipy.misc import imshow, toimage, imsave

FREQUENCY = 0.25

try:
    # Windows
    from msvcrt import getch
except ImportError:
    # Unix
    import sys
    import tty
    import termios

    def getch():
        fd = sys.stdin.fileno()
        old = termios.tcgetattr(fd)
        try:
            tty.setraw(fd)
            return sys.stdin.read(1)
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old)

if __name__ == '__main__':
    env = NewBrightEnv('192.168.42.1', 5000)
    time_init = time()
    env.reset()
    while True:
        action = [0.0, 0.0, 0.0, 0.0]
        c = getch()
        # c = None
        # sleep(0.2)
        if c == 'q':
            break
        if c == 'w':
            action[0] = 1.0
        if c == 's':
            action[1] = 1.0
        if c == 'd':
            action[2] = 1.0
        if c == 'a':
            action[3] = 1.0

        # TODO: Move the following code in env.step ?
        # TODO: Keep everything to 1.0 if key pressed ?
        # TODO: Deal with opposing directions ?
        env.step(action)
        # env.render()
    env.disconnect()


import matplotlib.pyplot as plt
import matplotlib as mpl
import scipy as sp
import matplotlib.cm as cm
import numpy as np
import pandas as pd
import seaborn as sns

with open("output.txt") as jurassicpark:
	jurassicpark_txt = jurassicpark.read()
	#jurassicpark.close()
	jurassicpark_tokens = jurassicpark_txt.split() #with no arguments, splits on whitespace
	print(len(jurassicpark_tokens))

lastword = jurassicpark_tokens[-1]
print(lastword)

strangestring = ",".join(jurassicpark_tokens)
strangestring[:100]

jurassicparklc_tokens = [word.lower() for word in jurassicpark_tokens]
count = jurassicparklc_tokens.count("had")
print(count)

#find unique set of words, count how often used with count method on lists
uniquelctokens = set(jurassicparklc_tokens)

tokendict = {}
for ut in uniquelctokens:
	tokendict[ut] = jurassicparklc_tokens.count(ut)

el = sorted(tokendict.items(), key= lambda args: args[1], reverse = True)[:100]
print(el)

topfreq = el[:20]
print(topfreq)
with sns.axes_style('whitegrid'):
	pos = np.arange(len(topfreq))
	plt.bar(pos, [e[1] for e in topfreq]);
	plt.xticks(pos+0.4, [e[0] for e in topfreq]);
	plt.xlabel("Words")
	plt.ylabel("Count")
	plt.title("Word Frequency in output text")

plt.show()


"""
@file
@brief  Function to solve the problem of the
`Route Inspection Problem <https://en.wikipedia.org/wiki/Route_inspection_problem>`_.
"""
import math
import random
import fractions

class SolutionException(Exception):
    """
    wrong solution
    """
    pass

def haversine_distance(lat1, lng1, lat2, lng2):
    """
    Computes `Haversine formula <http://en.wikipedia.org/wiki/Haversine_formula>`_.

    .. index:: Haversine

    @param      lat1    lattitude
    @param      lng1    longitude
    @param      lat2    lattitude
    @param      lng2    longitude
    @return             distance
    """
    radius = 6371
    dlat = math.radians(lat2 - lat1)
    dlon = math.radians(lng2 - lng1)
    a = math.sin(dlat / 2) * math.sin(dlat / 2) + math.cos(math.radians(lat1)) \
        * math.cos(math.radians(lat2)) * math.sin(dlon / 2) * math.sin(dlon / 2)
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
    d = radius * c
    return d

def distance_solution(edges_index, edges, distances, solution, exc=True):
    """
    Checks if a solution is really a solution and returns the distance of it,
    None if it is not a solution. The function does not case about the order.

    @param      edges_index     list of indices of edges (if None --> range(len(edges))
    @param      edges           list of tuple (vertex A, vertex B)
    @param      distances       list of distances of each edge
    @param      solution        proposed solutions (list of edge indices)
    @param      exc             raises an exception in case of error
    """
    if edges_index is None:
        edges_index = list(range(len(edges)))
    indices = set(edges_index)
    solset = set(solution)
    if len(indices) != len(solset):
        if exc:
            raise SolutionException("Different number of distinct edges:\nexpected={0} got={1}\n"
                                    "Did you cover all the edges?".format(len(indices), len(solset)))
        return None

    for s in solution:
        if s not in indices:
            raise SolutionException(
                "Index {0} is not in edges_index".format(s))

    doubles = {}
    for a, b in edges:
        if a > b:
            a, b = b, a
        if (a, b) in doubles:
            if exc:
                raise SolutionException(
                    "Edge {0} is duplicated in edges.".format((a, b)))
            return None
        doubles[a, b] = 1

    corres = {e: i for i, e in enumerate(edges_index)}
    degrees = {}
    for s in solution:
        a, b = edges[corres[s]]
        degrees[a] = degrees.get(a, 0) + 1
        degrees[b] = degrees.get(b, 0) + 1

    odd, even = 0, 0
    for d in degrees.values():
        if d % 2 == 0:
            even += 1
        else:
            odd += 1

    if odd > 2:
        if exc:
            red = list(sorted([(k, v)
                               for k, v in degrees.items() if v % 2 != 0]))
            if len(red) > 10:
                red = red[:10] + ["..."]
            raise SolutionException(
                "Are you sure? The path is inconsistent. Some help:\n" + str(red))
        return None
    else:
        return sum(distances[corres[s]] for s in solution)

def compute_degrees(edges):
    """
    Compute the degree of vertices.

    @param      edges       list of tuple
    @return                 dictionary {key: degree}
    """
    res = {}
    for a, b in edges:
        res[a] = res.get(a, 0) + 1
        res[b] = res.get(b, 0) + 1
    return res

def euler_path(edges_index, edges, solution):
    """
    Computes an Eulerian path.

    .. index:: Eulerian

    @param      edges_index     list of indices of edges (if None --> range(len(edges))
    @param      edges           list of tuple (vertex A, vertex B)
    @param      solution        proposed solutions (list of edge indices)
    @return                     path, list of edges indices

    The function assumes every vertex in the graph defined by *edges*
    has an even degree.
    """
    if edges_index is None:
        edges_index = range(len(edges))
    pos = {k: i for i, k in enumerate(edges_index)}
    indices = [pos[s] for s in solution]
    edges = [edges[i] for i in indices]

    degrees = compute_degrees(edges)
    odd = {k: v for k, v in degrees.items() if v % 2 == 1}
    if len(odd) not in (0, 2):
        odd = list(sorted((k, v) for k, v in odd.items()))
        if len(odd) > 10:
            odd = odd[:10] + ["..."]
        raise SolutionException(
            "Some vertices have an odd degree. This is not allowed.\n" + str(odd))

    if len(odd) == 2:
        # we add an extra edge which we remove later
        odd = list(odd.keys())
        remove = (odd[0], odd[1])
        edges.append(remove)
    else:
        remove = None

    res = _euler_path(edges)
    pathi = [_[1][0] for _ in res]
    if remove is not None:
        index = pathi.index(len(edges) - 1)
        if index == 0:
            pathi = pathi[1:]
        elif index == len(edges) - 1:
            pathi = pathi[:-1]
        else:
            pathi = pathi[index + 1:] + pathi[:index]
    pathi = [solution[i] for i in pathi]
    return pathi

def _euler_path(edges):
    """
    Computes an Eulerian path.

    @param      edges           edges
    @return                     path, list of (vertex,edge)

    The function assumes every vertex in the graph defined by *edges*
    has an even degree.
    """
    alledges = {}
    edges_from = {}
    for i, k in enumerate(edges):
        if isinstance(k, list):
            k = tuple(k)
        v = (i,) + k
        alledges[k] = v
        a, b = k
        alledges[b, a] = alledges[a, b]
        if a not in edges_from:
            edges_from[a] = []
        if b not in edges_from:
            edges_from[b] = []
        edges_from[a].append(alledges[a, b])
        edges_from[b].append(alledges[a, b])

    degre = {}
    for a, v in edges_from.items():
        t = len(v)
        degre[t] = degre.get(t, 0) + 1

    two = [a for a, v in edges_from.items() if len(v) == 2]
    odd = [a for a, v in edges_from.items() if len(v) % 2 == 1]
    if len(odd) not in (0, 2):
        add = "\n" + str(odd) if len(odd) < 10 else ""
        raise SolutionException(
            "Some vertices have an odd degree. This is not allowed." + add)

    begin = two[0]

    # checking
    for v, le in edges_from.items():
        for e in le:
            to = e[1] if v != e[1] else e[2]
            if to not in edges_from:
                raise SolutionException(
                    "Unable to find vertex {0} for edge ({0},{1}).".format(to, v))
            if to == v:
                raise SolutionException("Circular edge {0}.".format(to))

    # loop
    path = _explore_path(edges_from, begin)
    for p in path:
        if len(p) == 0:
            raise NotImplementedError("This exception should not happen.")
    while len(edges_from) > 0:
        start = None
        for i, p in enumerate(path):
            if p[0] in edges_from:
                start = i, p
                break
        sub = _explore_path(edges_from, start[1][0])
        i = start[0]
        path[i:i + 1] = path[i:i + 1] + sub
    return path

def _delete_edge(edges_from, n, to):
    """
    Removes an edge from the graph.

    @param      edges_from      structure which contains the edges (will be modified)
    @param      n               first vertex
    @param      to              second vertex
    @return                     the edge
    """
    le = edges_from[to]
    f = None
    for i, e in enumerate(le):
        if (e[1] == to and e[2] == n) or (e[2] == to and e[1] == n):
            f = i
            break

    assert f is not None
    del le[f]
    if len(le) == 0:
        del edges_from[to]

    le = edges_from[n]
    f = None
    for i, e in enumerate(le):
        if (e[1] == to and e[2] == n) or (e[2] == to and e[1] == n):
            f = i
            break

    assert f is not None
    keep = le[f]
    del le[f]
    if len(le) == 0:
        del edges_from[n]

    return keep

def _explore_path(edges_from, begin):
    """
    Explores an Eulerian path, remove used edges from *edges_from*.

    @param      edges_from      structure which contains the edges (will be modified)
    @param      begin           first vertex to use
    @return                     path
    """
    path = [(begin, None)]
    stay = True
    while stay and len(edges_from) > 0:

        n = path[-1][0]
        if n not in edges_from:
            # fin
            break
        le = edges_from[n]

        if len(le) == 1:
            h = 0
            e = le[h]
            to = e[1] if n != e[1] else e[2]
        else:
            to = None
            nb = 100
            while to is None or to == begin:
                h = random.randint(0, len(le) - 1) if len(le) > 1 else 0
                e = le[h]
                to = e[1] if n != e[1] else e[2]
                nb -= 1
                if nb < 0:
                    raise NotImplementedError(
                        "Algorithm issue {0}".format(len(path)))

        if len(edges_from[to]) == 1:
            if begin != to:
                raise NotImplementedError("Wrong algorithm.")
            stay = False

        keep = _delete_edge(edges_from, n, to)
        path.append((to, keep))

    return path[1:]

def distance_vertices(edges, vertices, distances):
    """
    Computes the length of edges if distances is None.

    @param      edges           list of tuple (vertex A, vertex B)
    @param      vertices        locations of the vertices
    @param      distances       distances (None or list of floats)
    @return                     distances (list of float)
    """
    if distances is None:
        distances = []
    while len(distances) < len(edges):
        distances.append(None)
    for i, edge in enumerate(edges):
        if distances[i] is not None:
            continue
        a, b = edge
        va = vertices[a]
        vb = vertices[b]
        d = haversine_distance(va[0], va[1], vb[0], vb[1])
        distances[i] = d
    return distances

def bellman_distances(edges, distances, fLOG=None):
    """
    Computes shortest distances between all vertices.
    It assumes edges are symmetric.

    @param      edges           list of tuple (vertex A, vertex B)
    @param      distances       distances (list of floats)
    @param      fLOG            logging function
    @return                     dictionary of distances

    This function could be implemented based on
    `shortest_path <https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csgraph.shortest_path.html>`_.
    """
    dist = {(a, b): d for d, (a, b) in zip(distances, edges)}
    dist.update({(b, a): d for d, (a, b) in zip(distances, edges)})
    dist0 = dist.copy()

    iter = 0  # pylint: disable=W0622
    up = dist
    while len(up) > 0:
        iter += 1
        up = {}
        for (a, b), d1 in dist.items():
            if (a, b) not in dist0 and (a, b) not in up:
                # no need to continue as no update during the last iteration
                continue
            for (aa, bb), d2 in dist.items():
                # not the most efficient
                if b == aa and a != bb:
                    d = d1 + d2
                    if (a, bb) not in dist or dist[a, bb] > d:
                        up[a, bb] = d
                        up[bb, a] = d
        dist.update(up)
        if fLOG:
            sum_values = sum(dist.values())
            mes = "iteration={0} modif={1} #dist={2} sum_values={3} avg={4}".format(
                iter, len(up), len(dist), sum_values, sum_values / len(dist))
            fLOG("[bellman_distances] " + mes)

    return dist

def dijkstra_path(edges, distances, va, vb):
    """
    Returns the best path between two vertices.
    Uses `Dikjstra <https://en.wikipedia.org/wiki/Dijkstra%27s_algorithm>`_ algorithm.

    @param      edges       list of edges.
    @param      distances   list of distances
    @param      va          first vertex
    @param      vb          last vertex
    @return                 list of edges

    This function could be implemented based on
    `shortest_path <https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csgraph.shortest_path.html>`_.
    """
    dist = {va: 0}
    prev = {va: None}
    modif = 1
    while modif > 0:
        modif = 0
        for (a, b), d in zip(edges, distances):
            if a in dist:
                d2 = dist[a] + d
                if b not in dist or dist[b] > d2:
                    dist[b] = d2
                    prev[b] = a
                    modif += 1
            if b in dist:
                d2 = dist[b] + d
                if a not in dist or dist[a] > d2:
                    dist[a] = d2
                    prev[a] = b
                    modif += 1
    rev = {(a, b): i for i, (a, b) in enumerate(edges)}
    rev.update({(b, a): i for i, (a, b) in enumerate(edges)})
    path = []
    v = vb
    while v is not None:
        path.append(v)
        v = prev[v]
    path.reverse()
    return [rev[a, b] for a, b in zip(path[:-1], path[1:])]

def matching_vertices(distances, algo="blossom"):
    """
    Finds the best match between vertices.

    @param      distances       result of function @bellman_distances but
                                only for odd vertices (odd = odd degree),
                                dictionary { (vi,vj) : distance}
    @param      algo            algorithm (see below)
    @return                     sequences of best matches.

    If ``algo=='hungarian'``,
    the function relies on `linear_sum_assignment
    <https://docs.scipy.org/doc/scipy-0.18.1/reference/generated/scipy.optimize.linear_sum_assignment.html>`_
    from `scipy <https://docs.scipy.org/doc/scipy-0.18.1/>`_ which uses the
    `Hungarian Algorithm <https://en.wikipedia.org/wiki/Hungarian_algorithm>`_.
    Vertex index do not have to start from zero and be continuous.
    The function will handle that particular case. However, the method
    works for a bi-partite matching and is not suitable here unless
    the algorithm is modified. Not done (yet?).

    If ``algo=='blossom'``, it uses the
    `Blossom algorithm <https://en.wikipedia.org/wiki/Blossom_algorithm>`_
    which is known to be in :math:`O(n^3)` and finds the optimal matching.

    If ``algo=='basic'``, the function sorts the distances
    by increasing order and builds new edges as they come in this list.
    It does not return an optimal solution but is much faster
    when the graph is big.
    """
    if not isinstance(distances, dict):
        raise TypeError(
            "Unexpected type for distances, this should be a dictionary.")
    from numpy import empty
    from scipy.optimize import linear_sum_assignment
    unique = set(_[0] for _ in distances) | set(_[1] for _ in distances)
    mapping = {}
    rev = {}
    for s in unique:
        n = len(mapping)
        mapping[s] = n
        rev[n] = s
    mx = len(mapping)

    # Hungarian algorithm
    if algo == "hungarian":
        cost = empty((mx, mx))
        mx = abs(max(distances.values())) * len(mapping) * 10
        cost[:] = mx
        for (i, j), v in distances.items():
            if i == j:
                raise ValueError(
                    "Unreasonable case: {0} == {1}, v={2}".format(i, j, v))
            cost[mapping[i], mapping[j]] = v

        row_ind, col_ind = linear_sum_assignment(   # pylint: disable=W0632
            cost)   # pylint: disable=W0632
        pairs = [(rev[i], rev[j]) for i, j in zip(row_ind, col_ind)]
        for a, b in pairs:
            if a == b:
                raise ValueError("Issue with one pair a == b.\n{0}".format(
                    "\n".join(str(_) for _ in pairs)))

        # we remove duplicates
        done = set()
        final = []
        for a, b in pairs:
            if (a, b) not in done:
                final.append((a, b))
                done.add((b, a))
        if len(final) * 2 != len(pairs):
            mes = "final={0}\n{3}={1}\ncost\n{2}".format(
                final, pairs, cost, algo)
            raise ValueError(
                "Did you use the tweak? The matching should be symmetric.\n" + mes)
        return final
    elif algo == "basic":
        # we sort pair by increasing order
        dists = [(v, i, j) for (i, j), v in distances.items()]
        dists.sort()

        # the graph is fully connected so we don't need to check
        # if selecting an edge leads to an impossible solution
        vdone = set()
        pairs = []
        for v, i, j in dists:
            if i not in vdone and j not in vdone:
                pairs.append((i, j))
                vdone.add(i)
                vdone.add(j)
        return pairs

    elif algo == "blossom":
        from .blossom import Vertex, StructureUpToDate, TreeStructureChanged, MaximumDualReached, INF

        vertices = {}
        for v in mapping:
            vertices[v] = Vertex(v)
        for (i, j), v in distances.items():
            f = fractions.Fraction(v)
            vertices[i].add_edge_to(vertices[j], f)

        def update_tree_structures(roots):
            try:
                while True:
                    try:
                        for root in roots:
                            root.alter_tree(roots)
                        raise StructureUpToDate()
                    except TreeStructureChanged:
                        pass
            except StructureUpToDate:
                pass

        def get_max_delta(roots):
            if len(roots) == 0:
                raise MaximumDualReached()
            delta = INF
            for root in roots:
                delta = min(delta, root.get_max_delta())
            assert delta >= 0
            return delta

        roots = set(vertices.values())

        try:
            while True:
                delta = get_max_delta(roots)
                for root in roots:
                    root.adjust_charge(delta)
                update_tree_structures(roots)
        except MaximumDualReached:
            # done
            pass

        M = set()
        for v in vertices.values():
            M.update(e for e in v.edges if e.selected)

        pairs = [e.extremities for e in M]
        return pairs
    else:
        raise NotImplementedError("Not recognized: {0}".format(algo))

def best_euler_path(edges, vertices, distances, edges_index=None, algo="blossom", fLOG=None):
    """
    Computes the final solution for the Eulerian path.

    @param      edges           edges (list of tuple)
    @param      vertices        location of the vertices (not needed if distances are filled)
    @param      distances       distances for each edge (list of distance)
    @param      edges_index     list of indices of edges (if None --> range(len(edges))
    @param      algo            algorithm to use to computes the matching
                                (see @see fn matching_vertices)
    @param      fLOG            logging function
    @return                     list of edges as tuple, list of edges as indices, distance
    """
    if fLOG:
        fLOG("[best_euler_path] distance_vertices #edges={0}".format(
            len(edges)))
    distances = distance_vertices(edges, vertices, distances)
    if fLOG:
        fLOG("[best_euler_path] bellman_distances")
    dist = bellman_distances(edges, distances, fLOG=fLOG)
    if fLOG:
        fLOG("[best_euler_path] degrees and distances between odd vertices")
    degrees = compute_degrees(edges)
    odd = {k: v for k, v in degrees.items() if v % 2 != 0}
    odd_dist = {k: v for k, v in dist.items() if k[0] in odd and k[1] in odd}
    if fLOG:
        fLOG("[best_euler_path] matching #odd={0}, #odd_dist={1}".format(
            len(odd), len(odd_dist)))
    pairs = matching_vertices(odd_dist, algo=algo)
    if fLOG:
        fLOG("[best_euler_path] build solution")
    solution = list(edges)
    for va, vb in pairs:
        dij = dijkstra_path(edges, distances, va, vb)
        solution.extend([edges[e] for e in dij])
    if fLOG:
        fLOG("[best_euler_path] order edges to get the path #edges={0}".format(
            len(solution)))
    if edges_index is None:
        edges_index = list(range(len(edges)))
    mapping = {}
    rev = {}
    for edge, index in zip(edges, edges_index):
        mapping[edge] = index
        rev[index] = edge
    sol_index = [mapping[e] for e in solution]
    euler_index = euler_path(edges_index, edges, sol_index)
    euler = [rev[i] for i in euler_index]
    if fLOG:
        fLOG("[best_euler_path] done.")
    d = distance_solution(edges_index, edges, distances, sol_index)
    return euler, euler_index, d


<reponame>Kyubyong/cross_vc
# -*- coding: utf-8 -*-
#/usr/bin/python2
'''
By <NAME>. <EMAIL>.
https://www.github.com/kyubyong/cross_vc
'''

from __future__ import print_function

from hparams import Hyperparams as hp
import librosa
import numpy as np
from scipy import signal
import copy
import tensorflow as tf

def get_mfcc_and_mag(fpath, trim=True):
    '''Returns normalized MFCCs and linear magnitude from `fpath`.
    Args:
      fpath: A string. The full path of a sound file.

    Returns:
      mfcc: A 2d array of shape (T, n_mfccs)
      mag: A 2d array of shape (T, 1+n_fft/2)
    '''
    # Loading sound file
    y, sr = librosa.load(fpath, sr=hp.sr)

    # Trimming
    if trim: y, _ = librosa.effects.trim(y)

    # Preemphasis
    y = np.append(y[0], y[1:] - hp.preemphasis * y[:-1])

    # stft
    linear = librosa.stft(y=y,
                          n_fft=hp.n_fft,
                          hop_length=int(hp.sr * hp.frame_shift),
                          win_length=int(hp.sr * hp.frame_length))

    # magnitude spectrogram
    mag = np.abs(linear)  # (1+n_fft//2, T)

    # mel spectrogram
    mel_basis = librosa.filters.mel(hp.sr, hp.n_fft, hp.n_mels)  # (n_mels, 1+n_fft//2)
    mel = np.dot(mel_basis, mag)  # (n_mels, t)

    # to decibel
    mel = 20 * np.log10(np.maximum(1e-5, mel))
    mag = 20 * np.log10(np.maximum(1e-5, mag))
    mag = np.clip((mag - hp.ref_db + hp.max_db) / hp.max_db, 1e-8, 1)

    # Get MFCCs
    mfcc = np.dot(librosa.filters.dct(hp.n_mfccs, mel.shape[0]), mel) # (n_mfcc, t)

    # Transpose for convenience
    mfcc = mfcc.T.astype(np.float32)  # (t, n_mfcc)
    mag = mag.T.astype(np.float32)
    return mfcc, mag

def spectrogram2wav(mag):
    '''# Generate wave file from magnitude spectrogram'''
    # transpose
    mag = mag.T

    # de-noramlize
    mag = (np.clip(mag, 0, 1) * hp.max_db) - hp.max_db + hp.ref_db

    # to amplitude
    mag = np.power(10.0, mag * 0.05)

    # wav reconstruction
    wav = griffin_lim(mag**hp.sharpening_factor)

    # de-preemphasis
    wav = signal.lfilter([1], [1, -hp.preemphasis], wav)

    # trim
    wav, _ = librosa.effects.trim(wav)

    return wav.astype(np.float32)

def griffin_lim(spectrogram):
    '''Applies Griffin-Lim's raw.
    '''
    hop_length = int(hp.sr*hp.frame_shift)
    win_length = int(hp.sr*hp.frame_length)
    X_best = copy.deepcopy(spectrogram)
    for i in range(hp.n_iter):
        X_t = invert_spectrogram(X_best)
        est = librosa.stft(X_t, hp.n_fft, hop_length, win_length=win_length)
        phase = est / np.maximum(1e-8, np.abs(est))
        X_best = spectrogram * phase
    X_t = invert_spectrogram(X_best)
    y = np.real(X_t)

    return y

def invert_spectrogram(spectrogram):
    '''
    spectrogram: [f, t]
    '''
    hop_length = int(hp.sr * hp.frame_shift)
    win_length = int(hp.sr * hp.frame_length)
    return librosa.istft(spectrogram, hop_length, win_length=win_length, window="hann")

def learning_rate_decay(init_lr, global_step, warmup_steps=4000.):
    '''Noam scheme from tensor2tensor'''
    step = tf.to_float(global_step + 1)
    return init_lr * warmup_steps**0.5 * tf.minimum(step * warmup_steps**-1.5, step**-0.5)


<gh_stars>0
## Calculation of acoustic mode contributions to internal energy (E*)

import math
import numpy as np
from scipy.integrate import quad

#import params.py for all input parameters
from params import *

##########################################################################

E_a = []

for i in ATEMP:
    X_array = [] # Nondimensionalized branches
    for j in U_array:
        X_array.append((132.32*j/(Vol**(1./3.)))*CONV/i)

    acoustic = []

    for z in X_array:
        def f_acoustic(x):
            return ((math.asin(x/z)**2.)*x)/(math.sqrt(z**2.-x**2.)*(math.exp(x)-1.))

        acoustic_quad, error = quad(f_acoustic, 0., z)
        acoustic.append(acoustic_quad)

    acoustic_E = (3.*AVO*BOLTZ*i*(acoustic[0]+acoustic[1]+acoustic[2])/(Natoms*Z))*(2./math.pi)*(2./math.pi)*(2./math.pi)
    E_a.append(acoustic_E)

##print(E_a)


# -*- coding: utf-8 -*-
"""
Created on Mon Feb 22 15:12:49 2021

@author: Administrator
"""

import math
import numpy as np
import pandas as pd
from scipy.io import arff
from sklearn.linear_model import LogisticRegression
from sklearn import svm, datasets
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
import sklearn.preprocessing as skpre
from matplotlib.colors import ListedColormap
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier

#DATA PROCESSING

data, meta = arff.loadarff('C:/Users/Administrator/Desktop/4year.arff')
df=pd.DataFrame(data)
df['bankruptcy'] = (df['class']==b'1')
df.drop(columns=['class'], inplace=True)
df.columns = ['X{0:02d}'.format(k) for k in range(1,65)] + ['bankruptcy']
df.fillna(df.mean(), inplace=True)
X_imp = df.values
X, y = X_imp[:, :-1], X_imp[:, -1]
X_train, X_test, y_train, y_test =\
    train_test_split(X, y, test_size=0.2, random_state=0, stratify=y)
scaler = skpre.StandardScaler()
X_train_std = scaler.fit_transform(X_train)
X_test_std = scaler.transform(X_test)

y_train=y_train*1
y_test=y_test*1
y_train=y_train.astype(int)
y_test=y_test.astype(int)
lr = LogisticRegression(penalty='l1',C=0.01, solver='liblinear')
lr.fit(X_train_std, y_train.astype(int))
lr.coef_[lr.coef_!=0].shape
X_train_std=X_train_std[:,lr.coef_[0]!=0]
X_test_std=X_test_std[:,lr.coef_[0]!=0]

#lR MODEL
lr = LogisticRegression(penalty='l1')
lr.fit(X_train_std, y_train.astype(int))
print('LR Training accuracy:', lr.score(X_train_std, y_train.astype(int)))
print('LR Test accuracy:', lr.score(X_test_std, y_test.astype(int)))

#models = (svm.SVC(kernel='linear', C=C),
#          svm.LinearSVC(C=C, max_iter=10000),
#          svm.SVC(kernel='rbf', gamma=0.7, C=C),
#          svm.SVC(kernel='poly', degree=3, gamma='auto', C=C))
##for clf in models:

clf=svm.SVC(C=1,kernel='rbf',gamma=10)
clf.fit(X_train_std,y_train)
print('SVM Training accuracy:', clf.score(X_train_std, y_train))
print('SVM Test accuracy:', clf.score(X_test_std, y_test))

tree= DecisionTreeClassifier(criterion='gini',
                                    max_depth=4,
                                    random_state=1)
tree.fit(X_train, y_train)

#############Codes below can't run, but I can't find the reason
#############Codes below can't run, but I can't find the reason
#############Codes below can't run, but I can't find the reason
#############print('DecTree Training accuracy:', tree.score(X_train_std,y_train))
#############print('DecTree Test accuracy:', tree.score(X_test_std,y_test))

def plot_decision_regions(X, y, classifier, test_idx=None, resolution=0.02):

    # setup marker generator and color map
    markers = ('s', 'x', 'o', '^', 'v')
    colors = ('red', 'blue', 'lightgreen', 'gray', 'cyan')
    cmap = ListedColormap(colors[:len(np.unique(y))])

    # plot the decision surface
    x1_min, x1_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    x2_min, x2_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    xx1, xx2 = np.meshgrid(np.arange(x1_min, x1_max, resolution),
                           np.arange(x2_min, x2_max, resolution))
    Z = classifier.predict(np.array([xx1.ravel(), xx2.ravel()]).T)
    Z = Z.reshape(xx1.shape)
    plt.contourf(xx1, xx2, Z, alpha=0.3, cmap=cmap)
    plt.xlim(xx1.min(), xx1.max())
    plt.ylim(xx2.min(), xx2.max())

    for idx, cl in enumerate(np.unique(y)):
        plt.scatter(x=X[y == cl, 0],
                    y=X[y == cl, 1],
                    alpha=0.8,
                    c=colors[idx],
                    marker=markers[idx],
                    label=cl,
                    edgecolor='black')

    # highlight test samples
    if test_idx:
        # plot all samples
        X_test, y_test = X[test_idx, :], y[test_idx]

        plt.scatter(X_test[:, 0],
                    X_test[:, 1],
                    c='',
                    edgecolor='black',
                    alpha=1.0,
                    linewidth=1,
                    marker='o',
                    s=100,
                    label='test set')

X_combined_std = np.vstack((X_train_std, X_test_std))
y_combined = np.hstack((y_train, y_test))

plot_decision_regions(X=X_combined_std, y=y_combined,
                      classifier=lr, test_idx=range(0, 50))
plt.xlabel('ratio1 [standardized]')
plt.ylabel('ratio2[standardized]')
plt.legend(loc='upper left')

plt.tight_layout()

plt.show()

#############Codes below can't run, but I can't find the reason
#############Codes below can't run, but I can't find the reason
#############Codes below can't run, but I can't find the reason

#plot_decision_regions(X_combined_std, y_combined,
#                      classifier=svm.SVC, test_idx=range(105, 150))
#plt.xlabel('petal length [standardized]')
#plt.ylabel('petal width [standardized]')
#plt.legend(loc='upper left')
#plt.tight_layout()
##plt.savefig('images/03_15.png', dpi=300)
#plt.show()
#
#plt.tight_layout()
##plt.savefig('images/03_01.png', dpi=300)
#plt.show()
#
#X_combined = np.vstack((X_train_std, X_test_std))
#y_combined = np.hstack((y_train, y_test))
#plot_decision_regions(X_combined, y_combined,
#                      classifier=tree_model,
#                      test_idx=range(105, 150))
#
#plt.xlabel('ratio1')
#plt.ylabel('ratio2')
#plt.legend(loc='upper left')
#plt.tight_layout()
##plt.savefig('images/03_20.png', dpi=300)
#plt.show()
##


import os
import math
import torch
import random
import scipy as sp
import scipy.stats
import numpy as np
import torch.nn as nn
from PIL import Image
from alisuretool.Tools import Tools
from torch.optim.lr_scheduler import StepLR
import torchvision.transforms as transforms
from torch.utils.data.sampler import Sampler
from torch.utils.data import DataLoader, Dataset

##############################################################################################################

class ClassBalancedSampler(Sampler):
    """ Samples 'num_inst' examples each from 'num_cl' pools of examples of size 'num_per_class' """

    def __init__(self, num_per_class, num_cl, num_inst, shuffle=True):
        self.num_per_class = num_per_class
        self.num_cl = num_cl
        self.num_inst = num_inst
        self.shuffle = shuffle
        pass

    def __iter__(self):
        # return a single list of indices, assuming that items will be grouped by class
        if self.shuffle:
            batch = [[i + j * self.num_inst for i in torch.randperm(self.num_inst)[:self.num_per_class]] for j in
                     range(self.num_cl)]
        else:
            batch = [[i + j * self.num_inst for i in range(self.num_inst)[:self.num_per_class]] for j in
                     range(self.num_cl)]
        batch = [item for sublist in batch for item in sublist]

        if self.shuffle:
            random.shuffle(batch)
        return iter(batch)

    def __len__(self):
        return 1

    pass

class ClassBalancedSamplerTest(Sampler):
    """ Samples 'num_inst' examples each from 'num_cl' pools of examples of size 'num_per_class' """

    def __init__(self, num_cl, num_inst, shuffle=True):
        self.num_cl = num_cl
        self.num_inst = num_inst
        self.shuffle = shuffle
        pass

    def __iter__(self):
        # return a single list of indices, assuming that items will be grouped by class
        if self.shuffle:
            batches = [[i + j * self.num_inst for i in torch.randperm(self.num_inst)] for j in range(self.num_cl)]
        else:
            batches = [[i + j * self.num_inst for i in range(self.num_inst)] for j in range(self.num_cl)]
        batches = [[batches[j][i] for j in range(self.num_cl)] for i in range(self.num_inst)]

        if self.shuffle:
            random.shuffle(batches)
            for sublist in batches:
                random.shuffle(sublist)
        batches = [item for sublist in batches for item in sublist]
        return iter(batches)

    def __len__(self):
        return 1

    pass

class MiniImageNetTask(object):

    def __init__(self, character_folders, num_classes, train_num, test_num):
        self.character_folders = character_folders
        self.num_classes = num_classes
        self.train_num = train_num
        self.test_num = test_num

        class_folders = random.sample(self.character_folders, self.num_classes)
        labels = dict(zip(class_folders, np.array(range(len(class_folders)))))

        samples = dict()
        self.train_roots = []
        self.test_roots = []
        for c in class_folders:
            temp = [os.path.join(c, x) for x in os.listdir(c)]
            samples[c] = random.sample(temp, len(temp))
            random.shuffle(samples[c])

            self.train_roots += samples[c][:train_num]
            self.test_roots += samples[c][train_num:train_num + test_num]
            pass

        self.train_labels = [labels[os.path.split(x)[0]] for x in self.train_roots]
        self.test_labels = [labels[os.path.split(x)[0]] for x in self.test_roots]
        pass

    pass

class MiniImageNet(Dataset):

    def __init__(self, task, split='train', transform=None, target_transform=None):
        self.transform = transform
        self.target_transform = target_transform
        self.task = task
        self.split = split
        self.image_roots = self.task.train_roots if self.split == 'train' else self.task.test_roots
        self.labels = self.task.train_labels if self.split == 'train' else self.task.test_labels
        pass

    def __len__(self):
        return len(self.image_roots)

    def __getitem__(self, idx):
        image_root = self.image_roots[idx]
        image = Image.open(image_root)
        image = image.convert('RGB')
        if self.transform is not None:
            image = self.transform(image)
        label = self.labels[idx]
        if self.target_transform is not None:
            label = self.target_transform(label)
        return image, label

    @staticmethod
    def mini_imagenet_folders(train_folder, val_folder, test_folder):
        metatrain_folders = [os.path.join(train_folder, label) for label in os.listdir(train_folder)
                             if os.path.isdir(os.path.join(train_folder, label))]
        metaval_folders = [os.path.join(val_folder, label) for label in os.listdir(val_folder)
                            if os.path.isdir(os.path.join(val_folder, label))]
        metatest_folders = [os.path.join(test_folder, label) for label in os.listdir(test_folder)
                            if os.path.isdir(os.path.join(test_folder, label))]
        random.seed(1)
        random.shuffle(metatrain_folders)
        random.shuffle(metaval_folders)
        random.shuffle(metatest_folders)
        return metatrain_folders, metaval_folders, metatest_folders

    @staticmethod
    def get_mini_imagenet_data_loader(task, num_per_class=1, split='train', shuffle=False):
        normalize = transforms.Normalize(mean=[0.92206, 0.92206, 0.92206], std=[0.08426, 0.08426, 0.08426])
        dataset = MiniImageNet(task, split=split, transform=transforms.Compose([transforms.ToTensor(), normalize]))

        if split == 'train':
            sampler = ClassBalancedSampler(num_per_class, task.num_classes, task.train_num, shuffle=shuffle)
        elif split == "val":
            sampler = ClassBalancedSampler(num_per_class, task.num_classes, task.test_num, shuffle=shuffle)
        else:  # test
            sampler = ClassBalancedSamplerTest(task.num_classes, task.test_num, shuffle=shuffle)
            pass

        return DataLoader(dataset, batch_size=num_per_class * task.num_classes, sampler=sampler)

    pass

##############################################################################################################

class CNNEncoder(nn.Module):

    def __init__(self):
        super().__init__()
        self.layer1 = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=0),
                                    nn.BatchNorm2d(64, momentum=1, affine=True), nn.ReLU(), nn.MaxPool2d(2))
        self.layer2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, padding=0),
                                    nn.BatchNorm2d(64, momentum=1, affine=True), nn.ReLU(), nn.MaxPool2d(2))
        self.layer3 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, padding=1),
                                    nn.BatchNorm2d(64, momentum=1, affine=True), nn.ReLU())
        self.layer4 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, padding=1),
                                    nn.BatchNorm2d(64, momentum=1, affine=True), nn.ReLU())
        pass

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        # out = out.view(out.size(0),-1)
        return out  # 64

    pass

class RelationNetwork(nn.Module):
    """docstring for RelationNetwork"""

    def __init__(self, input_size, hidden_size):
        super(RelationNetwork, self).__init__()
        self.layer1 = nn.Sequential(nn.Conv2d(128, 64, kernel_size=3, padding=0),
                                    nn.BatchNorm2d(64, momentum=1, affine=True), nn.ReLU(), nn.MaxPool2d(2))
        self.layer2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, padding=0),
                                    nn.BatchNorm2d(64, momentum=1, affine=True), nn.ReLU(), nn.MaxPool2d(2))
        self.fc1 = nn.Linear(input_size * 3 * 3, hidden_size)
        self.fc2 = nn.Linear(hidden_size, 1)
        pass

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.view(out.size(0), -1)
        out = torch.relu(self.fc1(out))
        out = torch.sigmoid(self.fc2(out))
        return out

    pass

##############################################################################################################

class Runner(object):

    def __init__(self):
        self.feature_dim = 64
        self.relation_dim = 8
        self.hidden_unit = 10

        self.class_num = 5
        self.sample_num_per_class = 1
        self.batch_num_per_class = 15

        self.train_episode = 1000  # 500000
        self.val_episode = 600
        self.test_avg_num = 10
        self.test_episode = 600

        self.learning_rate = 0.001

        self.print_freq = 100
        self.val_freq = 500  # 5000

        self.feature_encoder_dir = Tools.new_dir("../models/miniimagenet_feature_encoder_{}way_{}shot.pkl".format(
            self.class_num, self.sample_num_per_class))
        self.relation_network_dir = Tools.new_dir("../models/miniimagenet_relation_network_{}way_{}shot.pkl".format(
            self.class_num, self.sample_num_per_class))

        # data
        self.metatrain_folders, self.metaval_folders, self.metatest_folders = MiniImageNet.mini_imagenet_folders(
            train_folder='/mnt/4T/Data/miniImagenet/train',
            val_folder='/mnt/4T/Data/miniImagenet/val', test_folder='/mnt/4T/Data/miniImagenet/test')

        # model
        (self.feature_encoder, self.relation_network, self.feature_encoder_scheduler,
         self.relation_network_scheduler, self.feature_encoder_optim, self.relation_network_optim) = self._model()

        self.loss = self._loss()
        pass

    @staticmethod
    def _weights_init(m):
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            m.weight.data.normal_(0, math.sqrt(2. / n))
            if m.bias is not None:
                m.bias.data.zero_()
        elif classname.find('BatchNorm') != -1:
            m.weight.data.fill_(1)
            m.bias.data.zero_()
        elif classname.find('Linear') != -1:
            n = m.weight.size(1)
            m.weight.data.normal_(0, 0.01)
            m.bias.data = torch.ones(m.bias.data.size())
            pass
        pass

    def _model(self):
        Tools.print("Init neural networks")

        feature_encoder = CNNEncoder()
        relation_network = RelationNetwork(self.feature_dim, self.relation_dim)

        feature_encoder.apply(self._weights_init)
        relation_network.apply(self._weights_init)

        feature_encoder.cuda()
        relation_network.cuda()

        feature_encoder_optim = torch.optim.Adam(feature_encoder.parameters(), lr=self.learning_rate)
        feature_encoder_scheduler = StepLR(feature_encoder_optim, step_size=100000, gamma=0.5)
        relation_network_optim = torch.optim.Adam(relation_network.parameters(), lr=self.learning_rate)
        relation_network_scheduler = StepLR(relation_network_optim, step_size=100000, gamma=0.5)

        return (feature_encoder, relation_network, feature_encoder_scheduler,
                relation_network_scheduler, feature_encoder_optim, relation_network_optim)

    @staticmethod
    def _loss():
        mse = nn.MSELoss().cuda()
        return mse

    def _val(self, meta_folders, split, episode):

        def mean_confidence_interval(data, confidence=0.95):
            a = 1.0 * np.array(data)
            n = len(a)
            m, se = np.mean(a), scipy.stats.sem(a)
            h = se * sp.stats.t._ppf((1 + confidence) / 2., n - 1)
            return m, h

        accuracies = []
        for i in range(episode):
            total_rewards = 0
            counter = 0
            # 随机选5类，每类中取出1个作为训练样本，每类取出15个作为测试样本
            task = MiniImageNetTask(meta_folders, self.class_num, self.sample_num_per_class, self.batch_num_per_class)
            sample_dataloader = MiniImageNet.get_mini_imagenet_data_loader(task, 1, split="train", shuffle=False)
            batch_dataloader = MiniImageNet.get_mini_imagenet_data_loader(task, 3, split=split, shuffle=True)
            sample_images, sample_labels = sample_dataloader.__iter__().next()

            for batch_images, batch_labels in batch_dataloader:
                ###########################################################################
                # calculate features
                relations = self._compare_fsl(sample_images, batch_images)
                ###########################################################################

                _, predict_labels = torch.max(relations.data, 1)
                batch_size = batch_labels.shape[0]
                rewards = [1 if predict_labels[j].cpu() == batch_labels[j] else 0 for j in range(batch_size)]
                total_rewards += np.sum(rewards)

                counter += batch_size
                pass

            accuracy = total_rewards / 1.0 / counter
            accuracies.append(accuracy)
            pass

        accuracy, h = mean_confidence_interval(accuracies)
        return accuracy, h

    def load_model(self):
        if os.path.exists(self.feature_encoder_dir):
            self.feature_encoder.load_state_dict(torch.load(self.feature_encoder_dir))
            # self.feature_encoder.load_state_dict(torch.load(self.feature_encoder_dir, map_location='cuda:0'))
            Tools.print("load feature encoder success from {}".format(self.feature_encoder_dir))

        if os.path.exists(self.relation_network_dir):
            self.relation_network.load_state_dict(torch.load(self.relation_network_dir))
            # self.relation_network.load_state_dict(torch.load(self.relation_network_dir, map_location='cuda:0'))
            Tools.print("load relation network success from {}".format(self.relation_network_dir))
        pass

    def _compare_fsl(self, samples, batches):
        # calculate features
        sample_features = self.feature_encoder(samples.cuda())  # 5x64*19*19
        batch_features = self.feature_encoder(batches.cuda())  # 75x64*19*19
        batch_size, _, feature_width, feature_height = batch_features.shape

        # calculate relations
        sample_features_ext = sample_features.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1)
        batch_features_ext = batch_features.unsqueeze(0).repeat(
            self.sample_num_per_class * self.class_num, 1, 1, 1, 1)
        batch_features_ext = torch.transpose(batch_features_ext, 0, 1)
        relation_pairs = torch.cat((sample_features_ext, batch_features_ext),
                                   2).view(-1, self.feature_dim * 2, feature_width, feature_height)
        relations = self.relation_network(relation_pairs).view(-1, self.class_num * self.sample_num_per_class)
        return relations

    def train(self):
        Tools.print()
        Tools.print("Training...")

        last_accuracy = 0.0
        for episode in range(self.train_episode):
            # init dataset
            task = MiniImageNetTask(self.metatrain_folders, self.class_num,
                                    self.sample_num_per_class, self.batch_num_per_class)
            sample_dataloader = MiniImageNet.get_mini_imagenet_data_loader(
                task, num_per_class=self.sample_num_per_class, split="train", shuffle=False)
            batch_dataloader = MiniImageNet.get_mini_imagenet_data_loader(
                task, num_per_class=self.batch_num_per_class, split="val", shuffle=True)
            samples, sample_labels = sample_dataloader.__iter__().next()
            batches, batch_labels = batch_dataloader.__iter__().next()

            ###########################################################################
            # calculate features
            relations = self._compare_fsl(samples, batches)
            ###########################################################################

            one_hot_labels = torch.zeros(self.batch_num_per_class * self.class_num,
                                         self.class_num).scatter_(1, batch_labels.view(-1, 1), 1).cuda()
            loss = self.loss(relations, one_hot_labels)

            self.feature_encoder.zero_grad()
            self.relation_network.zero_grad()

            loss.backward()

            torch.nn.utils.clip_grad_norm_(self.feature_encoder.parameters(), 0.5)
            torch.nn.utils.clip_grad_norm_(self.relation_network.parameters(), 0.5)

            self.feature_encoder_optim.step()
            self.relation_network_optim.step()
            self.feature_encoder_scheduler.step(episode)
            self.relation_network_scheduler.step(episode)

            if (episode + 1) % self.print_freq == 0:
                Tools.print("Episode: {} loss: {}".format(episode + 1, loss.item()))
                pass

            if episode % self.val_freq == 0:
                Tools.print()
                Tools.print("Valing...")
                val_accuracy, h = self.val(is_print=True)
                if val_accuracy > last_accuracy:
                    # save networks
                    torch.save(self.feature_encoder.state_dict(), self.feature_encoder_dir)
                    torch.save(self.relation_network.state_dict(), self.relation_network_dir)
                    Tools.print("Save networks for episode: {}".format(episode))
                    last_accuracy = val_accuracy
                    pass

                pass

            pass

        pass

    def val(self, is_print=False):
        val_accuracy, h = self._val(self.metaval_folders, split="val", episode=self.val_episode)
        if is_print:
            Tools.print()
            Tools.print("Val Accuracy: {} h: {}".format(val_accuracy, h))
            pass
        return val_accuracy, h

    def test(self, is_print=True):
        Tools.print()
        Tools.print("Testing...")
        total_accuracy = 0.0
        for episode in range(self.test_avg_num):
            test_accuracy, h = self._val(self.metatest_folders, split="test", episode=self.test_episode)
            total_accuracy += test_accuracy
            Tools.print("episode={}, Test accuracy={}, H={}, Total accuracy={}".format(episode, test_accuracy,
                                                                                h, total_accuracy))
            pass

        final_accuracy = total_accuracy / self.test_avg_num
        if is_print:
            Tools.print("Final accuracy: {}".format(final_accuracy))
            pass
        return final_accuracy

    pass

##############################################################################################################

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
    runner = Runner()
    runner.load_model()
    runner.test(is_print=True)
    runner.val(is_print=True)
    # runner.train()
    runner.val(is_print=True)
    runner.test(is_print=True)